OneFlowSBI: One Model, Many Queries for Simulation-Based Inference
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Simulation-Based Inference | SBIBM Bernoulli GLM raw | MMD^20.041 | 12 | |
| Simulation-Based Inference | SBIBM Lotka–Volterra | MMD^20.557 | 12 | |
| Simulation-Based Inference | SBIBM Two Moons | MMD^20.22 | 12 | |
| Simulation-Based Inference | SBIBM Bernoulli GLM | MMD^29.05 | 12 | |
| Simulation-Based Inference | SBIBM SIR | MMD^20.344 | 12 | |
| Simulation-Based Inference | SBIBM SLCP | C2ST Score90.5 | 12 | |
| Simulation-Based Inference | SBIBM Bernoulli GLM | C2ST0.684 | 12 | |
| Simulation-Based Inference | SBIBM Gaussian Mixture | MMD^20.3 | 12 | |
| Simulation-Based Inference | SBIBM Gaussian Linear Uniform | C2ST0.63 | 12 | |
| Simulation-Based Inference | SBIBM Two Moons | C2ST0.512 | 12 |