Sensitivity-Aware Amortized Bayesian Inference

TL;DR

SA-ABI extends amortized Bayesian inference by embedding sensitivity analyses into the training and inference pipeline through context variables that cover likelihood, prior, data, and approximator configurations. By using weight sharing across contexts and deep ensembles to quantify approximator uncertainty, SA-ABI enables near-instant evaluation of many configurations without refitting, while still enabling rigorous quantitative and qualitative sensitivity assessments. The approach demonstrates robust performance in parameter estimation and model comparison across COVID-19 dynamics, climate trajectories, and hierarchical decision-making models, and it can detect simulation gaps via ensemble variability. Overall, SA-ABI provides a scalable framework for provenance-aware Bayesian inference that directly surfaces hidden sensitivity dimensions and supports robust scientific inference in complex systems.

Abstract

Sensitivity analyses reveal the influence of various modeling choices on the outcomes of statistical analyses. While theoretically appealing, they are overwhelmingly inefficient for complex Bayesian models. In this work, we propose sensitivity-aware amortized Bayesian inference (SA-ABI), a multifaceted approach to efficiently integrate sensitivity analyses into simulation-based inference with neural networks. First, we utilize weight sharing to encode the structural similarities between alternative likelihood and prior specifications in the training process with minimal computational overhead. Second, we leverage the rapid inference of neural networks to assess sensitivity to data perturbations and preprocessing steps. In contrast to most other Bayesian approaches, both steps circumvent the costly bottleneck of refitting the model for each choice of likelihood, prior, or data set. Finally, we propose to use deep ensembles to detect sensitivity arising from unreliable approximation (e.g., due to model misspecification). We demonstrate the effectiveness of our method in applied modeling problems, ranging from disease outbreak dynamics and global warming thresholds to human decision-making. Our results support sensitivity-aware inference as a default choice for amortized Bayesian workflows, automatically providing modelers with insights into otherwise hidden dimensions.

Figures (18)

• Figure 1: Our proposed approach for sensitivity-aware amortized Bayesian inference (SA-ABI). Stage 1: During training, a distribution $p(C_L, C_P)$ over plausible likelihood and prior choices is encoded via context variables $C_L$ and $C_P$ in a deep ensemble of neural approximators. Stage 2: During inference, we cast costly model refits as a near-instant neural network prediction task conditioned on user-specified context $C$. Our amortized neural approach unlocks fast large-scale sensitivity analyses of all components in a Bayesian model: likelihood ($C_L$), prior ($C_P$), data ($C_D$), and approximator ($C_A$). Experiment 3 uses $V = 8\,100$ variations in prior and data alongside $M = 20$ deep ensemble members. The resulting amortized sensitivity analysis encompassing $V\cdot M=162\,000$ approximate posteriors would have been infeasible with existing methods.
• Figure 2: Experiment 1. The bivariate posterior (\ref{['fig:covid_b']}) of best recoverable parameters $\lambda$ and $\psi$ indicates substantial sensitivity in terms of uncertainty reduction for $\psi$, but the posterior predictive distribution (\ref{['fig:covid_c']}) appears robust (largely overlapping median prediction lines and 90% CIs).
• Figure 3: Experiment 2. Two examples of simulated observations from the climate model ACCESS-ESM1-5 (SSP-3) with known time-to-threshold (training data; left, center) and the current empirical observation that we use for the forecasts (right).
• Figure 4: Experiment 2. Global warming forecasts are sensitive to the assumed climate model (rows) but not the emission scenario (SSP; groups of rows) or the prior (dotted: weakly informative prior; solid: informative prior).
• Figure 5: Experiment 3. (\ref{['fig:levy_sensitivity']}) Our sensitivity-aware posterior model probabilities indicate substantial approximator sensitivity but robustness to additional prior scaling and data perturbations. Dots represent the original results by elsemuller2023deep. (\ref{['fig:levy_prior_sensitivity']}) The posteriors are quantitatively sensitive to power-scaling of the prior location $\mu_\alpha$, as indexed by the ensemble-averaged probability for $\mathcal{M}_3$ (left) as well as the KL divergence between the original results by elsemuller2023deep vs. scaled model posteriors (right). Notable qualitative sensitivity is present mainly due to different $\mu_\alpha$ values.