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.
Related
- RECAP: Representation-Based Calibration with Amortized Pooling for Frequentist Simulation-Based Inference
- Whole-brain modeling of dynamic causal circuits in human cognition using amortized variational inference
- Amortized Bayesian Inference on Multilevel Models of Arbitrary Structure
- Amortized Variational Inference of Latent Intensities from Replicated Point Patterns
- Amortized Structured Stochastic Variational Inference for Gaussian Process Latent Variable Models
BibTeX
@misc{olischlager2026mitigating,
title = {{Mitigating Representation Gaps in Amortized Bayesian Inference with Auxiliary Supervision}},
author = {Hans Olischläger and Svenja Jedhoff and Šimon Kucharský and Aayush Mishra and Stefan T. Radev and Paul Bürkner},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.39525},
eprint = {2609.39525},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.39525},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Mitigating Representation Gaps in Amortized Bayesian Inference with Auxiliary Supervision
AU - Hans Olischläger
AU - Svenja Jedhoff
AU - Šimon Kucharský
AU - Aayush Mishra
AU - Stefan T. Radev
AU - Paul Bürkner
PY - 2026
DA - 2026/09/30
PB - arXiv
DO - 10.48550/arxiv.2609.39525
UR - https://arxiv.org/abs/2609.39525
AB - 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.
ER -