Mitigating Representation Gaps in Amortized Bayesian Inference with Auxiliary Supervision

Hans Olischläger; Svenja Jedhoff; Šimon Kucharský; Aayush Mishra; Stefan T. Radev; Paul Bürkner

stat.MLSep 30, 2026

Mitigating Representation Gaps in Amortized Bayesian Inference with Auxiliary Supervision

Authors: Hans Olischläger, Svenja Jedhoff, Šimon Kucharský, Aayush Mishra, Stefan T. Radev, Paul Bürkner

Organizations: Department of Statistics, TU Dortmund University, Germany · Rensselaer Polytechnic Institute

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.

Figures & tables

Figure 1: Location-scale models (Experiment 1) . Trajectories of posterior accuracy, measured by the expected log probability (y-axes), as well as summary informativeness, measured by the minimum canonical correlation (color) during training. The effectiveness of our guidance losses is evident in both metrics and across inference tasks, with parameter dimension $D=64,32$ and $16$ for the Normal , Uniform and Student-t likelihood respectively (rows). All models are trained until convergence but with simulation budgets varying from 10k to 100k (columns). For Normal and Uniform tasks guidance is applied directly to half of the summary statistics with varying weight factor $\tilde{\lambda}$ . For Student-t , we additionally include ensemble guidance, fixing $\tilde{\lambda}=4$ for both guidance losses.

Figure 2: Human decision making (Experiment 2) . Guidance removes a long plateau in parameter recovery on the joint EEG–drift-diffusion model at a 10k simulation budget. Overall posterior accuracy and summary informativeness (left), posterior-mean error on the prior scale for the four worst-recovered parameters (right). Unguided runs (black) leave $\sigma$ and $s_{\tau}$ near their prior spread for thousands of steps, each seed (line style) breaking at a different point, while guided runs ( $\tilde{\lambda}=4$ , red) recover both immediately and without seed spread.

Figure 3: Eye Movements (Experiment 3) . Parameter-specific metrics for two parameters, and global metrics. Metrics averaged over three independent test sets of 500 simulations each. Error bars represent standard deviation across the test sets. CF = Coupling flow, FM = Flow matching, P = Point estimation network (posterior mean).

Figure 4: Strong gravitational lensing (Experiment 4) . Rows show the two external-shear components, the least identified parameters. The right block shows global posterior metrics. Unguided networks suffer from a representation gap as indicated by the CCA diagnostic:

\rho_{\min}=0.61

for coupling flow (CF) and

\rho_{\min}=0.73

for flow matching (FM). Guidance through the summary loss (red) or an ensemble member (blue) clearly improves all metrics and raises

\rho_{\min}\geq 0.87

Figure 5: Social interactions between mice (Experiment 5) . Parameter-specific metrics for the two parameters network density and exchange factor, and global metrics. Metrics are grouped by the four different summary networks and averaged over three independent test sets of 500 simulations each. Error bars represent standard deviation across the test sets.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Decision making (Experiment 2): Complements Figure 2 by showing the trajectories of the last canonical correlations (right) instead of individual parameter’s posterior mean accuracy.

Figure 7: Experiment 3: Eye Movements . Parameter-specific metrics for all four parameters of the model.

Figure 8: Experiment 3: Eye Movements . Expected log probability of the parameters and minimum canonical correlation, across all combinations of inference and summary network architectures, and guidance type and weight. CF=Coupling flow, FM=Flow matching.

Figure 9: Experiment 3: Eye Movements . RMSE of the parameters, across all combinations of inference and summary network architectures, and guidance type and weight. CF=Coupling flow, FM=Flow matching, P=Point network.

Figure 10: Experiment 3: Eye Movements . Posterior contraction of the parameters, across all combinations of inference and summary network architectures, and guidance type and weight. CF=Coupling flow, FM=Flow matching, P=Point network.

Figure 11: Experiment 3: Eye Movements . Calibration error of the parameters, across all combinations of inference and summary network architectures, and guidance type and weight. CF=Coupling flow, FM=Flow matching, P=Point network.

Figure 12: Experiment 3: Eye Movements . Minimum and mean canonical correlations computed against the true parameter labels versus against the approximate posterior mean (from the set transformer summary network, 16 output dimensions), across all simulation settings and all three evaluation test sets. Canonical correlations computed against the estimated posterior mean are typically higher than those against the true labels, but the two targets yield strongly correlated diagnostics overall.

