Amortized Bayesian Workflow

TL;DR

The Amortized Bayesian Workflow tackles the scalability challenge of Bayesian inference across large collections of datasets by integrating fast amortized posterior approximations with gold-standard MCMC when necessary. It introduces a diagnostic-driven pipeline that reuses computations across steps, employing PSIS refinements and amortized initializations to accelerate convergence while preserving posterior quality. Across four diverse problems, the approach achieves substantial runtime reductions (5×–120×) with accuracy close to reference MCMC posteriors, even under distribution shifts. The modular design supports flexible component replacement and offers a practical path toward scalable, trustworthy Bayesian inference in real-world, data-rich settings.

Abstract

Bayesian inference often faces a trade-off between computational speed and sampling accuracy. We propose an adaptive workflow that integrates rapid amortized inference with gold-standard MCMC techniques to achieve a favorable combination of both speed and accuracy when performing inference on many observed datasets. Our approach uses principled diagnostics to guide the choice of inference method for each dataset, moving along the Pareto front from fast amortized sampling via generative neural networks to slower but guaranteed-accurate MCMC when needed. By reusing computations across steps, our workflow synergizes amortized and MCMC-based inference. We demonstrate the effectiveness of this integrated approach on several synthetic and real-world problems with tens of thousands of datasets, showing efficiency gains while maintaining high posterior quality.

Figures (11)

• Figure 1: Our workflow adaptively moves along the Pareto front and reuses previous computations.
• Figure 2: Our adaptive workflow leverages near-instant amortized posterior sampling when possible and gradually resorts to slower— but more accurate— sampling algorithms. As indicated by the blue dashed arrows, we reuse the $S$ draws from the amortized posterior in Step 1 for the subsequent steps in the form of PSIS proposals (Step 2) and initial values in ChEES-HMC (Step 3).
• Figure 3: Illustration of our sampling-based hypothesis test that flags atypical datasets.
• Figure 4: We initialize many ChEES-HMC chains with amortized draws.
• Figure 5: Evaluation of posterior draws across four problems based on two metrics: W1 distance (top row) and MMTV distance (bottom row). Lower values indicate better posterior approximation. ABI(✓) and ABI(✗) denote accepted and rejected draws, respectively, from amortized Bayesian inference in Step 1. PSIS denotes importance-weighted draws accepted in Step 2, and C-HMC denotes draws accepted via ChEES-HMC in Step 3.