PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.39525v1stat.MLcs.LGstat.COstat.ME

TL;DR

Amortized Bayesian inference can provide near-instant inference after training, but learned summaries may leave posterior directions underlearned. This paper adds auxiliary point-estimate guidance and reports faster convergence and improved recovery across toy and several real-world cases, though one real-world case shows no consistent improvement.

  • Problem

    Amortized Bayesian inference can retain representation gaps when summary networks fail to learn maximally informative statistics, leaving some posterior directions underlearned.

    Source: §1 · p. 2, §3.1 · p. 3

  • Method

    The paper applies auxiliary losses based on posterior point estimates to guide internal summary representations and introduces a diagnostic to distinguish summary failures from inference failures.

    Source: §1 · p. 2, §3.3 · p. 4

  • Results

    Guidance speeds convergence in all toy settings and substantially improves recovery of poorly estimated parameters across three real-world case studies, without sacrificing uncertainty calibration; a fourth real-world case shows no consistent improvement.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The proposed diagnostic confirms that the observed representation gaps trace back to insufficient summaries, helping localize their source.

    Source: §1 · p. 2, §6 · p. 10

    Faithful posterior inference requires a strictly proper objective and maximally informative learned summaries; convergence alone does not ensure the latter.

    Source: §2.1 · p. 2

Abstract

from arXiv · show

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.

1 Introduction

Amortized Bayesian inference offers fast posterior estimates after upfront simulation and training, but finite budgets leave practical systems with representation gaps. The paper proposes auxiliary supervision, formalizes how those gaps arise, and reports improved convergence and accuracy across challenging inference problems.

1 Introduction

  • Finite simulation, compute, and time budgets force practitioners to settle for solutions that can retain representation gaps, while problem-specific fixes demand expertise and transfer poorly.Amortized inference instead trains on labeled simulations to approximate posteriors for new data, yielding near-instantaneous inference after upfront costs.Source: §1 · pp. 1–2
  • Auxiliary losses on internal representations provide a generic way to improve amortized inference training dynamics, with faster convergence when data is abundant and better performance when it is scarce.The approach targets both convergence speed and final accuracy across challenging real-world inference problems.Source: §1 · p. 2
  • Representation gaps are formalized as local optima at the information bottleneck between summary-feature learning and conditional-distribution learning.The paper also introduces a generic diagnostic to distinguish summary failures from inference failures and localize a gap's origin.Source: §1 · p. 2

2 Background

Amortized Bayesian inference trains a conditional posterior network by minimizing expected scoring loss, with ideal convergence requiring informative learned summaries. The summary–inference decomposition acts as an information bottleneck that can produce representation gaps through capacity, data, or optimization limits.

2.1 Amortized Bayesian Inference

  • Amortized posterior approximation minimizes empirical Bayes risk over a parameterized family of conditional distributions using a strictly proper scoring rule.The posterior network is composed of a summary network and a conditional inference network.Source: §2.1 · p. 2
  • With a strictly proper score, the analytic posterior is the unique minimizer, but reaching the global optimum also requires maximally informative, sufficient summaries.Joint optimization of the summary and inference networks is motivated by this requirement.Source: §2.1 · p. 2

2.2 Information bottlenecks and representation gaps

  • The summary–inference decomposition forms an information bottleneck that discards unnecessary information while retaining task-relevant information.This bottleneck can control generalization error and inference speed.Source: §2.2 · pp. 2–3
  • Representation gaps arise from failure to learn maximally informative summaries, with approximation, estimation, and optimization errors corresponding to limited capacity, finite simulations, and suboptimal training dynamics.Under log-score and an idealized exact partial posterior, the excess over the full-data posterior is a mutual-information gap that vanishes if and only if the summary is sufficient.Source: §2.2 · p. 3
  • With finite-capacity inference networks, summary and inference learning depend on each other: summaries must be usable, while the inference network must provide signal to improve them.The inference network may ignore summaries it cannot immediately exploit, while the summary network needs its training signal to improve those representations.Source: §2.2 · p. 3

