1 Complete data
We first consider the case that the data are collected completely over a regular grid.
We develop a neural Bayes estimator for a spatial Gaussian process model, where \(\boldsymbol{Z} \equiv (Z_{1}, \dots, Z_{n})'\) are data collected at locations \(\{\boldsymbol{s}_{1}, \dots, \boldsymbol{s}_{n}\}\) in a spatial domain that is a subset of \(\mathbb{R}^2\). The data are assumed to be spatially-correlated mean-zero Gaussian random variables with exponential covariance function, \[ \textrm{cov}(Z_i, Z_j) = \textrm{exp}(-\|\boldsymbol{s}_i - \boldsymbol{s}_j\|/\theta), \] with unknown range parameter \(\theta > 0\). Here, we take the spatial domain be the unit square, we simulate data on a grid with \(16^2 = 256\) possible observation locations, and we adopt a uniform prior, \(\theta \sim \rm{Unif}(0, 0.4)\).
To begin, we define a function for sampling from the prior
distribution and a function for marginal simulation from the statistical
model. Note that our simulation code can be made faster (e.g., using
parallel versions of lapply, utilising GPUs, etc.), but we
provide a relatively simple implementation for pedagogical purposes:
# Sampling from the prior distribution
# K: number of samples to draw from the prior
prior <- function(K) {
theta <- 0.4 * runif(K) # draw from the prior distribution
theta <- t(theta) # reshape to matrix
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 = 1) {
# Spatial locations, a grid over the unit square, and spatial distance matrix
N <- 16 # number of locations along each dimension of the grid
S <- expand.grid(seq(1, 0, len = N), seq(0, 1, len = N))
D <- as.matrix(dist(S))
# Simulate conditionally independent replicates for each parameter vector
Z <- lapply(1:ncol(theta), function(k) {
Sigma <- exp(-D/theta[k]) # covariance matrix
L <- t(chol(Sigma)) # lower Cholesky factor of Sigma
n <- nrow(L) # number of observation locations
mm <- if (length(m) == 1) m else sample(m, 1) # allow for variable sample sizes
z <- matrix(rnorm(n*mm), nrow = n, ncol = mm) # standard normal variates
Z <- L %*% z # conditionally independent replicates from the model
Z <- array(Z, dim = c(N, N, 1, mm)) # reshape to multidimensional array
Z
})
return(Z)
}A note on data format: When working with the package
NeuralEstimators, the following general rules regarding
assumed data format apply:
- sets of parameter vectors are always stored as matrices, where each
column corresponds to one parameter vector \(\boldsymbol{\theta}\);
- simulated data are stored as a list, where each element of the list corresponds to a data set \(\boldsymbol{Z}\) (possibly containing independent replicates) simulated conditional on a parameter vector \(\boldsymbol{\theta}\).
When using CNNs, each data set must be stored as a multi-dimensional array. The penultimate dimension stores the so-called channels (this dimension is singleton for univariate processes, two for bivariate processes, etc.), while the final dimension stores conditionally independent replicates. For example, to store \(50\) conditionally independent replicates of a bivariate spatial process measured over a \(10\times15\) grid, one would construct an array of dimension \(10\times15\times2\times50\).
Now, let’s visualise a few realisations from our spatial Gaussian process model:
Next, we design our neural-network architecture.
The package NeuralEstimators is centred on the DeepSets
framework (Zaheer et al.,
2017), which allows for making inference with an arbitrary number of
independent replicates and for incorporating both neural (learned) and
user-defined summary statistics. Specifically, the neural Bayes
estimator is taken to be of the form, \[
\hat{\boldsymbol{\theta}}(\boldsymbol{Z}_1,\dots,\boldsymbol{Z}_m) =
\boldsymbol{\phi}\bigg\{\frac{1}{m}
\sum_{i=1}^m\boldsymbol{\psi}(\boldsymbol{Z}_i)\bigg\},
\] where \(\boldsymbol{Z}_1,\dots,\boldsymbol{Z}_m\)
are replicates of \(\boldsymbol{Z}\),
\(\boldsymbol{\psi}(\cdot)\) is a
conventional neural network whose architecture depends on the
multivariate structure of the data (e.g., a CNN for gridded data), and
\(\boldsymbol{\phi}(\cdot)\) is
(always) a multilayer perceptron (MLP). See Sainsbury-Dale
et al. (2024) for further details on the use of DeepSets in the
context of neural Bayes estimation and a discussion on the framework’s
connection to conventional estimators. See also Dumoulin and Visin (2016)
for a detailed description of CNNs and their construction, which is
beyond the scope of this vignette:
estimator <- juliaEval('
# Summary network
psi = Chain(
Conv((3, 3), 1 => 32, relu), # 3x3 convolutional filter, 1 input channel to 32 output channels, ReLU activation
MaxPool((2, 2)), # 2x2 max pooling layer for dimension reduction
Conv((3, 3), 32 => 64, relu), # 3x3 convolutional filter, 32 input channels to 64 output channels, ReLU activation
MaxPool((2, 2)), # 2x2 max pooling layer for dimension reduction
Flux.flatten # flatten the output to feed into a fully connected layer
)
# Inference network
phi = Chain(
Dense(256, 64, relu), # fully connected layer, 256 input neurons to 64 output neurons, ReLU activation
Dense(64, 1) # fully connected layer, 64 input neurons to 1 output neuron
)
# Construct DeepSet object and initialise a point estimator
deepset = DeepSet(psi, phi)
estimator = PointEstimator(deepset)
')Next, we train the estimator using parameter vectors sampled from the prior distribution, and data simulated from the statistical model conditional on these parameter vectors:
# Construct training and validation sets
K <- 25000 # size of the training set
theta_train <- prior(K) # parameter vectors used in stochastic-gradient descent during training
theta_val <- prior(K/10) # parameter vectors used to monitor performance during training
Z_train <- simulate(theta_train) # data used in stochastic-gradient descent during training
Z_val <- simulate(theta_val) # data used to monitor performance during training
# Train the estimator
estimator <- train(
estimator,
theta_train = theta_train,
theta_val = theta_val,
Z_train = Z_train,
Z_val = Z_val
)
#> Computing the initial validation risk... Initial validation risk = 0.27961352
#> Computing the initial training risk... Initial training risk = 0.27962294
#> Epoch: 1 Training risk: 0.075 Validation risk: 0.053 Run time of epoch: 30.346 seconds
#> Epoch: 2 Training risk: 0.043 Validation risk: 0.033 Run time of epoch: 6.594 seconds
#> Epoch: 3 Training risk: 0.031 Validation risk: 0.028 Run time of epoch: 6.537 seconds
#> Epoch: 4 Training risk: 0.027 Validation risk: 0.027 Run time of epoch: 6.498 seconds
#> Epoch: 5 Training risk: 0.024 Validation risk: 0.023 Run time of epoch: 6.63 seconds
#> Epoch: 6 Training risk: 0.023 Validation risk: 0.023 Run time of epoch: 6.551 seconds
#> Epoch: 7 Training risk: 0.022 Validation risk: 0.022 Run time of epoch: 6.513 seconds
#> Epoch: 8 Training risk: 0.022 Validation risk: 0.022 Run time of epoch: 6.526 seconds
#> Epoch: 9 Training risk: 0.021 Validation risk: 0.021 Run time of epoch: 6.815 seconds
#> Epoch: 10 Training risk: 0.021 Validation risk: 0.021 Run time of epoch: 6.791 seconds
#> Epoch: 11 Training risk: 0.02 Validation risk: 0.021 Run time of epoch: 6.583 seconds
#> Epoch: 12 Training risk: 0.02 Validation risk: 0.021 Run time of epoch: 6.705 seconds
#> Epoch: 13 Training risk: 0.02 Validation risk: 0.02 Run time of epoch: 6.775 seconds
#> Epoch: 14 Training risk: 0.019 Validation risk: 0.02 Run time of epoch: 6.652 seconds
#> Epoch: 15 Training risk: 0.019 Validation risk: 0.02 Run time of epoch: 6.586 seconds
#> Epoch: 16 Training risk: 0.019 Validation risk: 0.021 Run time of epoch: 6.554 seconds
#> Epoch: 17 Training risk: 0.019 Validation risk: 0.021 Run time of epoch: 6.73 seconds
#> Epoch: 18 Training risk: 0.019 Validation risk: 0.02 Run time of epoch: 6.554 seconds
#> Epoch: 19 Training risk: 0.018 Validation risk: 0.02 Run time of epoch: 6.556 seconds
#> Epoch: 20 Training risk: 0.018 Validation risk: 0.02 Run time of epoch: 6.57 seconds
#> Epoch: 21 Training risk: 0.018 Validation risk: 0.02 Run time of epoch: 6.613 seconds
#> Epoch: 22 Training risk: 0.018 Validation risk: 0.02 Run time of epoch: 6.571 seconds
#> Epoch: 23 Training risk: 0.018 Validation risk: 0.02 Run time of epoch: 6.54 seconds
#> Epoch: 24 Training risk: 0.017 Validation risk: 0.02 Run time of epoch: 6.648 seconds
#> Epoch: 25 Training risk: 0.017 Validation risk: 0.02 Run time of epoch: 6.695 seconds
#> Epoch: 26 Training risk: 0.017 Validation risk: 0.02 Run time of epoch: 6.587 seconds
#> Stopping early since the validation loss has not improved in 5 epochsIt is good practice to assess the performance of a trained estimator using out-of-sample test data:
# Test the estimator using out-of-sample data
theta <- prior(1000)
Z <- simulate(theta)
assessment <- assess(estimator, theta, Z, estimator_names = "NBE")
#> Running estimator NBE...
plotestimates(assessment)
Once the neural Bayes estimator has been trained and assessed, it can
be applied to observed data collected completely over a \(16\times 16\) grid using
estimate(). Below, we use simulated data as a surrogate for
observed data:
theta <- matrix(0.1) # true parameter
Z <- simulate(theta) # pretend that this is observed data
estimate(estimator, Z) # point estimates from the "observed" data
#> [,1]
#> [1,] 0.1008148
#> attr(,"JLDIM")
#> [1] 1 1In practice, data are often incomplete and, in the next section, we describe methods that facilitate neural Bayes estimation in these settings.
