Advanced usage
Backends
Flux.jl and Lux.jl are the primarily supported backends. These frameworks differ in a key way: Flux stores trainable parameters and states inside the network object, while Lux represents them explicitly as separate objects. Flux's stateful, object-oriented style will feel familiar to PyTorch users, while Lux's explicit, functional style will feel familiar to JAX/Flax users.
Despite these differences, the high-level API of NeuralEstimators.jl is largely consistent across backends. A minimal working example illustrating the typical workflows:
using NeuralEstimators, Flux
# Functions to sample from the model
d, n = 2, 100 # dimension of θ and number of replicates
sampler(K) = NamedMatrix(μ = randn(K), σ = rand(K))
simulator(θ::AbstractVector) = θ["μ"] .+ θ["σ"] .* sort(randn(n))
simulator(θ::AbstractMatrix) = reduce(hcat, map(simulator, eachcol(θ)))
# Neural network mapping n inputs into d outputs
network = Flux.Chain(Flux.Dense(n, 64, gelu), Flux.Dense(64, 64, gelu), Flux.Dense(64, d))
# Initialise a neural estimator
estimator = PointEstimator(network, d; num_summaries = d)
# Train the estimator
estimator = train(estimator, sampler, simulator)
# Assess the estimator
θ_test = sampler(250)
Z_test = simulator(θ_test);
assessment = assess(estimator, θ_test, Z_test)
bias(assessment)
rmse(assessment)
# Apply to observed data
θ = sampler(1) # ground truth (not known in practice)
Z = simulator(θ); # stand-in for real observations
infer(estimator, Z) # point estimateusing NeuralEstimators, Lux, Zygote
# Functions to sample from the model
d, n = 2, 100 # dimension of θ and number of replicates
sampler(K) = NamedMatrix(μ = randn(K), σ = rand(K))
simulator(θ::AbstractVector) = θ["μ"] .+ θ["σ"] .* sort(randn(n))
simulator(θ::AbstractMatrix) = reduce(hcat, map(simulator, eachcol(θ)))
# Neural network mapping n inputs into d outputs
network = Lux.Chain(Lux.Dense(n, 64, gelu), Lux.Dense(64, 64, gelu), Lux.Dense(64, d))
# Initialise a neural estimator
estimator = PointEstimator(network, d; num_summaries = d)
# Train the estimator
estimator = train(estimator, sampler, simulator)
# Assess the estimator
θ_test = sampler(250)
Z_test = simulator(θ_test);
assessment = assess(estimator, θ_test, Z_test)
bias(assessment)
rmse(assessment)
# Apply to observed data
θ = sampler(1) # ground truth (not known in practice)
Z = simulator(θ); # stand-in for real observations
infer(estimator, Z) # point estimateusing NeuralEstimators, Lux, Zygote, Random, Optimisers
# Functions to sample from the model
d, n = 2, 100 # dimension of θ and number of replicates
sampler(K) = NamedMatrix(μ = randn(K), σ = rand(K))
simulator(θ::AbstractVector) = θ["μ"] .+ θ["σ"] .* sort(randn(n))
simulator(θ::AbstractMatrix) = reduce(hcat, map(simulator, eachcol(θ)))
# Neural network mapping n inputs into d outputs
network = Lux.Chain(Lux.Dense(n, 64, gelu), Lux.Dense(64, 64, gelu), Lux.Dense(64, d))
# Initialise a neural estimator
estimator = PointEstimator(network, d; num_summaries = d)
# Initialise the parameters and states
rng = Random.default_rng()
ps, st = Lux.setup(rng, estimator)
# Train the estimator
optimiser = Adam(5e-4)
trainstate = Lux.Training.TrainState(estimator, ps, st, optimiser)
trainstate = train(trainstate, sampler, simulator)
ps = trainstate.parameters
st = trainstate.states
# Assess the estimator
θ_test = sampler(250)
Z_test = simulator(θ_test);
assessment = assess(estimator, θ_test, Z_test, ps, st)
bias(assessment)
rmse(assessment)
# Apply to observed data
θ = sampler(1) # ground truth (not known in practice)
Z = simulator(θ); # stand-in for real observations
infer(estimator, Z, ps, st) # point estimatePerformance tips with Lux.jl
Consider loading the optional dependencies for improved performance on CPUs.
If you plan to use the GPU via both CUDA and XLA/Reactant in the same session, ensure that CUDA.jl/cuDNN.jl are loaded before Reactant.jl:
using CUDA, cuDNN
using Reactant
Reactant.set_default_backend("gpu")For the most computationally efficient setup, use XLA/Reactant.jl during training by passing device = reactant_device() to train.
GPU acceleration
To improve computational efficiency, various GPU backends are supported. Once the relevant package is loaded and a compatible GPU is available, it will be used automatically:
using CUDA, cuDNNusing AMDGPUusing Metalusing oneAPISaving and loading estimators
The recommended approach is to let train save the neural network for you, by providing a savepath, and then to load it with loadestimator. Since the saved files store the parameters (and states) of the neural networks and not their architecture, loading involves constructing an estimator with the same architecture and then loading the saved parameters into it:
# Training: saves best_estimator.bson (and final_estimator.bson) to the given savepath
estimator = train(estimator, θ_train, θ_val, Z_train, Z_val, savepath = "path/to/folder")
# Loading (possibly in a later session): construct the same architecture, then load
estimator = PointEstimator(...) # same architecture as above
estimator = loadestimator(estimator, "path/to/folder")By default, the parameters corresponding to the best validation risk are loaded, that is, those of the estimator that was returned by train(); pass best = false to load the parameters from the final epoch instead.
Note that the same deep-learning backend must be loaded as when training. Note also that, when the estimator contains Lux networks, loadestimator() returns a LuxEstimator (which stores the neural-network parameters and states) irrespective of whether the estimator passed to it was wrapped in one, so the returned value should always be assigned, as above.
Alternatively, neural estimators can be saved and loaded manually, in the same way as regular Flux/Lux models (see the Flux documentation). For example, to save and load the model state of a Flux-based neural estimator:
using Flux
using BSON: @save, @load
# Save
model_state = Flux.state(estimator)
@save "estimator.bson" model_state
# Load (initialise an estimator with the same architecture, then load the state)
@load "estimator.bson" model_state
Flux.loadmodel!(estimator, model_state)For Lux users, we save the parameters/states directly:
using Lux
using BSON: @save, @load
# Save
@save "estimator.bson" parameters=estimator.ps states=estimator.st
# Load (initialise an estimator with the same architecture, then load the parameters/states)
@load "estimator.bson" parameters states
estimator = LuxEstimator(estimator.estimator, parameters, states)It is also straightforward to save the entire estimator including its architecture (see here for Flux), though saving the model state as above is recommended for long-term storage.
On-the-fly and just-in-time simulation
When data simulation is (relatively) computationally inexpensive, the training data can be simulated continuously during training, a technique known as "simulation-on-the-fly". This strategy prevents overfitting and facilitates the use of larger networks that are prone to overfitting when the training data are fixed. Further, it allows for data to be simulated "just-in-time", in the sense that data can be simulated in small batches, used to train the neural estimator, and then immediately removed from memory.
One may also regularly refresh the set of parameters (i.e., inferential targets) used during training, and doing so leads to similar benefits. However, fixing the parameters allows computationally expensive terms, such as Cholesky factors when working with Gaussian process models, to be reused throughout training, which can substantially reduce the training time for some models. Hybrid approaches are also possible, whereby the parameters (and possibly the data) are held fixed for several epochs (i.e., several passes through the training set when performing stochastic gradient descent) before being refreshed.
The above strategies are facilitated with various methods of train() and through user-defined subtypes of AbstractParameterSet.
Feature scaling
It is important to ensure that the data passed through the neural network are on a reasonable numerical scale, since values with very large absolute value can lead to numerical instability during training (e.g., exploding gradients).
A relatively simply way to achieve this is by including a transformation in the first layer of the neural network. For example, if the data have positive support, one could define the neural network with the first layer applying a log transformation:
network = Chain(z -> log.(1 + z), ...)If the data are not strictly positive, one may consider the following signed transformation:
network = Chain(z -> sign.(z) .* log.(1 .+ abs.(z)), ...)A simple preprocessing layer or transformation pipeline such as this can make a significant difference in performance and stability. See feature scaling for further discussion and possible approaches.
Regularisation
The term regularisation refers to a variety of techniques aimed to reduce overfitting when training a neural network, primarily by discouraging complex models.
Simulation on-the-fly
When the training data and parameters are simulated dynamically (i.e., "on the fly"), overfitting is generally not a concern.
One popular regularisation technique is known as dropout, implemented with Dropout (Flux/Lux). Dropout involves temporarily dropping ("turning off") a randomly selected set of neurons (along with their connections) at each iteration of the training stage, which results in a computationally-efficient form of model (neural-network) averaging (Srivastava et al., 2014).
Another class of regularisation techniques involve modifying the loss function. For instance, L₁ regularisation (sometimes called lasso regression) adds to the loss a penalty based on the absolute value of the neural-network parameters. Similarly, L₂ regularisation (sometimes called ridge regression) adds to the loss a penalty based on the square of the neural-network parameters. Note that these penalty terms are not functions of the data or of the statistical-model parameters that we are trying to infer. These regularisation techniques can be implemented straightforwardly by providing a custom optimiser rule to train that includes a SignDecay object for L₁ regularisation, or a WeightDecay object for L₂ regularisation. See the Optimisers.jl and Flux.jl documentation for further details.
For illustration, the following code constructs a neural Bayes estimator using dropout and L₁ regularisation with penalty coefficient
using NeuralEstimators, Flux
# Functions to simulate data Z|μ,σ ~ N(μ, σ²) with μ ~ N(0, 1) and σ ~ U(0, 1)
d, n = 2, 100 # number of parameters and number of replicates
sampler(K) = NamedMatrix(μ = randn(K), σ = rand(K))
simulator(θ::AbstractVector, n) = θ["μ"] .+ θ["σ"] .* sort(randn(n))
simulator(θ::AbstractMatrix, n) = reduce(hcat, simulator.(eachcol(θ), n))
# Fixed training/validation sets
K = 10000
θ_train = sampler(K)
θ_val = sampler(K)
Z_train = simulator(θ_train, n)
Z_val = simulator(θ_val, n)
# Neural network with dropout layers
network = Chain(
Dense(n, 128, relu),
Dropout(0.1),
Dense(128, 128, gelu),
Dropout(0.1),
Dense(128, d)
)
# Initialise estimator
estimator = PointEstimator(network)
# Optimiser with L₁ regularisation
optimiser = OptimiserChain(SignDecay(1e-4), Adam(5e-4))
# Train the estimator
train(estimator, θ_train, θ_val, Z_train, Z_val; optimiser = optimiser)