Mitigating Representation Gaps in Amortized Bayesian Inference with Auxiliary Supervision

Abstract

Casting Bayesian inference as a neural network optimization problem targeting an amortized posterior is attractive, as it extends to otherwise intractable statistical models and offers near instantaneous inference for new datasets after prepaying the training cost. Although theory guarantees faithfulness under ideal convergence, practical amortized inference still requires iterating over architectures and optimization choices and ultimately ``satisficing'' under finite simulation, compute, and time budgets. Even the best-performing solution may thus retain avoidable representation gaps that typically require problem-specific fixes. Here, we propose a generic alternative which improves training dynamics with auxiliary guidance losses applied to internal representations. Specifically, we show how such guidance leads to faster convergence when training data is abundant and to better performance when it is scarce. We formalize representation gaps as getting stuck in a local optimum at the information bottleneck between the parts of the network tasked with feature learning and those tasked with conditional distribution learning, and offer a generic diagnostic to separate summary failures from inference failures. Finally, we demonstrate that auxiliary supervision improves convergence speed and accuracy on a range of challenging real-world inference problems.

The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.

More in Machine Learning (Statistics)

• arXiv preprint · 30 September 2026

  Learn-Then-Differentiate Gradient Estimation

  Nifei Lin, Qingkai Zhang and L. Jeff Hong
• arXiv preprint · 1 October 2026

  On skew-symmetric distributions and their use in Monte Carlo sampling algorithms: coordinate-free, Gibbs-style and manifold versions of the Barker proposal

  Minh Vu, Samuel Livingstone and Pantelis Samartsidis
• arXiv preprint · 1 October 2026

  MCIR: A Feature Dependence-Aware Explainability Method with Reliability Guarantees

  Poushali Sengupta, Sabita Maharjan, Frank Eliassen and 2 others
• arXiv preprint · 30 September 2026

  Amortized ratio-estimation importance sampling and localized simulation-based calibration for intractable likelihoods

  Umberto Picchini
• arXiv preprint · 1 October 2026

  Flow Matching for Fast Posterior Sampling in Bayesian Inverse Problems

  Jan Blechschmidt, Oliver G. Ernst, Moritz Poguntke and 1 other

All new papers in Machine Learning (Statistics)