Hidden Axis of Uncertainty: Latent-Posterior Alignment in Graph Neural Networks with Bayesian Output Layers
Authors: Suk Hoon Choi, Damdae Park, Junhyuk Choi, Hyein Jung, Changsoo Kim, Ung Lee, Kyeongsu Kim
Organizations: Clean Energy Research Center, Korea Institute of Science and Technology, Seoul 02792, Republic of Korea. · Department of Chemical and Biological Engineering, Korea University, Seoul 02841, Republic of Korea. · Division of Energy & Environment Technology, Korea University of Science and Technology, 217 Gajeong-ro, Yuseong-gu, Daejeon 34113, Republic of Korea.
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.
Figures & tables
Figure 1: Uncertainty behavior of BGNN. (a) Schematic of the BGNN architecture. (b–g) Deviation from conventional expectation. As training data increases, predictive error (b) and uncertainty decrease (c), whereas weight posterior uncertainty expands (d), leading to broader distributions in the data-rich regime (f, red) compared to the data-poor regime (f, blue). The posterior means remain near zero on average (e), while both the NLL and KL terms decrease (g). Shaded areas represent standard deviation across n=30 independent runs. “Training data [%]” denotes the percentage of the available training set used to train each model.
Figure 2: Latent-Posterior Alignment (LPA) of BGNN. (a-c) Normalized latent vectors ( z~ , blue) progressively align with low-variance weight dimensions ( σ~w , orange) as data accumulates (1% → 80%). Grey bars indicate the relative uncertainty contribution. (d) High-uncertainty samples (green) exhibit amplified latent magnitudes compared to the full data distribution (blue) while maintaining the alignment. (e) t-SNE projection of latent vectors colored by predictive uncertainty. (f) Monotonic increase in the Latent-Posterior Alignment Score (LPAS), quantifying the alignment mechanism. Shaded areas represent standard deviation across n=30 independent runs.
Figure 3: Interventional evidence for a causal contribution of Latent–Posterior Alignment (a-f). Structurally disrupting the LPA via L1, L2, and anti-alignment regularization degrades accuracy (a) and amplifies predictive uncertainty (b), accompanied by the decrease in the alignment score (LPAS) (c) compared to baseline (dashed line). Visualizations of latent geometry (d–f) illustrate this structural breakdown: unlike the baseline, the regularized models fail to maintain LPA. Shaded areas represent standard deviation across n=30 independent runs.
Figure 4: Alignment-Guided Learning (AGL) and uncertainty calibration analysis. (a) Normalized performance metrics. For LPAS, normalized values represent the proportional reduction in 1−LPAS relative to the baseline; values above 100 indicate stronger alignment. AGL improves structural calibration (DUC) and increases the Latent-Posterior Alignment Score (LPAS) across both regimes without compromising accuracy (MAE) compared to baseline (dashed line). (b–e) Latent alignment. AGL enforces latent vectors ( z~ ) to shift toward low-variance posterior dimensions. This geometric reconfiguration is more pronounced in the data-rich regime (d, e) compared to the data-poor regime (b, c). (f-m) Qualitative visualization of density awareness. Each pair displays estimated density (left, KDE) and predictive uncertainty (right, σy ) on the same t-SNE coordinates. In the data-poor regime (f–i), the baseline shows no correlation between density (f) and uncertainty (g). AGL establishes this missing relationship, aligning low-density regions (h, blue in KDE) with high uncertainty (i, red in σy ). In the data-rich regime (j-m), AGL refines this structural correspondence, sharpening the boundary so that uncertainty (m) is precisely concentrated along the low-density edges of the data manifold (l), compared to the baseline (j, k).
Figure 1: Schematic of Bayesian GNN models developed in this study. (a) Backbone GNN model and (b) BGNN model.
Feature
Note
Implementation
Element type
Ten elemental types corresponding to atomic numbers (1, 5, 6, 7, 8, 15, 16, 17, 35, and 53) are represented using a one-hot encoding. For instance, a carbon atom is assigned the vector [0, 0, 1, 0, 0, 0, 0, 0, 0, 0].
atom_GetAtomicNum()
Formal charge
The formal charge of each atom is one-hot encoded over the set {-2, -1, 0, +1, +2}; any value outside this range (e.g., –3) is mapped to an all-zero vector.
atom_GetFormalCharge()
Hybridization
The hybridization state of each atom is encoded over the set {s, sp, sp 2 , sp 3 }, with each state mapped to its corresponding one-hot vector.
atom_GetHybridization()
Total degree
The number of bonded neighbors is encoded using a one-hot scheme over the values {1, 2, 3, 4}.
atom_GetTotalDegree()
Total number of neighboring hydrogen atoms
The total number of attached hydrogen atoms is represented by a one-hot vector defined over the set {0, 1, 2, 3}.
atom_GetTotalNumHs()
Aromatic
Whether an atom is aromatic is encoded as a binary feature taking values from {0, 1}.
atom_GetIsAromatic()
Table 1: Description of utilized atomic features.
Layer
Parameter
Value
Graph convolutions with GIN
Number of layers
5
Number of hidden channels
128
Dimensionality of fingerprint channel (latent dimension)
64
Dropout probability
0.3
Bayesian layer
Number of layers
1
Table 2: Detailed description of model architecture.
Parameter
Value
Input data scaling method
Normalization to (0, 1) range
Output data scaling method
Standardization to N(0,1)
Number of epochs
500
Batch size
64 for 0<N<104 ; 256 for 104≤N<105 ; 1024 for 105≤N
Optimizer
AdamW
Initial learning rate
10−3
Table 3: Detailed description of model training configuration.
Figure 2: Parity plots for the physicochemical properties examined in this study: (a) boiling point (bp), (b) melting point (mp), (c) vapor pressure at 25 ∘C (Pvap), (d) partition coefficient (P), (e) standard enthalpy of formation (H), and (f) standard Gibbs free energy of formation (G). The mean absolute error (MAE) for each property is annotated. Marker color denotes the amine type, and marker size represents molecular weight [g/mol]. Amine types are assigned priority in the order: tertiary > secondary > primary, when a compound belongs to multiple classes.
Property (abbreviation) [unit]
Data sources
Note
Total number of data points
Boiling point (bp) [K]
Sanchez-Lengeling et al. Sanchez-Lengeling et al. (2017) https://github.com/aspuru-guzik-group/ORGANIC/blob/master/data/datasets OPERA dataset v2.9 Mansouri et al. (2018) https://github.com/kmansouri/OPERA/tree/master
–
5,656
Melting point (mp) [K]
Jean-Claude Bradley Open Melting Point Dataset https://figshare.com/articles/dataset/Jean_Claude_Bradley_Open_Melting_Point_Dataset/1031637?file=1503990 OPERA dataset v2.9 Mansouri et al. (2018) https://github.com/kmansouri/OPERA/tree/master
–
21,414
Partition coefficient (P) [-]
OPERA dataset v2.9 Mansouri et al. (2018) https://github.com/kmansouri/OPERA/tree/master
Polarity of a compound: logP=log10([solute]water[solute]oct)
13,962
Vapor pressure (Pvap) [mmHg]
Sanchez-Lengeling et al. Sanchez-Lengeling et al. (2017) https://github.com/aspuru-guzik-group/ORGANIC/blob/master/data/datasets OPERA dataset v2.9 Mansouri et al. (2018) https://github.com/kmansouri/OPERA/tree/master
Vapor pressure at 25 ∘ C
2,925
Enthalpy (H) [kcal/mol]
QM9* star dataset Tang et al. (2024) https://github.com/gentlej1999/qm9star_query
Standard enthalpy of formation
2,008,806
Gibbs free energy (G) [kcal/mol]
Standard Gibbs free energy of formation
Table 4: Description of datasets and their properties.
Figure 3: Negligible influence of bias posteriors on predictive uncertainty. Analyses are presented for six different datasets (a-f), where each panel displays the evolution of bias posterior mean, posterior standard deviation, and relative uncertainty contribution (blue for bias, orange for weight) as a function of training data size.
Figure 4: Verification of the non-classical uncertainty behavior. Each row (a-e) corresponds to a distinct dataset, while columns show metrics as a function of training data size (from left to right): Mean Absolute Error (MAE), predictive uncertainty ( σy ), weight posterior uncertainty ( σw ), and weight posterior mean ( μw ). Consistent with the main analysis, predictive uncertainty (second column) decreases while posterior uncertainty (third column) increases across all tasks.
Figure 5: Cross-dataset consistency of Latent-Posterior Alignment. Results are presented for five different datasets (a-e). Columns 1–2: Comparison of latent geometry between data-poor (1%) and data-rich (80%) regimes. The plots overlay the normalized latent vector ( z~ , blue), posterior standard deviation ( σ~w , orange), and relative uncertainty contribution (grey bars). Column 3: The Latent-Posterior Alignment Score (LPAS) versus data size. Column 4: Comparison of latent structure between the overall data distribution (blue) and samples with the highest predictive uncertainty (green).
Dataset (Property)
λl1
λl2
λanti−align
γ
bp
0.05
0.2
0.1
0.05
mp
0.3
1.5
0.5
0.02
P
0.05
0.3
0.05
0.02
Pvap
0.02
0.1
0.01
0.005
H
0.02
0.2
0.1
0.05
G
0.02
0.1
0.05
0.02
Table 5: Hyperparameters for latent regularization experiments and Alignment-Guided Learning. The regularization and AGL strengths ( λ and γ ) were adjusted for each dataset.
Figure 6: Impact of latent regularization on predictive uncertainty and latent geometry. Rows (a-e) correspond to five different datasets. Columns 1-3: Comparison of MAE, predictive uncertainty, and LPAS between baseline (dashed line) and regularized models. Columns 4-6: Visualization of latent vector and posterior standard deviation under L1, L2, and anti-alignment penalties, respectively.
Figure 7: Robustness of latent-posterior alignment mechanism. (a–f) Robustness to KL weight. Reducing β leads to smaller σw (b). However, this variation has negligible impact on accuracy (a) and output uncertainty (c). This invariance arises because LPA serves as a robust geometric mechanism: data-rich models maintain a consistently high alignment score (d, purple) and a clearly aligned latent-posterior geometry (f) regardless of β . In contrast, data-poor models exhibit low LPAS (d, cyan) and fail to form such structure (e), confirming that predictive confidence is governed by geometric alignment rather than posterior magnitude. (g-l) Impact of prior variance constraints. Restricting the prior standard deviation (x-axis) forces posterior contraction (h) but paradoxically increases predictive error (g) and output uncertainty (i). This failure arises because tight priors restrict the parameter-space flexibility required to maintain LPA, as evidenced by the decay of alignment score (j) and disrupted geometry (k) compared to the relaxed prior (l). Shaded bands denote standard deviation across n=30 independent runs.
Figure 8: Influence of KL-divergence weight on Latent-Posterior Alignment and uncertainty dynamics. Rows (a-e) correspond to five different property datasets. Cyan and purple represent data-poor and data-rich regimes, respectively. Columns 1-4 show the KL-divergence weight impact on MAE, posterior uncertainty ( σw ), predictive uncertainty ( σy ), and LPAS. Columns 5-6 visualize the geometry of latent vectors and posterior standard deviations for data-poor (Column 5) and data-rich (Column 6) models.
Figure 9: Impact of prior distribution on Latent-Posterior Alignment and predictive uncertainty. Analysis is conducted on the data-rich regime (80% train data) across five physicochemical properties (a-e). Columns 1-4 show the evolution of MAE, posterior uncertainty ( σw ), predictive uncertainty ( σy ), and LPAS as a function of prior standard deviation. Columns 5-6 visualize latent vectors under restrictive prior (Column 5, σprior=0.0001 ) and relaxed prior (Column 6, σprior=1 ).
Figure 10: Performance evaluation of Alignment-Guided Learning (AGL) across diverse datasets. Results are shown for five different properties (a-e). The AGL model (colored bars) is compared with the baseline (gray dashed line) for six metrics, with values normalized relative to the baseline scaled to 100. Light and dark shaded bars represent the data-poor and data-rich regimes, respectively.
Figure 11: Impact of Alignment-Guided Learning (AGL) on latent geometry in the data-poor regime. Visualizations of normalized latent vector magnitudes ( z~ , blue line) and posterior standard deviations ( σ~w , orange line) across five physicochemical datasets (a–e). Each panel compares the unconstrained baseline model (left) with the AGL model (right). The arrows indicate the structural shift induced by the alignment objective. In this data-scarce regime, while AGL induces a tendency toward the alignment, the resulting structural reorganization is not clearly visible (or remains visually subtle) due to severe information scarcity. Shaded areas represent the standard deviation across n=30 independent runs.
Figure 12: Impact of Latent-Posterior Alignment via AGL in the data-rich regime. Comparison of latent geometry between Baseline (left) and AGL (right) models trained in a data-abundant setting across five datasets (a–e). Consistent with the data-poor regime, AGL enforces alignment; however, the extent of this reorganization is more pronounced. With sufficient data, AGL actively reconfigures the latent space, creating a sharp boundary where latent activity is strictly confined to stable, low-variance axes (forming a distinct “X” shape). This confirms that the availability of information facilitates more precise geometric calibration. Shaded areas denote standard deviation ( n=30 ).
Figure 13: Structural calibration in data-poor regime. t-SNE dimensional reduction of latent representations is used to visualize latent geometry and compare baseline models against AGL models for five different property datasets. Each sub-figure displays two maps: data density (left, estimated via KDE) and output uncertainty (right, σy ). Colors range from low (blue) to high (red). Left columns (a, c, e, g, i) and right columns (b, d, f, h, j) represent baseline and AGL models, respectively.
Figure 14: Structural calibration in data-rich regime. t-SNE dimensional reduction of latent representations is used to visualize latent geometry and compare baseline models against AGL models for five different property datasets. Each sub-figure displays two maps: data density (left, estimated via KDE) and output uncertainty (right, σy ). Colors range from low (blue) to high (red). Left columns (a, c, e, g, i) and right columns (b, d, f, h, j) represent baseline and AGL models, respectively.
On evolving graphs, node classifiers must satisfy two key requirements: inductive generalization to newly arriving nodes under distribution shift and calibrated uncertainty for safety-sensitive applications. Standard graph neural networks (GNNs) are typically trained once and address neither requirement. We adapt the Bayesian last-layer (BLL) model by placing random last-layer parameters on top of a deterministic GNN encoder for uncertainty quantification. The categorical softmax likelihood required for classification breaks Gaussian conjugacy, so neither the training posterior nor the test-time streaming update has a closed-form solution. To address both challenges, we introduce a variational Bayesian last-layer (VBLL) objective that jointly trains the encoder and an approximate last-layer posterior by maximizing an evidence lower bound with a Monte Carlo expected log-likelihood. At test time, we freeze the encoder and apply an online Laplace update to the last-layer posterior. This update corresponds to a power-prior Bayesian model with exponential forgetting and a Kullback-Leibler anchor to the training posterior. Across five node-classification benchmarks under distribution shift, online GVBLL is the only method to achieve the best accuracy and negative log-likelihood on every dataset. It improves accuracy by up to 17 percentage points on Cora and 14 percentage points on ogbn-arxiv over the strongest non-GVBLL baseline, while remaining competitive in calibration with MC Dropout, Deep Ensembles, Temperature Scaling, and Gaussian-process classifiers.
Jinwen Xu, Gonzalo Mateos Buckstein, Qin Lu
School of Electrical and Computer Engineering, University of Georgia, Athens, GA, USA · Department of Electrical and Computer Engineering, University of Rochester, Rochester, NY, USA
Gaussian neural networks (GaNNs) are proposed as a novel regularization mechanism for neural networks. From a Bayesian perspective standard regularization techniques can be viewed as imposing priors over weight-space. Assuming priors over activation-space remains a largely unexplored possibility. GaNNs assume such priors. They do this by treating activities from earlier layers like signals with Gaussian noise and predicting the properties of the noise distribution using an additional unsupervised loss. While training, the unsupervised loss acts as a penalty on unexpected activities, allowing greater weight updates in less surprising directions. The paper demonstrates the superiority of Gaussian neural networks over standard neural networks on a variety of classification and regression tasks. We also investigate the ability of GaNNs to quantify uncertainty.
Peter Kuhn, Victoria Heusinger-Heß
Fraunhofer Institute for High-Speed Dynamics, Ernst-Mach-Institut, EMI, Germany
Uncertainty quantification (UQ) in graph neural networks (GNNs) is crucial in high-stakes domains but remains a significant challenge. In graph settings, message passing often relies on strong assumptions such as exchangeability, which are rarely satisfied in practice, and achieving reliable UQ typically requires costly resampling or post-hoc calibration. To address these issues, we introduce Quantile-free Prediction Interval GNN (QpiGNN), a framework that builds on quantile regression (QR) to enable GNN-based UQ by directly optimizing coverage and interval width without requiring quantile inputs or post-processing. QpiGNN employs a dual-head architecture that decouples prediction and uncertainty, and is trained with label-only supervision through a quantile-free joint loss. This design allows efficient training and yields robust prediction intervals, with theoretical guarantees of asymptotic coverage and near-optimal width under mild assumptions. Experiments on 19 synthetic and real-world benchmarks show QpiGNN achieves average 22% higher coverage and 50% narrower intervals than baselines, while ensuring efficiency and robustness to noise and structural shifts.
Soyoung park, Hwanjun Song, Sungsu Lim
Department of Computer Science and Engineering, Chungnam National University, Daejeon, Korea · Department of Industrial and Systems Engineering, KAIST, Daejeon, Korea