Robust Amortized Bayesian Inference with Self-Consistency Losses on Unlabeled Data

TL;DR

The paper tackles brittleness in amortized Bayesian inference when test data lie outside the simulator distribution. It introduces a semi-supervised framework that combines a standard simulation-based objective with Bayesian self-consistency losses computed on unlabeled data, guaranteeing strict propriety under reasonable conditions. The authors provide theoretical results showing the semi-supervised loss is strictly proper and demonstrate robust posterior estimation across synthetic and real-world case studies, including high-dimensional time-series and image data, while preserving inference speed. The approach offers a practical path toward safe and reliable ABI in the presence of simulation gaps, with broad implications for applications requiring fast, calibrated probabilistic inference on unseen data.

Abstract

Amortized Bayesian inference (ABI) with neural networks can solve probabilistic inverse problems orders of magnitude faster than classical methods. However, ABI is not yet sufficiently robust for widespread and safe application. When performing inference on observations outside the scope of the simulated training data, posterior approximations are likely to become highly biased, which cannot be corrected by additional simulations due to the bad pre-asymptotic behavior of current neural posterior estimators. In this paper, we propose a semi-supervised approach that enables training not only on labeled simulated data generated from the model, but also on \textit{unlabeled} data originating from any source, including real data. To achieve this, we leverage Bayesian self-consistency properties that can be transformed into strictly proper losses that do not require knowledge of ground-truth parameters. We test our approach on several real-world case studies, including applications to high-dimensional time-series and image data. Our results show that semi-supervised learning with unlabeled data drastically improves the robustness of ABI in the out-of-simulation regime. Notably, inference remains accurate even when evaluated on observations far away from the labeled and unlabeled data seen during training.

Figures (15)

• Figure 1: Contour plot of the normal means problem using standard NPE (red) or our semi-supervised approach (NPE + SC, blue), with the analytic posterior in gray. Symbols indicate posterior mean estimates (red cross: NPE only; blue square: NPE + SC; gray triangle: reference). Each subplot shows posterior inference on observed data that are increasingly distant from the labeled training data ($\mu_{\rm prior} = 0$). Only the first two dimensions of the 10-dimensional posterior are shown. While standard NPE collapses to zero variance for $\mu_{\text{obs}} \geq 2$, adding the self-consistency loss preserves accurate posterior estimates even far beyond both training spaces ($\mu_{\text{obs}} > 3$). Training was performed using the default configuration (see Section \ref{['sec:case-study-normal-means']}).
• Figure 2: Posterior distance quantified by maximum mean discrepancy (MMD) to the analytic posterior for variations of the default configuration. Errorbars show $\pm 1$ SDs over 10 model refits.
• Figure 3: Comparison of posterior estimates for 15 countries (ISO 3166 alpha-2 codes) among standard NPE (red circles), NPE + self-consistency loss (blue squares), and Stan (reference; gray triangles). Central 50% (thick lines) and 95% (thin lines) posterior intervals of the autoregressive component $\beta$ are shown, sorted by lower 5% quantile as per Stan (i.e., established benchmark). The self-consistency loss was evaluated on data from $M=\mathbf{8}$ countries during training, greatly enhancing ABI's robustness in both no-misspecification scenarios and real-data evaluations.
• Figure 4: (a) Posterior predictive samples (gray) inferred from an out-of-simulation dataset (black). NPE only produces highly biased predictions while NPE+SC yields accurate results. (b) Histogram of the mean absolute bias (MAB) difference of posterior predictions computed for 1000 out-of-simulation datasets. NPE+SC has lower bias than NPE for almost all datasets.
• Figure 5: Example of denoising results for MNIST images of digit “0” in the held-out test set. The first row shows ten randomly selected MNIST images ($\theta$), the second row depicts the same images after applying the Gaussian blur ($x$), third and fourth rows depict mean and SD of 500 posterior samples estimated from the corresponding blurry observations using NPLE+SC, and the fifth and sixth rows depict the mean and SD of 500 posterior samples using NPLE only. Incorporating SC loss significantly improves denoising: reconstructed means become smoother, less pixelated, and closer to the ground truth. In the standard-deviation maps, darker regions indicate higher output variability; NPLE + SC approach produces coherent maps with variability confined to the inner and outer edges.

Theorems & Definitions (6)

• Proposition 1
• Proposition 2
• Proposition 3
• Proposition 4
• proof
• proof