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

Jonas Arruda; Jan Hasenauer; Stefan T. Radev

stat.MLOct 15, 2026

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

Authors: Jonas Arruda, Jan Hasenauer, Stefan T. Radev

Organizations: University of Bonn · Rensselaer Polytechnic Institute

Abstract

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.

Figures & tables

Figure 1: Learned two-moons posteriors as samplers and as densities. Top row: 5,000 draws per method for the canonical observation $\mathbf{y}=(0,0)$ ; bottom row: the learned log-density on a $512{\times}512$ grid with the exact 95% highest-density region dashed. ( A ) A normalizing flow; ( B ) a consistency model with $S{=}32$ steps, which has no tractable density; ( C ) a mean flow with $S{=}1$ . ( D ) The same mean flow with $S{=}32$ and ( E ) a separately trained flow matching model integrated with 32 Euler steps. ( F ) The velocity head of the mean flow and ( G ) the flow matching model, both integrated with the adaptive Tsit5 ODE solver.

Figure 2: Composing one learned trial likelihood over

N

trials of a cognitive model. ( A ) NRMSE of the composed posterior to the generating parameter. ( B ) Worst posterior predictive discrepancy across choice proportion, conditional response-time quantiles and timeout rate. ( C ) Posterior contraction relative to the prior. Lines show the mean over 10 independent datasets, bands the sample standard deviation.

Figure 3: Posterior predictive simulations of the renewal model on one dataset. Observations

\tilde{y}

simulated from the learned likelihood at posterior draws of ( A ) our conditional mean flow with the masked transformer and ( B ) the normalizing flow with sentinel values; median with 50 and 90% bands against the observed data (black). The gray strip is the seeded window, and the striped columns mark the missing time points.

	Mean flow	Normalizing flow
NRMSE	$\mathbf{0.187}$ (0.008)	0.199 (0.009)
RMSE $\tilde{y}_{t}$	$\mathbf{0.055}$ (0.007)	0.159 (0.036)
RMSE $y_{t}$	$\mathbf{0.060}$ (0.013)	0.125 (0.031)
ESS	$\mathbf{4{,}234}$ (247)	2,700 (629)
Leapfrog steps	$\mathbf{121}$ (4)	480 (57)
Wall time (s)	36,146 (6,157)	$\mathbf{1{,}624}$ (285)

Table 1: Likelihoods with missing observations. Mean (sample standard deviation) over 10 datasets.

\tilde{y}_{t}

is simulated with the learned likelihood.

Figure 4: Budget-matched posterior accuracy on g-and-k. ( A ) C2ST against the exact posterior over mean flow steps

S

, at the shared budget of Appendix D . ( B ) Minimum bulk ESS per 1,000 network evaluations. ( C ) Draws kept per chain under the shared budget. Lines show the mean over 10 independent datasets, bands the standard deviation; the delayed acceptance MCMC line is indexed by its coarse step count, and horizontal lines mark the normalizing flow and the flow matching model with 200 Euler steps outside the budget.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Experiment	Method	ATC
Two moons	Normalizing flow	$-0.008$ (0.019)
	Mean flow, $S{=}1$ / $S{=}2$ / $S{=}32$	$-0.001$ (0.010) / $+0.007$ (0.011) / $+0.008$ (0.011)
	Consistency model, $S{=}1$ / $S{=}2$ / $S{=}32$	$+0.102$ (0.011) † / $+0.113$ (0.011) / $+0.001$ (0.014)
	Mean flow ( $\mathbf{v}$ -head), Euler $S{=}1$ / $S{=}2$ / $S{=}32$	$-0.179$ (0.013) † / $-0.134$ (0.015) / $-0.007$ (0.012)
	Flow matching, Euler $S{=}1$ / $S{=}32$	$-0.176$ (0.011) † / $-0.004$ (0.014)
	Mean flow ( $\mathbf{v}$ -head) / flow matching, adaptive	$-0.007$ (0.012) / $\mathbf{+0.000}$ (0.011)

Appendix

Table S1: TARP calibration of every trained estimator. Area between the TARP curve and the diagonal (ATC; 0 is calibrated, positive is overdispersed), mean (standard deviation) over 10 folds of 200 prior-predictive validation pairs (2,000 in total) with 200 draws each. Two moons calibrates the posterior; g-and-k, the renewal model and the DDM calibrate the likelihood.

\dagger

marks methods whose Kolmogorov-Smirnov

p

-value is below 0.01; the ATC closest to zero per experiment is in bold.

