Accurate and Scalable Estimation of Epistemic Uncertainty for Graph
  Neural Networks
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.
Puja Trivedi, Mark Heimann, Rushil Anirudh, Danai Koutra, Jayaraman J. Thiagarajan
# Accurate and Scalable Estimation of Epistemic Uncertainty for Graph Neural Networks

###### 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.

## 1 Introduction

As graph neural networks (GNNs) are increasingly deployed in critical applications with test-time distribution shifts (Zhang and Chen, 2018; Gaudelet et al., 2020; Yang et al., 2018; Yan et al., 2019; Zhu et al., 2022), it becomes necessary to expand model evaluation to include safety-centric metrics, such as calibration errors (Guo et al., 2017), out-of-distribution (OOD) rejection rates (Hendrycks and Gimpel, 2017), and generalization gap estimates (Jiang et al., 2019), to holistically understand model performance in such shifted regimes (Hendrycks et al., 2022b; Trivedi et al., 2023b). Notably, such additional metrics often rely on _confidence indicators_ (CIs), such as maximum softmax or predictive entropy, which can be derived from prediction probabilities. Although there is a clear understanding in the computer vision literature that the quality of confidence indicators can noticeably deteriorate under distribution shifts (Wiles et al., 2022; Ovadia et al., 2019), and additional factors like model size or expressivity can exacerbate this deterioration (Minderer et al., 2021), the impact of these phenomena on graph neural networks (GNNs) remains under-explored.

Indeed, there is an expectation that adopting more advanced or expressive architectures (Chuang and Jegelka, 2022; Alon and Yahav, 2021; Topping et al., 2022; Rampasek et al., 2022; Zhao et al., 2022) would inherently improve CI calibration on graph classification tasks. Yet, we find that using graph transformers (GTrans) (Rampasek et al., 2022) or positional encodings (Dwivedi et al., 2022a; Wang et al., 2022b; Li et al., 2020) do not significantly improve CI calibration over vanilla message-passing GNNs (MPGNNs) even under controlled, label-preserving distribution shifts. Notably, when CIs are not well-calibrated, GNNs with high accuracy may perform poorly on the additional safety metrics,
leading to unforeseen risks during deployment. Given that using advanced architectures is not an immediately viable solution for improving CI calibration, we instead advocate for modulating CIs using epistemic _uncertainty estimates_.

Uncertainty quantification (UQ) methods (Gal and Ghahramani, 2016; Lakshminarayanan et al., 2017; Blundell et al., 2015) have been extensively studied for vision models (Guo et al., 2017; Minderer et al., 2021), and have been used to improve vision model CI performance under distribution shifts. Our work not only studies the effectiveness of such methods on improving GNN CIs, but also proposes a novel UQ method, G-\(\Delta\)UQ, which extends the recently proposed, state-of-the-art stochastic data-centering (or "anchoring") framework (Thiagarajan et al., 2022; Netanyah et al., 2023) to support partial stochasticity and structured data. In brief, stochastic centering provides a scalable alternative to highly effective, but prohibitively expensive, deep ensembles by efficiently sampling a model's hypothesis space, in lieu of training multiple, independently trained models. When using the uncertainty-modulated confidence estimates from G-\(\Delta\)UQ, we outperform other popular UQ methods, on not only improving the CI calibration under covariate, concept and graph size shifts, but also in improving generalization gap prediction and OOD detection performance.

**Proposed Work.** This work studies the effectiveness of GNN CIs on the graph classification tasks with distribution shifts, and proposes a novel uncertainty-based method for improving CI performance. Our contributions can be summarized as follows:

**Sec. 2: Case Study on CI Calibration.** We find that improving GNN expressivity does not mitigate CI quality degradation under distribution shifts.

**Sec. 4: (Partially) Stochastic Anchoring for GNNs.** We propose G-\(\Delta\)UQ, a novel UQ method based on stochastic centering for GNNs with support for partial stochasticity.

**Sec. 5: Evaluating Uncertainty-Modulated CIs under Distribution Shifts.** Across covariate, concept and graph-size shifts and a suite of evaluation protocols (calibration, OOD rejection, generalization gap prediction), we demonstrate the effectiveness of G-\(\Delta\)UQ.

## 2 Case Study on GNN CI Calibration

In this section, we demonstrate that GNNs struggle to provide calibrated confidence estimates under distribution shift (Dwivedi et al., 2020) despite, improvements in architectures (He et al., 2022; Corso et al., 2020; Zhao et al., 2022) and expressivity (Wang et al., 2022; Dwivedi et al., 2022). Since assessing calibration performance does not require any additional, potentially confounding, post-processing, we perform a direct assessment of GNN CI reliability and motivate why uncertainty-based CI modulation is needed.