3 Origins and Mitigation of Representation Gaps

Representation gaps can persist when training settles into partial posterior learning, because ignored summary directions receive little signal to become informative. The paper diagnoses those failures with canonical correlations and proposes point-estimate guidance losses to improve summary learning.

3.1 Failure to escape fixed points in training dynamics

  • When summaries inform only some parameter directions, the inference network can fall back to the prior on the rest, producing a partial posterior that is self-reinforcing.The signal that could make ignored summaries informative vanishes when the inference network ignores them; minibatch noise can mask the weak escape signal, leaving a metastable state.Source: §3.1 · p. 3

3.2 Diagnosing unresolved posterior directions

  • Canonical correlation analysis (CCA) measures whether linear combinations of learned summaries preserve parameter directions, by whitening summaries and parameters and examining their cross-covariance.Its singular values are canonical correlations that identify maximally aligned coordinates.Source: §3.2 · pp. 3–4
  • The minimum canonical correlation diagnoses the worst-preserved parameter direction: values near one indicate strong representation across directions, while small values flag a weak or missing direction.With ground-truth parameters, the minimum is bounded by a ceiling determined by posterior variance; with reference point estimates, CCA measures recovery of that point and can be applied to empirical data.Source: §3.2 · p. 4

3.3 Guidance losses in summary space

  • Guidance losses train designated summary dimensions to learn posterior point estimates, with weights controlling each loss's influence.For K point estimates across D parameters, the summary dimension must satisfy S ≥ KD; the authors recommend S ≥ (K + 1)D to leave at least D dimensions unguided.Source: §3.3 · p. 4
  • An ensemble variant applies guidance through learnable functions conditional on summaries, avoiding direct restriction of the summary space at the cost of slightly greater architectural complexity.The functions join the posterior estimator as an ensemble sharing a summary network.Source: §3.3 · p. 4
  • Posterior-mean guidance fixes representation gaps in all but one case study, though posterior means are not guaranteed sufficient outside conjugate exponential-family models.The posterior mean is sufficient for conjugate exponential-family models with known sample size and otherwise retains first-moment information.Source: §3.3 · p. 4

4 Related work

The paper distinguishes its use of point estimates to guide full-posterior training from prior uses as standalone estimators or ABC summaries. It also contrasts representation-gap guidance with calibration and domain-alignment losses.

4 Related work

  • Point estimates have been used as standalone estimators and ABC summaries, whereas this work uses them to guide training of a full posterior network.The paper builds on estimate-specific loss functions and posterior-mean summary work.Source: §4 · p. 5
  • Coverage-error and balancing losses address overconfidence, but calibration and balance alone cannot distinguish a correct posterior from one that ignores data and returns the prior.The paper's guidance instead targets representation gaps that can lead to that failure.Source: §4 · p. 5
  • Domain-alignment losses match summary distributions across simulated and empirical data for robustness to shifts, while the paper's guidance improves training dynamics without altering the inference target.The cited domain-alignment approach changes the target to a posterior on adjusted data.Source: §4 · p. 5

5 Experiments

Across one toy and four real-world case studies, estimators share summary-network configurations so performance differences test auxiliary-loss signals rather than architecture or capacity.

5.2 Experiment 2: Human decision making

  • At 10k simulations, guidance removes parameter-recovery plateaus: σ and sτ recover immediately without seed spread, whereas unguided runs linger near prior spread.Source: Figure 2 · p. 7, §5.2 · p. 7

5.1 Experiment 1: Simple location-scale models

  • Guidance closes CCA-detected gaps across Normal, Uniform, and Student-t; posterior-mean guidance helps even when insufficient and stabilizes Student-t training.Unguided Normal and Uniform gaps persist even at 100k simulations.Source: Figure 1 · p. 6, §5.1 · p. 7

