Neural-network building blocks

Any Flux.jl or Lux.jl model can be used to construct a neural network when using the package. In addition to the standard layers and architectures provided with these deep-learning packages (Flux/Lux), the following components can be useful.

Modules

The structures listed below are often useful when constructing neural estimators. In particular, DeepSet provides a convenient wrapper for embedding standard neural networks (e.g., MLPs, CNNs, GNNs) into a framework suited to making inference with an arbitrary number of replicates.

NeuralEstimators.DeepSet Type

DeepSet(ψ, ϕ, a = mean; condition_on_sample_size = false)
DeepSet(ψ; latent_dim, output_dim, a = mean, condition_on_sample_size = false, kwargs...)
(object::DeepSet)(Z::V) where V <: AbstractVector{A} where A
(object::DeepSet)(P::PackedReplicates)

The DeepSets representation (Zaheer et al., 2017; Sainsbury-Dale et al., 2024),

where 𝐙 ≡ (𝐙₁', …, 𝐙ₘ')' are exchangeable replicates of data, ψ and ϕ are neural networks, and a is a permutation-invariant aggregation function.

The function a must operate on arrays and have a keyword argument dims for specifying the dimension of aggregation (e.g., mean, sum, maximum, minimum, logsumexp).

The convenience constructor DeepSet(ψ; latent_dim, output_dim, ...) builds ϕ as an MLP, where latent_dim is the dimension of (i.e., the output dimension of ψ) and output_dim is the dimension of . Additional keyword arguments are passed to MLP.

DeepSet objects act on data of type Vector{A}, where each element of the vector is associated with one data set (i.e., one set of exchangeable replicates), and where A depends on the chosen architecture for ψ. Exchangeable replicates within each data set are stored in the batch dimension. For example, with data collected over a two-dimensional grid and with ψ chosen to be a CNN, A should be a 4-dimensional array, with replicates stored in the 4ᵗʰ dimension. Alternatively, a vector of arrays may be wrapped in PackedReplicates so that the data are stored as a single multidimensional array and device-transfer overhead is reduced when training with a GPU.

For computational efficiency, array data are first concatenated along their final dimension (i.e., the replicates dimension) before being passed into the inner network ψ, thereby ensuring that ψ is applied to a single large array rather than multiple small ones.

When data sets of varying sample size are envisaged, set condition_on_sample_size = true to concatenate with before it is passed to ϕ:

In this case, the input dimension of ϕ must be one greater than the dimension of .

Variable sample sizes and Reactant

When training with Reactant on data sets of varying sample size, data Z must be wrapped in PackedReplicates(Z; max_sample_size = M) with a fixed cap M so that every batch has the same array shape.

Graph data

Graph data via GNNSummary is currently supported only with the Flux backend. Data sets with differing numbers of replicates are supported: a batch of graphs is packed into a single supergraph on the host, padding the replicate dimension to a common length and masking the padded entries when needed. Padding assumes that no layer of the summary network mixes information across the replicate dimension, so normalisation layers (e.g., BatchNorm) in the propagation module are not supported with a varying number of replicates. As with padded PackedReplicates, the padded path supports only mean, sum, maximum, minimum, and logsumexp aggregation.

Examples

using NeuralEstimators, Flux

# Two data sets containing 3 and 4 replicates
d = 5  # number of parameters in the model
n = 10 # dimension of each replicate
Z = [rand32(n, m) for m ∈ (3, 4)]

# Construct DeepSet object
latent_dim = 16
output_dim = 8
ψ  = Chain(Dense(n, 32, relu), Dense(32, latent_dim, relu))
ϕ  = Chain(Dense(latent_dim, 32, relu), Dense(32, output_dim))
ds = DeepSet(ψ, ϕ)

# Convenience constructor
ds = DeepSet(ψ; latent_dim = latent_dim, output_dim = output_dim)

# Apply DeepSet object to data
ds(Z)
ds(PackedReplicates(Z))

source

NeuralEstimators.GNNSummary Type

GNNSummary(propagation, readout)

A graph neural network (GNN) module designed to serve as the inner network ψ in the DeepSet representation when the data are graphical (e.g., irregularly observed spatial data).

The propagation module transforms graph data into a set of hidden-feature graphs. The readout module aggregates these feature graphs into a single hidden feature vector of fixed length. The network ψ is then defined as the composition of the propagation and readout modules.