**Notations.** Let \(\mathcal{G}=(\mathbf{X},\mathbb{E},\mathbf{A},Y)\) be a graph with node features \(\mathbf{X}\in\mathbb{R}^{N\times d_{\ell}}\), (optional) edge features \(\mathbb{E}\in\mathbb{R}^{m\times d_{\ell}}\), adjacency matrix \(\mathbf{A}\in\mathbb{R}^{N\times N}\), and graph-level label \(Y\in\{0,1\}^{c}\), where \(N,m,d_{\ell},c\) denote the number of nodes, number of edges, feature dimension and number of classes, respectively. We use \(i\) to index a particular sample in the dataset, e.g. \(\mathcal{G}_{i},\mathbf{X}_{i}\).

Then, we can define a graph neural network consisting of \(\ell\) message passing layers \((\mathtt{MPNN})\), a graph-level readout function (READOUT), and classifier head (MLP) as follows:

\[\mathbf{X}_{M}^{\ell+1},\,\mathbb{E}^{\ell+1} = \mathtt{MPNN}_{e}^{\ell}\left(\mathbf{X}^{\ell},\mathbb{E}^{\ell},\mathbf{A}\right), \tag{1}\] \[\mathbf{G} = \mathtt{READOUT}\left(\mathbf{X}_{M}^{\ell+1}\right),\] (2) \[\hat{Y} = \mathtt{MLP}\left(\mathbf{G}\right), \tag{3}\]

where \(\mathbf{X}_{M}^{\ell+1},\mathbb{E}^{\ell+1}\) are intermediate node and edge representations, and \(\mathbf{G}\) is the graph representation. We focus on a graph classification setting throughout our paper.

**Experimental Set-up:** Our experimental set-up is as follows. All results are reported over three seeds.

_Models._ While improving the expressivity of GNNs is an active area of research, positional encodings (PEs) and graph-transformer (GTran) architectures (Muller et al., 2023) have proven to be particularly popular due to their effectiveness, and flexibility. Indeed, GTrans are known to not only help mitigate over-smoothing (a phenomenon where GNNs lose discriminative power) and over-squashing (a
phenomenon where GNNs collapse node representations) (Alon and Yahav, 2021; Topping et al., 2022) but also to better capture long-range dependencies in large graphs (Dwivedi et al., 2022). Critical to the success of any transformer architecture are well-designed PEs. Notably, graph PEs help improve GNN and GTran expressivity by distinguishing between isomorphic nodes, as well as capturing structural vs. proximity information (Dwivedi et al., 2022). Here, we ask if these enhancements translate to improved calibration under distribution shift with respect to simple MPNNs by: (i) incorporating equivariant and stable PEs (Wang et al., 2022); (ii) utilizing MPNN vs. GTran architectures; and, (iii) changing model depth and width. We utilize the state-of-the-art, flexible "general, powerful, scalable" (GPS) GTran (Rampasek et al., 2022) with the GatedGCN backbone. For fair comparison, we use a GatedGCN as the compared MPNN.

_Data._ Superpixel-MNIST (Dwivedi et al., 2020; Knyazev et al., 2019; Velickovic et al., 2018) is a popular graph classification benchmark that converts MNIST images into \(k\) nearest-neighbor graphs of superpixels (Achanta et al., 2012). We select this benchmark as it allows for (i) a diverse set of well-trained models without requiring independent, extensive hyper-parameter search and (ii) controlled, label preserving distribution shifts. Inspired by Ding et al. (2021), we create structurally distorted but valid graphs by rotating MNIST images by a fixed number of degrees and then creating the super-pixel graphs from these rotated images. (See Appendix, Fig. 8.) Since superpixel segmentation on these rotated images will yield different superpixel \(k\)-nn graphs without harming class information, we can emulate label-preserving structural distortion shifts. Note, the models are trained only using the original (\(0^{\circ}\) rotation) graphs.

_Evaluation._ Calibrated models are expected to produce confidence estimates that match the true probabilities of the classes being predicted (Naeini et al., 2015; Guo et al., 2017; Ovadia et al., 2019). While poorly calibrated CIs are over/under confident in their predictions, calibrated CIs are more trustworthy and can also improve performance on other safety-critical tasks which implicitly require reliable prediction probabilities (see Sec. 5). Here, we report the top-1 label expected calibration error (ECE) (Kumar et al., 2019; Detlefsen et al., 2022). Let \(p_{i}\) be the top-1 probability, \(c_{i}\) be the predicted confidence, \(b_{i}\) a uniformly sized bin in \([0,1]\). Then, \(ECE:=\sum_{i}^{N}b_{i}\|(p_{i}-c_{i})\|\).

**Observations.** In Fig. 1, we present our results and make the following observations.

