Accurate and Scalable Estimation of Epistemic Uncertainty for Graph Neural Networks

TL;DR

This paper introduces G-ΔUQ, a novel uncertainty quantification method for graph neural networks, improving confidence calibration under distribution shifts and enhancing safety-critical decision-making.

Contribution

It proposes G-ΔUQ, a flexible UQ method for GNNs, and provides insights into confidence indicator calibration under distribution shifts.

Findings

G-ΔUQ outperforms existing UQ methods in calibration accuracy.

G-ΔUQ improves OOD detection and generalization gap prediction.

Model size and expressivity do not guarantee better confidence calibration.

Abstract

Safe deployment of graph neural networks (GNNs) under distribution shift requires models to provide accurate confidence indicators (CI). However, while it is well-known in computer vision that CI quality diminishes under distribution shift, this behavior remains understudied for GNNs. Hence, we begin with a case study on CI calibration under controlled structural and feature distribution shifts and demonstrate that increased expressivity or model size do not always lead to improved CI performance. Consequently, we instead advocate for the use of epistemic uncertainty quantification (UQ) methods to modulate CIs. To this end, we propose G-$\Delta$UQ, a new single model UQ method that extends the recently proposed stochastic centering framework to support structured data and partial stochasticity. Evaluated across covariate, concept, and graph size shifts, G-$\Delta$UQ not only outperforms several popular UQ methods in obtaining calibrated CIs, but also outperforms alternatives when CIs are used for generalization gap prediction or OOD detection. Overall, our work not only introduces a new, flexible GNN UQ method, but also provides novel insights into GNN CIs on safety-critical tasks.