PaperlayerView exact revision on arXiv

Source-linked AI summary

Hidden Axis of Uncertainty: Latent-Posterior Alignment in Graph Neural Networks with Bayesian Output Layers

Suk Hoon Choi, Damdae Park, Junhyuk Choi, Hyein Jung, Changsoo Kim, Ung Lee, Kyeongsu Kim

arXiv:2608.20758v1cs.LG

TL;DR

Deep Bayesian models can reduce predictive uncertainty without posterior contraction, raising questions about how uncertainty is formed. The paper analyzes Bayesian output-layer GNNs, identifies LPA as a latent-geometric mechanism, and introduces AGL; AGL reduces uncertainty while preserving accuracy and improving density-aware calibration.

  • Problem

    The paper addresses limited understanding of predictive-uncertainty formation when deep Bayesian models violate classical posterior-contraction assumptions.

  • Method

    The paper combines analytical and empirical analysis of Bayesian output-layer GNNs with interventions that disrupt LPA and a training framework, AGL, that explicitly promotes it.

  • Results

    Predictive uncertainty decreases as latent representations align with low-variance posterior directions despite broader posteriors, while AGL preserves accuracy and improves structural calibration.

  • Takeaways & Limitations

    Uncertainty in mean-field Bayesian output-layer GNNs is governed by latent–posterior geometry as well as posterior magnitude, enabling confidence to better reflect data density.

  • Takeaways & Limitations

    The intervention studies do not isolate LPA's causal effect because latent norm, sparsity, representational capacity, and optimization dynamics are not held fixed.

Abstract

from arXiv · show

Bayesian Neural Networks (BNNs) with Bayesian output layers provide a principled and tractable framework for quantifying predictive uncertainty, yet the mechanisms shaping that uncertainty remain unclear. While conventional theory attributes uncertainty reduction to posterior contraction, the corresponding assumptions need not hold for deep models. In the Graph Neural Networks (GNNs) with Bayesian output layers studied here, we observe that predictive uncertainty decreases as latent representations shift toward lower-variance posterior directions, even though the posterior variance does not contract. We term this behavior Latent-Posterior Alignment (LPA) and conduct interventional experiments that support its functional role in shaping predictive uncertainty. Building on this insight, we propose Alignment-Guided Learning (AGL), which explicitly promotes this alignment during training. AGL effectively reduces predictive uncertainty while preserving accuracy and improves structural calibration, ensuring that the model confidence faithfully mirrors underlying data density. These findings provide a new perspective on uncertainty dynamics in GNNs with mean-field Bayesian output layers, shifting the focus from the magnitude of the posterior to the geometric interplay between latent and parameter spaces.

1 Introduction

The paper challenges posterior contraction as the explanation for decreasing predictive uncertainty in deep Bayesian models. It identifies Latent–Posterior Alignment (LPA), tests its functional role, and proposes Alignment-Guided Learning (AGL) to improve density-aware uncertainty.

  • Motivation: Classical posterior-contraction assumptions may fail in deep models because well-specification, finite-dimensional parameters, and positive-definite Fisher information are not guaranteed.Empirical studies also report complex, noncontracting posteriors and overconfident predictions in low-density regions.
  • Latent–Posterior Alignment: LPA describes latent representations shifting toward posterior directions with lower variance, allowing predictive variance to decrease without uniformly shrinking posterior variances.The alignment strengthens as training data increases and concentrates uncertainty contributions in lower-variance dimensions.
  • Alignment-Guided Learning: AGL explicitly promotes LPA, reducing predictive uncertainty while preserving accuracy and improving structural calibration between uncertainty and data density.AGL increases DUC and aligns high uncertainty with low-density regions, while posterior uncertainty changes little and ECE remains largely unchanged.
  • Observed uncertainty dynamics: Predictive uncertainty decreases with more data even as Bayesian output-layer weight uncertainty broadens, indicating a decoupling between posterior and predictive uncertainty.The posterior uncertainty increases from 0.05 with 1% data to 0.11 with 20% data, while predictive uncertainty decreases.
  • Latent geometry: High-uncertainty samples retain the global alignment pattern but have larger latent-vector magnitudes, showing that alignment does not eliminate local variability.The high-uncertainty subset comprises the top 5% of test samples ranked by posterior-predictive standard deviation.
  • Interventional evidence: Interventions that suppress or reverse LPA reduce alignment, degrade accuracy, and increase predictive uncertainty, providing convergent evidence for a causal contribution.The anti-alignment intervention directly pushes representations toward dimensions with larger posterior standard deviations, while L1 and L2 constrain latent sparsity or magnitude.

3 Discussion