Despite the aforementioned benefits in model expressivity, GPS performs noticeably worse compared to the MPGNN, despite having comparable accuracies. This is apparent particularly at severe shifts (\(60^{\circ}\), \(90^{\circ}\), \(180^{\circ}\) rotations). Furthermore, we find that PEs have minimal effects on both calibration and accuracy. This suggests that while these techniques may enhance theoretical and empirical expressivity, they do not necessarily transfer to the safety-critical task of obtaining calibrated predictions under distribution shifts. In addition, we investigate the impact of model depth and width on calibration performance, considering that model size has been known to affect both the calibration of vision models (Guo et al., 2017; Minderer et al., 2021) and the propensity for over

Figure 1: **Calibration on Structural Distortion Distribution Shifts.** On a controlled graph structure distortion shift, we evaluate models trained on the standard superpixel MNIST benchmark (Dwivedi et al., 2020) on super-pixel \(k\)-nn graphs created from rotated MNIST images. While accuracy is expected to decrease as distribution shift increases, we observe that the expected calibration error also grows significantly worse. Importantly, this trend is persistent when considering transformer architectural variants (GPS (Rampasek et al., 2022)), as well as different depths and widths. In contrast, our proposed G-\(\Delta\)UQ method achieves substantial improvement in ECE without significantly compromising on accuracy.
squashing in GNNs (Xu et al., 2021). We see that increasing the number of message passing layers (\(L=3\to L=5\)) can marginally improve accuracy, but it may also marginally decrease ECE. Moreover, we find that increasing the width of the model can lead to slightly worse calibration at high levels of shift (\(90^{\circ}\), \(180^{\circ}\)), although accuracy is not compromised.

Notably, when we apply our proposed method G-\(\Delta\)UQ, (see Sec. 4), to the MPGNN with no positional encodings, it significantly improves the calibration over more expressive variants (GPS, LPE), across all levels of distribution shifts, while maintaining comparable accuracy. We briefly note that we did not tune the hyper-parameters to our method to ensure a fair comparison, so we expect that accuracy could be further improved. Overall, our results emphasize that obtaining reliable GNN CIs remains a difficult problem that cannot be easily solved through advancements in architectures and expressivity. This motivates our uncertainty-modulated CIs as an architecture-agnostic solution.

## 3 Related Work

Here, we discuss techniques for improving CI reliability and the recently proposed stochastic centering paradigm, before introducing our proposed method in Sec. 4.

### Improving Confidence Indicators

It is well known in computer vision that CIs are often unreliable or uncalibrated directly out-of-the-box (Guo et al., 2017), especially under distribution shifts (Ovadia et al., 2019; Wiles et al., 2022; Hendrycks et al., 2019). Given that reliable CIs are necessary for a variety of safety-critical tasks, including generalization error prediction (GEP) (Jiang et al., 2019) and out-of-distribution (OOD) detection (Hendrycks and Gimpel, 2017), many strategies have been proposed to improve CI calibration (Lakshminarayanan et al., 2017; Guo et al., 2017; Gal and Ghahramani, 2016; Blundell et al., 2015). One particularly effective strategy is to create a deep ensemble (DEns) (Lakshminarayanan et al., 2017) by training a set of independent models (e.g., different hyper-parameters, random-seeds, data order) where the mean predictions over the set is noticeably better calibrated. However, since DEns requires training multiple models, in practice, it can be prohibitively expensive to use. To this end, we focus on single-model strategies.

Single-model UQ techniques attempt to scalably and reliably provide uncertainty estimates, which can then optionally be used to modulate the prediction probabilities. Here, the intuition is that when the epistemic uncertainties are large in a data regime, confidence estimates can be tempered so that they better reflect the accuracy degradation during extrapolation (e.g., training on small-sized graphs but testing on large-sized graphs). Some popular strategies include: Monte Carlo dropout (MCD) (Gal and Ghahramani, 2016) which performs Monte Carlo dropout at inference time and takes the average prediction to improve calibration, temperature scaling (Temp) (Guo et al., 2017) which rescales logits using a temperature parameter computed from a validation set, and SVI (Blundell et al., 2015) which proposes a stochastic variational inference method for estimating uncertainty. While such methods are more scalable than DeepEns, in many cases, they struggle to match its performance (Ovadia et al., 2019). We note that while some recent works have studied GNN calibration, they focus on node classification settings (Hsu et al., 2022; Wang et al., 2021; Kang et al., 2022) and are not directly relevant to this work as they make assumptions that are only applicable to node classification tasks (e.g., proposing regularizers that rely upon on similarity to training nodes or neighbourhood similarity).

### Stochastic Centering for Uncertainty Quantification

