PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jonas Arruda, Jan Hasenauer, Stefan T. Radev

arXiv:2610.21657v1stat.MLcs.LG

TL;DR

Mechanistic simulators often lack closed-form likelihoods, and repeated density evaluation can make MCMC costly. The paper develops conditional mean flows for posterior and likelihood approximation, including masked marginal estimation. Across its benchmarks, single-step mean-flow likelihood estimation has error comparable to 200-step flow matching while enabling about 200-fold faster downstream MCMC sampling.

  • Problem

    Mechanistic simulators often have unavailable likelihoods, while SBI approaches can restrict architecture or make repeated density evaluation expensive.

    Source: §1 · pp. 1–2

  • Method

    Conditional mean flows learn average velocities to define finite-time maps for posterior and likelihood approximation, with masking for marginal densities.

    Source: §1 · p. 2

  • Results

    Across four NLE and NPE tasks, one mean-flow step matches 200-step flow-matching error in a single-trial likelihood comparison with about 200-fold faster sampling.

    Source: §1 · p. 2

  • Takeaways & Limitations

    In low-dimensional benchmarks, one or two steps matched flow-matching accuracy while reducing MCMC time by up to two orders of magnitude.

    Source: §5 · p. 8

    High-dimensional density evaluation is costly because each step requires a full Jacobian and log-determinant, whose cost grows cubically with dimension; NUTS also differentiates the log-determinant.

    Source: §5 · p. 8

Abstract

from arXiv · show

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.

1 INTRODUCTION

SBI addresses inference when mechanistic simulators have unavailable likelihoods, but existing estimators can restrict architecture or make density evaluation expensive. The paper introduces conditional mean flows for posterior and likelihood approximation, with masked marginal-density estimation and benchmarks across four tasks.

1 INTRODUCTION

  • SBI uses simulated parameter-observation pairs to infer unknown parameters when mechanistic simulators lack closed-form likelihoods.Amortized methods reuse the trained network across datasets.Source: §1 · p. 1
  • Repeated likelihood evaluation is costly for NLE's MCMC, while existing SBI methods either restrict network architecture or make density evaluation expensive.Other density-based approaches can require hundreds of ODE network evaluations, and few-step consistency models lack tractable densities.Source: §1 · pp. 1–2
  • Conditional mean flows learn average velocity between times, so one network evaluation defines a finite-time map for posterior and likelihood approximation.The paper also proves the mean-flow objective identifies conditional flow maps and develops masked marginal-density estimation.Source: §1 · p. 2
  • Across four NLE and NPE tasks, one mean-flow step matches 200-step flow-matching error in a single-trial likelihood comparison with about 200-fold faster sampling.The benchmark includes normalizing flows and regular flow matching.Source: §1 · p. 2

2 BACKGROUND AND RELATED WORK

SBI methods address inference when likelihoods are unavailable, with NPE and NLE offering different trade-offs in amortization and posterior sampling. Flow-based and few-step generative methods provide relevant density and sampling approaches, but each has constraints.

2 BACKGROUND AND RELATED WORK

  • NPE learns a posterior and samples directly, whereas NLE learns a likelihood and uses MCMC to sample the posterior.NPE is usually tied to the training prior and data format; NLE permits changing the prior after training and composes exchangeable likelihoods across i.i.d. observations.Source: §2 · p. 2
  • Flow matching learns a velocity field and uses ODE integration for sampling and density evaluation, so repeated integration can dominate NLE costs.Density evaluation integrates velocity divergence through the continuous change-of-variables formula.Source: §2 · p. 2
  • Consistency models enable one- or few-step generation and fast SBI posterior sampling, but their standard formulation lacks a tractable posterior density.Mean flows parameterize finite-time ODE maps through average velocities and can be trained without distillation.Source: §2 · p. 2
  • Delayed-acceptance MCMC screens proposals with a cheap approximation while preserving the fine target, and multilevel variants organize several approximations.Related SBI approaches also use simulators at different resolutions.Source: §2 · p. 2

3 METHOD

Conditional mean flows parameterize conditional finite-time maps through average velocities, supporting few-step sampling and tractable densities. The method composes stepwise maps for density evaluation, extends them to MCMC and delayed acceptance, and uses masking to estimate marginal densities.

