Flow Matching for Scalable Simulation-Based Inference

1 Introduction

• FMPE uses flexible architectures to emphasize complex, high-dimensional conditioning data, while its proposed design uses a pyramid from x to v_t and gated linear units for (θ, t).The authors also propose an alternative time weighting that improves performance in many tasks.Source: Figure 1 · p. 2, §1 · p. 3
• Discrete normalizing flows can be difficult to scale to the large networks needed for big-data experiments.NPE commonly uses these flows as conditional density estimators.Source: §1 · p. 2
• FMPE adapts flow matching to Bayesian inference, representing the posterior through a continuous flow defined by a vector field rather than a discrete sequence of transforms.Its deterministic trajectories follow ordinary differential equations, and the method provides direct access to the density.Source: §1 · pp. 2–3
• FMPE performs comparably or better than NPE across most SBI benchmarks and substantially outperforms it in gravitational-wave inference on training time, posterior accuracy, and scaling to larger networks.The benchmark evaluation also finds mass-covering posteriors in all cases.Source: §1 · p. 3

2 Preliminaries

• A discrete normalizing flow defines q(θ|x) by an invertible map from a simple base distribution, with composed transforms enabling expressive densities and efficient density evaluation.The transforms use triangular Jacobians and are interspersed with shuffling of θ.Source: §2 · p. 3
• A continuous flow evolves a base distribution into q(θ|x) over time through a vector field; sampling follows its ODE and density evolves according to the transport equation.Although the density is tractable in principle, sampling and density estimation can require many network passes to solve the ODE numerically.Source: §2 · pp. 3–4
• Flow matching trains a continuous flow by regressing its vector field on one that generates a target probability path, avoiding ODE integration during training.With sample-conditional paths, the objective becomes simple; minimizing it is equivalent to regression on the marginal vector field.Source: §2 · p. 4
• The sample-conditional Gaussian paths can be freely specified subject to boundary conditions; the experiments focus on the optimal transport path.The path should begin at the base distribution and end concentrated around the target sample within a region of size σ_min.Source: §2 · p. 4, Related work · p. 4

Related work

• Prior score-based SBI methods use stochastic SDE or Langevin formulations, which are not directly suited to posterior-density evaluation; FMPE has a tractable density and can support sequential NPE approaches.That sequential extension would no longer be amortized and is left for future work.Source: Related work · pp. 4–5

3 Flow matching posterior estimation

• FMPE replaces the intractable posterior expectation in flow matching with prior-and-simulator samples, training q(θ|x) by empirical risk minimization.Training samples θ from the prior and then simulates x corresponding to θ.Source: §3 · p. 5

3.1 Probability mass coverage

• For a bimodal target, FMPE’s Gaussian fit is mass-covering, paralleling NPE’s forward-KL fit, which captures both modes unlike the backward-KL fit.Mass coverage means containing the full posterior support, reducing the risk of ruling out possible explanations.Source: §3.1 · p. 5
• Under regularity assumptions on p0, ut, and vt, a forward-KL bound can be obtained; without such control, vector-field MSE alone cannot generally bound it.The theorem assumes square-integrable p0 with bounded growth of |∇ln p0(θ)| and vector fields with bounded second derivatives; neural-network parameterizations are not guaranteed to satisfy the regularity assumptions.Source: §3.1 · p. 6

3.2 Network architecture

• For SBI with complex observations x and typically low-dimensional θ, the architecture maps x to vector-field features and uses gated linear units for (t, θ) conditioning.The design starts from a feature extractor for x and adjusts its output dimension to dim(θ).Source: §3.2 · p. 6

3.3 Re-scaling the time prior

• FMPE samples t from a power-law prior that gives greater weight to larger t when α > 0, and this improved learning for distributions with sharp bounds in experiments.α = 0 recovers uniform time sampling.Source: §3.3 · p. 6

4 SBI benchmark