Recently, it was found that applying a (random) constant bias to vector-valued (and image) data leads to non-trivial changes in the resulting solution of a DNN (Thiagarajan et al., 2022). This behavior was attributed to the lack of shift-invariance in the neural tangent kernel (NTK) induced by conventional neural networks such as MLPs and CNNs. Building upon this observation, Thiagarajan et al. proposed a single model uncertainty estimation method, \(\Delta\)-UQ, based on the principle of _anchoring_. Conceptually, anchoring is the process of creating a relative representation for an input sample \(x\) in terms of a random anchor \(c\) (which is used to perform the _stochastic centering_), \([x-c,c]\). By choosing different anchors randomly in each training iteration, \(\Delta\)-UQ emulates the process of sampling different solutions from the hypothesis space (akin to an ensemble). During inference, \(\Delta\)
UQ aggregates multiple predictions obtained via different random anchors and produces uncertainty estimates. Formally, given a trained stochastically centered model, \(f_{\theta}:[\mathbf{X}-\mathbf{C},\mathbf{C}]\rightarrow\hat{\mathbf{Y}}\), let \(\mathbf{C}:=\mathbf{X}_{train}\) be the anchor distribution, \(x\in\mathbf{X}_{test}\) be a test sample, and anchor \(c\in\mathbf{C}\) be anchor. Then, the mean target class prediction, \(\mathbf{\mu}(y|\mathrm{x})\), and corresponding variance, \(\mathbf{\sigma}(y|\mathrm{x})\) over \(K\) random anchors are computed as:

\[\mathbf{\mu}(y|\mathrm{x}) =\frac{1}{K}\sum_{k=1}^{K}f_{\theta}([\mathrm{x}-\mathrm{c}_{k}, \mathrm{c}_{k}]) \tag{4}\] \[\mathbf{\sigma}(y|\mathrm{x}) =\sqrt{\frac{1}{K-1}\sum_{k=1}^{K}(f_{\theta}([\mathrm{x}-\mathrm{ c}_{k},\mathrm{c}_{k}])-\mathbf{\mu})^{2}} \tag{5}\]

Since the variance over \(K\) anchors captures epistemic uncertainty by sampling different hypotheses, these estimates can be used to modulate the predictions: \(\mathbf{\mu}_{\text{calib}}=\mathbf{\mu}(1-\mathbf{\sigma})\). The resulting calibrated predictions and uncertainty estimates have led to state-of-the-art performance on image outlier rejection and calibration tasks, while still only requiring a single model. Furthermore, it was separately shown that anchoring can also be used to improve the extrapolation behavior of DNNs (Netanyahu et al., 2023). However, while an attractive paradigm, there are several challenges to using stochastic centering with GNNs and graph data. We discuss and remedy these below in Sec. 4.

## 4 Graph-\(\Delta\)UQ: Uncertainty-based Prediction Calibration

In this section, we introduce, G-\(\Delta\)UQ, a novel single-model UQ method that helps improve the performance of CIs without sacrificing computational efficiency or accuracy by extending the recently proposed stochastic centering paradigm to graph data. (See Fig. 2 for an overview.)

As discussed in Sec. 3, the stochastic centering and anchoring paradigm has demonstrated significant promise in computer vision, yet there are several challenges that must be addressed prior to applying it to GNNs and graph data. Notably, previous research on stochastic centering has focused on traditional vision models (CNNs, ResNets, ViT) and relied on straightforward input space transformations (e.g., subtraction and channel-wise concatenation: \([\mathbf{X}-\mathbf{C},\mathbf{C}]\)) to construct anchored representations. However, graph datasets are structured, discrete, and variable-sized, where such trivial transformations do not exist. Moreover, the distribution shifts encountered in graph datasets exhibit distinct characteristics compared to those typically examined in the vision literature that must be accounted for when sampling the underlying GNN hypothesis space. Therefore, it is non-trivial to design anchors that are capable of appropriately capturing epistemic uncertainty.

Below, we discuss not only how to extend stochastic centering to GNNs and structured data, (G-\(\Delta\)UQ), but also propose partially stochastic and pretrained variants that further improve the capabilities of anchored GNNs. In Section 5, we empirically demonstrate the advantages of our approach.

Figure 2: **Overview of G-\(\Delta\)UQ. We propose three different stochastic centering variants that induce varying levels of stochasticity in the underlying GNN. Notably, READOUT stochastic centering allows for using pretrained models with G-\(\Delta\)UQ.**
### Node Feature Anchoring