The data should be stored as a Vector{A} where each element is associated with one parameter vector. For spatial data collected over irregular locations, A is typically a GNNGraph, where independent replicates (possibly with differing spatial locations) are stored as subgraphs. See spatialgraph for constructing these graphs from matrices of spatial locations and data.

Examples

using NeuralEstimators, Flux, GraphNeuralNetworks
using Statistics: mean

# Spatial data
n = 100                # number of spatial locations
m = 50                 # number of replicates
S = rand(n, 2)         # spatial locations
Z = rand(n, m)         # observed data
g = spatialgraph(S, Z) # construct the graph

# Propagation module
nₕ = 32    # dimension of node feature vectors
propagation = Chain(SpatialGraphConv(1 => nₕ), SpatialGraphConv(nₕ => nₕ))

# Readout module
readout = GlobalPool(mean)

# Inner network
ψ = GNNSummary(propagation, readout)

# Outer network
d = 3     # number of parameters
w = 64    # width of hidden layer
ϕ = Chain(Dense(nₕ, w, relu), Dense(w, d))

# DeepSet object 
ds = DeepSet(ψ, ϕ)

# Apply to data 
ds(g)        # single graph with subgraphs corresponding to independent replicates
ds([g, g])   # vector of graphs, corresponding to multiple data sets

source

NeuralEstimators.MLP Function

MLP(in::Integer, out::Integer; kwargs...)

A traditional fully-connected multilayer perceptron (MLP) with input dimension in and output dimension out.

Keyword arguments

• depth::Integer = 2: the number of hidden layers.

• width::Integer = 128: the width of each hidden layer.

• activation = relu: the activation function used in each hidden layer.

• output_activation = identity: the activation function used in the output layer.

• backend::Union{Nothing, Module} = nothing: the backend to use for constructing the network (e.g., Lux or Flux). If nothing, the backend is resolved automatically.

source

NeuralEstimators.MultiHeadMLP Function

MultiHeadMLP(in::Integer, out::Integer, num_heads::Integer; growing::Bool = false, kwargs...)

A multi-head MLP consisting of num_heads independent MLP heads, each with input dimension in and output dimension out. The outputs of all heads are concatenated to give a final output of dimension num_heads * out.

Keyword arguments

• growing::Bool = false: if true, the input dimension of the i-th head is in + i, allowing each head to receive an incrementally larger input.

• kwargs: keyword arguments passed to MLP.

Examples

using NeuralEstimators, Lux, Random
rng = Random.default_rng()
Random.seed!(rng, 0)

# Dummy data
batchsize = 16
num_summaries = 5
num_parameters = 3
s = rand(Float32, num_summaries, batchsize)
θ = rand(Float32, num_parameters, batchsize)

m = MultiHeadMLP(num_summaries, 1, num_parameters)
ps, st = Lux.setup(rng, m)
m(s, ps, st)[1]

m = MultiHeadMLP(num_summaries + num_parameters, 1, num_parameters)
ps, st = Lux.setup(rng, m)
m(vcat(s, θ), ps, st)[1]

m = MultiHeadMLP(num_summaries, 1, num_parameters; growing = true)
sθ_split = Tuple(vcat(s, θ[1:i, :]) for i in 1:num_parameters)
ps, st = Lux.setup(rng, m)
m(sθ_split, ps, st)[1]

source

Layers

In addition to the built-in layers provided by Flux and Lux, the following layers may be used when building a neural-network architecture.

NeuralEstimators.ResidualBlock Function

ResidualBlock(filter, in => out; stride = 1, backend = nothing)

Basic residual block (see here), consisting of two sequential convolutional layers and a skip (shortcut) connection that connects the input of the block directly to the output, facilitating the training of deep networks.

Examples

using NeuralEstimators, Flux
z = rand(16, 16, 1, 1)
b = ResidualBlock((3, 3), 1 => 32)
b(z)

source

NeuralEstimators.SpatialGraphConv Type

SpatialGraphConv(in => out, g=relu; args...)

Implements a spatial graph convolution for isotropic spatial processes (Sainsbury-Dale et al., 2025),

where is the hidden feature vector at location at layer ,   is a non-linear activation function applied elementwise, and are trainable parameter matrices, is a trainable bias vector, denotes the indices of neighbours of ,   is a (learnable) spatial weighting function, denotes elementwise multiplication, and    is a (learnable) function.

By default, the function    is modelled using a PowerDifference function. One may alternatively employ a nonlearnable function, for example, f = (hᵢ, hⱼ) -> (hᵢ - hⱼ).^2, specified through the keyword argument f.