• The benchmark covers ten tasks spanning Gaussian toy models and epidemiology and ecology problems, with parameter dimensions 2–10 and observation dimensions 2–100.Each task is evaluated at simulation budgets N ∈ {10^3, 10^4, 10^5}.Source: §4 · p. 6
• FMPE is comparable to NPE on most benchmark tasks and outperforms it on several by C2ST; MMD also favors FMPE, though it can be hyperparameter-sensitive.NPE is among the benchmark’s highest-ranking methods on many tasks.Source: §4 · p. 7
• On SLCP with a large simulation budget, FMPE greatly improves over NPE on a task designed to have a simple likelihood but complex posterior.The authors relate this result to the enhanced flexibility of the FMPE density estimator.Source: §4 · p. 7
• In the displayed evaluation, all reference posterior samples fall within FMPE’s support, although some posterior regions may still be undersampled.The appendix reports mass coverage for the other benchmark tasks as well.Source: §4 · p. 7

5.1 Background

• Under stationary Gaussian detector-noise assumptions, GW data comprise 15,744 real values for two LIGO detectors, and conventional analyses can require up to ∼10^8 likelihood evaluations.Conventional analyses may take hours to months, depending on the event and signal-model complexity.Source: §5.1 · pp. 7–8

5.2 Experiments

• FMPE trains in ≈2 days versus ≈3 days for NPE on 5 · 10^6 simulations, despite using a larger network.The FMPE network has 1.9 · 10^8 learnable parameters, compared with 1.3 · 10^8 for NPE; both train for 400 epochs on an A100 GPU.Source: §5.2 · p. 8
• The GNPE comparison is only rough because results use different data conditioning; evaluation also omits three phase-related parameters under phase marginalization.The reference uses a different basis for the omitted parameters.Source: §5.2 · p. 8
• FMPE achieves mean JSD 0.5 mnat (NPE: 3.6 mnat) and max JSD < 2.0 mnat, while performing similarly to symmetry-enhanced GNPE.The maximum JSD is described as an indistinguishability criterion for GW posteriors.Source: §5.2 · p. 8, Figure 5 · p. 9
• Consistent log evidences—−7667.969±0.005 for FMPE and −7667.958±0.006 for NPE—provide another indication that FMPE posteriors cover the reference posterior.The passage notes that correct importance-sampling evidence requires the inferred posterior to cover the entire posterior.Source: §5.2 · p. 9

5.3 Discussion

• FMPE can devote network capacity to interpreting high-dimensional input rather than discrete-flow splines, but typically requires hundreds of forward passes at inference.The authors suggest this capacity allocation helps explain FMPE's higher accuracy on GW150914.Source: §5.3 · p. 9
• GW150914 is a loud event with good data quality; FMPE does not integrate physical symmetries and may need further enhancements to fully compete with GNPE.The authors mention a symmetry-aware architecture or larger networks as possible enhancements.Source: §5.3 · p. 9

6 Conclusions

• FMPE parameterizes a conditional vector field rather than a full distribution, allowing flexible path specifications, direct posterior-density access, and greater scalability.The conclusion attributes scalability to needing fewer parameters to learn the vector field, leaving capacity for larger networks.Source: §6 · p. 10
• Across 10 benchmark tasks FMPE performs competitively or better than other SBI methods, and on gravitational-wave inference it outperforms comparable discrete flows with samples on par with a symmetry-based method.The authors also note that optimal-transport paths yield more naturally structured latent spaces than discrete flows.Source: §6 · p. 10

A Gaussian flow

• The Gaussian-flow derivation uses the optimal-transport path μ_t(θ)=tθ_1 and σ_t=1−(1−σ_min)t.These specify the time-dependent mean and standard deviation in the stated path.Source: A Gaussian flow · p. 16

Marginal vector field

• The appendix derives a marginalized vector field from the sample-conditional field and marginal path, choosing its form so the continuous flow is restricted to a one-dimensional Gaussian.The chosen field has learnable mean and variance parameters.Source: Marginal vector field · pp. 16–17