Figure 13: Experiment 4: Strong gravitational lensing. Eight independent draws from the prior, each rendered as a single noisy cutout as per the simulation pipeline.

Explore similar work

Jan 12, 2026stat.ML

Neural Architectures for Amortized Bayesian Inference: Statistical Foundations and Empirical Assessments

Since the turn of the century, approximate Bayesian inference has steadily evolved as new computational techniques have been incorporated to handle increasingly complex, large-scale predictive problems. The recent success of deep neural networks and foundation models has now given rise to a new paradigm in statistical modeling, in which Bayesian inference can be amortized through large-scale learned predictors. In amortized inference, substantial computation is required at the beginning to train a neural network, but it can subsequently produce approximate posteriors or predictions at much lower computational cost across a wide range of tasks. While the typical Bayesian inference procedures are computationally expensive due to repeated likelihood calculations and Monte Carlo steps for each new dataset, amortized inference provides a much lower computational cost at deployment. Despite the growing popularity of amortized inference, its statistical interpretation and position within Bayesian inference remain poorly explored. In this paper, we present a statistical perspective on several major neural architectures, including feedforward networks, Deep Sets, and Transformers, and examine how they naturally support amortized Bayesian inference. We explore how these models perform structured approximation and also probabilistic reasoning in ways that yield controlled generalization error throughout a wide range of deployment scenarios, and how these properties can be harnessed for Bayesian computation. Via simulation studies, we evaluate the accuracy, robustness, and uncertainty quantification of amortized inference across varying sample sizes, varying noise distributional families, varying sparsity levels, and multimodality, highlighting its strengths and limitations.

Roy Shivam Ram Shreshtth, Arnab Hazra, Gourab Mukherjee

Department of Statistics and Data Science, Indian Institute of Technology Kanpur, Kanpur, India 208016 · Marshall School of Business, University of Southern California, Los Angeles, United States 90089-0809

Bayesian Inference Neural Network Inference Approximation

Ask PDF

May 26, 2026cs.LG

Amortized Factor Inference Networks for Posterior Inference

Amortized inference promises fast test-time Bayesian inference, but existing methods are inherently tied to fixed models. Extending amortization to unseen models typically requires retraining or costly test-time finetuning. In this paper, we ask: is it possible to build a single inference network capable of generalizing across varying priors, likelihoods, and dimensionality? We introduce Amortized Factor Inference Networks (AFINs), a family of encode-merge-decode inference networks built on dimension-independent modules that map a model specification and its observations to the parameters of a variational posterior. Experimentally, a single trained AFIN achieves posterior accuracy comparable to NUTS and several variational inference methods, while requiring 2 to 4 orders of magnitude less test-time compute. Code is available at https://github.com/joohwanko/AFINs.

Joohwan Ko, Justin Domke

Manning College of Information and Computer Sciences University of Massachusetts Amherst

Bayesian Inference Posterior Neural Network Inference

Ask PDF

Sep 30, 2026stat.ML

Amortized Bayesian Inference on Multilevel Models of Arbitrary Structure

We develop a general method for amortized Bayesian inference on multilevel models of arbitrary structure. Given a generative model specified as a directed acyclic graph, our method automatically derives valid factorizations of the joint posterior and matching neural network architectures. The key steps, graph expansion and graph inversion, yield an inverse graph that determines how inference networks are stacked and conditioned, producing factorizations that amortize over the number of groups and the number of observations within each group. Unlike approaches that simplify the dependency structure to speed up learning or inference, our method preserves all conditional independence and exchangeability assumptions of the generative model. Across three case studies, it closely matches gold-standard samplers on models with more than 6,500 parameters while reducing inference to a near-instant forward pass once trained.

Daniel Habermann, Andreas Bulling, Stefan T. Radev +1

Department of Statistics TU Dortmund University, Germany · Institute for Visualisation and Interactive Systems University of Stuttgart, Germany · Department of Cognitive Science Rensselaer Polytechnic Institute, USA

Bayesian Inference Generalized Linear Model Generative Models Graph Representations

Ask PDF