Recall that in \(\Delta\)-UQ, input samples are transformed into an anchored representation, by directly subtracting the input and anchor, and then concatenating them channel-wise, where the first DNN layer is correspondingly modified to accommodate the additional channels. While this is reasonable for vector-valued data or images, due to the variability graph size and discrete nature, performing a structural residual operation, \((\mathbf{A}-\mathbf{A}_{c},\mathbf{A}_{c})\) with respect to a graph sample, \(\mathcal{G}=(\mathbf{X},\mathbb{E},\mathbf{A},Y)\), and another anchor graph, \(\mathcal{G}_{c}=(\mathbf{X}_{c},\mathbb{E}_{c},\mathbf{A}_{c},Y_{c})\), would introduce artificial edge weights and connectivity artifacts that can harm convergence. Likewise, we cannot _directly_ anchor using the node features, \(\mathbf{X}\), since the underlying graphs are different sizes, and a set of node features cannot be considered IID. To this end, we first create a distribution over the training dataset node features and sample anchors from this distribution as follows.

We first fit a Gaussian distribution (\(\mathcal{N}(\mu,\sigma)\)) to the training node features. Then, during training, we randomly sample an anchor for each node. Mathematically, given the anchor \(\mathbf{C}^{N\times d}\sim\mathcal{N}(\mu,\sigma)\), we create the anchor/query node feature pair \([\mathbf{X}_{i}-\mathbf{C}||\mathbf{X}_{i}]\), where \(||\) denotes concatenation, and \(i\) is the node index. During inference, we sample a fixed set of \(K\) anchors and compute residuals for all nodes with respect to the same anchor, e.g., \(\mathbf{c}^{1\times d}{}_{k}\sim\mathcal{N}(\mu,\sigma)\) (\([\mathbf{X}_{i}-c_{k}||\mathbf{X}_{i}]\)), with appropriate broadcasting. For datasets with categorical node features, it is more beneficial to perform the anchoring operation after embedding the node features in a continuous space. Alternatively, considering the advantages of PEs in enhancing model expressivity (Wang et al., 2022), one can compute positional information for each node and perform anchoring based on these encodings. While performing anchoring with respect to the node features is perhaps the most direct extension of \(\Delta\)-UQ to graphs, as it results in a fully stochastically centered GNN, only using node features for anchoring neglects direct information about the underlying structure, which may lead to less diversity when sampling from the hypothesis space. Below, we introduce hidden layer variants that create partially stochastic GNNs that exploit message-passing to capture both feature and structural information during hypothesis sampling.

### Hidden Layer Anchoring

While performing anchoring in the input space creates a fully stochastic neural network as all parameters are learned using the randomized input, it was recently demonstrated with respect to Bayesian neural networks that relaxing the assumption of fully stochastic to partially stochastic neural networks not only leads to strong computational benefits, but also may improve calibration (Sharma et al., 2023). Motivated by this observation, we extend G-\(\Delta\)UQ to support anchoring in intermediate layers, in lieu of the input layer. This allows for _partially stochastic_ GNNs, wherein the layers prior to the anchoring step are deterministic. Moreover, intermediate layer anchoring has the additional benefit that anchors will be able to sample hypotheses that consider both topological and node feature information due to MPNN steps, and supports using pretrained GNNs. We introduce these variants below. (See Fig. 2 for a visual representation.)

_Intermediate MPNN Anchoring:_ Given a GNN containing \(\ell\) MPNN layers, let \(r\leq\ell\) be the layer at which we perform node feature anchoring. We obtain the anchor/sample pair by computing the intermediate node representations from the first \(r\) MPNN layers. We then randomly shuffle the node features over the entire _batch_, \((\mathbf{C}=\text{SHUFFLE}(\mathbf{X}_{i}^{r+1}))\), concatenate the residuals, and proceed with the READOUT and MLP layers as with the standard \(\Delta\)-UQ model. Note that, we do not consider the gradients of the query sample when updating the parameters, and the MPNN\({}^{r+1}\) layer is modified to accept inputs of dimension \(d_{r}\times 2\) (to take in anchored representations as inputs). Another difference from the input space implementation is that we fix the set of anchors and subtract a single anchor from all node representations in an iteration (instead of sampling uniquely), e.g., \(\mathbf{c}^{1\times d}=\mathbf{X}_{c}^{r+1}[n,:]\) and \([\mathbf{X}_{i,n}^{r+1}-\mathbf{c}||\mathbf{c}]\). This process is shown below, assume appropriate broadcasting:

\[\mathbf{X}^{r+1} =\textsc{MPNN}^{1\ldots r}\] \[\mathbf{X}^{r+1} =\textsc{MPNN}^{r+1\ldots\ell}\left([\mathbf{X}^{r+1}-\mathbf{C}, \mathbf{X}^{r+1}],\mathbf{A}\right)\] \[\hat{Y} =\textsc{MLP}(\textsc{READOUT}\left(\mathbf{X}^{\ell+1}\right))\]
_Intermediate Read Out Anchoring:_ While READOUT anchoring is conceptually similar to intermediate MPNN anchoring, we now only obtain a different anchor for each hidden graph representation, instead of individual nodes. This allows us to sample hypotheses after all node information has been aggregated over \(\ell\) hops. This is demonstrated below:

