Amortized Bayesian Workflow

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.

💡 Research Summary

Bayesian inference traditionally balances two competing goals: computational speed and sampling accuracy. Markov chain Monte Carlo (MCMC) offers strong theoretical guarantees and diagnostic tools but becomes prohibitively expensive when applied independently to thousands or tens of thousands of datasets, as each dataset requires a fresh run of the sampler. Amortized Bayesian inference (ABI), by contrast, learns a neural network mapping from observations to posterior distributions using simulated data, enabling near‑instant inference at test time. However, ABI lacks the rigorous diagnostics of MCMC and can fail under distributional shift.

The authors propose an adaptive workflow that moves along a Pareto front between these two extremes, reusing computations wherever possible. The workflow consists of a training phase and an inference phase. In the training phase, a conditional normalizing flow qϕ(θ|y) is trained by minimizing the forward KL divergence on a large simulated dataset (θ, y) ∼ p(θ)p(y|θ). Calibration is assessed with simulation‑based calibration (SBC) and parameter‑recovery plots; if diagnostics fail, hyper‑parameters or the simulation budget are adjusted and training is repeated.

During inference, each observed dataset y(k) is processed through three sequential steps. Step 1 draws S posterior samples from the amortized estimator. An out‑of‑distribution (OOD) test based on Mahalanobis distance of low‑dimensional summary statistics flags datasets whose summary lies in the extreme α‑tail (default α = 0.05) of the empirical training distribution. Datasets that pass are accepted as reliable amortized draws.

If a dataset is flagged OOD, Step 2 applies Pareto‑smoothed importance sampling (PSIS) to re‑weight the amortized samples. When PSIS efficiency exceeds a preset threshold (e.g., 0.7), the re‑weighted draws are accepted. Otherwise, the workflow proceeds to Step 3, where ChEES‑HMC—a GPU‑accelerated Hamiltonian Monte Carlo implementation—is initialized with the amortized samples and run to obtain gold‑standard draws. Importantly, the draws from earlier steps are reused as proposals for PSIS and as initial states for HMC, dramatically reducing overall computational effort.

The authors also discuss optional, more expensive diagnostics such as posterior SBC and the local classifier two‑sample test (L‑C2ST) for datasets where higher confidence is required. These can be employed depending on the number of datasets, simulation cost, and dimensionality of the observations.

Empirical evaluation on synthetic benchmarks and real‑world problems involving up to tens of thousands of datasets demonstrates substantial speed‑ups (often >5×) while preserving or improving effective sample size (ESS) and R̂ diagnostics relative to pure MCMC. The adaptive workflow automatically selects the cheapest method that meets the diagnostic criteria for each dataset, thereby achieving a scalable, principled solution for large‑scale Bayesian analysis.

In summary, the paper contributes (1) a modular, theoretically motivated adaptive Bayesian workflow that integrates amortized inference, importance sampling, and MCMC; (2) a suite of per‑dataset diagnostics that guide method selection and ensure posterior quality; and (3) extensive experimental validation showing that the approach delivers both efficiency and accuracy at scale, offering a practical pathway for deploying Bayesian methods in high‑throughput scientific and industrial settings.