3.1 Flow-Based Density Evaluation

  • A conditional mean flow predicts average velocity along a trajectory, and one evaluation applies the finite-time map Φψ(z; s, t, x) = z + (t − s)uψ(z, s, t, x).The conditioning variable x remains fixed along the trajectory; at s=t, the mean flow corresponds to instantaneous velocity.Source: §3.1 · p. 3
  • A cycle loss encourages map invertibility, while an auxiliary velocity head allows the trained mean flow to be used as a flow-matching model.The loss also replaces some conditioning inputs with another simulation's condition to encourage invertibility across conditions.Source: §3.1 · p. 4
  • The method composes scheduled maps toward noise and evaluates density with the base Gaussian log-density plus the stepwise Jacobian log-determinants.The chain-rule expression is a pullback of the base density when the composed map is a diffeomorphism; Jacobians are computed by autodiff.Source: §3.1 · p. 4
  • Every population minimizer of the unweighted mean-flow loss identifies the conditional flow map, making the composed-step density equal to log p(z | x) for every schedule under the stated assumptions.The proposition characterizes the average velocity using the conditional expectation of the flow-matching velocity.Source: §3.1 · p. 4

3.2 Sampling and Posterior Inference

  • Sampling draws from a Gaussian base and applies the flow maps backward; samples follow the density only if the backward maps invert the forward maps, which the cycle loss encourages.In posterior estimation the draws are parameters conditional on data; in likelihood estimation they are observations conditional on parameters.Source: §3.2 · p. 4
  • Two-level delayed-acceptance MCMC uses a coarse likelihood to propose through a reversible subchain, then corrects proposals to leave the fine posterior invariant.The fine likelihood is evaluated once per subchain rather than at every proposal.Source: §3.2 · p. 4

3.3 Masked Transformers for Marginal Densities

  • Masked attention excludes omitted target tokens, and the losses operate only on observed target coordinates to estimate marginal densities.The proposition states that the masked-loss minimizer yields log p(z_O | x) for every training mask under its assumptions.Source: §3.3 · p. 5
  • Choosing parameters as the target gives marginal posterior densities, while choosing observations gives marginal likelihoods for missing observations.The paper evaluates marginal likelihoods for missing observations.Source: §3.3 · p. 5

4 EXPERIMENTS

Across four SBI benchmarks, conditional mean flows support accurate few-step sampling and density evaluation, with results varying by task and inference setting. The experiments compare posterior and likelihood quality, missing-data handling, and MCMC cost.

4.1 Few Steps Suffice for High Posterior Sample and Density Quality

  • On two moons, one mean-flow step matches the sampling accuracy across step counts, outperforming the normalizing flow at 0.523 against 0.543.S=1 sampling is about 23 times faster than flow matching with 32 Euler steps and 140 times faster than with the adaptive solver.Source: §4.1 · p. 5
  • In two moons, additional mean-flow steps improve density accuracy, while S=1 places mass between the two modes but still outperforms the normalizing flow.Mean-flow density is close to normalized; Euler densities overshoot unit mass by many orders of magnitude at S=1 and remain far from normalized at S=32.Source: §4.1 · pp. 5–6

4.2 Single-Step Likelihoods Yield Error Comparable to 200-Step Flow Matching

  • With S=1, cognitive-model NRMSE and posterior contraction are comparable to 200-step flow matching at every N, while sampling is 182 to 216× faster.The comparison covers N ∈ {4, 16, 64, 256} composed trial likelihoods; at N=256, additional discretization error further lowers coverage.Source: §4.2 · p. 6, Figure 2 · p. 6
  • In the cognitive-model experiments, 8.5% of MCMC runs failed to converge, and every failed run used S ≥4.Additional steps increased cost and convergence failures without clearly reducing parameter error; the authors attribute this to poorly trained intermediate maps.Source: §4.2 · p. 6