\[\mathbf{G} =\texttt{READOUT}\left(\mathbf{X}\right),\mathbf{G}_{c}=\texttt{ READOUT}\left(\mathbf{X}_{c}\right)\] \[\hat{Y} =\texttt{MLP}\left(\left[\mathbf{G}-\mathbf{G}_{c},\mathbf{G}_{ c}\right]\right)\]

_Pretrained Anchoring:_ Lastly, we note that in order to be compatible with the stochastic centering framework (the input layer or chosen intermediate layer), the network architecture must be modified and retrained from scratch. To circumvent this, we consider a variant of READOUT anchoring using a pretrained GNN backbone. Here, the final MLP layer of a pretrained model is discarded, and reinitialized to accommodate query/anchor pairs. We then freeze the MPNN, and only train the anchored classifier head. This allows for an inexpensive, limited stochasticity GNN.

While all G-\(\Delta\)UQ variants are able to sample from the underlying hypothesis space (see Fig. 10), each variant will provide somewhat different uncertainty estimates. Through our extensive evaluation, we show that the complexity of the task and the nature of the distribution shift will determine which of the variants is best suited and make some recommendations on which variants to use.

## 5 Uncertainty-based Prediction Calibration under Distribution Shift using G-\(\Delta\)Uq

In this section, we demonstrate the effectiveness of G-\(\Delta\)UQ in improving the reliability of CIs on various tasks (calibration, generalization gap prediction and OOD detection) as well as various distribution shifts (size, covariate and concept).

### Size Generalization

While GNNs are well-known to struggle when generalizing to larger size graphs (Buffelli et al., 2022; Yehudai et al., 2021; Chen et al., 2022), their predictive uncertainty behavior with respect to such shifts remains under studied. Given that such shifts can be expected at deployment, reliable uncertainty estimates under this setting are important for safety critical applications. We note that while sophisticated training strategies can be used to improve size generalization (Buffelli et al., 2022; Bevilacqua et al., 2021), our focus is primarily on the quality of uncertainity estimates, so we do not consider such techniques. However, we note that G-\(\Delta\)UQ can be used in conjunction with such techniques.

**Experimental Set-up.** Following the procedure of (Buffelli et al., 2022; Yehudai et al., 2021), we create a size distribution shift by taking the smallest 50%-quantile of graph size for the training set, and reserving the larger quantiles (>50%) for evaluation. Unless, otherwise noted, we report results on the largest 10% quantile to capture performance on the largest shift. We utilize this splitting procedure on four well-known benchmark binary graph classification datasets from the TUDataset repository (Morris et al., 2020): D&D, NCI1, NCI09, and PROTEINS. (See App. A.3 for dataset statistics.) We further consider three different backbone GNN models, GCN (Kipf & Welling, 2017), GIN (Xu et al., 2019), and PNA (Corso et al., 2020). All models contain three message passing layers and the same sized hidden representation. The accuracy and expected calibration error on the larger-graph test set are reported for models trained with and without stochastic anchoring.

Figure 3: **Impact of Layer Selection on G-\(\Delta\)UQ.** Performing anchoring at different layers leads the sampling of different hypothesis spaces. On D&D, we see that later layer anchoring corresponds to a better inductive bias and can lead to dramatically improved performance.
**Results.** As noted in Sec. 4, stochastic anchoring can be applied at different layers, leading to the sampling of different hypothesis spaces and inductive biases. In order to empirically understand this behavior, we compare the performance of stochastic centering when applied at different layers on the D&D dataset, which comprises the most severe size shift from training to test set (see Fig. 3). We observe that applying stochastic anchoring after the READOUT layer (L3) dramatically improves both accuracy and calibration as the depth increases. While this behavior is less pronounced on other datasets (see Fig. 11), we find overall that applying stochastic anchoring at the last layer yields competitive performance on size generalization benchmarks and better convergence compared to stochastic centering performed at earlier layers.

Indeed, in Fig. 4, we compare the performance of last-layer anchoring against a non-anchored model on four datasets. We observe that G-\(\Delta\)UQ improves calibration performance on most datasets, while generally maintaining or even improving the accuracy. Indeed, improvement is most pronounced on the largest shift (D&D), further emphasizing the benefits of stochastic centering.

### Evaluation under Concept and Covariate Shifts

Here, we seek to understand the behavior of GNN CIs under controlled covariate and concept shifts, as well demonstrate the benefits of G-\(\Delta\)UQ in providing reliable estimates under such shifts. Notably, we expand our evaluation beyond calibration error to include the safety-critical tasks of OOD detection (Hendrycks and Gimpel, 2017; Hendrycks et al., 2019) and generalization gap prediction tasks (Guillory et al., 2021; Ng et al., 2022; Trivedi et al., 2023; Garg et al., 2022). We begin by introducing our data and additional tasks, and then present our results.

