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

TL;DR

OneFlowSBI addresses the lack of a generalizable SBI framework by learning a single flow-matching model over the joint distribution $p(\boldsymbol{\theta}, \mathbf{y})$, enabling flexible queries through a query-aware masking scheme. It introduces a masked linear interpolant and a mask-conditioned flow objective $\mathcal{L}_{\text{OneFlowSBI}}(\phi)$ that transports probability mass only along unobserved coordinates, allowing posterior sampling, likelihood estimation, and arbitrary conditionals from a single model. Across ten SBIBM benchmarks and two high-dimensional real-world problems, OneFlowSBI achieves competitive accuracy with state-of-the-art solvers while maintaining robustness to noise and missing data and achieving fast sampling with few ODE steps. The framework offers a versatile, amortized SBI tool suitable for exploratory workflows and multi-modal observations across domains.

Abstract

We introduce \textit{OneFlowSBI}, a unified framework for simulation-based inference that learns a single flow-matching generative model over the joint distribution of parameters and observations. Leveraging a query-aware masking distribution during training, the same model supports multiple inference tasks, including posterior sampling, likelihood estimation, and arbitrary conditional distributions, without task-specific retraining. We evaluate \textit{OneFlowSBI} on ten benchmark inference problems and two high-dimensional real-world inverse problems across multiple simulation budgets. \textit{OneFlowSBI} is shown to deliver competitive performance against state-of-the-art generalized inference solvers and specialized posterior estimators, while enabling efficient sampling with few ODE integration steps and remaining robust under noisy and partially observed data.

Figures (22)

• Figure 1: OneFlowSBI architecture. The flow network operates on the joint state $\mathbf{z}_t=(\bm{\theta}_t,\mathbf{y}_t)$, where observed and latent components are controlled by a binary mask $\mathbf{m}$. Conditioned on time $t$, the network learns a masked velocity field, yielding a single model of the joint distribution $p(\bm{\theta}, \mathbf{y})$.
• Figure 2: SBIBM benchmark results. C2ST across ten benchmark tasks and three simulation budgets ($10{,}000$; $20{,}000$; $30{,}000$), comparing OneFlowSBI with Simformer (Dense), Simformer (Undirected), and NPE. Lower C2ST indicates higher posterior fidelity.
• Figure 3: Multi-query inference on Two Moons. Using a single OneFlowSBI model, we target diverse densities $p(\cdot|\cdot)$ solely by varying the inference mask. Panels show the (a) posterior, (b) prior, (c) likelihood, and (d) evidence, alongside (e–g) arbitrary partial conditionals.
• Figure 4: Posterior inference for Bayesian image deblurring. We visualize the ability of OneFlowSBI to recover the posterior $p(\bm{\theta}\mid\mathbf{y})$ from a noisy, blurred observation $\mathbf{y}$. The posterior samples $\hat{\bm{\theta}}_{1:4}$ demonstrate high-frequency detail recovery, while the uncertainty map $\sigma$ (right) effectively captures the combined effects of pixel-level noise and information loss due to the blur kernel, particularly at object boundaries.
• Figure 5: Posterior predictive checks for shallow water inference. (a) Posterior samples of the inferred depth profile compared to the ground truth. (b–d) Wave amplitudes at selected time steps comparing observed and posterior predictive wavefields. (e–f) Full spatiotemporal wavefields illustrating posterior predictive reconstructions over time.

Theorems & Definitions (2)

• Theorem 1.1: Conditional continuity equation
• proof