Neural Network Model Domains

DeepPumas extends the standard Pumas.jl by adding neural networks (mainly Multilayer Perceptrons - MLPs) into the NLME modeling workflow. Here we describe the basis for this extension: the underlying neural network types, practically useful functions that construct model components from said types, setting of neural network parameter priors/regularization and interfacing with the neural network domain objects outside of a PumasModel.

Constructing Neural Network Domains

The main workhorse function when dealing with PumasModel embeddable neural networks is MLPDomain which creates the neural network domain. Here we will focus on its practical usage and postpone the technical details about the specific types created by MLPDomain to Advanced Features.

DeepPumas.MLPDomain Function

MLPDomain(input, layers...; kwargs...)

Construct a, possibly regularized, MLP.

Arguments

• input::Int - the number of inputs to the neural network.

• layers... - Each additional argument after input specifies a new neural network layer. Layers can either be specified by an Int specifying the length of the layer or by a Tuple (length, [act,] [bias]) which allows you to optionally specify an activation function, act, and whether the layer should have bias parameters, bias. Such a tuple of inputs overrides the default values that are passed as keyword arguments to MLPDomain. The output layer should typically be different from the rest - the identity activation function often makes sense and bias may be undesirable.

Keyword arguments

• act::Function = tanh - the default activation function for the layers. Can be overridden by tuples in layers.... tanh has proven effective for small and shallow MLPDomains but relu and leakyrelu are the go-to activation functions for anything deeper.

• bias::Bool = true - the default specification of whether layers should have bias parameters. Can be overridden by tuples in layers....

• reg::Union{Nothing, L1, L2, Prior} = nothing - specify what, if any, regularization to use. Available options are L1, L2, a Prior, or nothing for no regularization. When the MLPDomain is embedded in a PumasModel then the regularization is equivalent to a prior on the neural network weights and will only be applied during fitting when using an algorithm that supports parameter priors, such as MAP or JointMAP.

• backend::Symbol=:flux - specify which machine learning package to use in the backend. Available options are (:flux, :simplechains, :staticflux). :flux is flexible and stable. :staticflux is much faster for small networks on a CPU but can take a long time to compile for anything too large. :simplechains is fast but should be considered experimental and is not fully featured. Defaults to :flux unless the default has been modified using set_mlp_backend.

Examples

MLPDomain(3, 5, 1)

3 inputs, a hidden layer of length 5 and an output layer of length 1. All using tanh activation, all with bias parameters, and none of the parameters regularized.

MLPDomain(3, 5, 5, (1, identity, false); reg = L1(1.))

3 inputs, two hidden layers of length 5 with tanh activation and bias parameters, an output layer with a single output with no activation function and no bias. All weights have an L1 regularization when fitted with a prior supporting fitting function such as Pumas.MAP.

MLPDomain(3, (5, tanh, true), 5, (1, identity); act = softplus, bias=false, backend=:staticflux, reg=Prior(Normal(0., 1.), output=false))

Construct a multi-layer perceptron that is equivalent to a Flux model, but backed by StaticArrays.MArrays with 3 inputs, one hidden layer of length 5 with tanh activation and bias parameters (as specified in the tuple), one hidden layer of length 5 with softplus activation and no bias parameters (as defined in act and bias), and one output layer of length 1 with no activation function (identity in the tuple) and no bias (bias kwarg). All weights except those of the output layer have a Normal(0, 1) prior when fitted with an algorithm that supports parameter priors.

target = preprocess(fpm)
nn = MLPDomain(numinputs(target), 4, (numoutputs(target), identity))

Construct an MLPDomain to match the input-output dimension of a target for the augment workflow. See numinputs, numoutputs.

source

Setting up Priors & Regularization

DeepPumas.Prior Type

Prior(distribution; input=true, output=true, bias=true)

Specify a prior parameter distribution for a neural network embedded in a PumasModel.

Arguments:

• distribution::UnivariateDistribution - the prior distribution according to which the neural network weights will be regularized

• input::Bool = true - toggle regularization of the input layer

• output::Bool = true - toggle regularization of the output layer

• bias::Bool = true - toggle regularization of bias parameters

For now, the input, output and bias specifications are only respected during the fitting of a full PumasModel, not during hyperoptimization or fitting in the augment workflow where everything is regularized.

Examples:

using Distributions

d1 = Normal(0.0, 1.0)
prior_1 = Prior(d1)

d2 = Uniform(0.0, 1.0)
prior_2 = Prior(d2; output=false, bias=false)

source

