Uncertainty-Aware Surrogate-based Amortized Bayesian Inference for Computationally Expensive Models

Bayesian inference typically relies on a large number of model evaluations to estimate posterior distributions. Established methods like Markov Chain Monte Carlo (MCMC) and Amortized Bayesian Inference (ABI) can become computationally challenging. While ABI enables fast inference after training, generating sufficient training data still requires thousands of model simulations, which is infeasible for expensive models. Surrogate models offer a solution by providing approximate simulations at a lower computational cost, allowing the generation of large data sets for training. However, the introduced approximation errors and uncertainties can lead to overconfident posterior estimates. To address this, we propose Uncertainty-Aware Surrogate-based Amortized Bayesian Inference (UA-SABI) – a framework that combines surrogate modeling and ABI while explicitly quantifying and propagating surrogate uncertainties through the inference pipeline. Our experiments show that this approach enables reliable, fast, and repeated Bayesian inference for computationally expensive models, even under tight time constraints.

💡 Research Summary

This paper addresses the fundamental bottleneck of Bayesian inference for computationally expensive simulators: the need for a large number of model evaluations. Classical approaches such as Markov Chain Monte Carlo (MCMC) provide high‑quality posterior samples but are prohibitively slow, while Amortized Bayesian Inference (ABI) offers near‑instant inference after a costly training phase that itself requires thousands of simulator runs. The authors propose a novel framework called Uncertainty‑Aware Surrogate‑based Amortized Bayesian Inference (UA‑SABI) that combines surrogate modelling with ABI while explicitly quantifying and propagating surrogate uncertainties throughout the inference pipeline.

The methodology consists of two main stages. In the first stage, a surrogate model (f_{\mathcal{M}}^{c}) is trained on a sparse set of expensive simulator outputs (D_T={(x^{(i)},\omega^{(i)},y^{(i)})}_{i=1}^{N_T}) with (N_T \ll N_B). Rather than producing a point estimate of the surrogate parameters (c), the authors adopt a Bayesian treatment: they infer a joint posterior (p(c,\sigma\mid D_T)) over the surrogate coefficients and an error‑scale parameter (\sigma) that governs a distribution (p(\varepsilon\mid\sigma)) for the approximation error (\varepsilon). This captures both epistemic uncertainty (limited training data) and aleatoric error (model misspecification). Sampling from this posterior can be performed with MCMC because evaluating the surrogate likelihood is cheap.

In the second stage, ABI is performed using data generated from the uncertainty‑aware surrogate. For each synthetic training pair, the algorithm draws a sample ((c_i,\sigma_i)) from the surrogate posterior, computes the deterministic surrogate output (e_y^{(i)} = f_{\mathcal{M}}^{c_i}(x^{(i)},\omega^{(i)})), and then perturbs it with a draw (\varepsilon_i \sim p(\varepsilon\mid\sigma_i)) to obtain an error‑adjusted observation (e_y^{\varepsilon,(i)}). This pair ((x^{(i)}, e_y^{\varepsilon,(i)})) is fed into a summary network (S_\theta) that produces a fixed‑size embedding (s_i). An inference network (I_\phi) (implemented as a normalizing flow) then learns to map (s_i) to posterior samples of the original parameters (\omega). The training loss remains the standard negative log‑likelihood (-\sum_i \log q_\phi(\omega^{(i)}\mid S_\theta(x^{(i)},e_y^{\varepsilon,(i)}))), but the crucial difference is that the inputs now reflect the full predictive distribution of the surrogate, not a single point estimate.

Key contributions include:

1. A Bayesian surrogate modelling step that yields a full posterior over surrogate parameters and error scales, enabling principled uncertainty quantification.
2. A systematic propagation of this uncertainty into the ABI training data, preventing the over‑confidence that plagues naïve surrogate‑based ABI.
3. Demonstration that replacing MCMC with ABI after surrogate‑augmented training dramatically improves sampling efficiency (orders of magnitude speed‑up) while retaining accurate posterior estimates.

Experimental validation spans three settings: a low‑dimensional toy problem, a groundwater flow model governed by nonlinear PDEs, and a high‑dimensional atmospheric pollution simulator. In the toy case, UA‑SABI’s posterior matches the ground‑truth distribution, whereas a baseline surrogate‑ABI (ignoring uncertainty) yields biased, overly narrow posteriors. In the real‑world case studies, the method achieves accurate parameter posteriors with only a few hundred expensive simulator runs for surrogate training, and inference times on the order of milliseconds—three to four orders of magnitude faster than MCMC.

Limitations are acknowledged: the surrogate posterior still requires an initial MCMC run, and the current error model assumes i.i.d. errors. Future work is suggested on variational approximations for the surrogate posterior, ensemble surrogate strategies to reduce initial cost, and heteroscedastic or non‑Gaussian error models.

In summary, UA‑SABI provides a practical pathway to perform reliable Bayesian inference for costly simulators by leveraging cheap surrogate evaluations, rigorously accounting for their uncertainties, and exploiting the amortization benefits of deep‑learning‑based inference. This opens the door to real‑time, repeated inference in fields such as climate modeling, hydrology, and biomedical device control, where both accuracy and speed are essential.