Uncertainty Quantification in Graph Neural Networks with Shallow Ensembles

TL;DR

This work addresses the challenge of unreliable GNN predictions for out-of-domain data in materials modeling by introducing Direct Propagation of Shallow Ensembles (DPOSE) integrated with SchNet. DPOSE uses lightweight, weight-sharing shallow ensembles to deliver both predictive means and uncertainties without the heavy cost of deep ensembles, trained via a Negative Log-Likelihood objective. Evaluations on QM9, OC20, and Gold MD datasets demonstrate that DPOSE can distinguish in-domain from out-of-domain configurations and captures uncertainty trends related to molecular size, distortions, and compositional changes, with dataset-specific strengths and limitations. The approach offers a scalable path toward robust uncertainty-aware materials discovery and can be integrated with active learning to guide efficient exploration of design spaces.

Abstract

Machine-learned potentials (MLPs) have revolutionized materials discovery by providing accurate and efficient predictions of molecular and material properties. Graph Neural Networks (GNNs) have emerged as a state-of-the-art approach due to their ability to capture complex atomic interactions. However, GNNs often produce unreliable predictions when encountering out-of-domain data and it is difficult to identify when that happens. To address this challenge, we explore Uncertainty Quantification (UQ) techniques, focusing on Direct Propagation of Shallow Ensembles (DPOSE) as a computationally efficient alternative to deep ensembles. By integrating DPOSE into the SchNet model, we assess its ability to provide reliable uncertainty estimates across diverse Density Functional Theory datasets, including QM9, OC20, and Gold Molecular Dynamics. Our findings often demonstrate that DPOSE successfully distinguishes between in-domain and out-of-domain samples, exhibiting higher uncertainty for unobserved molecule and material classes. This work highlights the potential of lightweight UQ methods in improving the robustness of GNN-based materials modeling and lays the foundation for future integration with active learning strategies.

Figures (9)

• Figure 1: Comparison of the Original Architecture and the DPOSE Architecture. The Original Architecture produces a single output ($y$) using a deep neural network. The DPOSE Architecture consists of multiple shallow ensembles in the last layer predict outputs ($y_1, y_2, \dots, y_n$), with the final output ($\bar{y}$) computed as their average and the uncertainty ($\sigma$) estimated as their variance.
• Figure 2: Internal energy predictions and uncertainty estimates for C$_2$N$_2$, C$_2$H$_6$, CH$_3$CN, and C$_2$H$_2$O$_2$ as a function of bond length between two carbon atoms. Solid lines represent predicted internal energies, and shaded regions indicate prediction variance. Minimum variance is observed at equilibrium bond lengths from the training dataset.
• Figure 3: (a) Parity plot for inter-metallic slabs showing high accuracy. (b) Parity plot for non-metal slabs showing lower accuracy. (c) Variance estimates show higher uncertainty for non-metal slabs than inter-metallic slabs.
• Figure 4: Change in volume per atom for three test cases in inter-metallic slabs: (a) Ag$_36$Cd$_12$, (b) Bi$_16$Ni$_16$Zr$_16$, and (c) Nb$_16$Ru$_32$Sn$_16$. The uncertainty region (shaded in blue) is negligible and overlaps with the relative energy curve (black dotted line), with no discernible trend in uncertainty. The difference between maximum and minimum uncertainty is approximately 0.002 eV/atom, indicating consistently low uncertainty across all cases. The red dotted line is the equilibrium volume per atom for the respective systems. The equilibrium volume per atom, indicated by the red dotted line, aligns with the point of zero relative energy, demonstrating strong model accuracy.
• Figure 5: Change in volume per atom for three nonmetal slab systems: (a) Pb$_16$S$_32$, (b) Cl$_24$S$_48$Sn$_60$, and (c) Pb$_24$S$_66$Sb$_24$. The equilibrium volume per atom, indicated by the red dotted line, is not accurately predicted by the model. The relative energy curve (red line) is non-zero at equilibrium, indicating suboptimal model performance. The uncertainty (shaded in blue) follows the expected trend, being lower near equilibrium states and higher away from them.