OneFlowSBI: One Model, Many Queries for Simulation-Based Inference
Authors: Mayank Nautiyal, Li Ju, Melker Ernfors, Klara Hagland, Ville Holma, Maximilian Werkö Söderholm, Andreas Hellander, Prashant Singh
Organizations: Department of Information Technology, Uppsala University, Uppsala, Sweden · Science for Life Laboratory, Uppsala University, Uppsala, Sweden
Simulation-based inference (SBI) methods typically amortize a single query, usually the posterior, so each further conditional or marginal needs its own model. Existing generalists amortize many queries in one transformer with one token per variable, so their attention cost grows quadratically with the number of variables. We introduce OneFlowSBI, a masked flow matching framework that trains one network on the joint distribution of parameters and observations. A ternary mask marks each coordinate as observed, generated or absent, so any conditional or marginal is sampled by changing the mask. The mask is an input, rather than a tokenization, so the backbone can follow the data's structure and scale to many variables. On high-dimensional image denoising and shallow water inference, OneFlowSBI achieves lower posterior error than the baselines, samples far faster than a transformer-based generalist, and handles missing regions and partial observation windows without retraining. On most of the ten SBIBM tasks at the largest simulation budget, it matches or exceeds neural posterior estimation and is comparable to transformer-based generalists, while also providing the likelihood and marginals.
Figures & tables
Figure 1: OneFlowSBI overview. Training paths under the masked linear interpolant for joint, posterior, likelihood, and marginal queries. Trajectories induced by the learned velocity field need not be straight. Absent coordinates use zero placeholders, not conditioning values.
Task
Target density
Mask m
Joint
p(θ,y)
[0dθ;0dy]
Posterior
p(θ∣y)
[0dθ;1dy]
Likelihood
p(y∣θ)
[1dθ;0dy]
Marginal
p(y)
[−1dθ;0dy]
Marginal
p(θ)
[0dθ;−1dy]
Mixed conditional
p(zG∣zO)
mG=0,mO=1
Table 1: Masks for SBI queries. For mixed conditionals, G and O partition the coordinates into generated and observed subsets.
Figure 2: Posterior C2ST versus simulation budget (lower is better, 0.5 is optimal). Bars show ± one standard deviation: over five training runs and over the ten reference observations.
Figure 3: Lotka–Volterra. All panels use one trained model. Bands in (b) are central 90% intervals from 10,000 samples each; the gray band in (c) marks the 95% range of a calibrated model.
Figure 4: Bayesian denoising. One query per row, all from one trained OneFlowSBI model except row 2, which shows the posterior of NPE; revealed clean pixels are in yellow and the missing observation block in black. For the likelihood (row 3), the error is taken against the mean of simulator runs.
Denoising
Shallow water
Method
RMSE
Time (s)
RMSE (m)
Time (s)
OneFlowSBI
0.111
6.0
0.329
16.6
Simformer (all)
0.326
447
0.795
403
Simformer (post.)
0.245
443
0.537
395
NPE
0.247
0.08
1.265
0.04
Table 2: Case studies. RMSE of the posterior mean and sampling time for 1,000 posterior samples, with 30,000 simulations (Appendix F.1 – F.3 ).
Figure 5: Shallow water inference. All panels use one trained model. (a) Posterior standard deviation at each position and cutoff T , with the wave arrival time τ(x) dashed. (b) Depth posteriors at three cutoffs.
Query
Specialist
OneFlowSBI
p(θ∣y)
0.697±0.013
0.720±0.011
p(y∣θ)
0.604±0.013
0.758±0.017
p(θ)
0.493±0.009
0.496±0.006
p(y)
0.508±0.017
0.505±0.007
Number of models
4
1
Table 3: Cost of generality on Lotka–Volterra. C2ST (mean ± SD over five runs), with one specialist per query.
Figure 7: Sensitivity to the mask distribution on Lotka–Volterra (mean ± SD over five runs). Weights on joint, posterior, likelihood, p(y) , p(θ) and partial masks: default (0.15, 0.2, 0.2, 0.15, 0.1, 0.2), balanced (all 1/6), no partial (0.2, 0.2, 0.2, 0.2, 0.2, 0) and partial heavy (0.1, 0.1, 0.1, 0.1, 0.1, 0.5).
Appendix figures & tables25 assets
Supplementary material from the paper’s appendix.
Appendix
Query
Mask
Initial state
Output
Joint
[0θ;0y]
[εθ;εy]
(θ1,y1)
Posterior
[0θ;1y]
[εθ;yobs]
θ1
Likelihood
[1θ;0y]
[θobs;εy]
y1
Marginal p(y)
[−1θ;0y]
[0θ;εy]
y1
Marginal p(θ)
[0θ;−1y]
[εθ;0y]
θ1
Appendix
Table 4: Sampling under canonical query masks.
Figure 8: Gaussian Mixture.
Figure 9: Two Moons.
Figure 10: Gaussian Linear.
Figure 11: Gaussian Linear Uniform.
Figure 12: SLCP.
Figure 13: SLCP Distractors.
Figure 14: Bernoulli GLM.
Figure 15: Bernoulli GLM (raw).
Figure 16: SIR.
Figure 17: Lotka–Volterra.
Figure 18: C2ST of each learned distribution at 30,000 simulations (lower is better, 0.5 is optimal). Bars show ± one standard deviation over the ten reference observations, or over ten sample sets for the prior and prior predictive.
Table 5: Posterior C2ST across simulation budgets; closer to 0.5 is better. OneFlowSBI entries report the mean ± standard deviation of five per-run scores, each averaged over ten observations. Results are rounded to three decimal places. Best means are bold and second-best means are underlined , with ties sharing the same formatting.
Figure 19: Posterior calibration. Coverage of central posterior intervals against the credibility level, one curve per parameter, for one model per task at 30,000 simulations. The gray band marks the 95% range of a calibrated model.
Figure 20: Likelihood calibration. As Figure 19 , for central intervals of p(y∣θ∗) , one curve per observation coordinate.
Posterior (s)
Likelihood (s)
Task
dθ+dy
Ours
Simformer
ACE
Ours
Simformer
ACE
Gaussian mixture
2+2
1.15
2.96
0.03
1.07
2.91
0.03
Two moons
2+2
1.45
3.00
0.03
1.30
2.94
0.02
SIR
2+10
1.22
4.02
0.03
1.19
4.05
0.09
SLCP
5+8
1.43
4.20
0.08
1.33
4.42
0.08
Bernoulli GLM
10+10
1.03
5.52
0.18
0.96
5.34
0.12
Appendix
Table 7: Sampling cost of the three generalists at 30,000 simulations. Seconds per 10,000 posterior and likelihood samples, with tasks sorted by the number of variables.
Figure 21: Bayesian denoising: comparison with baselines. Posterior mean ( μ ), with its RMSE, and standard deviation ( σ ) for four test images. The colour bar gives the posterior standard deviation.
Method
RMSE
Energy score
OneFlowSBI
0.111±0.0002
2.20±0.005
OneFlowSBI (MLP)
0.181±0.006
3.59±0.12
Simformer (all conditionals)
0.326±0.006
6.48±0.12
Simformer (posterior)
0.245±0.007
4.88±0.14
NPE
0.247±0.007
5.42±0.13
Appendix
Table 8: Bayesian denoising: posterior accuracy on 100 test images. RMSE of the posterior mean and energy score ( Gneiting and Raftery,, 2007 ) , which also accounts for the posterior spread, from 100 posterior samples per image (lower is better). For OneFlowSBI with the convolutional network, mean ± standard deviation over five training runs; for the other rows, one run, mean ± standard error over images.
Figure 22: Bayesian denoising: calibration. TARP coverage of the posterior and likelihood against the credibility level, over 1,000 test pairs; the gray band marks the 95% range of a calibrated model.
Figure 23: Bayesian denoising: missing observation block. Posterior mean, with its RMSE, and standard deviation ( σ ) given an observation whose central 8×8 block is missing, for four test images.
Figure 24: Shallow water: comparison with baselines. Posterior of the depth profile given the full observation, with the RMSE of the posterior mean, for the test pair of Figure 6 .
Method
RMSE (m)
Energy score (m)
OneFlowSBI
0.329±0.021
2.41±0.18
Simformer (all conditionals)
0.795±0.088
5.52±0.61
Simformer (posterior)
0.537±0.080
3.75±0.55
NPE
1.265±0.059
9.25±0.40
Appendix
Table 9: Shallow water: posterior accuracy on 100 test pairs. RMSE and energy score in metres, from 100 samples per pair (lower is better); mean ± SD over five runs for OneFlowSBI , mean ± SE over pairs otherwise.
RMSE (m)
Energy score (m)
90% coverage
T (h)
OneFlowSBI
Simformer
OneFlowSBI
Simformer
OneFlowSBI
Simformer
10
3.18±0.11
4.30±0.13
22.44±0.85
31.80±1.07
0.87
0.57
30
0.93±0.12
1.61±0.10
6.57±0.86
11.91±0.78
0.94
0.66
60
0.43±0.06
0.98±0.09
3.08±0.42
6.93±0.67
0.93
0.81
Appendix
Table 10: Shallow water: window posteriors p(θ∣yobs≤T) on 100 test pairs. RMSE and energy score in metres, and 90% coverage, from 100 samples per pair; mean ± SD over five runs for OneFlowSBI , mean ± SE over pairs for Simformer.
Figure 25: Shallow water: window posteriors p(θ∣yobs≤T) for the test pair of Figure 6 . Posterior mean and central 90% interval of the depth profile; dotted lines mark the central 90% interval of the prior, and the vertical line the wave front at T=10 h. Simformer, trained on all conditionals, generates the later hours and discards them.
Bayesian denoising
Shallow water
Method
Training (min)
Sampling (s)
Training (min)
Sampling (s)
OneFlowSBI
39.3±0.2
6.0±0.4
68.7±0.4
16.6±0.7
Simformer (all conditionals)
416.5
447
98.6
403.2
Simformer (posterior)
414.2
443
97.5
395.3
NPE
33.7
0.08
22.7
0.04
Appendix
Table 11: Training and sampling cost. Training time on 30,000 simulations and time for 1,000 posterior samples for one observation.
The cost of simulator evaluations is a key practical bottleneck for Simulation Based Inference (SBI). In hierarchical settings with shared global parameters and exchangeable site-level parameters and observations, this structure can be exploited to improve simulation efficiency. Existing hierarchical SBI approaches factorise the posterior yet still simulate across multiple sites per training sample; We instead explore likelihood factorisation (LF) to train from single-site simulations. In LF sampling we learn a per-site neural surrogate of the simulator and then assemble synthetic multi-site observations to amortise inference for the full hierarchical posterior. Building on this, we propose Tokenised Flow Matching for Posterior Estimation (TFMPE), a tokenised flow matching approach that supports function-valued observations through likelihood factorisation. To enable systematic evaluation, we introduce a benchmark for hierarchical SBI. We validate TFMPE on this benchmark and on realistic infectious disease and computational fluid dynamics models, finding well-calibrated posteriors while reducing computational cost.
Giovanni Charles, Cosmo Santoni, Seth Flaxman +1
Imperial College London, London, United Kingdom · University of Oxford, Oxford, United Kingdom
Flow and diffusion generative models have established themselves as widely adopted density estimators for simulation-based inference (SBI), extending naturally from neural posterior estimation to likelihood and joint density estimation. Their principled optimization objectives and freedom from architectural constraints have driven rapid adoption across the natural sciences. Yet the most widely used SBI libraries remain PyTorch-based, leaving researchers who develop their forward models and analysis pipelines in JAX without a native option. We present GenSBI, an open-source library that implements flow matching, score matching, and denoising diffusion entirely in JAX. The library offers three transformer-based architectures - SimFormer, Flux1, and a novel Flux1Joint that extends gate-modulated transformer blocks to joint density estimation - all interchangeable through a unified interface that decouples generative method, neural backbone, and inference mode. GenSBI provides an end-to-end workflow from training through posterior calibration (SBC, TARP, LC2ST) and supports custom architectures with domain-specific embedding networks. We validate the framework on standard SBI benchmarks, achieving near-ideal mean C2ST scores (0.50-0.56, where 0.50 is ideal) on SBIBM tasks with minimal per-task tuning and well-calibrated posterior coverage across all tested configurations. The code is publicly available at https://github.com/aurelio-amerio/GenSBI.
Aurelio Amerio
Instituto de Física Corpuscular (IFIC) Universitat de València & CSIC València, Spain
Flow matching has emerged as a powerful generative family for simulation-based inference (SBI). However, sampling and density evaluation require numerical integration at inference time. This makes posterior sampling in neural posterior estimation costly. More importantly, it can make repeated likelihood evaluation prohibitively expensive and impractical for Markov chain Monte Carlo (MCMC) sampling, which may require thousands of repeated density evaluations. Here, we develop conditional mean flows for posterior and likelihood approximation, combining few-step sampling with tractable density evaluation. We show that the mean flow objective identifies conditional flow maps and that independent masking enables direct estimation of marginal likelihoods. We evaluate both posterior and likelihood estimation across several tasks and show that conditional mean flows retain the flexibility of flow-based density estimators while significantly reducing inference time. In particular, single-step conditional mean flow density evaluation yields estimation error comparable to 200-step flow matching while enabling about 200-fold faster downstream MCMC sampling.
Jonas Arruda, Jan Hasenauer, Stefan T. Radev
University of Bonn · Rensselaer Polytechnic Institute