On Estimating the Epistemic Uncertainty of Graph Neural Networks using Stochastic Centering

Graph neural networks (GNNs) are known to have limited expressivity (poor size generalization; over-smoothing; over-squashing). However, at test time, GNNs may encounter distributions where such factors are present. For example, test datasets may have larger sizes than those used for training. In such settings, to ensure safe deployment, it is necessary that GNNs provide accurate confidence indicators that can then be utilized in a variety of downstream safety tasks (generalization gap prediction; calibration; OOD detection). Here, we assess the ability of several baseline uncertainty estimators (Monte Carlo Dropout, Deep Ensembles, Temperature Scaling) in producing well-calibrated confidence estimates under covariate and concept shifts, and study the impact of architecture and model size on the quality of these estimates. Moreover, we adapt a recently proposed stochastic centering framework to graph datasets/GNNs, identifying several graph-specific challenges in the process. Overall, our work not only rigorously studies UQ under challenging graph distribution shifts, but also provides multiple insights into designing effective UQ estimators on graphs that are effective on a variety of safety-critical tasks.