NeuralEstimators.jlEfficient simulation-based inference using neural networks

Introduction ​

NeuralEstimators.jl facilitates neural methods for parameter inference in scenarios where simulation from the model is feasible. These methods are likelihood-free and amortised, in the sense that, once the neural networks are trained on simulated data, they enable rapid inference across arbitrarily many observed data sets in a fraction of the time required by conventional approaches.

The package supports neural posterior estimators (NPEs), which approximate the full posterior distribution using generative neural networks; neural Bayes estimators (NBEs), which transform data into functionals (e.g., point summaries) of the posterior distribution; and neural ratio estimators (NREs), which approximate the likelihood-to-evidence ratio and thereby enable flexible Bayesian or frequentist inference.

User-friendliness is a central focus of the package, which is designed to minimise "boilerplate" code while preserving complete flexibility in the neural-network architecture and other workflow components. The package accommodates any model for which simulation is feasible by allowing users to define their model implicitly through simulated data.

Installation ​

To install the package, first install the current stable release of Julia. Then, install the current stable version of the package using the following command inside Julia:

using Pkg; Pkg.add("NeuralEstimators")

Or install the current development version using the command:

using Pkg; Pkg.add(url = "https://github.com/msainsburydale/NeuralEstimators.jl")

Quick start ​

In the following minimal example, we develop a neural estimator for   from data  , where  , and we use the priors   and  .

using NeuralEstimators, Lux, Enzyme

# Functions to sample from the prior and simulate data
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 = Chain(Dense(n, 64, gelu), Dense(64, 64, gelu), 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 (here, simulated as a stand-in)
θ = sampler(1)                   # ground truth (not known in practice)
Z = simulator(θ)                 # stand-in for real observations
estimate(estimator, Z)           # point estimate

using NeuralEstimators, Lux, Enzyme

# Functions to sample from the prior and simulate data
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 = Chain(Dense(n, 64, gelu), Dense(64, 64, gelu), Dense(64, d))

# Initialise a neural estimator
estimator = PosteriorEstimator(network, d; num_summaries = d, q = GaussianMixture)

# 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 (here, simulated as a stand-in)
θ = sampler(1)                   # ground truth (not known in practice)
Z = simulator(θ)                 # stand-in for real observations
sampleposterior(estimator, Z)    # approximate posterior sample

using NeuralEstimators, Lux, Reactant

# Functions to sample from the prior and simulate data
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 = Chain(Dense(n, 64, gelu), Dense(64, 64, gelu), Dense(64, d))

# Initialise a neural estimator
estimator = RatioEstimator(network, d; num_summaries = d)

# Train the estimator
estimator = train(estimator, sampler, simulator, device = reactant_device())

# Assess the estimator
θ_test = sampler(250)
Z_test = simulator(θ_test)
assessment = assess(estimator, θ_test, Z_test)
bias(assessment)
rmse(assessment)

# Apply to observed data (here, simulated as a stand-in)
θ = sampler(1)                   # ground truth (not known in practice)
Z = simulator(θ)                 # stand-in for real observations
sampleposterior(estimator, Z)    # approximate posterior sample

Contributing ​

We welcome contributions of all sizes. To get started, see CONTRIBUTING.md for an overview of the code structure, development workflow, and how to submit contributions. Specific instructions for contributing to the documentation can be found in docs/README.md.

If you encounter a bug or have a suggestion, please feel free to open an issue or submit a pull request.

Supporting and citing ​

This software was developed as part of academic research. If you would like to support it, please star the repository. If you use it in your research or other activities, please also use the citation given here.

Papers using NeuralEstimators ​

• Likelihood-free parameter estimation with neural Bayes estimators [paper] [code]

• Neural methods for amortized inference [paper][code]

• Neural Bayes estimators for irregular spatial data using graph neural networks [paper][code]

• Neural Bayes estimators for censored inference with peaks-over-threshold models [paper] [code]

• Neural parameter estimation with incomplete data [paper][code]

• Neural Bayes inference for complex bivariate extremal dependence models [paper][code]

• Fast likelihood-free parameter estimation for Lévy processes [paper]

• Joint modeling of low and high extremes using a multivariate extended generalized Pareto distribution [paper]