Accurate and Scalable Estimation of Epistemic Uncertainty for Graph Neural Networks
This work addresses safety-critical deployment of GNNs by improving uncertainty estimation, though it is incremental as it extends an existing framework to structured data.
The paper tackles the problem of accurate confidence indicators for graph neural networks under distribution shift by proposing G-ΔUQ, a new epistemic uncertainty quantification method that outperforms existing approaches in calibration, generalization gap prediction, and out-of-distribution detection across various shifts.
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-$Δ$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-$Δ$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.