• Without regularity, arbitrarily small mean squared error can still produce a non-mass-covering distribution with infinite forward KL divergence.In the example, the target is uniform on [−1, 1], while the learned distribution has support only on [−1, 0] ∪ [ε, 1].Source: B Mass covering properties of flows · p. 18
• Even for Lipschitz-continuous vector fields, the forward KL divergence cannot in general be bounded by mean squared error.The analysis establishes a weaker bound in this setting.Source: B Mass covering properties of flows · pp. 18–19
• Under additional regularity assumptions on the base distribution and vector fields, the forward KL divergence can be bounded by mean squared error, recovering mass coverage, though with weaker scaling than for SDEs.The theorem assumes regularity of p0 and bounded second derivatives of the vector fields.Source: B Mass covering properties of flows · pp. 19–20
• The stated KL bound depends on second-derivative regularity; the authors note that weaker sufficient assumptions may be possible.Source: B Mass covering properties of flows · p. 20

• The section gathers further details and results for the SBI benchmark.Source: C SBI Benchmark · p. 22

• For each task and budget, the search varies batch size, learning rate, network size, and the time-prior parameter, selecting the configuration with the best validation loss.Five percent of each simulation budget is reserved for validation.Source: §C.1 · p. 22

C.2 Additional results

• FMPE outperforms NPE on most benchmark tasks and simulation budgets under MMD, with larger improvements than those shown by C2ST.MMD is an alternative metric; the paper uses C2ST as its primary metric because MMD can be hyperparameter-sensitive.Source: §C.2 · p. 22, Figure 6 · p. 23
• FMPE with the optimal transport path clearly outperforms the comparable score-matching model on almost all tested configurations.The score-matching model uses a variance-preserving diffusion path [15].Source: §C.2 · p. 22, Figure 7 · p. 24
• For the two displayed tasks with dim(x) = 100 and dim(θ) ∈ {5, 10}, gated linear unit conditioning improves performance over simple concatenation.The authors interpret the gain as reflecting difficulty learning the lower-dimensional (t, θ) when concatenated with x.Source: §C.2 · p. 22, Figure 8 · p. 24
• FMPE covers the support of reference posterior samples, though some may have low learned density when model-reference deviation is high, especially for Lotka–Volterra.The histograms show learned densities for reference samples across simulation budgets.Source: §C.2 · p. 23, Figure 9 · p. 25
• On Two Moons, the proposed time-prior re-weighting improves performance by assigning greater training weight to samples with t closer to 1.The figure caption reports improvement across all simulation budgets.Source: §C.2 · p. 23, Figure 11 · p. 26

D Gravitational-wave inference

• Appendix D collects missing details and additional results for the gravitational-wave inference problem analyzed in Section 5.Source: D Gravitational-wave inference · p. 27

D.1 Network architecture and hyperparameters

• FMPE outputs a vector field, using additional residual blocks, gated (t, θ)-conditioning at each block, and a small residual network to embed (t, θ).The network builds on the embedding network developed in [7].Source: §D.1 · p. 27
• The ablation reuses the NPE embedding network and tests applying (t, θ)-conditioning from different layers of the residual network.Source: §D.1 · p. 27

D.2 Data settings

• The GW setup uses IMRPhenomPv2 and a 5,000,000-sample dataset generated in about 1 hour on 64 CPUs, with a reduced frequency range and no PSD conditioning.The changes were applied to FMPE and NPE alike; preliminary tests found similar performance with PSD conditioning.Source: §D.2 · p. 27

D.3 Additional results

• FMPE requires many ODE-solver passes, and computing probabilities adds vector-field divergence calculations that use memory and limit inference batch size.The reported FMPE inference setup uses dopri5 with absolute and relative tolerances of 1e-7.Source: §D.3 · p. 27, Table 5 · p. 28
• The ablation network matches the NPE embedding network, and the GW models infer masses in (Mc, q) while marginalizing over phase ϕc.The priors are uniform over specified ranges unless indicated otherwise.Source: Table 3 · p. 27, Table 4 · p. 28
• FMPE achieves performance comparable to NPE with ≈1/3 of its network size; delaying (t, θ)-conditioning impairs performance but reduces inference FLOPs.The comparison uses networks with the same hidden dimensions as the NPE embedding network.Source: Figure 12 · p. 28, §D.3 · p. 28