Unsupervised Continual Learning for Amortized Bayesian Inference
| Authors | Aayush Mishra et al. |
| Year | 2026 |
| Field | Statistics / ML |
| arXiv | 2602.22884 |
| Download | |
| Categories | stat.ML, cs.LG |
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.
Engineering Breakdown
Plain English
This paper tackles a critical bottleneck in amortized Bayesian inference (ABI)—a technique that trains neural networks on simulated data to quickly estimate probability distributions from real observations. The core problem is that these pre-trained networks degrade when the simulator doesn't perfectly match reality (model misspecification), and existing fixes only work on static datasets. The authors propose a continual learning framework that first pre-trains on simulation, then continuously refines the network on real-world data streams using self-consistency training without labeled data, while preventing catastrophic forgetting through episodic replay and elastic weight consolidation.
Core Technical Contribution
The key innovation is decoupling amortized Bayesian inference into two explicit phases: simulation-based pre-training followed by unsupervised continual fine-tuning on sequential real data. Rather than treating self-consistency (SC) training as a one-shot improvement step, the authors reformulated it as an ongoing adaptation mechanism that can handle distribution shifts and sequentially arriving observations. They introduced two complementary strategies to prevent catastrophic forgetting—episodic replay (maintaining a buffer of past observations) and elastic weight constraints (protecting learned parameters)—enabling the network to simultaneously improve robustness to model misspecification while retaining prior knowledge. This is fundamentally different from prior work that assumed static empirical datasets or required labeled data for fine-tuning.
How It Works
The system operates in three stages. First, a neural network posterior estimator is pre-trained on simulated data from a potentially misspecified simulator, learning to map observations to approximate posterior distributions. Second, when real data arrives sequentially, the network undergoes unsupervised self-consistency fine-tuning: the network generates its own predictions, computes divergence between the prior and posterior estimates given those predictions, and updates weights to minimize this divergence without requiring ground truth labels. Third, to prevent catastrophic forgetting (where learning on new data destroys performance on old data), two mechanisms activate in parallel: an episodic replay buffer stores representative past observations and interleaves them during training batches, ensuring the network rehearses old tasks; and elastic weight consolidation penalizes large changes to weights that were important for prior observations, identified through Fisher information matrices. The combination allows the network to adapt to new distribution shifts while maintaining historical performance.
Production Impact
For production systems using Bayesian inference at scale, this approach directly addresses the painful gap between simulator accuracy and real-world data. Engineers deploying amortized inference (common in scientific computing, anomaly detection, or parameter estimation pipelines) currently face a choice: either maintain expensive simulators with high fidelity, or accept degraded posterior estimates when reality diverges from training assumptions. This framework eliminates that choice by enabling passive refinement on production data streams without requiring labels or simulator retraining—the network improves itself through self-supervision. The trade-off is computational overhead: you must maintain a replay buffer in memory (tunable, but adds storage proportional to data volume), compute Fisher information matrices for elastic weights (expensive, scales with parameter count), and run SC fine-tuning on every mini-batch of new data (adds ~10-20% latency to inference depending on batch size). For systems where data arrives continuously (sensor networks, finance, real-time observational systems), this cost is often acceptable given the accuracy gains.
Limitations and When Not to Use This
The paper assumes access to the original simulator or at least the ability to sample from it, which breaks down if the simulator is proprietary, deprecated, or too expensive to query during deployment. Self-consistency training works best when the prior is well-specified; if both prior and likelihood are severely misspecified, the self-supervised signal becomes noisy and training can diverge. The episodic replay buffer introduces a critical hyperparameter (buffer size and composition strategy) whose optimal setting likely varies per application, and the paper doesn't provide principled guidance for tuning it in new domains. The elastic weight consolidation assumes the Fisher information matrix accurately captures parameter importance, but this assumes local quadratic loss landscape geometry that may not hold after distribution shifts; pathological cases exist where protected weights become irrelevant or irrelevant weights become critical. Finally, the framework is demonstrated only in relatively low-dimensional settings; scaling to high-dimensional observations (images, sequences) and ensuring stable training over long data streams remains an open question.
Research Context
This work builds on two established research areas: amortized Bayesian inference (which itself emerged from variational autoencoders and neural density estimation) and continual learning (also called online learning or lifelong learning). The motivation for self-consistency training comes from prior work showing that fine-tuning pre-trained inference networks on unlabeled empirical data improves robustness to model misspecification. The use of episodic replay is standard in continual learning (popularized by experience replay in RL and later adapted for supervised continual learning), while elastic weight consolidation is a classic approach from the continual learning literature for preventing catastrophic forgetting. The contribution is bringing these pieces together specifically for the Bayesian inference setting where neither labels nor ground truth parameters are available during fine-tuning—a constraint that rules out many standard continual learning methods. This opens directions for future work: adaptive buffer management strategies, theoretical analysis of when self-consistency fine-tuning converges, and extensions to hierarchical or meta-learned priors.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::