4.3 Masked Transformer Backbone Reduces Error under Missing Observations

  • For the renewal task, the masked transformer has six times higher ESS per network evaluation, although each transformer evaluation costs more computationally.The model uses a masked transformer for missing coordinates, while the normalizing flow is trained and evaluated with −1 in missing entries.Source: §4.3 · p. 7
  • On the renewal task, a masked-transformer mean flow has lower posterior error than a sentinel-value normalizing flow, but sampling finishes 22× sooner with the normalizing flow.The mean-flow predictive bands contract around observed data, whereas the normalizing flow frequently generates the sentinel value at observed times.Source: §4.3 · pp. 6–7

4.4 More Posterior Draws Are Better Than More Flow Steps

  • For g-and-k, every mean-flow setting converges on every dataset with similar accuracy, while the normalizing flow and regular flow matching have higher C2ST.The normalizing flow requires more than four times the wall time of mean flow at S=2; mean flow is about 50× faster than regular flow matching.Source: §4.4 · p. 8, Figure 4 · p. 8
  • For g-and-k, more posterior draws outperform more mean-flow steps beyond S=4 under the shared budget, and S=2 is recommended.The minimum bulk ESS per network evaluation drops by an order of magnitude from S=1 to S=8; delayed acceptance is comparable to NUTS without gradients.Source: §4.4 · p. 8, Figure 4 · p. 8

5 DISCUSSION

Conditional mean flows let users choose the number of inference steps after training, and few steps retain accuracy in low-dimensional benchmarks. High-dimensional density costs and other limitations constrain their practical performance.

5 DISCUSSION

  • In low-dimensional benchmarks, one or two steps matched flow-matching accuracy while reducing MCMC time by up to two orders of magnitude.Few steps also yielded accurate posterior samples, and under a fixed MCMC budget more draws were more valuable than more steps.Source: §5 · p. 8
  • In high dimensions, full Jacobian log-determinants scale cubically and NUTS requires second-order derivatives; renewal-task MCMC was slower than with normalizing flows.This occurred despite using fewer network evaluations and alongside higher transformer-evaluation costs.Source: §5 · p. 8
  • The cycle loss encourages but does not guarantee invertibility, so learned densities may not be exactly normalized.Source: §5 · p. 8
  • Coverage decreased for every method, including 200-step flow matching, at large trial counts, consistent with error accumulation.Source: §5 · p. 8
  • Fast mean-flow densities could support self-consistency objectives and evidence estimation, and may make flexible flow-based estimators practical in SBI settings requiring repeated density evaluation.Source: §5 · p. 8

A.1 Conditional Mean Flows

Conditional mean flow training identifies the conditional velocity and average-velocity maps, yielding exact conditional densities for any step schedule under the stated assumptions. The result has regularity and weighting qualifications, including for targets with discontinuous densities.

A.1 Conditional Mean Flows

  • For compact-support or discontinuous targets, the velocity may fail to be Lipschitz near time zero, and the argument applies from δ>0 to a slightly smoothed target.In practice, the schedule is not truncated at zero because time zero is almost surely never sampled.Source: §A.1 · p. 13
  • Every population minimizer recovers the conditional velocity and average-velocity maps, and their composed map transports the conditional target to the Gaussian base with exact log density for any schedule.This result assumes the unweighted objective and assumptions (A1)–(A4).Source: §A.1 · p. 13
  • The proof uses conditional-expectation decomposition and the MeanFlow identity: the exact velocity and average-velocity pair makes the auxiliary loss terms zero.Continuity and full support at interior times extend the minimizer equalities to all states and times.Source: §A.1 · p. 13
  • Residual-adaptive weighting with p>0 generally changes the learned velocity because the weighted regression decomposition no longer applies, though it empirically seems to stabilize training.Adding the cycle loss does not change the population minimizers under the stated exact-pair condition.Source: §A.1 · p. 14

A.2 Marginalization by Masking

Independent coordinate masks let conditional mean flows estimate marginal densities from retained target coordinates. The result requires representable marginal flows and applies to masks with positive training probability.