The paper argues that predictive uncertainty is governed by latent–posterior geometry rather than posterior variance alone. Alignment-Guided Learning promotes this structure, improving density awareness while preserving calibration, but the findings remain bounded to the studied model class and evaluation setting.

  • Predictive uncertainty can decrease while the Bayesian output-layer posterior broadens, because latent representations shift predictive information toward lower-variance posterior directions.The authors interpret this joint representation–posterior behavior as Latent-Posterior Alignment rather than posterior contraction.
  • AGL strengthens density awareness without explicitly including data density in its objective, redistributing uncertainty toward sparse or boundary regions rather than lowering it uniformly.After AGL, higher uncertainty more often coincides with lower-density latent regions, while well-supported regions benefit from reduced uncertainty.
  • Density-aware uncertainty may support decisions to defer weakly supported MLIP predictions for additional DFT calculations, although this use was not prospectively evaluated.DUC complements ECE by relating uncertainty to latent support before the true label is available.
  • The conclusions are restricted to deterministic GNN feature extractors with mean-field Bayesian output layers and fixed-distribution molecular-property benchmarks.The authors identify fully Bayesian architectures, alternative posterior approximations, distribution shifts, and prospective active learning as unresolved evaluation directions.

4 Methods

The study uses deterministic GNN backbones with Bayesian output layers, repeated multi-run training, and metrics that quantify accuracy, uncertainty, latent–posterior alignment, calibration, and density awareness.

  • Model and training: The model uses five sequential Graph Isomorphism Networks, pooling, and a fully connected Bayesian layer to process molecular graphs.Nodes encode 36 atomic features and edges represent covalent bonds.
  • Model and training: Training minimizes the negative ELBO with AdamW, and reported metrics aggregate 30 independent runs with distinct random seeds.The datasets use an 80% training and 20% testing split after curation from public sources.
  • Predictive metrics: MAE and average output uncertainty σy are computed across all test samples and independent runs.The uncertainty is the predicted standard deviation averaged over the evaluation samples and runs.
  • Posterior and alignment analysis: Global posterior mean and uncertainty aggregate weight-distribution statistics across 64 hidden dimensions and independent runs.For each run, the weight posterior statistics use the mean and standard deviation of each parameter's distribution.
  • Posterior and alignment analysis: Latent–posterior alignment is assessed by normalizing latent-component magnitudes and posterior standard deviations, then sorting latent values by posterior scale.The visualization distinguishes alignment with high-variance axes, low-variance axes, or no systematic structure.
  • Calibration and density metrics: Regression ECE measures the discrepancy between target confidence intervals and empirical coverage under a Gaussian predictive distribution.The calibration gap is evaluated across confidence levels from 0.1 through 0.9 and 0.99, then integrated using the trapezoidal rule.
  • Calibration and density metrics: DUC measures statistical dependence between latent-data scarcity and predictive uncertainty, with latent density estimated by Gaussian-kernel KDE.Higher DUC indicates stronger dependence, while the expected low-density/high-uncertainty direction is checked separately using paired KDE–uncertainty maps.

1 Experimental details

The study uses deterministic GNN feature extraction followed by a Bayesian fully connected output layer, trained with variational inference and negative-ELBO optimization. Molecular datasets are curated, split into training and testing sets, and evaluated across repeated runs.

  • Model architecture: The backbone uses sequential graph-convolutional layers, pooling, and a fully connected layer to transform atomic and bond features into molecular representations.The described implementations use GIN layers and atomic features encoded on graph nodes.
  • Bayesian output layer: The Bayesian GNN reparameterizes final-layer weights as Gaussian variables, sampling them during forward passes to produce stochastic predictions.This creates an implicit ensemble effect within a single model.
  • Model training: Training minimizes a negative ELBO using mean-field variational inference, with KL regularization balancing fit and posterior constraints.The backbone is pretrained and its weights transferred to the Bayesian model for a warm start.
  • Data and evaluation: Experiments use curated public datasets with 80% allocated for training and 20% for testing, alongside property-specific parity plots and repeated runs.The study reports 30 independent runs for each configuration.

2 Theoretical context for posterior contraction in the present setting