The spatial distances between locations must be stored as an edge feature, as facilitated by spatialgraph(). The input to   is a   matrix (i.e., a row vector) of spatial distances. The output of   must be either a scalar; a vector of the same dimension as the feature vectors of the previous layer; or, if the features vectors of the previous layer are scalars, a vector of arbitrary dimension. To promote identifiability, the weights are normalised to sum to one (row-wise) within each neighbourhood set. By default,   is taken to be a multilayer perceptron with a single hidden layer, although a custom choice for this function can be provided using the keyword argument w. The hidden layer of the default   uses the activation g, while its output layer uses softplus, so that the weights are strictly positive. A custom w should likewise return non-negative weights; note that an output activation which can return exactly zero for every edge (e.g., relu) risks being identically zero with no gradient available to recover from it.

GPU memory and the choice of batch size

The messages    are formed on the edges of the batched graph, so every intermediate array in the layer is of size (w_out, m, E), where m is the number of independent replicates and E is the total number of edges in the batch, that is, the batch size multiplied by the number of edges per data set. Peak memory during training therefore grows in proportion to m × (batch size), and it is the peak that matters: the total volume of memory allocated over an epoch is independent of the batch size.

The practical consequence is that, with many replicates, increasing the batch size can make training slower rather than faster, because memory pressure causes the CUDA.jl allocator to run the garbage collector inside the allocation path. The transition is abrupt rather than gradual, and once it is crossed the run time is dominated by garbage collection rather than by computation. If throughput degrades as the batch size is raised, reduce the batch size; if a large effective batch is needed for optimisation reasons, accumulate gradients over several smaller sub-batches instead.

Arguments

• in: dimension of input features.

• out: dimension of output features.

• g = relu: activation function.

• bias = true: add learnable bias?

• init = glorot_uniform: initialiser for , , and .

• f = nothing

• w = nothing

• w_width = 128 (applicable only if w = nothing): the width of the hidden layer in the MLP used to model   .

• w_out = in (applicable only if w = nothing): the output dimension of   .

Examples

using NeuralEstimators, Flux, GraphNeuralNetworks

# Toy spatial data
n = 250                # number of spatial locations
m = 5                  # number of replicates
S = rand(n, 2)         # spatial locations
Z = rand(n, m)         # data
g = spatialgraph(S, Z) # construct the graph

# Construct and apply spatial graph convolution layer
l = SpatialGraphConv(1 => 10)
l(g)

source

Output layers

In addition to the standard activation functions provided by NNlib.jl (e.g., relu, gelu, softplus), the following layers can be used at the end of a neural network to ensure the outputs satisfy certain constraints. Note that Parallel (Flux/Lux) can be useful for applying several parameter constraints simultaneously.

Layers vs. activation functions

The following types should be incorporated as separate layers in the final stage of a Chain. In particular, they cannot be used as the activation function of a Dense layer.

NeuralEstimators.Compress Type

Compress(a, b, k = 1)

Layer that compresses its input to be within the range a and b, where each element of a is less than the corresponding element of b.

The layer uses a logistic function,

where the arguments a and b together combine to shift and scale the logistic function to the range (a, b), and the growth rate k controls the steepness of the curve.

The logistic function given here contains an additional parameter, θ₀, which is the input value corresponding to the functions midpoint. In Compress, we fix θ₀ = 0, since the output of a randomly initialised neural network is typically around zero.

Examples

using NeuralEstimators, Flux

a = [25, 0.5, -pi/2]
b = [500, 2.5, 0]
p = length(a)
K = 100
θ = randn(p, K)
l = Compress(a, b)
l(θ)

n = 20
θ̂ = Chain(Dense(n, p), l)
Z = randn(n, K)
θ̂(Z)

source

NeuralEstimators.CorrelationMatrix Type

CorrelationMatrix(d)
(object::CorrelationMatrix)(x::Matrix, cholesky::Bool = false)

Transforms unconstrained input into the parameters of a d×d correlation matrix or, if cholesky = true, the lower Cholesky factor of a d×d correlation matrix.

The expected input is a Matrix with T(d-1) = (d-1)d÷2 rows, where T(d-1) is the (d-1)th triangular number (the number of free parameters in an unconstrained d×d correlation matrix), and the output is a Matrix of the same dimension. The columns of the input and output matrices correspond to independent parameter configurations (i.e., different correlation matrices).

