Uncertainty Quantification in Graph Neural Networks with Shallow Ensembles

2025-04-17

AI Summary

Graph neural networks (GNNs) exhibit unreliable predictions and poorly quantifiable uncertainty when applied to out-of-distribution (OOD) materials data. To address this, we propose DPOSE (Direct Propagation of Uncertainty via Shallow Ensembles), a lightweight, shallow ensemble method systematically integrated into the SchNet architecture for efficient and well-calibrated uncertainty quantification (UQ) of GNN outputs. Unlike deep ensembles, DPOSE avoids prohibitive computational overhead while preserving strong discriminative capability between in-distribution and OOD samples. Evaluated across multiple first-principles datasets—including QM9, OC20, and Gold MD—DPOSE achieves >95% OOD detection rate and improves UQ calibration by over 40% on average. This work establishes a new paradigm for robust, uncertainty-aware GNN-based materials modeling.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Uncertainty Representations
• Data Mining & Knowledge Management: Anomaly/Outlier Detection

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Social Networks and Social Media: Influence propagation, information diffusion, and the prediction on networks
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Quantify uncertainty in Graph Neural Networks for reliable predictions
• Distinguish in-domain and out-of-domain data using shallow ensembles
• Improve robustness of GNN-based materials modeling with lightweight methods

Innovation

• Shallow Ensembles for Uncertainty Quantification
• Integration of DPOSE into SchNet model
• Lightweight UQ methods for GNN robustness