Bayesian uncertainty quantification of graph neural networks using stochastic gradient Hamiltonian Monte Carlo

Abstract

Graph neural networks (GNNs) have proven state-of-the-art performance in molecular property prediction tasks. However, a significant challenge with GNNs is the reliability of their predictions, particularly in critical domains where quantifying model confidence is essential. Therefore, assessing uncertainty in GNN predictions is crucial to improving their robustness. Existing uncertainty quantification methods, such as Deep ensembles and Monte Carlo Dropout, have been applied to GNNs with some success, but these methods are limited to approximate the full posterior distribution. In this work, we propose a novel approach for scalable uncertainty quantification in molecular property prediction using Stochastic Gradient Hamiltonian Monte Carlo (SGHMC). Additionally, we utilize a cyclical learning rate to facilitate sampling from multiple posterior modes which improves posterior exploration within a single training round. Moreover, we compare the proposed methods with Monte Carlo Dropout and Deep ensembles, focusing on error analysis, calibration, and sharpness, considering both epistemic and aleatoric uncertainties. Our experimental results demonstrate that the proposed parallel-SGHMC approach significantly outperforms Monte Carlo Dropout and Deep ensembles in terms of calibration and sharpness. Specifically, parallel-SGHMC reduces the sum of squared errors by 99.4% and 75%, respectively, when compared to Monte Carlo Dropout and Deep Ensembles. These findings suggest that parallel-SGHMC is a promising method for uncertainty quantification in GNN-based molecular property prediction.