Internally, the layer constructs a valid Cholesky factor 𝐋 for a correlation matrix, and then extracts the strict lower triangle from the correlation matrix 𝐑 = 𝐋𝐋'. The lower triangle is extracted and vectorised in line with Julia's column-major ordering: for example, when modelling the correlation matrix

the rows of the matrix returned by a CorrelationMatrix layer are ordered as

which means that the output can easily be transformed into the implied correlation matrices using vectotril and Symmetric.

See also CovarianceMatrix.

Examples

using NeuralEstimators, LinearAlgebra

d  = 4
l  = CorrelationMatrix(d)
p  = (d-1)*d÷2
x  = randn(p, 100)

# Returns a matrix of parameters, which can be converted to correlation matrices
R = l(x)
R = map(eachcol(R)) do r
	R = Symmetric(vectotril(r, strict = true), :L)
	R[diagind(R)] .= 1
	R
end

# Obtain the Cholesky factor directly
L = l(x, true)
L = map(eachcol(L)) do x
	# Only the strict lower diagonal elements are returned
	L = LowerTriangular(vectotril(x, strict = true))

	# Diagonal elements are determined under the constraint diag(L*L') = 𝟏
	L[diagind(L)] .= sqrt.(1 .- rowwisenorm(L).^2)
	L
end
L[1] * L[1]'

source

NeuralEstimators.CovarianceMatrix Type

CovarianceMatrix(d)
(object::CovarianceMatrix)(x::Matrix, cholesky::Bool = false)

Transforms unconstrained input into the parameters of a d×d covariance matrix or, if cholesky = true, the lower Cholesky factor of a d×d covariance matrix.

The expected input is a Matrix with T(d) = d(d+1)÷2 rows, where T(d) is the dth triangular number (the number of free parameters in an unconstrained d×d covariance matrix), and the output is a Matrix of the same dimension. The columns of the input and output matrices correspond to independent parameter configurations (i.e., different covariance matrices).

Internally, the layer constructs a valid Cholesky factor 𝐋 and then extracts the lower triangle from the positive-definite covariance matrix 𝚺 = 𝐋𝐋'. The lower triangle is extracted and vectorised in line with Julia's column-major ordering: for example, when modelling the covariance matrix

the rows of the matrix returned by a CovarianceMatrix are ordered as

which means that the output can easily be transformed into the implied covariance matrices using vectotril and Symmetric.

See also CorrelationMatrix.

Examples

using NeuralEstimators, LinearAlgebra

d = 4
l = CovarianceMatrix(d)
p = d*(d+1)÷2
x = randn(p, 50)

# Returns a matrix of parameters, which can be converted to covariance matrices
Σ = l(x)
Σ = [Symmetric(vectotril(x), :L) for x ∈ eachcol(Σ)]

# Obtain the Cholesky factor directly
L = l(x, true)
L = [LowerTriangular(vectotril(x)) for x ∈ eachcol(L)]
L[1] * L[1]'

source

Miscellaneous

NeuralEstimators.IndicatorWeights Type

IndicatorWeights(h_max, n_bins::Integer)
(w::IndicatorWeights)(h::Matrix)

For spatial locations and , creates a spatial weight function defined as

where   denotes the indicator function,    is the spatial distance between and , and    is a set of  n_bins equally-sized distance bins covering the spatial distances between 0 and h_max.

Examples

using NeuralEstimators, GraphNeuralNetworks

h_max = 1
n_bins = 10
w = IndicatorWeights(h_max, n_bins)
h = rand(1, 30) # distances between 30 pairs of spatial locations 
w(h)

source

NeuralEstimators.KernelWeights Type

KernelWeights(h_max, n_bins::Integer)
(w::KernelWeights)(h::Matrix)

For spatial locations and , creates a spatial weight function defined as

where    is the spatial distance between and , and and are the means and standard deviations of the Gaussian kernels for each bin, covering the spatial distances between 0 and h_max.

Examples

using NeuralEstimators, GraphNeuralNetworks

h_max = 1
n_bins = 10
w = KernelWeights(h_max, n_bins)
h = rand(1, 30) # distances between 30 pairs of spatial locations 
w(h)

source

NeuralEstimators.PowerDifference Type

PowerDifference(a, b)

Function    , where    and    are constrained transformations of the trainable parameters a and b.

Examples

using NeuralEstimators

X = rand(5, 100)
Y = rand(5, 100)
f = PowerDifference(0, 1.55)
f(X, Y)   # two arg method
f((X, Y)) # tuple method

source