5.3 Experiment 3: Eye movements

  • Ensemble guidance gives the best RMSE and contraction; summary guidance also improves flow matching more than coupling flows, while calibration stays stable and global log probability and CCA improve.Source: Figure 3 · p. 8, §5.3 · p. 8

5.4 Experiment 4: Strong gravitational lensing

  • ρmin rose from 0.61 (CF) and 0.73 (FM) to ≥0.87 with guidance, while posterior contraction improved and calibration error decreased.Source: §5.4 · pp. 8–9, Figure 4 · p. 9

5.5 Experiment 5: Social interactions between mice

  • For GCN, guidance yields only small gains and does not match other summaries; other networks gain nothing, with high unguided CCA indicating little room for representation-gap fixes.The point-estimation network also does not improve RMSE over unguided inference for any summary network.Source: §5.5 · p. 9, Figure 5 · p. 10

6 Conclusion

The paper uses posterior point-estimate guidance to address representation gaps in amortized Bayesian inference. Guided training improves recovery in most evaluated settings, but not when architectural representation gaps cannot be fixed through training dynamics alone.

6 Conclusion

  • Across three real-world case studies, guided training substantially improves recovery of poorly estimated parameters without sacrificing uncertainty calibration.In a fourth case study, guidance yields no consistent improvement.Source: §6 · p. 10
  • The method uses auxiliary losses based on posterior point estimates, in summary-guided and ensemble-guided forms, across one toy and four real-world case studies.Source: §6 · p. 10
  • The diagnostic attributes the observed representation gaps to insufficient summaries in the studied cases.Source: §6 · p. 10
  • Guidance has no room to improve performance when neural point estimation cannot recover parameters better than the unguided inference network.The method focuses on posterior-mean guidance; other point estimates may suit gaps in uncertainty calibration, but their effectiveness remains an open question.Source: §6 · p. 10

A Appendix

The appendix expands the experiments described in Section 5, first covering shared evaluation and training details, then presenting five experiments.

A Appendix

  • The appendix provides further details on the experiments in Section 5.Source: A Appendix · p. 14
  • It first describes the evaluation metrics and then outlines network architectures and training concepts shared across experiments.Source: A Appendix · p. 14
  • It presents five experiments in turn, including their simulation setups, experiment-specific architecture and training choices, and additional results where applicable.Source: A Appendix · p. 14

A.1 Metrics

The appendix reports parameter-wise and global metrics for evaluating posterior estimates, calibration, information retention, and summary quality.

A.1 Metrics

  • The RMSE measures performance of the posterior mean estimate across test datasets and parameters.Source: §A.1 · p. 14
  • Posterior contraction measures how much the approximate posterior narrows relative to the prior; values near 1 indicate strong concentration, while values near 0 indicate no information gain from the data.Source: §A.1 · p. 14
  • Simulation-based calibration coverage measures how often true parameters fall within nominal credible intervals, and calibration error aggregates deviations from nominal coverage across interval levels.A calibration error of 0 corresponds to perfect calibration.Source: §A.1 · p. 15
  • Expected log probability scores the log probability assigned to the true parameters given test data conditions, with higher values favorable.Source: §A.1 · p. 15
  • CCA between summaries and parameters indicates how much parameter information the summaries retain; against a posterior-mean target, a minimum canonical correlation near 1 is desirable.Source: §A.1 · p. 15

A.2 Properties of the CCA diagnostic

The CCA diagnostic differs depending on whether it is computed against ground-truth parameters or the posterior mean. The posterior uncertainty remaining after conditioning on data constrains the ground-truth comparison.

