Sbi Reloaded: a Toolkit for Simulation-based Inference Workflows

2024-11-26

Journal of Open Source Software

AI Summary

Likelihood-free and gradient-free parameter calibration in black-box simulators poses significant challenges for Bayesian inference. Method: This paper introduces the first simulation-based, fully amortized, gradient-free, and parallelizable neural Bayesian inference framework, accompanied by the open-source PyTorch package SBI. The framework unifies neural posterior estimation (NPE), neural likelihood estimation (NLE), neural ratio estimation (NRE), and mixture density networks (MDNs), integrating Monte Carlo sampling, Bayesian optimization, and simulation scheduling into a modular, end-to-end workflow with production-ready defaults and comprehensive diagnostic tools. Contribution/Results: Evaluated across physics, biology, and astronomy, SBI substantially lowers the barrier to simulation-based inference, accelerates posterior estimation by multiple-fold, and achieves state-of-the-art reusability and scalability.

Technology Category

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

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Scientists and engineers use simulators to model empirically observed phenomena. However, tuning the parameters of a simulator to ensure its outputs match observed data presents a significant challenge. Simulation-based inference (SBI) addresses this by enabling Bayesian inference for simulators, identifying parameters that match observed data and align with prior knowledge. Unlike traditional Bayesian inference, SBI only needs access to simulations from the model and does not require evaluations of the likelihood-function. In addition, SBI algorithms do not require gradients through the simulator, allow for massive parallelization of simulations, and can perform inference for different observations without further simulations or training, thereby amortizing inference. Over the past years, we have developed, maintained, and extended $ exttt{sbi}$, a PyTorch-based package that implements Bayesian SBI algorithms based on neural networks. The $ exttt{sbi}$ toolkit implements a wide range of inference methods, neural network architectures, sampling methods, and diagnostic tools. In addition, it provides well-tested default settings but also offers flexibility to fully customize every step of the simulation-based inference workflow. Taken together, the $ exttt{sbi}$ toolkit enables scientists and engineers to apply state-of-the-art SBI methods to black-box simulators, opening up new possibilities for aligning simulations with empirically observed data.

Problem

• Tuning simulator parameters to match observed data
• Enabling Bayesian inference without likelihood evaluations
• Providing flexible tools for simulation-based inference workflows

Innovation

• Neural network-based Bayesian SBI algorithms
• No likelihood function or gradients required
• Amortized inference for diverse observations