**Experimental Set-up.** In brief, concept shift corresponds to a change in the conditional distribution of labels given input from the training to evaluation datasets, while covariate shift corresponds to

Figure 4: **Predictive Uncertainty under Size Distribution Shifts. When evaluating the accuracy and calibration error of models trained with and without stochastic anchoring on dataset with a graph size distribution shift, we observe that stochastic centering decreases calibration error while improving or maintaining accuracy across datasets and different GNNs.**

Figure 5: **Predictive Uncertainty under Concept and Covariate Shifts. Stochastic anchoring leads to competitive in-distribution and out-of distribution test accuracy while improving calibration, across domains and shifts. This is particularly true when comparing to other single-model UQ methods.**
change in the input distribution. We use the recently proposed Graph Out-Of Distribution (GOOD) benchmark (Gui et al., 2022) to obtain four different datasets (GOODCMNIST, GOODMotif-basis, GOODMotif-size, GOODSST2) with their corresponding in-/out- of distribution concept and covariate splits. To ensure fair comparison, we use the architectures and hyper-parameters suggested by the benchmark when training. Please see the supplementary for more details.

We consider the following baseline UQ methods in our analysis: Deep Ensembles (Lakshminarayanan et al., 2017), Monte Carlo Dropout (MCD) (Gal and Ghahramani, 2016), and our proposed G-\(\Delta\)UQ, including the pretrained variant. DeepEns is well known to be a highly performative baseline on uncertainty estimation tasks, but we emphasize that it requires training multiple models. This is in contrast to single model estimators, such as MCD and G-\(\Delta\)UQ. We note that while MCD and G-\(\Delta\)UQ can be applied at intermediate layers; we present results on the best performing layer but include the full results in the supplementary.

**Using Confidence Estimates in Safety Critical Tasks.** The safe deployment of graph machine learning models in critical applications requires that GNNs not only generalize to ID and OOD datasets, but that they do so safely. To this end, recent works (Hendrycks et al., 2022; Hentryedi et al., 2023; Trivedi et al., 2023) have expanded model evaluation to include additional robustness metrics to provide a holistic view of model performance. Notably, while reliable confidence indicators are critical to success on these metrics, the impact of distributions shift on GNN confidence estimates remains under-explored. We introduce these additional tasks below.

_Generalization Error Prediction:_ Accurate estimation of the expected generalization error on unlabeled datasets allows models with unacceptable performance to be pulled from production. To this end, generalization error predictors (GEPs) (Garg et al., 2022; Ng et al., 2022; Jiang et al., 2019; Trivedi et al., 2023; Guillory et al., 2021) which assign sample-level scores, \(S(x_{i})\) which are then aggregated into dataset-level error estimates, have become popular. We use maximum softmax probability and a simple thresholding mechanism as the GEP (since we are interested in understanding the behavior of confidence indicators), and report the error between the predicted and true target dataset accuracy:

\[GEPError:=||\text{Acc}_{target}-\frac{1}{|X|}\sum_{i}\mathbb{I}(\mathrm{S}( \tilde{x}_{i};\mathrm{F})>\tau)||\]

where \(\tau\) is tuned by minimizing GEP error on the validation dataset. We use the confidences obtained by the different baselines as sample-level scores, \(\mathrm{S}(x_{i})\) corresponding to the model's expectation that a sample is correct. The MAE between the estimated error and true error is reported on both in- and out-of -distribution test splits provided by the GOOD benchmark.

Figure 6: **Generalization Gap Prediction.The mean absolute error when using scores obtained from different baselines in the challenging (and to the best of our knowledge, yet unexplored for graphs) task of generalization error prediction are reported. While there is not a dominant method, stochastic anchoring is very competitive, and yields among the lowest MAE of single-model UQ estimators. Notably, pretrained G-\(\Delta\)UQ is particularly effective and outperforms the end-to-end variant.**
Out-of-Distribution DetectionBy reliably detecting OOD samples and abstaining from making predictions, models can avoid over extrapolating to distributions which are not relevant. While many scores have been proposed for detection (Hendrycks et al., 2019, 2022; Lee et al., 2018; Wang et al., 2022; Liu et al., 2020), flexible, popular baselines, such as maximum softmax probability and predictive entropy (Hendrycks and Gimpel, 2017), can be derived from confidence indicators relying upon prediction probabilities. Here, we report the AUROC for the binary classification task of detecting OOD samples using the maximum softmax probability (Kirchheim et al., 2022).

We briefly note that while more sophisticated scores can be used, our focus is on the reliability of GNN confidence indicators and thus we choose scores directly related to those estimates. Moreover, since sophisticated scores can often be derived from prediction probabilities, we expect their performance would also be improved with better estimates.

