Table of Contents
Fetching ...

BayesFlow 2.0: Multi-Backend Amortized Bayesian Inference in Python

Lars Kühmichel, Jerry M. Huang, Valentin Pratz, Jonas Arruda, Hans Olischläger, Daniel Habermann, Simon Kucharsky, Lasse Elsemüller, Aayush Mishra, Niels Bracher, Svenja Jedhoff, Marvin Schmitt, Paul-Christian Bürkner, Stefan T. Radev

TL;DR

BayesFlow 2.0 tackles the computational bottleneck of Bayesian inference by providing a backend-agnostic ABI framework that learns estimators from simulations to produce rapid posterior, likelihood, and model-quantity inferences. The approach emphasizes a modular Python-based workflow with components for simulation, data preparation, neural architectures, and diagnostics, enabling end-to-end analyses from simulated data to inference and model checking. The Lotka-Volterra case study demonstrates end-to-end posterior estimation, diagnostic tooling, and point estimation with summary statistics, illustrating robust calibration and posterior predictive checks. Compared with existing SBI tools, BayesFlow extends functionality with multiple backends, hierarchical models, rich generative architectures, ensemble training, and integrated diagnostics, positioning it for broader adoption and future probabilistic programming integrations.

Abstract

Modern Bayesian inference involves a mixture of computational methods for estimating, validating, and drawing conclusions from probabilistic models as part of principled workflows. An overarching motif of many Bayesian methods is that they are relatively slow, which often becomes prohibitive when fitting complex models to large data sets. Amortized Bayesian inference (ABI) offers a path to solving the computational challenges of Bayes. ABI trains neural networks on model simulations, rewarding users with rapid inference of any model-implied quantity, such as point estimates, likelihoods, or full posterior distributions. In this work, we present the Python library BayesFlow, Version 2.0, for general-purpose ABI. Along with direct posterior, likelihood, and ratio estimation, the software includes support for multiple popular deep learning backends, a rich collection of generative networks for sampling and density estimation, complete customization and high-level interfaces, as well as new capabilities for hyperparameter optimization, design optimization, and hierarchical modeling. Using a case study on dynamical system parameter estimation, combined with comparisons to similar software, we show that our streamlined, user-friendly workflow has strong potential to support broad adoption.

BayesFlow 2.0: Multi-Backend Amortized Bayesian Inference in Python

TL;DR

BayesFlow 2.0 tackles the computational bottleneck of Bayesian inference by providing a backend-agnostic ABI framework that learns estimators from simulations to produce rapid posterior, likelihood, and model-quantity inferences. The approach emphasizes a modular Python-based workflow with components for simulation, data preparation, neural architectures, and diagnostics, enabling end-to-end analyses from simulated data to inference and model checking. The Lotka-Volterra case study demonstrates end-to-end posterior estimation, diagnostic tooling, and point estimation with summary statistics, illustrating robust calibration and posterior predictive checks. Compared with existing SBI tools, BayesFlow extends functionality with multiple backends, hierarchical models, rich generative architectures, ensemble training, and integrated diagnostics, positioning it for broader adoption and future probabilistic programming integrations.

Abstract

Modern Bayesian inference involves a mixture of computational methods for estimating, validating, and drawing conclusions from probabilistic models as part of principled workflows. An overarching motif of many Bayesian methods is that they are relatively slow, which often becomes prohibitive when fitting complex models to large data sets. Amortized Bayesian inference (ABI) offers a path to solving the computational challenges of Bayes. ABI trains neural networks on model simulations, rewarding users with rapid inference of any model-implied quantity, such as point estimates, likelihoods, or full posterior distributions. In this work, we present the Python library BayesFlow, Version 2.0, for general-purpose ABI. Along with direct posterior, likelihood, and ratio estimation, the software includes support for multiple popular deep learning backends, a rich collection of generative networks for sampling and density estimation, complete customization and high-level interfaces, as well as new capabilities for hyperparameter optimization, design optimization, and hierarchical modeling. Using a case study on dynamical system parameter estimation, combined with comparisons to similar software, we show that our streamlined, user-friendly workflow has strong potential to support broad adoption.
Paper Structure (38 sections, 16 equations, 7 figures, 3 tables)

This paper contains 38 sections, 16 equations, 7 figures, 3 tables.

Figures (7)

  • Figure 1: Side-by-side comparison of likelihood-based (explicit) and simulation-based (implicit) models from a Bayesian perspective. In likelihood-based models, the data model $p(\mathcal{D} \mid \theta)$ and prior $p(\theta)$ can be explicitly sampled and evaluated. In simulation-based models, the data model $p_{\text{implicit}}(\mathcal{D} \mid \theta)$ and prior $p(\theta)$ can be sampled through stochastic simulation, but cannot be (easily) evaluated.
  • Figure 2: Overview of the basic amortized Bayesian workflow provided by BayesFlow. The Simulator interface provides users with automated grouping of data-generating functions and simulated data. The Adapter interface transforms the raw simulator outputs for training the neural approximator. The neural approximator typically consists of a SummaryNetwork that encodes observables into a latent summary vector, and an InferenceNetwork conditioned on this vector to approximate target distributions. The resulting estimates are then used for model validation through BayesFlow's diagnostics module, with a wide array of graphical and numerical diagnostics for checking computational faithfulness and model sensitivity. The entire workflow is encapsulated at the high level via the Workflow object, which allows users to rapidly iterate through the Bayesian workflow.
  • Figure 3: Simulated trajectories aggregates from the LV model with sampled priors.
  • Figure 4: Automatically generated visual diagnostics from the BasicWorkflow object. From top to bottom: training and validation loss, calibration ECDF, parameter recovery, and posterior $z$-score contraction.
  • Figure 5: Estimated posterior distribution of the LV model parameters for a single dataset, showing contraction towards the estimation targets compared to the priors.
  • ...and 2 more figures