Hyperparameter	Search space	Selected
Cycle weight $\lambda$ (log)	$[1,100]$	$2.5$
Probability of replacing the condition in $\mathcal{L}_{\mathrm{cyc}}$	$[0,0.8]$	$0.6$
Fraction of time pairs with $t{=}s$	$[0.2,0.8]$	$0.7$
Exponent $p$ of the adaptive weight	$\{0,0.25,0.5,0.75,1\}$	$0.25$
Offset $c$ of the adaptive weight (log)	$[10^{-5},1]$	$0.1$
Scale of the Fourier time embedding (log)	$[0.02,30]$	$0.3$

Appendix

Table S2: Hyperparameters of the conditional mean flow. Search space of the Optuna search and the selected configuration used in all experiments.

		$S{=}2$	$S{=}4$	$S{=}8$	$S{=}16$	$S{=}32$
Consistency model	$\rho$	1	2	3	3.5	3.5
	C2ST, selected	0.918	0.542	0.555	0.539	0.535
	C2ST, worst	0.973	0.975	0.946	0.700	0.666
Mean flow	$\rho$	4.54	3.82	4.12	8.34	7.42
	$\sigma_{\min}$	$1.1{\times}10^{-4}$	$4.9{\times}10^{-2}$	$5.7{\times}10^{-5}$	$7.0{\times}10^{-2}$	$2.0{\times}10^{-3}$
	C2ST, selected	0.502	0.504	0.509	0.505	0.507

Appendix

Table S3: Sampling schedules selected on two moons. Discretization exponent

\rho

of the consistency model and EDM schedule

(\rho,\sigma_{\min})

of the conditional mean flow for each step count

S

, tuned on held-out observations; at

S{=}1

, the schedule has no free parameter. To show how much the selection matters, we report the tuning C2ST (3 held-out observations, 2,000 draws each) of the selected and of the worst of the 7 candidates.

Figure S1: Sample and density quality against step count on two moons. ( A ) C2ST against the reference posterior over mean flow steps

S

(

0.5

indicates indistinguishable samples). ( B ) Hellinger distance of the learned density from the exact posterior on a

512{\times}512

grid. Lines show the mean over 10 independent datasets, bands the standard deviation.

Figure S2: Wall time of every method across experiments. Warm wall times by method (log scale), mean over 10 datasets, for ( A ) two-moons sampling (5,000 draws), ( B ) g-and-k MCMC under the shared budget, with the flow matching baseline outside it, ( C ) renewal MCMC at equal draws and ( D ) DDM MCMC per trial count.

Method	Steps	Mass	KL	Hellinger
Normalizing flow	$S{=}1$	0.977 (0.055)	0.446 (0.073)	0.383 (0.023)
Mean flow	$S{=}1$	1.066 (0.017)	0.347 (0.137)	0.322 (0.052)
Mean flow	$S{=}2$	1.045 (0.030)	0.179 (0.077)	0.230 (0.042)
Mean flow	$S{=}8$	1.085 (0.093)	0.165 (0.053)	0.219 (0.050)
Mean flow	$S{=}32$	$\mathbf{1.003}$ (0.025)	0.160 (0.052)	0.169 (0.027)
Mean flow ( $\mathbf{v}$ -head)	Euler, $S{=}1$	$7.2{\times}10^{39}$ $(2.3{\times}10^{40})$	73.2 (19.0)	1.000 (0.000)

Appendix

Table S4: Density metrics on two moons. Total mass, KL divergence and Hellinger distance of each learned posterior density against the exact one, on a

512{\times}512

grid over the prior; mean (standard deviation) over 10 independent datasets. The best value per column is in bold, for the mass the one closest to 1.

Method	$N$	NRMSE	Cov. 95%	Min ESS	Div.	$\max_{k}\|z_{k}\|$
Mean flow $S{=}1$	4	0.268 (0.079)	0.90 (0.16)	1,396 (312)	1.0 (1.9)	0.73 (0.14)
Mean flow $S{=}1$	16	0.161 (0.033)	0.90 (0.16)	1,174 (378)	6.6 (11.5)	0.90 (0.30)
Mean flow $S{=}1$	64	0.102 (0.028)	0.90 (0.16)	1,205 (186)	2.4 (4.0)	1.13 (0.47)
Mean flow $S{=}1$	256	0.070 (0.019)	0.60 (0.31)	1,146 (335)	7.8 (21.7)	2.15 (1.03)
Mean flow $S{=}2$	256	0.061 (0.016)	0.73 (0.21)	1,094 (335)	0.6 (1.6)	1.93 (0.77)
Mean flow $S{=}4$	256	0.066 (0.018)	0.77 (0.27)	503 (620)	50.7 (50.1)	2.02 (0.77)

Appendix

Table S5: Composing the DDM trial likelihood, per method and trial count. NRMSE to the generating parameter, coverage of the 95% central interval, minimum bulk ESS, divergent transitions, and worst predictive discrepancy; mean (standard deviation) over 10 independent datasets.

