Accurate and Scalable Estimation of Epistemic Uncertainty for Graph Neural Networks

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 & Gimpel, 2017), many strategies have been proposed to improve CI calibration (Lakshminarayanan et al., 2017; Guo et al., 2017; Gal & 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 & 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 neighbhorhood similarity).

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, $\bm{\mu}(y|\mathrm{x})$, and corresponding variance, $\bm{\sigma}(y|\mathrm{x})$ over $K$ random anchors are computed as:

Since the variance over $K$ anchors captures epistemic uncertainty by sampling different hypotheses, these estimates can be used to modulate the predictions: $\bm{\mu}_{\text{calib.}}=\bm{\mu}(1-\bm{\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.

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., 2022b), 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.

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 $\texttt{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:

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:

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.

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.

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.

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 & Gimpel, 2017; Hendrycks et al., 2019) and generalization gap prediction tasks  (Guillory et al., 2021; Ng et al., 2022; Trivedi et al., 2023a; 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 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 & 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., 2022b; 2021; Trivedi et al., 2023b) 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., 2023a; 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:

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}(\mathrm{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.

Out-of-Distribution Detection: By 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; 2022a; Lee et al., 2018; Wang et al., 2022a; Liu et al., 2020), flexible, popular baselines, such as maximum softmax probability and predictive entropy (Hendrycks & 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 varaints 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 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.