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

TL;DR

This paper introduces a semi-supervised Bayesian inference method that uses unlabeled data and self-consistency losses to improve robustness and accuracy, especially on out-of-distribution observations.

Contribution

It proposes a novel semi-supervised approach leveraging self-consistency losses for robust amortized Bayesian inference on unlabeled and real-world data.

Findings

Significantly improves robustness of Bayesian inference on out-of-distribution data

Maintains high accuracy on real-world high-dimensional time-series and image data

Outperforms traditional methods in safety-critical applications

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.