Figure S3: Composed DDM posteriors on one dataset, over the trial count. Rows are

N\in\{4,16,64,256\}

; columns are the marginals of ( A ) the drift

\nu

, ( B ) the boundary separation

a

and ( C ) the non-decision time

t_{0}

for

S{=}1

S{=}8

and the 200-step baseline, with the generating parameter dashed and the sampled

t_{0}

bound dotted. ( D ) Posterior predictive CDF of the signed response time (bands: mean

\pm

one standard deviation over 200 replicates) against the observed trials. Shown for one of 10 datasets.

Figure S4: Effect of step count on DDM inference. A NRMSE over mean flow steps

S

, one line per trial count

N

, with the flow matching baseline with 200 Euler steps dashed. B Step count selected within each dataset by NRMSE among converged settings, over the trial count

N

. C Minimum bulk ESS over

S

, with the baseline dash-dotted. Lines show the mean over 10 independent datasets, bands the standard deviation.

Figure S5: Renewal-model observations and what each method is shown of them. ( A ) The observation: four log-incidence curves at one parameter, with the latent

\log I_{t}

of the first behind them and the seeded steps shaded, giving 123 observed values on a 128-step grid. ( B ) The reproduction number that generated them, against the

R{=}1

line and the prior’s 90% interval, for comparison with the generating trajectory. ( C ) The same curve under one 30% missingness pattern, as each method receives it: the observed coordinates alone for the masked transformer, against the vector with sentinel values for the normalizing flow.

Method	Steps	Samples	Min ESS	C2ST
Mean flow $S{=}1$	1	32,000	$\mathbf{10{,}498}$ (1,552)	0.604 (0.049)
Mean flow $S{=}2$	2	16,000	4,856 (939)	0.587 (0.046)
Mean flow $S{=}4$	4	8,000	2,436 (708)	$\mathbf{0.577}$ (0.039)
Mean flow $S{=}8$	8	4,000	1,143 (476)	0.599 (0.030)
Normalizing flow	1	32,000	8,023 (3,169)	0.679 (0.078)
Flow matching, Euler $S{=}200$	200	4,000	1,270 (331)	0.614 (0.048)

Appendix

Table S6: Budget-matched MCMC on g-and-k. Every method except the flow matching baseline shares the same budget of requested draws times network evaluations per density. The realized totals in forward passes are

0.9

2.0\times 10^{7}

for the NUTS arms and

1.0

1.4\times 10^{8}

for delayed acceptance MCMC. Mean (standard deviation) over 10 independent datasets; the best minimum ESS and C2ST are in bold.

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

Flow Matching for Fast Posterior Sampling in Bayesian Inverse Problems
Oct 1, 2026Jan Blechschmidt, Oliver G. Ernst, Moritz Poguntke +1Posterior Sampling Bayesian Inverse Problems
OneFlowSBI: One Model, Many Queries for Simulation-Based Inference
Jan 30, 2026Mayank Nautiyal, Li Ju, Melker Ernfors +5Simulation-Based Inference Amortized Inference
Department of Information Technology, Uppsala University, Uppsala, Sweden · Science for Life Laboratory, Uppsala University, Uppsala, Sweden
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

Oct 1, 2026J/K move · Enter open · S save

Flow Matching for Fast Posterior Sampling in Bayesian Inverse Problems

Jan Blechschmidt, Oliver G. Ernst, Moritz Poguntke, Björn Sprungk

Posterior Sampling Bayesian Inverse Problems Amortized Inference Conditional Flow Matching PDE Inverse Problems

Open paper AskPDF arXiv

Sampling from the posterior is the central task of computational Bayesian inverse problems. The standard workhorse in Bayesian inference - Markov chain Monte Carlo (MCMC) - is sequential, yields correlated samples, and must be rerun for each observation. Conditional flow matching offers an amortized alternative: a transport map, trained once on joint samples of parameter and data, that yields independent approximate posterior samples for any observation at negligible online cost, without new likelihood evaluations. We give a careful, MCMC-literate assessment of flow matching for PDE-based inverse problems with function-valued parameters. Exploiting the flow's tractable density, we derive computable accuracy estimates of the underlying approximate posterior in total-variation distance and Kullback-Leibler divergence and, moreover, propose a hybrid sampler that is asymptotically exact by Metropolization. We validate the accuracy estimates and demonstrate the amortization in several numerical examples, including electrical impedance tomography and a likelihood-free Lotka-Volterra model.

Abstract

Simulation-Based Inference

Explore similar work

Flow Matching for Fast Posterior Sampling in Bayesian Inverse Problems

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

Tokenised Flow Matching for Hierarchical Simulation Based Inference