A.2 Marginalization by Masking

  • For every mask O with positive probability, the unweighted masked objective identifies the marginal velocity and average-velocity maps, giving the log marginal density in |O| dimensions for any schedule.The proposition assumes the mask is independent of the simulation variables and that network outputs on O depend on the target only through zO.Source: §A.2 · p. 14
  • The masked loss is a positively weighted sum of conditional mean flow objectives for the marginal targets, so simultaneous realizability lets each component attain its lower bound.The proposition is applied separately to each marginal target z0,O.Source: §A.2 · p. 14
  • The synthetic training-mask independence is imposed by construction and does not assume a missingness mechanism in observed data; only observed coordinates enter the marginal density.Additional attention restrictions are covered only if they preserve representation of the marginal velocity and its average velocity.Source: §A.2 · p. 14

B EVALUATION METRICS

The evaluation metrics assess calibration, distributional agreement, parameter recovery, predictive checks, and sampler convergence across the paper’s experiments. Several diagnostics use repeated validation folds or posterior draws, and failed convergence runs remain included in reported summaries.

B EVALUATION METRICS

  • C2ST held-out accuracy measures distinguishability between approximate and reference posterior draws, with 0.5 indicating indistinguishable distributions and 1 indicating perfect separation.Accuracy is averaged over five-fold cross-validation.Source: B EVALUATION METRICS · p. 14
  • On two moons, grid evaluation reports learned-density total mass and compares normalized grid shapes using Hellinger distance, which remains finite when the learned density vanishes where the exact posterior has mass.The exact posterior and learned densities are evaluated on a 512×512 midpoint grid.Source: B EVALUATION METRICS · p. 15
  • Posterior-draw metrics include normalized RMSE relative to prior parameter ranges and coverage, the fraction of generating parameter values inside central 95% credible intervals.Posterior contraction is defined separately as the relative reduction in variance from prior to posterior.Source: B EVALUATION METRICS · p. 15
  • TARP reports area between the calibration curve and diagonal, where zero is calibrated, and a Kolmogorov–Smirnov p-value across 2,000 validation pairs and 200 draws each.It assesses posterior calibration on two moons and likelihood-generated data calibration on the other experiments.Source: B EVALUATION METRICS · p. 15, Table S1 · p. 15
  • For the DDM, posterior predictive checks compare observed choice proportions, response-time quantiles, and timeout rate with replicate distributions, reporting the largest standardized discrepancy.The checks use 200 posterior draws.Source: B EVALUATION METRICS · p. 16
  • A run passes convergence criteria at max R-hat ≤1.05, minimum bulk ESS of 100, and at least 500 retained draws; reported means and standard deviations include failed runs.All metrics are summarized over 10 independent datasets.Source: B EVALUATION METRICS · pp. 14, 16

C ARCHITECTURES, TRAINING AND HYPERPARAMETER SEARCH

The experiments compare conditional mean flows with several generative architectures using specified network configurations, shared training procedures, and selected hyperparameters. The schedule comparison on two moons evaluates both the selected schedule and the worst candidate at each step count.

C ARCHITECTURES, TRAINING AND HYPERPARAMETER SEARCH

  • The mean-flow network is a five-layer width-256 MLP with a separate auxiliary velocity head, while flow matching uses the same depth and width without that head.The consistency model is a six-layer width-256 MLP; the normalizing flow uses spline coupling layers.Source: C ARCHITECTURES, TRAINING AND… · p. 16
  • The masked transformer uses five width-128 blocks with four attention heads and independently drops each observation during training with probability 0.3.Dropped tokens are excluded from attention and loss; a causal mask restricts each observation token to tokens up to its time step.Source: C ARCHITECTURES, TRAINING AND… · p. 16
  • Training used 50,176 simulations for two moons and the DDM, 100,000 for g-and-k, and 8,192 for renewal, with 1,024 additional validation simulations.Networks were trained offline with batch size 256 and AdamW; mean flow, consistency, and flow matching ran for 1,000 epochs, while the normalizing flow ran for 200 epochs with early stopping.Source: C ARCHITECTURES, TRAINING AND… · p. 16
  • Mean-flow hyperparameters were selected by Optuna search on two moons and a three-dimensional Gaussian linear task, and the selected configuration was used in all experiments.The search evaluated 200 trials using a score combining accuracy, calibration error, and grid Hellinger distance.Source: C ARCHITECTURES, TRAINING AND… · pp. 16–17, Table S2 · p. 17
  • On two moons, schedules were tuned separately by step count on held-out observations, and Table S3 compares the selected schedule with the worst of seven candidates using tuning C2ST.At one step the schedule has no free parameter; other experiments use ρ=7 and σmin=0.001.Source: C ARCHITECTURES, TRAINING AND… · p. 17, Table S3 · p. 17

