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

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.

💡 Research Summary

Amortized Bayesian inference (ABI) has emerged as a powerful tool for fast Bayesian inference by training neural networks on simulated data to map observations directly to posterior distributions. Despite its speed, ABI suffers from a critical robustness problem: when the test observations lie outside the support of the simulated training data, the learned posterior can become severely biased. Existing remedies—such as simulation‑gap detection, post‑hoc corrections, or regularization—either require additional costly simulations, rely on MCMC corrections, or alter the target distribution, thereby compromising the original Bayesian goal.

The paper introduces a novel semi‑supervised framework that leverages Bayesian self‑consistency to create a strictly proper loss function that can be applied to unlabeled real observations. The key insight is that Bayes’ rule implies the marginal likelihood p(x) is constant across all parameter values:

 p(x) = p(x|θ) p(θ) / p(θ|x) for any θ.

If we replace the true posterior p(θ|x) with a neural estimator q(θ|x) (and optionally the likelihood with q(x|θ)), the ratio becomes non‑constant. Minimizing the variability of this ratio therefore forces q(θ|x) toward the true posterior. The authors formalize this by defining a self‑consistency loss C(x) = Var_{θ∼p_C}