3 Examples
We now consider examples where the data are univariate (Section 3.1), multivariate and unstructured (Section
3.2), or collected over a regular grid
(Section 3.3). These examples focus on
neural Bayes estimation using a PointEstimator:
for a list of other estimator types available in the package (e.g., for
neural posterior or ratio estimation), see here.
Before proceeding, we load the required packages: see the README
for instructions on installing Julia and the Julia packages
NeuralEstimators and Flux. If you have access
to an NVIDIA graphics processing unit (GPU), you can utilise it by
simply adding CUDA to the list of Julia packages in the
final line below:
library("NeuralEstimators")
library("JuliaConnectoR")
library("ggplot2")
library("ggpubr")
library("latex2exp")
juliaEval('using NeuralEstimators, Flux')
#> Starting Julia ...3.1 Univariate data
Here, we develop a neural Bayes estimator for \(\boldsymbol{\theta} \equiv (\mu, \sigma)'\) from data \(\boldsymbol{Z} \equiv (Z_1, \dots, Z_m)'\), where each \(Z_i \overset{\mathrm{iid}}\sim \mathrm{Gau}(\mu, \sigma^2)\).
First, we sample parameters from the prior to construct parameter sets used for training and validating the estimator. Here, we use the priors \(\mu \sim \rm{Gau}(0, 1)\) and \(\sigma \sim \rm{Gamma}(1, 1)\), and we assume that the parameters are independent a priori. The sampled parameters are stored as \(d \times K\) matrices, with \(d\) the number of parameters in the model and \(K\) the number of sampled parameter vectors.
# Sampling from the prior
# K: number of samples to draw from the prior
sampler <- function(K) {
mu <- rnorm(K)
sigma <- rgamma(K, 1)
Theta <- matrix(c(mu, sigma), byrow = TRUE, ncol = K)
return(Theta)
}
K <- 10000
theta_train <- sampler(K)
theta_val <- sampler(K/10)Next, we define the statistical model implicitly through data
simulation. The simulated data are stored as a list, where
each element of the list corresponds to a data set simulated
conditionally on one parameter vector. The storage type of each data set
depends on the multivariate structure of the data (e.g., a matrix for
unstructured multivariate data, a multidimensional array for gridded
data). In this example, each replicate \(Z_1,
\dots, Z_m\) is univariate, so the data are stored as a matrix
with \(n = 1\) row and \(m\) columns:
# Marginal simulation from the statistical model
# theta: a matrix of parameters drawn from the prior
# m: number of conditionally independent replicates for each parameter vector
simulate <- function(Theta, m) {
apply(Theta, 2, function(theta) t(rnorm(m, theta[1], theta[2])), simplify = FALSE)
}
m <- 30
Z_train <- simulate(theta_train, m)
Z_val <- simulate(theta_val, m)Let’s visualise a few elements of the training set:
mu <- theta_train[1, 1:4]
sigma <- theta_train[2, 1:4]
Z <- Z_train[1:4]
df <- Map(function(z, m, s) {
data.frame(Z = c(z), mu = m, sigma = s)
}, Z, mu, sigma)
df <- do.call(rbind, df)
df$theta <- paste0("$\\mu$ = ", round(df$mu, 3), ", $\\sigma$ = ", round(df$sigma, 3))
df$theta <- as.factor(df$theta)
levels(df$theta) <- sapply(levels(df$theta), TeX)
ggplot(df) +
geom_histogram(aes(x = Z), bins = 30, fill = "blue", alpha = 0.5, color = "black") +
facet_grid(~theta, labeller = label_parsed) +
theme_bw()
We now design our neural network.
As we are constructing a neural Bayes estimator, the neural network is a mapping \(\mathscr{Z}\to\Theta\), and the dimensionality of the neural-network output is therefore \(d \equiv \textrm{dim}(\Theta) = 2\).
Since our data are replicated, we adopt the DeepSets framework,
implemented in NeuralEstimators with the type DeepSet.
DeepSets consist of two neural networks: an inner network and an outer
network. The inner network extracts summary statistics from the data,
and its architecture depends on the multivariate structure of the data.
For unstructured data (i.e., data without spatial or temporal
correlation within each replicate), we use a multilayer perceptron
(MLP). The input dimension matches the dimensionality of each data
replicate, while the output dimension corresponds to the number of
summary statistics appropriate for the model (a common choice is \(d\)). The outer network maps the learned
summary statistics to the output space (here, the parameter space, \(\Theta\)). The outer network is always an
MLP.
Below is an example of a DeepSets architecture for neural Bayes estimation in this example. Note that many models have parameter constraints (e.g., variance and range parameters that must be strictly positive). These constraints can be incorporated in the final layer of the neural network by choosing appropriate activation functions for each parameter. Here, we enforce the constraint \(\sigma > 0\) by applying the softplus activation function in the final layer of the outer network, ensuring that all parameter estimates are valid. For some additional ways to constrain parameter estimates, see Output layers.
# Initialise the estimator
estimator <- juliaEval('
n = 1 # dimension of each data replicate (univariate)
d = 2 # dimension of the parameter vector θ
w = 128 # width of each hidden layer
# Final layer has output dimension d and enforces parameter constraints
final_layer = Parallel(
vcat,
Dense(w, 1, identity), # mean parameter: no constraints
Dense(w, 1, softplus) # standard-deviation parameter: strictly positive
)
# Construct inner and outer networks and combine into DeepSet
psi = Chain(Dense(n, w, relu), Dense(w, w, relu))
phi = Chain(Dense(w, w, relu), final_layer)
network = DeepSet(psi, phi)
# Wrap the neural network as a PointEstimator
estimator = PointEstimator(network)
')Next, we train the estimator using train(), here using
the default mean-absolute-error loss, so that the resulting neural Bayes
estimator approximates the marginal posterior medians. We use the sets
of parameters and simulated data constructed earlier; one may
alternatively use the arguments sampler and
simulator to pass functions for sampling from the prior and
simulating from the model, respectively, which facilitates the technique
known as “simulation on-the-fly”. These functions can be defined in R
(as we have done in this example) or in Julia using the package
JuliaConnectoR (see the help file for
train()).
# Train the estimator
estimator <- train(
estimator,
theta_train = theta_train,
theta_val = theta_val,
Z_train = Z_train,
Z_val = Z_val
)
#> Constructing the training set...
#> Constructing the validation set...
#> Computing the initial validation risk... Initial validation risk = 0.74764067
#> Epoch: 1 Training risk: 0.282 Validation risk: 0.199 Learning rate: 5.00E-04 Epoch time: 20.288 seconds
#> Epoch: 2 Training risk: 0.184 Validation risk: 0.166 Learning rate: 5.00E-04 Epoch time: 2.816 seconds
#> Epoch: 3 Training risk: 0.165 Validation risk: 0.155 Learning rate: 5.00E-04 Epoch time: 2.814 seconds
#> Epoch: 4 Training risk: 0.156 Validation risk: 0.15 Learning rate: 5.00E-04 Epoch time: 2.828 seconds
#> Epoch: 5 Training risk: 0.151 Validation risk: 0.153 Learning rate: 4.99E-04 Epoch time: 2.779 seconds
#> Epoch: 6 Training risk: 0.148 Validation risk: 0.14 Learning rate: 4.98E-04 Epoch time: 2.782 seconds
#> Epoch: 7 Training risk: 0.146 Validation risk: 0.14 Learning rate: 4.97E-04 Epoch time: 2.937 seconds
#> Epoch: 8 Training risk: 0.143 Validation risk: 0.138 Learning rate: 4.96E-04 Epoch time: 2.838 seconds
#> Epoch: 9 Training risk: 0.141 Validation risk: 0.135 Learning rate: 4.94E-04 Epoch time: 2.99 seconds
#> Epoch: 10 Training risk: 0.14 Validation risk: 0.145 Learning rate: 4.92E-04 Epoch time: 2.795 seconds
#> Epoch: 11 Training risk: 0.139 Validation risk: 0.135 Learning rate: 4.90E-04 Epoch time: 2.765 seconds
#> Epoch: 12 Training risk: 0.138 Validation risk: 0.135 Learning rate: 4.88E-04 Epoch time: 2.797 seconds
#> Epoch: 13 Training risk: 0.138 Validation risk: 0.134 Learning rate: 4.85E-04 Epoch time: 2.756 seconds
#> Epoch: 14 Training risk: 0.137 Validation risk: 0.133 Learning rate: 4.82E-04 Epoch time: 2.599 seconds
#> Epoch: 15 Training risk: 0.136 Validation risk: 0.134 Learning rate: 4.79E-04 Epoch time: 2.771 seconds
#> Epoch: 16 Training risk: 0.135 Validation risk: 0.137 Learning rate: 4.76E-04 Epoch time: 2.731 seconds
#> Epoch: 17 Training risk: 0.136 Validation risk: 0.132 Learning rate: 4.73E-04 Epoch time: 2.789 seconds
#> Epoch: 18 Training risk: 0.134 Validation risk: 0.131 Learning rate: 4.69E-04 Epoch time: 2.72 seconds
#> Epoch: 19 Training risk: 0.135 Validation risk: 0.132 Learning rate: 4.65E-04 Epoch time: 2.528 seconds
#> Epoch: 20 Training risk: 0.134 Validation risk: 0.131 Learning rate: 4.61E-04 Epoch time: 2.708 seconds
#> Epoch: 21 Training risk: 0.134 Validation risk: 0.128 Learning rate: 4.57E-04 Epoch time: 2.776 seconds
#> Epoch: 22 Training risk: 0.133 Validation risk: 0.131 Learning rate: 4.52E-04 Epoch time: 2.735 seconds
#> Epoch: 23 Training risk: 0.133 Validation risk: 0.129 Learning rate: 4.48E-04 Epoch time: 2.789 seconds
#> Epoch: 24 Training risk: 0.132 Validation risk: 0.131 Learning rate: 4.43E-04 Epoch time: 2.707 seconds
#> Epoch: 25 Training risk: 0.131 Validation risk: 0.127 Learning rate: 4.38E-04 Epoch time: 2.589 seconds
#> Epoch: 26 Training risk: 0.132 Validation risk: 0.127 Learning rate: 4.32E-04 Epoch time: 2.707 seconds
#> Epoch: 27 Training risk: 0.131 Validation risk: 0.13 Learning rate: 4.27E-04 Epoch time: 2.748 seconds
#> Epoch: 28 Training risk: 0.132 Validation risk: 0.127 Learning rate: 4.21E-04 Epoch time: 2.752 seconds
#> Epoch: 29 Training risk: 0.131 Validation risk: 0.127 Learning rate: 4.15E-04 Epoch time: 3.204 seconds
#> Epoch: 30 Training risk: 0.13 Validation risk: 0.13 Learning rate: 4.09E-04 Epoch time: 3.007 seconds
#> Epoch: 31 Training risk: 0.13 Validation risk: 0.126 Learning rate: 4.03E-04 Epoch time: 2.824 seconds
#> Epoch: 32 Training risk: 0.131 Validation risk: 0.125 Learning rate: 3.97E-04 Epoch time: 2.646 seconds
#> Epoch: 33 Training risk: 0.13 Validation risk: 0.128 Learning rate: 3.91E-04 Epoch time: 2.844 seconds
#> Epoch: 34 Training risk: 0.129 Validation risk: 0.127 Learning rate: 3.84E-04 Epoch time: 2.778 seconds
#> Epoch: 35 Training risk: 0.13 Validation risk: 0.126 Learning rate: 3.77E-04 Epoch time: 2.75 seconds
#> Epoch: 36 Training risk: 0.129 Validation risk: 0.126 Learning rate: 3.70E-04 Epoch time: 2.876 seconds
#> Epoch: 37 Training risk: 0.129 Validation risk: 0.126 Learning rate: 3.63E-04 Epoch time: 2.973 seconds
#> Epoch: 38 Training risk: 0.129 Validation risk: 0.127 Learning rate: 3.56E-04 Epoch time: 2.833 seconds
#> Stopping early since the validation loss has not improved in 5 epochs
#> Finished training in 124.2119135 secondsIf the argument savepath in train() is
specified, the neural-network state (e.g., the weights and biases) will
be saved during training as bson files, and the best state
(as measured by validation risk) will be saved as
best_network.bson. To load these saved states into a neural
network in later R sessions, one may use the function
loadstate(). Note that one may also manually save a trained
estimator using savestate().
Once a neural Bayes estimator has been trained, its performance should be assessed. Simulation-based (empirical) methods for evaluating the performance of a neural Bayes estimator are ideal, since simulation is already required for constructing the estimator, and because the estimator can be applied to thousands of simulated data sets at almost no computational cost.
To assess the accuracy of the resulting neural Bayes estimator, one
may use the function assess(). The resulting object can
then be passed to the functions risk(),
bias(), and rmse() to compute a Monte Carlo
approximation of the Bayes risk, the bias, and the RMSE, or passed to
the function plotestimates() to visualise the estimates
against the corresponding true values:
# Assess the estimator
theta_test <- sampler(1000)
Z_test <- simulate(theta_test, m)
assessment <- assess(estimator, theta_test, Z_test, estimator_names = "NBE")
plotestimates(assessment, parameter_labels = c("θ1" = expression(mu), "θ2" = expression(sigma)))
Once the neural Bayes estimator has been trained and assessed, it can
be applied to observed data using the function estimate(),
and non-parametric bootstrap estimates can be computed using the
function bootstrap(). Below, we use simulated data as a
surrogate for observed data:
theta <- as.matrix(c(0, 0.5)) # true parameters
Z <- simulate(theta, m) # pretend that this is observed data
estimate(estimator, Z) # point estimate from the "observed data"
#> [,1]
#> [1,] -0.04679317
#> [2,] 0.48535904
bootstrap(estimator, Z)[, 1:6] # non-parametric bootstrap estimates
#> [,1] [,2] [,3] [,4] [,5] [,6]
#> [1,] -0.1178385 0.1237657 -0.1555784 -0.1538627 0.09082093 0.005188575
#> [2,] 0.5093642 0.4116983 0.4775196 0.4497076 0.49067265 0.4815713173.2 Multivariate data
Suppose now that our data consists of \(m\) replicates of a \(d\)-dimensional multivariate distribution.
For unstructured \(d\)-dimensional data, the estimator is based on a multilayer perceptron (MLP), and each data set is stored as a two-dimensional array (a matrix), with the second dimension storing the independent replicates. That is, in this case, we store the data as a list of \(d \times m\) matrices, and the inner network of the DeepSets representation has \(d\) neurons in the input layer.
For example, consider the task of estimating \(\boldsymbol{\theta} \equiv (\mu_1, \mu_2, \sigma, \rho)'\) from data \(\boldsymbol{Z}_1, \dots, \boldsymbol{Z}_m\) that are independent and identically distributed according to the bivariate Gaussian distribution, \[ \rm{Gau}\left(\begin{bmatrix}\mu_1 \\ \mu_2\end{bmatrix}, \sigma^2\begin{bmatrix} 1 & \rho \\ \rho & 1\end{bmatrix}\right). \] Then, to construct a neural Bayes estimator for this model, one may use the following code for defining a prior, the data simulator, and the neural-network architecture:
# Sampling from the prior
# K: number of samples to draw from the prior
sampler <- function(K) {
mu1 <- rnorm(K)
mu2 <- rnorm(K)
sigma <- rgamma(K, 1)
rho <- runif(K, -1, 1)
theta <- matrix(c(mu1, mu2, sigma, rho), byrow = TRUE, ncol = K)
return(theta)
}
# Marginal simulation from the statistical model
# theta: a matrix of parameters drawn from the prior
# m: number of conditionally independent replicates for each parameter vector
simulate <- function(Theta, m) {
apply(Theta, 2, function(theta) {
mu <- c(theta[1], theta[2])
sigma <- theta[3]
rho <- theta[4]
Sigma <- sigma^2 * matrix(c(1, rho, rho, 1), 2, 2)
Z <- MASS::mvrnorm(m, mu, Sigma)
t(Z)
}, simplify = FALSE)
}
# Initialise the estimator
estimator <- juliaEval('
d = 2 # dimension of each replicate
w = 32 # number of neurons in each hidden layer
# Layer to ensure valid estimates
final_layer = Parallel(
vcat,
Dense(w, 2, identity), # mean parameters
Dense(w, 1, softplus), # variance parameter
Dense(w, 1, tanh) # correlation parameter
)
psi = Chain(Dense(d, w, relu), Dense(w, w, relu), Dense(w, w, relu))
phi = Chain(Dense(w, w, relu), Dense(w, w, relu), final_layer)
deepset = DeepSet(psi, phi)
estimator = PointEstimator(deepset)
')Let’s visualise a few realisations from the current model:
theta <- sampler(3)
Z <- simulate(theta, m = 250)
mu1 <- theta[1, 1:3]
mu2 <- theta[2, 1:3]
sigma <- theta[3, 1:3]
rho <- theta[4, 1:3]
df <- Map(function(z, m1, m2, s, r) {
data.frame(Z1 = z[1, ], Z2 = z[2, ], mu1 = m1, mu2 = m2, sigma = s, rho = r)
}, Z, mu1, mu2, sigma, rho)
df <- do.call(rbind, df)
df$theta <- paste0(
"$\\mu_1$ = ", round(df$mu1, 2), ", ",
"$\\mu_2$ = ", round(df$mu2, 2), ", ",
"$\\sigma$ = ", round(df$sigma, 2), ", ",
"$\\rho$ = ", round(df$rho, 2)
)
df$theta <- as.factor(df$theta)
levels(df$theta) <- sapply(levels(df$theta), TeX)
ggplot(df) +
geom_point(aes(x = Z1, y = Z2), alpha = 0.5) +
facet_grid(~theta, labeller = label_parsed) +
theme_bw()
The training, assessment, and application of the neural Bayes estimator then remains as in the example above:
# Construct training and validation data sets
K <- 5000
m <- 250
theta_train <- sampler(K)
theta_val <- sampler(K/10)
Z_train <- simulate(theta_train, m)
Z_val <- simulate(theta_val, m)
# Train the estimator
estimator <- train(
estimator,
theta_train = theta_train,
theta_val = theta_val,
Z_train = Z_train,
Z_val = Z_val,
epochs = 20
)
#> Constructing the training set...
#> Constructing the validation set...
#> Computing the initial validation risk... Initial validation risk = 0.73066986
#> Epoch: 1 Training risk: 0.544 Validation risk: 0.328 Learning rate: 5.00E-04 Epoch time: 5.964 seconds
#> Epoch: 2 Training risk: 0.253 Validation risk: 0.212 Learning rate: 5.00E-04 Epoch time: 2.672 seconds
#> Epoch: 3 Training risk: 0.201 Validation risk: 0.195 Learning rate: 4.97E-04 Epoch time: 2.688 seconds
#> Epoch: 4 Training risk: 0.189 Validation risk: 0.187 Learning rate: 4.88E-04 Epoch time: 2.941 seconds
#> Epoch: 5 Training risk: 0.181 Validation risk: 0.176 Learning rate: 4.73E-04 Epoch time: 2.683 seconds
#> Epoch: 6 Training risk: 0.172 Validation risk: 0.17 Learning rate: 4.52E-04 Epoch time: 3.037 seconds
#> Epoch: 7 Training risk: 0.167 Validation risk: 0.166 Learning rate: 4.27E-04 Epoch time: 2.818 seconds
#> Epoch: 8 Training risk: 0.161 Validation risk: 0.159 Learning rate: 3.97E-04 Epoch time: 2.813 seconds
#> Epoch: 9 Training risk: 0.157 Validation risk: 0.157 Learning rate: 3.63E-04 Epoch time: 2.579 seconds
#> Epoch: 10 Training risk: 0.153 Validation risk: 0.154 Learning rate: 3.27E-04 Epoch time: 2.626 seconds
#> Epoch: 11 Training risk: 0.151 Validation risk: 0.149 Learning rate: 2.89E-04 Epoch time: 2.787 seconds
#> Epoch: 12 Training risk: 0.149 Validation risk: 0.15 Learning rate: 2.50E-04 Epoch time: 2.892 seconds
#> Epoch: 13 Training risk: 0.147 Validation risk: 0.149 Learning rate: 2.11E-04 Epoch time: 2.68 seconds
#> Epoch: 14 Training risk: 0.145 Validation risk: 0.147 Learning rate: 1.73E-04 Epoch time: 2.724 seconds
#> Epoch: 15 Training risk: 0.144 Validation risk: 0.147 Learning rate: 1.37E-04 Epoch time: 2.61 seconds
#> Epoch: 16 Training risk: 0.143 Validation risk: 0.145 Learning rate: 1.03E-04 Epoch time: 2.728 seconds
#> Epoch: 17 Training risk: 0.143 Validation risk: 0.144 Learning rate: 7.32E-05 Epoch time: 2.611 seconds
#> Epoch: 18 Training risk: 0.142 Validation risk: 0.144 Learning rate: 4.77E-05 Epoch time: 2.82 seconds
#> Epoch: 19 Training risk: 0.142 Validation risk: 0.143 Learning rate: 2.72E-05 Epoch time: 2.741 seconds
#> Epoch: 20 Training risk: 0.142 Validation risk: 0.143 Learning rate: 1.22E-05 Epoch time: 2.8 seconds
#> Finished training in 58.464254333 seconds
# Assess the estimator
theta_test <- sampler(1000)
Z_test <- simulate(theta_test, m)
assessment <- assess(estimator, theta_test, Z_test,
estimator_names = "NBE",
parameter_names = c("mu1", "mu2", "sigma", "rho"))
plotestimates(assessment)
3.3 Gridded data
For data collected over a regular grid, neural estimators are typically based on a convolutional neural network (CNN). We give specific attention to this case in a separate vignette. There, we also outline how to perform neural Bayes estimation with incomplete/missing data based on methods described by Wang et al. (2024) and Sainsbury-Dale et al. (2025b).