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

2026-01-30

AI Summary

This work proposes a unified flow-matching generative model for simulation-based inference that overcomes the inefficiency of traditional methods, which require separate model training for distinct query tasks such as posterior sampling or likelihood estimation. By introducing a query-aware masking distribution, the framework supports diverse conditional inference tasks within a single model without task-specific retraining. Leveraging an ordinary differential equation (ODE) solver, the approach enables efficient sampling and, for the first time, achieves general-purpose, efficient, and robust multi-task simulation-based inference within one unified architecture. Evaluated on ten benchmark problems and two high-dimensional real-world inverse problems, the method matches or exceeds state-of-the-art performance while significantly improving sampling efficiency and demonstrating strong robustness to missing data and noise.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Probabilistic Inference

Application Category

• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Large language models for search

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.

Problem

• simulation-based inference
• posterior sampling
• likelihood estimation
• conditional distributions
• generative modeling

Innovation

• simulation-based inference
• flow matching
• query-aware masking
• unified generative model
• conditional inference