OneFlowSBI: One Model, Many Queries for Simulation-Based Inference

Mayank Nautiyal; Li Ju; Melker Ernfors; Klara Hagland; Ville Holma; Maximilian Werkö Söderholm; Andreas Hellander; Prashant Singh

stat.MLJan 30, 2026

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

Abstract

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 $\mathbf{m}$
Joint	$p(\bm{\theta},\mathbf{y})$	$[\mathbf{0}_{d_{\theta}};\mathbf{0}_{d_{y}}]$
Posterior	$p(\bm{\theta}\mid\mathbf{y})$	$[\mathbf{0}_{d_{\theta}};\mathbf{1}_{d_{y}}]$
Likelihood	$p(\mathbf{y}\mid\bm{\theta})$	$[\mathbf{1}_{d_{\theta}};\mathbf{0}_{d_{y}}]$
Marginal	$p(\mathbf{y})$	$[-\mathbf{1}_{d_{\theta}};\mathbf{0}_{d_{y}}]$
Marginal	$p(\bm{\theta})$	$[\mathbf{0}_{d_{\theta}};-\mathbf{1}_{d_{y}}]$
Mixed conditional	$p(\mathbf{z}_{G}\mid\mathbf{z}_{O})$	$\mathbf{m}_{G}=0,\;\mathbf{m}_{O}=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

\pm

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	$\mathbf{0.111}$	$6.0$	$\mathbf{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 ).

Panel of Figure 5 — 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 $\tau(x)$ dashed. (b) Depth posteriors at three cutoffs.

Query	Specialist	OneFlowSBI
$p(\bm{\theta}\mid\mathbf{y})$	$0.697\pm 0.013$	$0.720\pm 0.011$
$p(\mathbf{y}\mid\bm{\theta})$	$0.604\pm 0.013$	$0.758\pm 0.017$
$p(\bm{\theta})$	$0.493\pm 0.009$	$0.496\pm 0.006$
$p(\mathbf{y})$	$0.508\pm 0.017$	$0.505\pm 0.007$
Number of models	4	1

Table 3: Cost of generality on Lotka–Volterra. C2ST (mean

\pm

SD over five runs), with one specialist per query.

Figure 7: Sensitivity to the mask distribution on Lotka–Volterra (mean

\pm

SD over five runs). Weights on joint, posterior, likelihood,

p(\mathbf{y})