Classical posterior-contraction theory predicts shrinking, increasingly concentrated posteriors under Bernstein–von Mises assumptions. The present model differs because it combines a mean-field output-layer approximation with a jointly learned, data-dependent representation, limiting direct application of that theory.

  • Classical theory: The Bernstein–von Mises theorem states that posteriors asymptotically approach Gaussian distributions under standard regularity conditions.Its covariance is linked to the inverse Fisher Information Matrix and sample size.
  • Classical theory: In the classical regime, posterior variance scales as O(1/n), so posterior mass concentrates around the true parameter as data increase.This underlies the conventional expectation that more data produce greater confidence.
  • Applicability: The theorem assumes a well-specified model, a fixed finite-dimensional parameter space, and regularity with a positive-definite Fisher Information Matrix.These assumptions are not automatically inherited by the studied training formulation.
  • Applicability: Misspecification changes the target of concentration to a pseudo-true parameter but does not alone imply persistent posterior broadening.This distinguishes altered concentration conditions from a demonstrated failure of concentration.
  • Applicability: The output-layer variational posterior is coupled to a data-dependent learned representation rather than a pre-specified fixed design.Full-network over-parameterization alone does not establish failure of contraction in the output layer.
  • Applicability: Posterior broadening is interpreted as an empirical property of the tested architecture and objective, not as a violation of classical posterior-contraction theory.The mean-field approximation and jointly learned representation differ from the exact-posterior setting addressed by classical theory.

3 Derivation of the Evidence Lower Bound (ELBO) objective

The Bayesian output layer uses variational inference with factorized weight and bias distributions. Its negative-ELBO objective combines expected negative log-likelihood with KL regularization weighted by β.

  • Probabilistic model: Weights and biases are modeled as random variables with prior distributions conditioned on the dataset.The resulting joint distribution forms the Bayesian output-layer model.
  • Variational approximation: The intractable joint posterior is approximated by a mean-field factorization qΦ(w,b) = qΦ(w)qΦ(b).This separates the variational distributions for weights and biases.
  • Training objective: The negative ELBO combines expected negative log-likelihood with KL-divergence regularization.The objective balances predictive fit against deviation from the priors.
  • Training objective: β weights the KL term to control the strength of variational regularization during optimization.The same weighting applies to the regularization component of the negative-ELBO objective.

4 Negligible impact of bias posterior distributions

Bias posterior distributions contribute negligibly to predictive uncertainty across the examined datasets and training sizes. The analysis therefore focuses on the weight posterior and its coupling with latent representations.

  • Bias posterior behavior: Bias posterior means converge near zero, while bias standard deviations remain small and fluctuate inconsistently with dataset size.The bias distribution shows no clear relationship with predictive uncertainty.
  • Uncertainty decomposition: The weight contribution consistently dominates the bias contribution to output uncertainty across datasets and training sizes.The decomposition identifies the weight posterior coupled with latent representations as the primary uncertainty contribution.
  • Uncertainty decomposition: Figure 3 compares bias and weight posterior means, standard deviations, and relative uncertainty contributions as training data increase across six datasets.The panels provide the cross-dataset view supporting the relative contribution analysis.

5 Cross-dataset consistency of posterior–predictive decoupling

Across six molecular-regression datasets, predictive uncertainty decreases as data grows even while weight-posterior uncertainty increases. The analysis formalizes this decoupling through latent–posterior pairing, showing that alignment minimizes predictive variance under fixed variance constraints.

  • Cross-dataset evidence: Across six datasets, predictive uncertainty and MAE decrease with training-set growth while average weight-posterior standard deviation increases.Posterior means remain stable near zero, making the decoupling reproducible across the tested architecture and protocol.
  • Cross-dataset evidence: Predictive uncertainty decreases while posterior uncertainty increases across all tasks as training data size grows.
  • Latent–posterior formulation: Under fixed posterior-variance spectra, predictive variance is minimized by pairing large latent energies with small posterior variances.This alignment follows from the trace-inequality formulation of the uncertainty coupling term.
  • Constrained solutions: With fixed total variance or standard deviation, optimal solutions concentrate posterior uncertainty along latent directions with the smallest energy.The fixed-standard-deviation problem is strictly convex and yields an inverse relationship between latent energy and posterior standard deviation.
  • Constrained solutions: The fixed-spectrum and budget-constrained analyses both identify inverse alignment between posterior variance and latent energy as the uncertainty-minimizing configuration.The constraints isolate latent–posterior pairing as the mechanism rather than enforcing it exactly during training.

7 Cross-dataset consistency of Latent-Posterior Alignment

Across property datasets, data-rich models reorganize latent representations toward posterior directions with lower variance. This shift raises LPAS, reallocates uncertainty toward reliable dimensions, and preserves aligned directions while increasing latent magnitudes for high-uncertainty samples.

  • Latent geometry: Across all datasets, data-rich models shift normalized latent vectors toward dimensions with minimized posterior standard deviation.The comparison uses 1% data for data-poor and 80% data for data-rich regimes.
  • Uncertainty composition: As more data are used, the uncertainty budget shifts toward low-variance dimensions despite residual contributions from high-variance dimensions.
  • Alignment score: LPAS increases with training-set size across all tested datasets, quantitatively validating the emergence of latent–posterior alignment.
  • High-uncertainty samples: High-uncertainty samples retain the global aligned directionality but exhibit substantially larger latent magnitudes.The results support a two-tiered pattern in which alignment sets a baseline and magnitude captures instance-specific epistemic risk.
  • Cross-dataset analysis: Figure 5 compares data-poor and data-rich latent geometry, LPAS versus data size, relative uncertainty contributions, and high-uncertainty sample structure.