D MCMC SAMPLERS

The section describes the NUTS setup and a two-level delayed-acceptance sampler that uses the same trained network at coarse and fine levels.

D MCMC SAMPLERS

  • Four NUTS chains use a diagonal mass matrix, maximum tree depth 10, and nutpie’s default target acceptance rate.The implementation runs through PyMC on the JAX backend.Source: D MCMC SAMPLERS · p. 17
  • NUTS uses 400 warmup iterations by default, while g-and-k uses warmup equal to 25% of each method’s requested draws.Source: D MCMC SAMPLERS · p. 18
  • The g-and-k budget equalizes requested draws times network evaluations per density, but ignores leapfrog steps and excludes automatic-differentiation costs from the pass count.Delayed acceptance adjusts requested iterations for five coarse evaluations and one fine evaluation; realized pass totals can therefore vary.Source: D MCMC SAMPLERS · p. 18
  • Delayed acceptance uses a mean-flow coarse posterior and a velocity-head fine posterior, with five coarse random-walk steps before the fine-level proposal.The coarse flow uses 1, 2, 4, or 8 steps; the fine level uses 100 fixed-size Tsit5 steps and the exact divergence.Source: D MCMC SAMPLERS · p. 18

E PER-EXPERIMENT DETAILS

The experiments describe four inference tasks and their data-generating or evaluation setups, alongside supplementary figures and tables covering density quality, sampling, and MCMC diagnostics.

E.1 Two Moons

  • The two-moons posterior has two crescent-shaped modes, and each dataset is evaluated against 32,000 exact reference samples.The first dataset uses y = (0, 0); the other nine use prior-predictive observations.Source: §E.1 · p. 18

E.2 g-and-k

  • The g-and-k task uses a Gaussian copula to couple two g-and-k marginals and compares inference against exact-log-posterior reference draws.Each dataset has 50 observations, with parameters drawn from the stated priors; the marginal density is obtained by inverting the strictly increasing quantile function for k ≥ 0.Source: §E.2 · p. 18

E.3 Epidemic Renewal

  • The renewal task uses 123 observed values on a 128-step grid, with five seeded but unobserved initial incidences and a causal attention mask.Observations are missing completely at random; the masked transformer learns marginal likelihoods for independently drawn missingness patterns, while the normalizing flow receives sentinel values.Source: §E.3 · p. 19

E.4 Cognitive Model

  • The cognitive task models signed response times from a Wiener process, with drift, boundary separation, and non-decision time as parameters.Non-decision time is sampled below the fastest observed response time, a bound that tightens as trials are added; trial counts are nested.Source: §E.4 · p. 19

F ADDITIONAL FIGURES AND TABLES

  • Supplementary tables report two-moons density metrics and DDM trial-likelihood diagnostics, each averaged over 10 independent datasets.The two-moons metrics are total mass, KL divergence, and Hellinger distance; DDM metrics include NRMSE, coverage, minimum bulk ESS, divergences, and predictive discrepancy.Source: Table S4 · p. 19, Table S5 · p. 20
  • Two-moons plots show C2ST and Hellinger distance as mean-flow step count varies; DDM plots show NRMSE, selected step count, and minimum bulk ESS against trial count.The plots summarize means and standard deviations across 10 datasets.Source: Figure S1 · p. 20, Figure S4 · p. 23
  • The g-and-k table compares methods under a matched requested-draws-times-evaluations budget, while supplementary plots report method wall times across experiments.The realized forward-pass totals differ by method and MCMC arm.Source: Table S6 · p. 20, Figure S2 · p. 21
  • Further figures show DDM posterior marginals and predictive CDFs across trial counts, plus renewal observations and the masked versus sentinel-filled inputs used by the methods.The DDM posterior and predictive displays are for one of 10 datasets; the renewal example includes one 30% missingness pattern.Source: Figure S3 · p. 22, Figure S5 · p. 23
Loading 2610.21657v1…