cs.LGAug 21, 2026
SaveHidden Axis of Uncertainty: Latent-Posterior Alignment in Graph Neural Networks with Bayesian Output Layers
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.
Abstract
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 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 ( , blue) progressively align with low-variance weight dimensions ( , 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 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 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 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 ( ) 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, ) 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 ). 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 |
| Number of epochs | 500 |
| Batch size | 64 for ; 256 for ; 1024 for |
| Optimizer | AdamW |
| Initial learning rate |
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 (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: | 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 ( ), weight posterior uncertainty ( ), and weight posterior mean ( ). 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 ( , blue), posterior standard deviation ( , 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) | ||||
|---|---|---|---|---|
| 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 (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 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 ( ), predictive uncertainty ( ), 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 ( ), predictive uncertainty ( ), and LPAS as a function of prior standard deviation. Columns 5-6 visualize latent vectors under restrictive prior (Column 5, ) and relaxed prior (Column 6, ).
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 ( , blue line) and posterior standard deviations ( , 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 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 ( ).
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, ). 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, ). 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.