8 Formulation of latent regularization objectives

The interventional framework adds latent regularization terms to the standard ELBO objective to suppress or reverse latent–posterior alignment. Regularization strengths are adjusted by dataset, while AGL uses a separate weight γ.

  • Regularization objectives: L1, L2, and anti-alignment penalties are introduced directly into the standard ELBO objective to inhibit LPA formation.
  • Regularization objectives: L1 enforces sparsity, L2 penalizes latent-vector magnitude, and anti-alignment pushes representations toward dimensions with larger posterior standard deviations.The interventions target directional concentration, latent magnitude, or direct reversal of the alignment score.
  • Hyperparameter selection: Regularization strengths λ are tuned separately for each dataset using a prespecified grid and validation data, without test-set selection.Table 5 lists the hyperparameters for latent regularization and AGL.

9 Extended interventional evidence across datasets

Interventions and robustness tests across additional datasets support LPA as a functional determinant of predictive uncertainty. Disrupting alignment worsens uncertainty and error, whereas changes to posterior regularization alone leave confidence largely stable unless restrictive priors break the alignment structure.

  • Interventional evidence: Across five additional datasets, L1, L2, and anti-alignment regularization degrade MAE and substantially increase predictive uncertainty relative to unconstrained baselines.
  • Interventional evidence: The degradation hierarchy is generally anti-alignment > L1 > L2, with anti-alignment and L1 producing greater uncertainty and MAE than L2.L2 tends to produce diffuse latent vectors, whereas L1 and anti-alignment induce sharper or sparser representations.
  • Interventional evidence: Figure 6 compares MAE, predictive uncertainty, LPAS, and latent-vector/posterior-standard-deviation geometry for baseline and regularized models across five datasets.
  • Prior robustness: When prior standard deviation falls below 0.01, coherent alignment fails, causing a sharp increase in predictive uncertainty and error.The experiments identify a need for sufficient exploratory variance to discover effective alignment pathways.
  • Overall robustness: Across the examined deterministic GNN feature extractors and mean-field Bayesian output layers, LPA remains associated with predictive uncertainty under posterior-regularization and prior variations.
  • KL-weight robustness: Lowering β reduces posterior uncertainty, but predictive accuracy and output uncertainty remain largely invariant across data-poor and data-rich regimes.The same pattern appears across five additional physicochemical property datasets.
  • KL-weight robustness: Data-rich models maintain higher LPAS under varying β, while data-poor models retain low LPAS, preserving the distinction between their uncertainty regimes.
  • Prior robustness: Narrower priors reduce posterior standard deviations but increase MAE and output uncertainty as LPAS decreases and latent vectors enter high-variance dimensions.

11 Cross-dataset evaluation of AGL and density awareness

Across five additional physicochemical datasets, AGL strengthens Latent-Posterior Alignment in both data-poor and data-rich regimes, with larger alignment changes in data-rich settings. It reduces predictive uncertainty without degrading MAE and improves structural calibration by better matching uncertainty to data density, despite mixed probabilistic-calibration results.

  • Evaluation design: The evaluation compares AGL with the baseline across five properties, six metrics, and data-poor versus data-rich regimes using baseline-normalized values.The reported metrics include LPAS, MAE, σw, σy, ECE, and DUC.
  • Cross-dataset evaluation: AGL consistently increases LPAS relative to the baseline across both data-poor and data-rich regimes.The increase is more pronounced in the data-rich regime.
  • Uncertainty and accuracy: Predictive uncertainty σy decreases substantially while MAE remains stable and posterior uncertainty σw remains unchanged or slightly increases.This pattern is reported across the evaluated datasets and regimes.
  • Calibration: Structural calibration DUC improves consistently and significantly, whereas probabilistic calibration ECE shows mixed results or slight degradation in some cases.The results indicate stronger gains for geometric calibration than for probability matching.
  • Density awareness: In data-poor settings, AGL corrects the baseline’s disjoint, quasi-random relationship between uncertainty and data density.The comparison uses t-SNE maps of latent representations, data density, and output uncertainty.
  • Density awareness: In data-rich settings, AGL sharpens an already slight baseline alignment, particularly in boundary regions.The data-rich comparison is visualized through t-SNE maps of density and output uncertainty.
Loading 2608.20758v1…