p(\bm{\theta})

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	$[\mathbf{0}_{\theta};\mathbf{0}_{y}]$	$[\bm{\varepsilon}_{\theta};\bm{\varepsilon}_{y}]$	$(\bm{\theta}_{1},\mathbf{y}_{1})$
Posterior	$[\mathbf{0}_{\theta};\mathbf{1}_{y}]$	$[\bm{\varepsilon}_{\theta};\mathbf{y}_{\rm obs}]$	$\bm{\theta}_{1}$
Likelihood	$[\mathbf{1}_{\theta};\mathbf{0}_{y}]$	$[\bm{\theta}_{\rm obs};\bm{\varepsilon}_{y}]$	$\mathbf{y}_{1}$
Marginal $p(\mathbf{y})$	$[-\mathbf{1}_{\theta};\mathbf{0}_{y}]$	$[\mathbf{0}_{\theta};\bm{\varepsilon}_{y}]$	$\mathbf{y}_{1}$
Marginal $p(\bm{\theta})$	$[\mathbf{0}_{\theta};-\mathbf{1}_{y}]$	$[\bm{\varepsilon}_{\theta};\mathbf{0}_{y}]$	$\bm{\theta}_{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

\pm

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

\pm

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(\mathbf{y}\mid\bm{\theta}^{*})

, one curve per observation coordinate.

		Posterior (s)			Likelihood (s)
Task	$d_{\theta}+d_{y}$	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 ( $\mu$ ), with its RMSE, and standard deviation ( $\sigma$ ) for four test images. The colour bar gives the posterior standard deviation.

Method	RMSE	Energy score
OneFlowSBI	$\mathbf{0.111\pm 0.0002}$	$\mathbf{2.20\pm 0.005}$
OneFlowSBI (MLP)	$0.181\pm 0.006$	$3.59\pm 0.12$
Simformer (all conditionals)	$0.326\pm 0.006$	$6.48\pm 0.12$
Simformer (posterior)	$0.245\pm 0.007$	$4.88\pm 0.14$
NPE	$0.247\pm 0.007$	$5.42\pm 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

\pm

standard deviation over five training runs; for the other rows, one run, mean

\pm

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 ( $\sigma$ ) given an observation whose central $8\times 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	$\mathbf{0.329\pm 0.021}$	$\mathbf{2.41\pm 0.18}$
Simformer (all conditionals)	$0.795\pm 0.088$	$5.52\pm 0.61$
Simformer (posterior)	$0.537\pm 0.080$	$3.75\pm 0.55$
NPE	$1.265\pm 0.059$	$9.25\pm 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

\pm

SD over five runs for OneFlowSBI , mean

\pm

SE over pairs otherwise.

	RMSE (m)		Energy score (m)		$90\%$ coverage
$T$ (h)	OneFlowSBI	Simformer	OneFlowSBI	Simformer	OneFlowSBI	Simformer
$10$	$\mathbf{3.18\pm 0.11}$	$4.30\pm 0.13$	$\mathbf{22.44\pm 0.85}$	$31.80\pm 1.07$	$0.87$	$0.57$
$30$	$\mathbf{0.93\pm 0.12}$	$1.61\pm 0.10$	$\mathbf{6.57\pm 0.86}$	$11.91\pm 0.78$	$0.94$	$0.66$
$60$	$\mathbf{0.43\pm 0.06}$	$0.98\pm 0.09$	$\mathbf{3.08\pm 0.42}$	$6.93\pm 0.67$	$0.93$	$0.81$

Appendix

Table 10: Shallow water: window posteriors

p(\bm{\theta}\mid\mathbf{y}_{\mathrm{obs}}^{\leq T})

100

test pairs. RMSE and energy score in metres, and

90\%

coverage, from

100

samples per pair; mean

\pm

SD over five runs for OneFlowSBI , mean

\pm

SE over pairs for Simformer.

Figure 25: Shallow water: window posteriors

p(\bm{\theta}\mid\mathbf{y}_{\mathrm{obs}}^{\leq 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\pm 0.2$	$6.0\pm 0.4$	$68.7\pm 0.4$	$16.6\pm 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.

Simulation-Based Inference

Simulation-based inference estimates parameters and uncertainty from simulators when evaluating a tractable likelihood is difficult or impossible.

See all 86 papers

Tokenised Flow Matching for Hierarchical Simulation Based Inference
Apr 22, 2026Giovanni Charles, Cosmo Santoni, Seth Flaxman +1Flow Matching Simulation-Based Inference
Imperial College London, London, United Kingdom · University of Oxford, Oxford, United Kingdom
GenSBI: Generative Methods for Simulation-Based Inference in JAX
May 26, 2026Aurelio AmerioFlow Matching Simulation-Based Inference
Instituto de Física Corpuscular (IFIC) Universitat de València & CSIC València, Spain
Conditional Mean Flows: Few-Step Densities for Simulation-Based Inference
Oct 15, 2026Jonas Arruda, Jan Hasenauer, Stefan T. RadevSimulation-Based Inference Flow-Based Generative Modeling
University of Bonn · Rensselaer Polytechnic Institute

Apr 22, 2026J/K move · Enter open · S save

Tokenised Flow Matching for Hierarchical Simulation Based Inference

Giovanni Charles, Cosmo Santoni, Seth Flaxman, Elizaveta Semenova

Imperial College London, London, United Kingdom · University of Oxford, Oxford, United Kingdom

Flow Matching Simulation-Based Inference Bayesian Inference Hierarchical Bayesian Modeling

Open paper AskPDF arXiv

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.

Abstract

Simulation-Based Inference

Explore similar work

Tokenised Flow Matching for Hierarchical Simulation Based Inference

GenSBI: Generative Methods for Simulation-Based Inference in JAX

Conditional Mean Flows: Few-Step Densities for Simulation-Based Inference