Unsupervised Continual Learning for Amortized Bayesian Inference

TL;DR

This paper introduces a continual learning framework for Amortized Bayesian Inference that improves robustness and mitigates forgetting when adapting to sequential real-world data, outperforming traditional methods.

Contribution

It proposes a novel continual learning approach with two adaptation strategies to enhance ABI's performance on sequential data and address catastrophic forgetting.

Findings

Methods significantly reduce forgetting across case studies.

Posterior estimates outperform standard simulation-based training.

Approaches achieve estimates closer to MCMC references.

Abstract

Amortized Bayesian Inference (ABI) enables efficient posterior estimation using generative neural networks trained on simulated data, but often suffers from performance degradation under model misspecification. While self-consistency (SC) training on unlabeled empirical data can enhance network robustness, current approaches are limited to static, single-task settings and fail to handle sequentially arriving data or distribution shifts. We propose a continual learning framework for ABI that decouples simulation-based pre-training from unsupervised sequential SC fine-tuning on real-world data. To address the challenge of catastrophic forgetting, we introduce two adaptation strategies: (1) SC with episodic replay, utilizing a memory buffer of past observations, and (2) SC with elastic weight consolidation, which regularizes updates to preserve task-critical parameters. Across three diverse case studies, our methods significantly mitigate forgetting and yield posterior estimates that outperform standard simulation-based training, achieving estimates closer to MCMC reference, providing a viable path for trustworthy ABI across a range of different tasks.