A.2 Properties of the CCA diagnostic

  • The decomposition θ = µ + ε separates the posterior mean, which captures the data-revealed first moment, from ε, the remaining posterior uncertainty.Source: §A.2 · p. 15
  • Ground-truth parameter CCA cannot exceed a bound set by residual posterior variance, even for a sufficient summary.The posterior mean is the most a summary can recover from the data; canonical correlations against it are never lower than those against ground truth.Source: §A.2 · p. 15
  • For a posterior-mean target, ρmin near 1 means the summary linearly recovers the mean in every direction, so conditioning on the summary loses almost no information about it.Source: §A.2 · p. 15
  • A low posterior-mean CCA does not necessarily indicate lost information, because the mean may depend nonlinearly on the summary.Source: §A.2 · p. 15

A.3 Network architectures and training

The experiments compare unguided, summary-guided, and ensemble-guided training across inference and summary network architectures, using experiment-specific defaults unless noted.

A.3 Network architectures and training

  • Inference networks use coupling flows, flow matching, or a point network, paired with summary architectures suited to each experiment.The point network is trained with a scoring rule to estimate the posterior mean.Source: §A.3 · p. 16
  • Three training regimes vary whether inference receives no auxiliary signal, a point-estimation loss on the summary, or a jointly trained point-estimation network.The guided losses are combined with the distributional loss using a weighting term.Source: §A.3 · p. 16
  • Training uses Adam or AdamW with cosine-decay scheduling, online or offline depending on simulator cost.Networks and training use default hyperparameters unless an experiment specifies otherwise.Source: §A.3 · p. 16

A.4 Choosing the auxiliary loss weight

The auxiliary-loss weight is scaled relative to the initial magnitudes of the generative and guidance losses, then held fixed during training.

A.4 Choosing the auxiliary loss weight

  • The guidance weight λ is set so the auxiliary loss is comparable in scale to the generative loss on a freshly initialized network and simulated batch.The authors report this as a robust policy yielding comparable results across architectures.Source: §A.4 · p. 16
  • An architecture-specific λ is computed as λ = ˜λ |J|/|ℓ|, using the initial-batch loss magnitudes.For ensemble guidance, the summary-network term sd(x) is replaced by f(s(x))d.Source: §A.4 · p. 16
  • Guidance was similarly effective across ˜λ values from 1 to 16 in the reported experiments.Source: §A.4 · p. 16

A.5 Experiment 1: Simple location-scale models

The appendices describe five inference experiments, with the supplied passages detailing their setups and reporting additional eye-movement results. In that case, guided training generally matches or outperforms unguided training, but its effects vary by architecture and weight.

A.5 Experiment 1: Simple location-scale models

  • The location-scale experiments cover Normal, Uniform, and Student-t models, using analytic or near-analytic posterior references for evaluation.Normal has a closed-form Gaussian posterior; Uniform and Student-t posteriors factorize by component and are integrated on fixed grids.Source: §A.5 · p. 16

A.6 Experiment 2: Human decision making

  • In human decision making, the least informative learned-summary directions become informative in distinct steps rather than gradually.The experiment used an older BayesFlow version because the later set-transformer modification largely fixed the representation gap under default configurations.Source: §A.6 · pp. 17–18

A.7 Experiment 3: Eye movements

  • Across eye-movement settings, guided training tends to match or outperform unguided training, but outcomes vary by architecture and weight.At λ = 100, RMSE can improve while expected joint log probability worsens; some settings also fail to improve τnonlabile RMSE and contraction.Source: §A.7 · p. 19, Figure 8 · p. 20, Figure 9 · p. 20

A.8 Experiment 4: Strong gravitational lensing

  • The lensing experiment infers 17 parameters from simulated 64 × 64-pixel images rendered with Euclid VIS settings and instrumental noise.Training is online, with fresh simulations and approximately 107 unique simulated images per run.Source: §A.8 · p. 23

A.9 Experiment 5: Social Interactions between mice

  • The mouse-interaction experiment infers network density δ and exchange factor α from an interaction network and final-day microbial composition after five simulated days.The simulator represents mice as a weighted interaction graph and transfers microbiota according to edge weights and α.Source: §A.9 · p. 24
Loading 2609.39525v1…