ENFR: Evidence Networks with Flow-Matched References for Amortized Simulation-Based Model Inference
Authors: Kai Lehman, Jameson Dong, Carolina Cuesta-Lazaro, Niall Jeffrey, Benjamin D. Wandelt, Francisco Villaescusa-Navarro, Sven Krippendorf, Jochen Weller
Organizations: Universitäts-Sternwarte München, Fakultät für Physik, Ludwig-Maximilians-Universität, Scheinerstr. 1, 81679 München, Germany · Center for Computational Astrophysics, Flatiron Institute, 162 5th Avenue, New York, NY, 10010, USA · Excellence Cluster ORIGINS, Boltzmannstr. 2, 85748 Garching, Germany · Department of Physics and Astronomy, Johns Hopkins University, 3400 North Charles Street, Baltimore, MD, 21218, USA · Center for Cosmology and Particle Physics, Department of Physics, New York University, 726 Broadway, New York, NY 10003, USA · Department of Physics & King’s Institute for Artificial Intelligence, King’s College London, Strand, London WC2R 2LS, United Kingdom · Department of Physics & Astronomy, University College London, Gower Street, London, WC1E 6BT, UK · Department of Applied Mathematics and Statistics, Johns Hopkins University, 3400 North Charles Street, Baltimore, MD, 21218, USA · Department of Astrophysical Sciences, Princeton University, Princeton, NJ 08544, USA · Cavendish Laboratory and Department of Applied Mathematics and Theoretical Physics, University of Cambridge, Cambridge CB3 0WA, United Kingdom · Max-Planck-Institut für extraterrestrische Physik, Giessenbachstr. 1, 85748 Garching, Germany
The evidence, Z=p(x∣M), of a data-generating model M is a powerful tool in Bayesian frameworks: it enables model inference, model misspecification tests, and goodness-of-fit assessment as an alternative to standard discrepancy statistics, e.g. χ2. While versatile, the evidence is notoriously difficult to compute, even when the likelihood is available. In both likelihood-based and simulation-based inference, estimating Z provides an immediate solution to high-dimensional goodness-of-fit or model misspecification tests: using Z(x) as the test statistic induces a minimum-volume level-α test. Such tests are essential to assess model fidelity, which is necessary to trust simulation-based parameter inference. Existing learned approaches largely target Bayes factors, which only measure evidence ratios. Instead, we provide accurate, fast, and amortized estimates of the absolute evidence. We present ENFR: Evidence Networks with Flow-matched References. Building on the Evidence Network framework of Jeffrey & Wandelt (2024), the proposed method first learns Bayes factors with respect to a flow-matched reference, from which the absolute evidence is recovered. We show improvement in estimation of Z over the flow-based baseline in all applications considered and the best performance of all density estimation methods considered in our benchmarks. We furthermore introduce and provide the proof for an architecture that jointly learns the individual evidences of a set of models {Mi}. We demonstrate the performance and scaling behavior of these approaches on data from a variety of scientific applications.
Figures & tables
Figure 1: A visualization of our proposed SingleENFR and MultiENFR architectures. Given simulated data from one or multiple models, we train flow matching once over all models. In a second step, we use the simulated data and samples from the flow to train an evidence network classifier, which learns all Bayes factors. Those can be combined with the evaluated flow density to obtain the evidence under a given model.
Figure 2: Scaling of density-estimation accuracy with (a) data dimensionality and (b) number of training simulations on the Gaussian mixture model of Appendix A.1 , for which the true density is known analytically. We report the RMSE between estimated and true log-density over 200 held-out test points. In (a) the simulation budget is fixed at 5×106 ; in (b) the dimensionality is fixed at 64. The line shows the mean across three trained models and the errorbars show the estimated standard deviation. SingleENFR (orange) improves on the flow it is built on (blue) at every point and is the most accurate of all estimators considered, while CTSM-VP (yellow) is consistently the least accurate.
Simulation
flow
SingleENFR (ours)
MultiENFR (ours)
CTSM-SB + FlowRef
Λ CDM
1.34
0.53
0.52
2.74
w CDM
1.40
0.76
0.67
2.46
w0wa CDM
1.16
0.83
0.81
3.79
ν CDM
1.41
0.47
0.46
1.77
Ωk CDM
1.52
0.54
0.78
2.29
Table 1: RMSE of our learned log-evidence estimates against the likelihood-based log-evidence estimate over 100 held out BAO test simulations (lower is better). Our Models SingleENFR and MultiENFR consistently outperform both the flow and CTSM-SB with a flow reference.
Model
Likelihood-based analysis
MultiENFR
Δ
MultiENFR logK
ΛCDM
-13.99 ± 0.16
-14.07 ± 0.34
-0.08
0.00
wCDM
-16.41 ± 0.18
-16.00 ± 0.35
0.40
-1.94
w0waCDM
-15.63 ± 0.18
-15.46 ± 0.40
0.17
-1.40
νCDM
-14.06 ± 0.16
-14.01 ± 0.30
0.04
0.06
ΩkCDM
-15.88 ± 0.17
-15.72 ± 0.23
0.16
-1.65
Table 2: Evidence estimates using MultiENFR on the DESI DR2 observed data. For MultiENFR we also quote the estimate of the standard deviation over the ensemble of 5 members. We can reproduce the likelihood-based evidence estimates of the DESI DR2 data with our simulation-based framework within the ensemble scatter.
Figure 3: (a) Evidence estimation for noisy observations of an ODE driven time series using flow matching and our ENFR models. MultiENFR outperforms the other models. (b) Diffusion drift application of standard flow matching and our ENFR networks. MultiENFR shows the best performance and good evidence estimation over a large dynamical range.
Figure 4: Blind coverage testing for the simulation-based applications for (a) the Hodgkin-Huxley application, and (b) the Galaxy assembly bias application. As we do not have access to a likelihood-based evidence estimate for our simulation based models we test whether the resulting Bayes factors are well calibrated. We find that our MultiENFR architecture is well calibrated in both cases. The error bars show 2σ binomial errors.
Figure 5: Evidence estimates of the MultiENFR model on prior predictive simulations and the DESI DR2 data under the Λ CDM model. While 46.1% of simulations have lower evidence than the data, no simulation has lower evidence than the toy w(z) model, and we can clearly pick up on the model misspecification.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Parameter
Λ CDM
w CDM
w0wa CDM
ν CDM
Ωk CDM
Base parameters (shared)
H0
[20,100]
Ωbh2
[0.005,0.1]
Ωch2
[0.001,0.99]
τ
[0.01,0.8]
log(1010As)
[1.61,3.91]
Appendix
Table 3: Boundaries of the uniform priors for the parameters of the five cosmological models used in our BAO example. The six base parameters of the Λ CDM model are shared by every model. Each extension introduces either one or two more parameters.
Parameter
Prior
Harmonic
Damped
Duffing
x0
U(−2,2)
✓
✓
✓
v0
U(−2,2)
✓
✓
✓
ω
U(0.4,2.5)
✓
✓
✓
γ
U(0,0.4)
✓
✓
β
U(−0.8,0.8)
✓
Appendix
Table 4: Prior distributions for the ODE oscillator models. All priors are independent uniforms; parameters shared between models have identical priors.
Method
harmonic
damped
Duffing
All
flow
6.13
3.69
3.91
4.71
SingleENFR
2.54
2.44
2.60
2.53
MultiENFR
2.36
0.94
2.17
1.93
flow + CTSM-SB (ens.)
135.47
26.23
7.96
79.80
Appendix
Table 5: RMSE of the learned evidence against nested sampling (dynesty) on the ODE oscillator benchmark, over 10 held-out simulations per model.
Parameter
Description
Prior
M0
M1
v
drift rate
U(0,5)
✓
✓
λ
leak rate
U(−20,−5)
✓
B
initial boundary height
U(0.3,2)
✓
✓
κ
boundary collapse rate
U(0.5,1.5)
✓
✓
tnd
non-decision time [s]
U(0.1,0.3)
✓
✓
Appendix
Table 6: Prior distributions for the drift-diffusion models.
Model
RMSE under M0
RMSE under M1
flow
4.0
5.5
SingleENFR
2.4
3.5
MultiENFR
1.8
1.7
flow + CTSM-SB
20.3
115.4
Appendix
Table 7: Estimated evidence against nested sampling (dynesty) on the DDM benchmark, over 50 held-out datasets from each model.
Parameter
Prior
Unit
gL
U(10−6,3×10−4)
Scm−2
EL
U(−80,−60)
mV
gK
U(1.5×10−3,1.5×10−2)
Scm−2
Vt
U(−70,−50)
mV
gNa
U(8×10−3,8×10−2)
Scm−2
σ
U(10−4,0.15)
μAcm−2ms1/2
Appendix
Table 8: Prior distributions for the Hodgkin–Huxley models.
Quantity
Value
Unit
ENa
50
mV
EK
−90
mV
ECa
120
mV
Cm
1
μFcm−2
Compartment length / radius
10 / 1
μm
Step current density (225–1225 ms)
2
μAcm−2
Appendix
Table 9: Fixed parameters for the Hodgkin–Huxley simulations.
Parameter
M0
M1
Ωm
U(0.1,0.5)
U(0.1,0.5)
Ωb
U(0.03,0.07)
U(0.03,0.07)
h
U(0.5,0.9)
U(0.5,0.9)
ns
U(0.8,1.2)
U(0.8,1.2)
σ8
U(0.6,1.0)
U(0.6,1.0)
log10Mmin
NGC empirical prior
NGC empirical prior
Appendix
Table 10: Priors for the Quijote galaxy-bias comparison. U(a,b) denotes a uniform distribution and TN(0,0.22;[−1,1]) a zero-mean Gaussian of standard deviation 0.2 , truncated to [−1,1] . Mass logarithms are base ten with masses in h−1M⊙ . Numerical constants denote fixed values. The HOD mass M1 is distinct from the model label M1 .
Example
Flow
CTSM-VP
SingleENFR
CTSM-SB + FlowRef
GMM 64D
0.76
2.14
0.45
1.06
GMM 128D
2.30
7.70
1.32
2.46
GMM 256D
6.61
27.87
3.04
6.57
GMM 512D
15.79
65.36
14.44
15.87
GMM 1024D
34.99
163.21
33.44
35.01
Appendix
Table 11: RMSE values of the estimated density logq(x) with respect to the true density logp(x) over 200 held out test draws (lower is better) averaged over the three seeds. The SingleENFR combination consistently outperforms the other models. The flow is shared at each seed.
Figure 6: The performance on test points in 128 dimensions for the flow and SingleENFR.
Model
Likelihood Based Analysis
MultiENFR
Δ
wCDM
-2.42 ± 0.24
-1.94
+0.48
w0waCDM
-1.64 ± 0.24
-1.40
+0.25
νCDM
-0.07 ± 0.22
0.06
+0.13
ΩkCDM
-1.89 ± 0.23
-1.65
+0.24
Appendix
Table 12: Bayes factors w.r.t. Λ CDM on the DESI DR2 data. The error of our simulation-based estimate is largely comparable to the error of the likelihood-based estimate in three of the four models.
Deep learning models are being used for the analysis of parametric statistical models based on simulation-only frameworks. Bayesian models using normalizing flows simulate data from a prior distribution and are composed of two deep neural networks: a summary network that learns a sufficient statistic for the parameter and a normalizing flow that conditional on the summary network can approximate the posterior distribution. Here, we explore frequentist models that are based on a single summary network. During training, input of the network is a simulated data set based on a parameter and the loss function minimizes the mean-square error between learned summary and parameter. The network thereby solves the inverse problem of parameter estimation. We propose a branched network structure that contains collapsing layers that reduce a data set to summary statistics that are further mapped through fully connected layers to approximate the parameter estimate. We motivate our choice of network structure by theoretical considerations. In simulations we demonstrate three desirable properties of parameter estimates: finite sample exactness, robustness to data contamination, and algorithm approximation. These properties are achieved offering the the network varying sample size, contaminated data, and data needing algorithmic reconstruction during the training phase. In our simulations an EM-algorithm for genetic data is automatically approximated by the network. Simulation only approaches seem to offer practical advantages in complex modeling tasks where the simpler data simulation part is left to the researcher and the more complex problem of solving the inverse problem is left to the neural network. Challenging future work includes offering pre-trained models that can be used in a wide variety of applications.
Stefan Böhringer
Department of Biomedical Data Sciences Leiden University Medical Center Einthovenweg 20, 2333 ZC Leiden, The Netherlands
Amortized inference promises fast test-time Bayesian inference, but existing methods typically target a fixed model or a restricted model family. Extending amortization to new model families often requires retraining or test-time fine-tuning. In this paper, we ask: is it possible to build a single inference network capable of generalizing across varying priors, likelihoods, latent dimensions, and observation counts? We introduce Amortized Factor Inference Networks (AFINs), which describe a model as a list of typed prior and likelihood factors and map this description, together with the observations, to a Gaussian posterior approximation through an encode--merge--decode network whose learned parameters do not depend on the latent dimension or the number of observations. Experiments on synthetic and real-data benchmarks show that a single trained AFIN provides competitive posterior approximations at substantially lower test-time cost than NUTS and per-task variational inference. Optional importance sampling can further improve accuracy.
Joohwan Ko, Justin Domke
Manning College of Information and Computer Sciences University of Massachusetts Amherst
Simulation plays a central role in scientific discovery. In many applications, the bottleneck is no longer running a simulator; it is choosing among large families of plausible simulators, each corresponding to different forward models/hypotheses consistent with observations. Over large model families, classical Bayesian workflows for model selection are impractical. Furthermore, amortized model selection methods typically hard-code a fixed model prior or complexity penalty at training time, requiring users to commit to a particular parsimony assumption before seeing the data. We introduce PRISM, a simulation-based encoder-decoder that infers a joint posterior over both discrete model structures and associated continuous parameters, while enabling test-time control of model complexity via a tunable model prior that the network is conditioned on. We show that PRISM scales to families with combinatorially many (up to billions) of model instantiations on a synthetic symbolic regression task. As a scientific application, we evaluate PRISM on biophysical modeling for diffusion MRI data, showing the ability to perform model selection across several multi-compartment models, on both synthetic and in vivo neuroimaging data.
Manuel Gloeckler, J. P. Manzano-Patrón, Stamatios N. Sotiropoulos +2
Machine Learning in Science, University of Tübingen and Tübingen AI Center, Tübingen, Germany · Sir Peter Mansfield Imaging Centre, School of Medicine, University of Nottingham, UK · National Institute for Health Research (NIHR) Nottingham Biomedical Research Centre, Queens Medical Centre, Nottingham, United Kingdom +1