DeepPumas.L1 Type

L1(λ=1; input=true, output=true, bias=true)

Construct an L1 regularization type.

Adds a regularization of λ⋅∑ⱼ|βⱼ| where βⱼ is a regularized parameter and λ is a regularization constant. Such a regularization is proportional to having a Laplace(0., 1/λ) prior.

• λ::Number - the regularization constant.

• input::Bool - toggle regularization of the input layer.

• output::Bool - toggle regularization of the output layer.

• bias::Bool - toggle regularization of bias parameters.

The regularization of the input and output layer bias parameters follow an AND rule where they're only regularized if input==true AND bias==true.

For now, the input, output and bias specifications are only respected during the fitting of a full PumasModel, not during hyperoptimization or fitting in the augment workflow where everything is regularized.

source

DeepPumas.L2 Type

L2(λ=1; input=true, output=true, bias=true)

Construct an L2 regularization type.

Adds a regularization of λ⋅∑ⱼ(βⱼ²) where βⱼ is a regularized parameter and λ is a regularization constant. Such a regularization is proportional to having a Normal(0., 1/√(2λ)) prior.

• λ::Number - the regularization constant.

• input::Bool = true - toggle regularization of the input layer.

• output::Bool = true - toggle regularization of the output layer.

• bias::Bool = false - toggle regularization of bias parameters.

The regularization of the input and output layer bias parameters follow an AND rule where they're only regularized if input==true AND bias==true.

For now, the input, output and bias specifications are only respected during the fitting of a full PumasModel, not during hyperoptimization or fitting in the augment workflow where everything is regularized.

source

Interface

In this subsection, we discuss the basic interface with the neural network domain types outside of a PumasModel. An important point to keep in mind is that a MLPDomain call returns a domain rather than an instance of a callable neural network. Therefore, you need to use Pumas.init to initialize from the domain:

Pumas.init Function

init(d)

Return an initialized object from d, where d can be a domain, Distribution, ParamSet or Formula.

Examples:

using Pumas
using DeepPumas

init(RealDomain(;))  
0.0

init(ParamSet((;a=RealDomain(;), b=RealDomain(;))
(a=0.0, 
b=0.0,)

NN = init(MLPDomain(1, 2, 1))
DeepPumas.FluxModelParam{Flux.Chain{Tuple{Flux.Dense{typeof(tanh), Matrix{Float64}, Vector{Float64}}, Flux.Dense{typeof(tanh), Matrix{Float64}, 
Vector{Float64}}}}, Vector{Float64}}(Chain(Dense(1 => 2, tanh), Dense(2 => 1, tanh)), [-0.6472278596494233, -1.052911126562834, 0.0, 0.0, 
-0.0563053509375365, 0.5064983064035643, 0.0])

typeof(NN(1.0))
Vector{Float64}

source

Having obtained an initialized neural network via NN = init(MLPDomain(1, 2, 1)) a few other useful functions are:

Pumas.init_params Function

init_params(ml::Union{FluxDomain, FluxModel})

Get a Vector of initial parameters for the machine learning model ml.

Uses a Glorot Normal initialization of the weights and sets biases to 0. A seeded random number generator is used for reproducibility.

source

Pumas.sample_params Function

sample_params(ml::Union{FluxDomain, FluxModel}; rng)

Get a vector of random initial parameters for the machine learning model ml.

Uses a Glorot Normal initialization of the weights and sets biases to 0. The optional rng keyword argument lets you pass a random number generator to be used.

See also init_params.

source

It is possible to fit and hyper-optimize NN by using functions fit and ho with optimization algorithms from a preferred backend - we present that as a part of a more involved workflow, see The Augment Workflow

Advanced Features

Here we present the types used by MLPDomain and on which the PumasModel embeddable neural networks are based and may be of interest to power users.

DeepPumas.FluxDomain Type

FluxDomain(model::FluxModel, prior::Prior)

Create a PumasModel-embeddable Flux.jl model with a prior specification for parameter regularisation. Currently only MLPs (DeepPumas.FluxMLP) and ResNets (DeepPumas.FluxStaticResNet) are supported.

Primarily used in the @param block of a @model in order to embed the model and make it visible/usable in the @pre block (from which it can be passed on further).

See the Flux.jl documentation for information about constructing models. See Prior for information about the prior specification.

source

Constructing models via FluxDomain can be useful if you want to include more elaborate layers (e.g. Flux.BatchNorm or Flux.Conv, see Built-in Layer Types) that are not used in MLPDomain, especially to use in The Augment Workflow.