**Results.** We report the results in Figs. 5, 6, and 7. Our observations are below. Results are reported over three seeds.

_Accuracy & Calibration._ In Fig. 5, we observe that using stochastic anchoring via G-\(\Delta\)UQ yields competitive accuracy, especially in comparison to other single-model methods such as MCD, temperature scaling, or the base GNN model: in-distribution accuracy is higher on 6 out of 8 dataset/shift combinations, and out-of-distribution accuracy is higher on 5 out of 8 combinations. While Deep Ensembles is the most accurate method on a majority of datasets, they are known to be computationally expensive. Moreover, the simpler stochastic anchoring procedure generally comes close to the accuracy of Deep Ensembles, and in a few cases (covariate shift on GOODCMNIST and GOODMotif-size datasets), can noticeably outperform it. Stochastic anchoring also excels in improving calibration, improving in-distribution calibration compared to all baselines on 4 out of 8 combinations. _Most importantly, out-of-distribution calibration error is decreased by stochastic anchoring on 7 of 8 dataset/shift combinations compared to all other methods (single-model or ensemble)._

_Generalization Gap Prediction._ Next, we study all of our methods on the GOOD benchmarks for the task of generalization gap prediction, and report the results in Fig. 6. On this challenging task, there is no clear winner across all benchmarks. However, G-\(\Delta\)UQ variants are consistently competitive in MAE, and yield among the lowest MAE (across the board lower than other single-model UQ methods). In particular, _the pretrained G-\(\Delta\)UQ variant produces on average the lowest MAE for generalization gap estimation._

_OOD Detection._ Finally, we consider the task of detecting out-of-distribution samples. In Fig. 7, we see that the performance of stochastic anchoring methods under concept shift is generally very competitive with other UQ methods. For covariate shifts, except for the GOODMotif-basis dataset, stochastic anchoring produces high AUROC scores. In particular, on the GOODCMNIST-color, GOODSST2-length and GOODMotif-size benchmarks, the pretrained variant of G-\(\Delta\)UQ produces significantly improved AUROC scores. Finally, on GOODMotif-basis, however, both have lower

Figure 7: **OOD Detection.** The AUROC is reported for the task of detecting out-of-distribution samples. Under concept shift, the proposed G-\(\Delta\)UQ variants are very competitive with other baselines, including DeepEns. Under covariate shifts, except for GOODMotif-basis, pretrained G-\(\Delta\)UQ produces significant improvements over all baselines, including end-to-end G-\(\Delta\)UQ training.
AUROC than other baselines; we suspect the reason for this to be the inherent simplicity of this dataset and that G-\(\Delta\)UQ was prone to shortcuts.

Overall, we find that G-\(\Delta\)UQ performs competitively across several tasks and distributions shifts, validating our approach as an effective mechanism for producing reliable confidence indicators.

## 6 Conclusion

In this work, we take a closer look at confidence estimation under distribution shifts in the context of graph neural networks. We begin by demonstrating that techniques for improving GNN expressivity, such as transformer architectures and using positional encodings, do not necessarily improve the estimation performance on a simple structural distortion shift benchmark. To this end, we seek to improve the uncertainty estimation of GNNs by adapting the principle of stochastic anchoring for discrete, structured settings. We propose several G-\(\Delta\)UQ variants, and demonstrate the benefits of partial stochasticity when estimating uncertainty. Our evaluation is extensive, spanning multiple types of distribution shift (size, concept, covariate) while considering multiple safety critical tasks that require reliable estimates (calibration, generalization gap prediction, and OOD detection.) The proposed G-\(\Delta\)UQ improves estimation performance on a number of tasks, while remaining scalable. Overall, our paper rigorously studies uncertainty estimation for GNNs, identifies several shortcomings in existing approaches and proposes a flexible framework for reliable estimation. In future work, we will extend our framework to support link prediction and node classification tasks, as well as provide an automated mechanism for creating partially stochastic GNNs.

## 7 Acknowledgements

This work was performed under the auspices of the U.S. Department of Energy by the Lawrence Livermore National Laboratory under Contract No. DE-AC52-07NA27344, Lawrence Livermore National Security, LLC and is partially supported by the LLNL-LDRD Program under Project No. 2-ERD-006. This work is also partially supported by the National Science Foundation under CAREER Grant No. IIS 1845491, Army Young Investigator Award No. W9-11NF1810397, and Adobe, Amazon, Facebook, and Google faculty awards. Any opinions, findings, and conclusions or recommendations expressed here are those of the author(s) and do not reflect the views of funding parties. PT thanks Ekdeep Singh Lubana and Vivek Sivaraman for useful discussions during the course of this project.