Accurate and Scalable Estimation of Epistemic Uncertainty for Graph Neural Networks
Authors: Puja Trivedi
Mark Heimann
Rushil Anirudh
Danai Koutra
Jayaraman J. Thiagarajan
Venue: ICLR 2024 poster (2024)
Keywords: Graph Neural Networks
Uncertainity Quantification
Calibration
Primary area: general machine learning (i.e., none of the above)
TL;DR: We propose a new method for performing uncertainity quantification for GNNs.

Abstract: While graph neural networks (GNNs) are widely used for node and graph representation learning tasks, the reliability of GNN uncertainty estimates under distribution shifts remains relatively under-explored. Indeed, while post-hoc calibration strategies can be used to improve in-distribution calibration, they need not also improve calibration under distribution shift. However, techniques which produce GNNs with better intrinsic uncertainty estimates are particularly valuable, as they can always be combined with post-hoc strategies later. Therefore, in this work, we propose G-$\Delta$UQ, a novel training framework designed to improve intrinsic GNN uncertainty estimates. Our framework adapts the principle of stochastic data centering to graph data through novel graph anchoring strategies, and is able to support partially stochastic GNNs. While, the prevalent wisdom is that fully stochastic networks are necessary to obtain reliable estimates, we find that the functional diversity induced by our anchoring strategies when sampling hypotheses renders this unnecessary and allows us to support G-$\Delta$UQ on pretrained models. Indeed, through extensive evaluation under covariate, concept and graph size shifts, we show that G-$\Delta$UQ leads to better calibrated GNNs for node and graph classification. Further, it also improves performance on the uncertainty-based tasks of out-of-distribution detection and generalization gap estimation. Overall, our work provides insights into uncertainty estimation for GNNs, and demonstrates the utility of G-$\Delta$UQ in obtaining reliable estimates.

Published as a conference paper at ICLR 2024
ACCURATE AND SCALABLE ESTIMATION OF EPISTEMIC UNCERTAINTY FOR GRAPH NEURAL NETWORKS
Puja Trivedi', Mark Heimann', Rushil Anirudh 23 Danai Koutra', Jayaraman J. Thiagarajan"
University of Michigan, CSE Department, PLawrence Livermore National Laboratory, *Amazon
ABSTRACT
While graph neural networks (GNNs) are widely used for node and graph represen- tation learning tasks, the reliability of GNN uncertainty estimates under distribution shifts remains relatively under-explored. Indeed, while post-hoc calibration strate- gies can be used to improve in-distribution calibration, they need not also improve calibration under distribution shift. However, techniques which produce GNNs with better intrinsic uncertainty estimates are particularly valuable, as they can always be combined with post-hoc strategies later. Therefore, in this work, we propose G-AUQ, a novel training framework designed to improve intrinsic GNN uncertainty estimates. Our framework adapts the principle of stochastic data center- ing to graph data through novel graph anchoring strategies, and is able to support partially stochastic GNNs. While, the prevalent wisdom is that fully stochastic networks are necessary to obtain reliable estimates, find that the functional diversity induced by our anchoring strategies when sampling hypotheses renders this unnecessary and allows to support G-AUQ on pretrained models. Indeed, through extensive evaluation under covariate, concept and graph size shifts, we show that G-AUQ leads better calibrated GNNs for node and graph classifi- cation. Further, also improves performance on the uncertainty-be out-of-distributtinn detection and generalization gap estimation. Overall, our work provides insights into uncertainty estimation for GNNs, and demonstrates the utility of G-AUQ in obtaining reliable estimates.
tasks of
INTRODUCTION
As graph neural networks (GNNs) are increasingly deployed in critical applications with test-time distribution shifts (Zhang& & Chen, 2018; Gaudelet et al., 2020; Yang et al., 2018; Yan et al., 2019; Zhu al., 2022), it becomes necessary expand model evaluation include safety-centric metrics, such as calibration errors (Guo et al., 2017) 7), out-of-distribution (OOD) rejection rates (Hendrycks & Gimpel, 2017), and generalization error predictions (GEP) (Jiang al., 2019), holistically understand model performance in such shifted regimes (Hendrycks et al., 2022b; Trivedi et al., 2023b). Notably, improving on these additional metrics often requires reliable uncertainty estimates, such as maximum softmax or predictive entropy, which can be derived from prediction probabilities. Although there is clear understanding in the computer vision literature that the quality uncertainty estimates can noticeably deteriorate under distribution shifts (Wiles et al., 2022; Ovadia et al., 2019), the impact of such shifts on graph neural networks (GNNs) remains relatively under-explored.
Post-hoc calibration methods (Guo et al., 2017; Gupta et al., 2021; Kull et al., 2019; Zhang et al., 2020), which use validation datasets to rescale logits to obtain better calibrated models, are an effective, accuracy-preservie strategy for improving uncertainty estimates and model trust-worthiness. Indeed, several post-hoe calibration strategies (Hsu et al., 2022; Wang et al., 2021) have been recently proposed to explicitly account for the non-IID nature of code-classification datasets. However, while these methods are effective at improving uncertainty estimate reliability on in-distribution (ID) data, they have not been evaluated OOD data, where they may become unreliable. To this end, training strategies which produce models with better intrinsic uncertainty estimates are valuable as they will provide better out-of-tha box ID and OOD estimates, which can then be further combined with post-hoc calibration strategies if desired.
"Correspondence pujjate umi ch. edu.

---

Published as a conference paper at ICLR 2024
The A-UQ training framework (Thiagarajan et al., 2022) was recently proposed scalable, single model alternative for vision models ensembles and has achieved the--o--thh-rr performance on calibration and OOD detection tasks. Central to A-UQ's success is the concept of anchored training, where models are trained on stochastic, relative representations of input samples in order to simulate sampling from different functional modes test time (Sec. 2.) While, on the surface, A-UQ also appears as a potentially attractive framework for obtaining reliable, intrinsic uncertainty estimates on graph-based tasks, there are several challenges that arise from the structured, discrete, and variable- sized nature of graph data that must be resolved first. Namely, the anchoring procedure used by A-UQ is not applicable for graph datasets, and it is unclear how design alternative anchoring strategies such that sufficiently diverse functional modes are sampled at inference to provide reliable epistemic uncertainty estimates.
Proposed Work. Thus, our work proposes G-AUQ, novel training paradigm which provides better intrinsic uncertainty estimates for both graph and node classification tasks through the use of newly introduced graph-specific, anchoring strategies. Our contributions can be summarized as follows:
(Partially) Stochastic Anchoring for GNNs. We propose G-AUQ, a novel training paradigm that improves the reliability of uncertainty estimates GNN-based tasks. Our novel graph-anchoring strategies support partial stochasticity GNNs as well as training with pretrained models. (Sec. 3).
CIs under Distribution Shifts. Across covariate, concept and graph-size shifts, we demonstrate that G-AUQ leads to better calibration. Moreover, G-AUQ's performance is further improved when combined with post-hoc calibration strategies on several node and graph-level tasks, including new safety-critical tasks (Sec. 5).
Evaluating Uncertainty-Mooduatee
GrainGr Analysis of G-AUQ. We study the calibration of architectures of varying expressivity and G-AUQ 's ability to improve them under varying distribution shift. We further demonstrate its utility as lightweight strategy for improving the calibration of pretrained GNNs (Sec. 6).
2 RELATED WORK & BACKGROUND
While uncertainty estimates are useful for a variety of safety-critical tasks (Hendrycks & Gimpel, 2017; Jiang et al., 2019; Guo et al., 2017), DNNs are well-known to provide poor uncertainty estimates directly out of the box (Guo et al., 2017). To this end, there has been considerable interest in building calibrated models, where the confidence of a prediction matches the probability of the prediction being correct. Notably, since GEP and OOD detection methods often rely upon transformations of a model's logits, improving calibration can in turn improve performance on these tasks as well. Due to their
properties, post-hoc calibration strategies, which rescale confidences after training using a validation dataset, are particularly popular. Indeed, several methods (Guo et al., 2017; Gupta et al., 2021; Kull et al., 2019; Zhang al., 2020) have been proposed for DNNs in general and, more recently, dedicated node-classifier calibration methods (Hsu et al., 2022; Wang et al., 2021) have also been proposed to accommodate the non-IID nature of graph data. (See App. A.9 for more details.) Notably, however, such post-hoe methods do not lead to reliable estimates under distribution shifts, enforcing calibration ID validation data does not directly lead reliable estimates on OOD data (Ovadia et al., 2019; Wiles et al.. 2022; Hendrycks et al., 2019).
Alternatively, Bayesian methods have been proposed for DNNs (Hernández-Lobato & Adams, 2015; Blundell et al., 2015), and more recently GNNs (Zhang et al., 2019; Hasanzadeh et al., 2020), as inherently "uncertainty-aware"" strategies. However, not only do such methods often lead to performance loss, require complicated architectures and additional training time, they often struggle to outperform the simple Deep Ensembles (DEns) baseline (Lakshminarayanan et al., 2017). By training a collection of independent models, DEns is able to sample different functional modes of the hypothesis space, and thus, capture epistemic variability to perform uncertainty quantification (Wilson & Izmailov. 2020). Given that DEns requires training and storing multiple models, the SoTA A- UQ framework (Thiagarajan et al., 2022) was recently proposed to sample different functional modes using only single model, based on the principle of anchoring.
Background on Anchoring. Conceptually, anchoring is the process of creating a relative representa- tion for an input sample in terms of a random "anchor." By randomizing anchors throughout training (c.g., stochastically centering samples with respect to different anchors), A-UQ emulates the process of sampling and learning different solutions from the hypothesis space.
2

---

Published as a conference paper at ICLR 2024
In detail, let Derain be the training distribution, Diest be the testing distribution, and D Danchor Dirain be the anchoring distribution. Existing research on stochastic centering has focused on vision models (CNNs, ResNets, ViT) and used imput space transformations to construct anchored representations. Specifically, given an image sample with corresponding label, (I, y), and anchor C E D Duccorr anchored samples were created by subtracting and then channel-wise concatenat- ing two images: -I- C||C] Given the anchored representation, a corresponding stochastically centered model can be defined as fo- [I- C|C] + D. and can be trained as shown in Fig. 1. At inference. similar to ensembles, predictions and uncertainties are aggregated over differ- ent hypotheses. Namely, given K random anchors, the mean target class prediction, μ [T), and the corresponding variance, o [y/ I) are computed as: p(y(T) = R E fo C22 C+)) and
(vll) -00-VK= VKL -- [x=r(fo((I Ck2 Cx]) p2).. Since the variance over K anchors captures epis- temic uncertainty by sampling different hypotheses, these estimates can be used to modulate the predictions: Mealib. =u(1- ).. Notably, the rescaled logits and uncertainty estimates have led to state-of-the-art performance on image outlier rejection, calibration, and extrapolation (Anirudh & Thiagarajan, 2022; Netanyahu al., 2023).
3 GRAPH-AUQ: UNCERTAINTY-AWAN PREDICTIONS
Given A-UQ's success in improving calibration and generalization (Netanyahu et al., 2023) under distribution shifts on computer vision tasks and the limitations of existing post-hoc strategies, stochastic centering appears as a potentially attractive framework for obtaining reliable uncertainty estimates when performing GNN-based classi- fication tasks. However, there are several challenges that must be addressed before to applying it to graph data. Namely, while input space transformations, which induce fully stochastic models, were sufficient for sampling diverse functional hypotheses from vision mod- els, it is (i) non-trivial to define such transformations when working with variable sized, discrete graph data and (ii) unclear whether full stochasticity is in fact needed when working with message passing models. Below, we explore these issues through novel graph anchor- ing strategies. However, we begin with a conceptual discussion of the role of anchoring strategies in generating reliable uncertainty estimates.
#training for (I,y) in trainloader:
C C= create =ccrste__nchhrre anchors (n:Batchit:) In CONCAT( (I-n,C),din-1)
preds nodel[I_Anc]
loss = criterion(preds,y)
#inference
anc_preds =
testAnc create_anchors(n=k]
testAncs:
Lanc = (I-,
preds rodel(I_anc)
anc_ ptt.prdddd...edpprr...
UiicATlan__predsddie
nu MEAN ddte() var STDOEV(P, dinn=))
Figure 1: Training/Infer- ence with Anchoring.
What are the goals of anchoring?: As discussed in Sec. 2, epistemic uncertainty can be estimated by aggregating the variability over different functional hypotheses (Hüllermeier & Wacgeman, 2021). Indeed, the prevalent wisdom behind the success of DeepEns is its ability to sample diverse f'unctional hypotheses. Since these hypotheses are more likely to differ on OOD inputs, aggregating them can lead to better generalization and uncertainity estimates. Insofar as stochastic centering seeks to simulate an ensemble through single model, a key goal of the anchoring distribution/strategy is then to ensure that sampled hypotheses are also diverse. Thiagarajan et al. (2022) obtained sufficient diversity by using input space anchoring to sample a fully stochastic network. However, in the context of Bayesian neural networks (BNNs), it was recently shown that partial stochasticity can perform equally well with respect to fully stochastic BNNs at significantly less cost (Sharma et al., 2023). This suggests that in addition to the "amount" of diversity, the "effective" or functional diversity is also important for performance. However, in practice, it is difficult to control this balance, so existing methods default to heuristics that only promote diverse hypotheses. For example, DeepEns uses different random seeds or shuffles the batch order when creating ensemble members, and AUUQ relies upon fully stochastic models. To this end, we propose three different anchoring strategies that only handle the difficulties of working with graph data and GNNs, but also induce different scales of the aforementioned balance. At a high-level, our strategies trade-off the amount of stochasticity (i.e., amount of diversity) and the semantic expressivity of the anchoring distribution to accomplish this.
Notations. Let G = (x", Y, be graph with node features X E R Nxd, adjacency matrix A E [0, N and labels Y, where N, d, q denote the number of nodes, feature dimension and number of classes, respectively. When performing graph classification, Y E {0, 139; for node classification, let Y E (0. 1) Nx,, We define a graph classification GNN consisting of é message
For example, channel wise concatenating two RGB images creates channel sample.
3

---

Published as a conference paper at ICLR 2024
passing layers (MPNN), a graph-level readout function (READOUT), and classifier head (MLP) as follows: X MPNN (X', A). G READOUT (x++1, ) and Ý MLP (G) where X NNdd is the intermediate node representation layer l 1, G E R gxdeia is the graph representation, and y (0. 139 the predicted label. When performing node classification, we do not include the READOUT layer, and instead output node-level predictions: y = MLP (X++1). We use subscript to indicate indexing and || to indicate concatenation.
3.1 NODE FEATURE ANCHORING
We begin by introducing& graph anchoring strategy for inducing fully stochastic GNNs. Due size variability and discreteness, performing a structural residual operation by subtracting two adjacency matrices would be ineffective at inducing an anchored GNN. Indeed, such a transform would introduce artificial edge weights and connectivity artifacts. Likewise, when performing graph classification, we cannot directly anchor over node features, since graphs are different sizes. Taking arbitrary subsets of node features is also inadvisable as node features cannot be considered IID. Further, due to iterative message passing, the network may not be able to converge after aggregating hops of stochastic node representations (see A. for details). Furthermore, there is risk of exploding stochasticity when anchoring MPNNs. Namely, after / rounds of message passing, a node's 's representations will have aggregated information from its I hop neighborhood. However, since anchors unique to individual nodes, these representations are not only stochastic due to their own anchors but also those of their neighbors.
To address both these challenges, we instead fit a d-dimensional Gaussian distribution over the training dataset's input node features which is then used as the anchoring distribution (see Fig. 2). While simple solution, the fitted distribution allows to easily sample anchors for arbitrarily sized graphs, and helps manage stochasticity by reducing the complexity of the anchoring distribution, ensuring that overall stochasticity is manageable, even after aggregating the h-hop neighbhorhood. (See A. for details.) We emphasize that this distribution is only used for anchoring and does not assume that the dataset's s node features are normally distributed. During training, we randomly sample unique anchor for each node. Mathematically, given anchors C.N xd ~ N(H,7), we create the anchored node features as: [X C| During inference, we sample fixed set of K anchors and compute residuals for all nodes with respect to the same anchor after performing appropriate broadcasting, e.g., clxd ~ "Nu, r), where C == REPEAT( N) and [X - Cx " the kth anchored sample. For datasets with categorical node features, anchoring can be performed after embedding the node features into a continuous space. If node features are not available, anchoring can still be performed via positional encodings (Wang et al., 2022b), which are known to improve the expressivity and performance of GNNs (Dwivedi et al., 2022a). Lastly, note that performing node feature anchoring (NFA) is the most analogous extension of A-UQ to graphs as it results in fully stochastic GNNs. This is particularly true on node classification tasks where each node can be viewed an individual sample, similar to image sample original A-UQ formulation.
Mode Classificatiin
Miraon Classifiiation
nitoral Nun Hodes, Inpondia)
OSCATI loox for (ettt))
#Creste Anchoring Dist nu K.stiidiam)) trpurdin))
finundia dieencional. Goucstar
#Creste TRADINN Anchers Lnoutiil
C.edde AndDist sampla
awdesigh inputiial
C_Grash wwdllst..ammllll.....o (8))
C_Hade Andiiet. www.paalll..aa........ (innk))
C.Grach AndDlat, hwwpet/t.....ttg...... (2000)
Mlschored Repressentatt Lono
Anddraah OINATI D C.rape),,ane
Figure 2: Node Feature An- choring Pseudocode.
3.2 HIDDEN LAYER ANCHORING FOR GRAPH CLASSIFICATION
While NFA can conceptually be used for graph classification tasks, there are several nuances that may limit its effectiveness. Notably, since each sample (and label) is at a graph-level, NFA not only effectively induces multiple anchors per sample, also ignores structural information that may be useful in sampling more functionally diverse hypotheses, e.g., hypotheses which capture functional modes that rely upon different high-level semantic, non-linear features. To improve the quality of hypothesis sampling,
introduce hidden layer anchoring below, which incorporates structural information into anchors at the expense of full stochasticity in the network (See Fig. 1):
Hidden Layer and Readout Anchoring: Given GNN containing e MPNN layers, let S E be the layer at which we perform anchoring. Then, given the intermediate node representations X -P (Xx-2 A), we randomly shuffle the node features over the entire batch, (C
4

---

Published as a conference paper at ICLR 2024
SHUFFLE(X" dim 0)), concatenate the residuals - and proceed with the READOUT and MLP layers as usual. (See A.1 for corresponding pseudocode.) Note the gradients of the query sample are not considered when updating parameters, and MPNN" is modified to accept inputs of dimension d, x 2 (to take in anchored representations as inputs). At inference, we subtract single anchor from all node representations using broadcasting. Hidden layer anchoring induces the following GNN: X MPNN (x--2, A), X = MPNN" ([X--1- - CI and X MP.N*+ (X", A), and Ý (READOOT (X(+)))
Not only do hidden layer anchors aggregate structural information over r hops, they induce a GNN that now partially stochastic, as layers ....r are deterministic. Indeed, by reducing network stochasticity, it is naturally expected that hidden layer anchoring will reduce the diversity of the hypotheses, but by sampling more functionally diverse hypotheses through deeper, semantically expressive anchors, it is possible that naively maximizing diversity is in fact not required for reliable uncertainty estimation. To validate this hypothesis, we thus propose the final variant, READOUT anchoring for graph classification tasks. While conceptually similar to hidden layer anchoring, here, we simultancously minimize GNN stochasticity (only the classifier is stochastic) and maximize anchor expressivity (anchors are graph representations pooled after l' rounds of message passing), Notably, READOUT anchoring is also compatible with pretrained GNN backbones, as the final MLP layer of pretrained model is discarded (if necessary), and reinitialized to accommodate query/anchor pairs. Given the frozen MPNN NN backbone, only the anchored classifier head is trained.
In Sec. 5, we empirically verify the effectiveness of our proposed G-AUQ variants and demon- strate that fully stochastic GNNs are, in fact, unnecessary to obtain highly generalizable solutions, meaningful uncertainties and improved calibration on graph classification tasks.
4 NODE CLASSIFICATION EXPERIMENTS: G-AUQ IMPROVES CALIBRATION
In this section, we demonstrate that G-AUQ improves uncertainty estimation in GNNs, particularly when evaluating node classifiers under distribution shifts. To the best of our knowledge, GNN calibration has not been extensively evaluated under this challenging setting, where uncertainty estimates are known to be unreliable (Ovadia et al., 2019). We demonstrate that G-AUQ not only directly provides better estimates, but also that combining G-UUQ with existing post-hoc calibration methods further improves performance.
Experimental Setup. We the concept and covariate shifts for WebKB, Cora and CBAS datasets provided by Gui et al. (2022), and follow the recommended hyperparameters for training. In our implementation of node feature anchoring, we use 10 random anchors obtain predictions with G-AUQ. All our results are averaged over 5 seeds and post-hoc calibration methods (described further in App. A.9) are fitted on the -idistriiuttinn validation dataset. The expected calibration error and accuracy on the unobserved "OOD test" split are reported.
Results. From Table (and expanded in Table. 12), we observe that across 4 datasets and shifts that G-AUQ, without any post-hoc calibration (X), is superior to the vanilla model on nearly every benchmark for better or same accuracy (8/8 benchmarks) and better calibration error (7/8), often with a significant gain in calibration performance. Moreover, we note that combining G- AUQ with a particular posthoc calibration method improves performance relative to using the same posthoc method with vanilla model. Indeed, on WebKB, across 9 posthoc strategies, "G- AUQ +<calibration method>" improves or maintains the calibration performance of the corresponding "no G- G-AUQ +<calibration method>" in 7/9 (concept) and 6/9 (covariate) cases. (See App. A.8 for more discussion.) Overall, across post hoc methods and evaluation sets, G-UUQ variants are very performant achieving (best accuracy: 8/8), best calibration (6/8) or second best calibration (2/8).
5 GRAPH CLASSIFICATION UNCERTAINTY EXPERIMENTS WITH G-AUQ
While applying G-AUQ to node classification tasks was relatively straightforward, performing stochastic centering with graph classification tasks is more nuanced. As discussed in Sec. 3. different anchoring strategies can introduce varying levels of stochasticity, and it is unknown how these strategies affect uncertainty estimate reliability. Therefore, we begin by demonstrating that fully stochastic GNNs are not necessary for producing reliable estimates (Sec. 5.1). We then extensively evaluate the calibration of partially stochastic GNNs on covariate and concept shifts with and without post-hoc calibration strategies (Sec. 5.2), as well as for different UQ tasks (Sec. 5.3).
5

---

Published as a conference paper at ICLR 2024
NCI1 GIN
NC1109
PROTEINS
GCN
W 02
0.0
NA
Layer
Figure 3: Effect of Anchoring Layer. Anchoring at different layers (L1, L2, L3) induces different hy- potheses spaces. Variations of stochastic anchoring outperform models without it, and the lightweight READOUT anchoring in particular generally performs well across datasets and architectures.
Lastly, we demonstrate that G-AUQ's uncertainty estimates remain reliable when used with different architectures and pretrained backbones (Sec. 6).
5.1 Is FULL STOCHASTICITY NECESSARY FOR G-AUQ?
By changing the anchoring strategy and intermediate anchoring layer, we can induce varying levels of stochasticity in the resulting GNNs. As discussed in Sec. 3, hypothesize that the decreased stochasticity incurred by performing anchoring at deeper network layers will lead to more functionally diverse hypotheses, and consequently more reliable uncertainty estimates. We verify this hypothesis here, by studying the effect of anchoring layer on calibration under graph-size distribution shift. Namely, we find that READOUT anchoring sufficiently balances stochasticity and functional diversity.
Experimental Setup. We study the effect of different anchoring strategies on graph classification calibration under graph-size size shift. 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 evaluate on the largest 10% quantile. Following (Buffelli et al., 2022), we apply this splitting procedure to NCI1, NC109, and PROTEINS (Morris et al., 2020), consider 3 GNN backbones (GCN (Kipf& Welling, 2017), GIN (Xu et al., 2019), and PNA (Corso et al., 2020)) and use the same architectures/prraned expected calibration error over 10 seeds on the grrpph test set are reported for models trained with and without stochastic anchoring.
(See Appendix A.6 for dataset statistics.) The accuracy and
Results. We compare the performance of anchoring at different layers in Fig. 3. While there is no find there
clear winner across datasets and architectures for which layer to perform anchoring, consistent trend across all datasets and architectures the best accuracy and ECE is obtained by G-AUQ variant. Overall, our results clearly indicate that partial stochasticity can yield substantial benefits when estimating uncertainty (though suboptimal layers selections are generally not too harmful). Insofar, we are the first to focus on partially stochastic anchored GNNs, automatically selecting the anchoring layer is an interesting direction of future work. However, in subsequent experiments, we use READOUT anchoring, unless otherwise noted, as is faster train (see App. A.13), and allow our methods to support pretrained models. Indeed, READOUT anchoring (L3) yields top performance for some datasets and architectures such as PNA on PROTEINS, compared earlier (L.1, L2) and, as we discuss below, is very performative on a variety of tasks and shifts.
5.2
CALIBRATION UNDER CONCEPT AND COVARIATE SHIFTS
Next, we assess the ability of G-AUQ to produce well-calibrated models under covariate and concept shift in graph classification tasks. We find that G-AUQ not only provides better calibration out of the box, its performance is further improved when combined with post-hoe calibration techniques.
Experimental Setup. We use three different datasets (GOODCMNIST GOODMotif-basiss. GOODSST2) with their corresponding splits and shifts from the recently proposed Graph Out- Of Distribution (GOOD) benchmark (Gui et al., 2022). The architectures and hyperparameters suggested by the benchmark are used for training. G-AU uses READOUT anchoring and 10 random anchors (see App. A.7 for more details). We report accuracy and expected calibration error for the OOD test dataset, taken over three seeds.
Results. As shown in Table 1, we observe that G- leads to inherently better calibrated models, as the ECE from G-AUQ without additional post-hoc calibration (X) is better than the vanilla ("No G-AUQ") counterparts on 5/6 datasets. Moreover, find that combining G-AUQ with

---

Published as a conference paper at ICLR 2024
Table I: Calibration under Covariate and Concept shifts. G-AUQ leads to better calibrated models for (GOOD(a)) and graph-level prediction tasks under different kinds of distribution shifts. Notably, G-AUQ can be combined with post-hoe calibration techniques to further improve calibration. The expected calibration error (ECE) is reported. Best, Second.
Shet: Concept
Shift Conatiate
Accuncy
BCE
(C)
Dutasct
Domain
Cullbeation
No UV
G-A UQ
G. UQ
GO
U-
UQ
G-G UQ
UQ
x
0.581:00:00
0.360-5008
0:307:40.00
011340011
0.47:6.000
0.51540011
0.308-0.009
0.141::0005
CAGCN
0.007e0.00
0.130:50003
0.12
0:4=00002
0:236200:8
0,231=0005
Drichlor
0.534:00007
0.501400004
0.32:0.004
0.10
0.33660.007
0.48:661
0.161:0.000
0.356:000 01
ETS
0.58120000
0.200550004
0,30120007
0.116
04770.802
023
0:31:0007
0,14
GOODC
Degne
GATS
0.581211443
0.18521001
0.47-6/002
0:211:00004
IRM
0.582:10022
0.507=0.000
0.120:5000
0.102200000
0.400
0.52243104
0.10
0.10
Onkeinoniin
0.581100443
0.502 6,000
0 02226.10024
0
0.47_.0002
0.318_0.002
Spline
0.57):10141
0.806
0.000
0.400s0.005
0.5260004
0.158a001
0.098
ViS
0.58 10001
(.5W) du00
0.9161000
0. 1000
0.67 6.00
1.592
0.345 0.006
0.497 -cace
0.391 0063910075
0:88810000
0.36511111
0.361_.0.11
Denchlir
1s:::1 000
0:510:5:006
0.305:0012
0.30
0:30000003
0.335
0.342:0.001
0. 30G-0 076
ETS
GOODCMNIST
Culor
IKM
0.000:50000
0:500:0.010
0:255:00011
200
0308:00010
0 0
0.416:0.064
D. 02555 093
Onkeinowient
1000
0.500 6.008
0557910099
0.388 1002
031810006
19
0.415 0.007
0.5 101
Spline
0.805:00000
0.857=6:000
0.21:00007
0000
0
0.414:0.05
0.455
Vi
0:500-4.122 00
0:330:000
0
0309:0007
0.33611007
0.319-0.0m
O. 108000000 050
Ersembiing
0,000e00001
0.509=0.004
0.
034
0:39760.006
0.43:-0.077
0.905:00001
0.966560008
0.00Ge00144
0.07
0.01=00.00
0 002
0.32
Diichliee
0.925:40:011
09220-60110
0.161-0013
0. 0:00000000
0886600.000
00081:0000
0.37
3 017
ETS
0,525:20.00
0.927
0
0.691:0011
0
0.314:0.041
GOODMont
Basis
Ordesioooriazt
IRM
0.525:10044
0.025
0.928
0,00-0000
0.067
0
6.007
0.611000011
0.0110001
0.008100
0.00
0.121:0.000
0.366.02011
3.305
0.
Syline
0.927
00001100000
0.S1S
VS
0,925:10000
0.027=0.012
0.
0
0.30
Ersemblin
0.582-10002
0.915 (0.006
0.0861005
0.714:0012
100091000
0.298 0.383
0.521 195
X
Déichles
0.65022002
0.683
0.16
0.138
0.798
D8820022
0.162000
ETS
0.853
0.31
0
0.991100005
081001
0:267
GOORSST2
Length
IBM
0.685a1400
0.680..000
0,200100012
0.79612000
0,80110011
0.2611106
2.26
Ondeinooienn
0-633-4.001
0.275:10000
0.321
0.79146.005
08:-2011
0:200-00000
0.23:10 007
Spline
0.654:3.02
0.68155.000
0:2350005
0.71
0
0:200
0.2620000122
VS
0:680:10.019
0.00340
0. 0:311:0011
0.374=00000
0777700006
018:3001
0.307:61116
0.31--311
Essersbling
0,705ein002
0.700
0,276620m6
0.24
0.88850.000
0.8421000
0.154:0.322
01112500119
particular post-hoc calibration methods further elevates its performance relative to combining the same strategy with vanilla models. Indeed, for a fixed post-hoc calibration strategy, G-AUQ improves the calibration, while maintaining comparable if not better accuracy on the vast majority of the methods and datasets. There are some settings where combining G-AUQ or the vanilla model with post-hoc method leads decreases performance (for example, GOODSST2, covariate, ETS, calibration) but we emphasize that this is not short-coming of G-AUQ. Posthoe strategies, which rely upon ID calibration datasets, may not be effective on shifted data. This further emphasizes the importance of our OOD evaluation and G-AUQ as intrinsic method for improving uncertainty estimation.
5.3 USING CONFIDENCE ESTIMATES IN SAFETY-CRITICAL TASKS
While post-hoc calibration strategies rely upon an additional calibration dataset to provide meaningful uncertainty estimates, such calibration datasets are not always available and may not necessarily improve OOD performance (Ovadia et al., 2019). Thus, we also evaluate the quality of the uncertainty estimates directly provided by G-AUQ on two additional UQ-based, safety-critical tasks (Hendrycks et al., 2022b; 2021: Trivedi et al., 2023b): (i) OOD detection (Hendrycks et al., 2019), which attempts to classify samples as in- or nut-of-distribution and (ii) generalization error prediction (GEP) (Jiang et al., 2019), which attempts to predict the generalization on unlabeled test datasets (to the best of our knowledge, are the first to study GEP of graph classifiers). In the interest of space, we present the results on GEP in the appendix.
OOD Detection Experimental Setup. By reliably detecting OOD samples and abstaining from making predictions on them, models can avoid over-extrapolating to irrelevant distributions. While many scores have been proposed for detection (Hendrycks et al., 2019; 2022a; Lee al., 2018; Wang et al., 2022a: Liu et al., 2020), popular scores, such as maximum softmax probability and predictive entropy (Hendrycks & Gimpel, 2017), are derived from uncertainty estimates. Here, we report the
7

---

Published as a conference paper at ICLR 2024
Table 2: GOOD-Datasets, OOD Detection Performance. The AUROC of the binary classification task of classifying OOD samples is reported. G-AUQ variants outperform the vanilla models on 6/8 datasets. [We further note that end-to-end G-AUQ does in fact lose performance relative to the vanilla model on 4 datasets. Investigating why pretrained G-AUQ is able to increase performance on those datasets is an interesting direction of future work. does not appear that a particular shift is more difficult for this task: concept shift is easier for GOODCMNIST and COODMotif(Basis) while covariate shift is easier for 2OODMotif(Size) and GOODSST2. Combining G-AU with more sophisticated, uncertainty confidence based OOD scores may further improve performance.]
CMNIST (Coloe)
MaciLPP (llacis)
ModifL (Siae)
S812
Method
Conceput))
(ovariate(
ConcepaiT)
Covariate
Concepil
(ovariate(t)
Canceplit)
(ovariate(
Varilla
0. 0.006
0.405 + 0.002
0.730 0.021
0. 1666
0 001
0. 0.050 + 0.003
0.755 + 0.074
0 0,350
0.014
0.545 = 0.0%0
AUQ
0.771 0,002
0,470 1 0.043
0.758 0.006
0.328
0.022
0.005
0.691 1.0067
0,338
0,023
0,351 0.042
Pretr. AUQ
0.774 0.016
0.543 1 0.152
0.709 0.029
0.272
0.025
686 004
829 0.113
0.324
9.055
0.446 0.049
Table 3: RoLMNIST-Calliraate
Here, we report expanded results (calibration) on the Rotated MNIST dataset, including a variant that combines G-AUQ with Deep Ens. Notably, we see that anchored ensembles outperform basic ensembles in both accuracy and calibration.
Architecture
LPE?
G-AUQ
Calibration
Avg.BCE ())
BCE (10( (1)
(CE((15) (1)
ECE (25) (1)
(38) (L)
(40) (L)
X
0.039 10.001
0.068 10.340
0.195
0.245 10.011
GatedGiCN
X
X
0.018 +0.005
0,023 +0.03
0.033 +60
0.000 +0.033
0.117 +0.048
0.162 +0.067
X
basembbling
0.026 10.000
0.088 +0.001
0.042 10.001
0.084 +0.002
0.135 +0.001
0.186 #0.005
X
Enscmbbling
0.014 1.0.000 003
0. 0.018 1..005
0.021 10.005
0.036
10.012
0.069
10.032
0.114 10.058
0.036 00.000
0.039 #0002
0.068
#0.340
0.125 40.005
0.191
+0J007
0,240
=0.008
GaldGGCN
0,022 10.007
0. 0.028 10.014
0.031 10.169
060 10.022
0.109
10.019
0.141
10.019
Easernbling
0.024 10.001
0.038 10.011
0.043
10.002
0.083
10.001
0.139 20.004
0.181 20,002
Ensernbling
0.017 #0.002
0.024 #0.00
0.027 20.00s
0.030
20.004
0. 0.036 =0.012
0.059
20000
0.026 +0.001
0.044 +0.00
0.002 +-0.156
0.108 +0.005
0. 0.197 +0.012
0:273 +0.005
0.022 #0.001
0.037 #0.05
0.044 +0.133
0.091 +11.08
0.165 +:1.18
0.239 #0.018
Eineembling
0.016 +0.001
0,026 +0.02
0.030 +0.000
0.066 +:0000
0.123 +0.000
0.195 +0.000
Enserabling
0.014 +40000
0.023 40.002
0.027 10,000
0.055 #00004
0. 0.103 400006
0.164 20.006
AUROC for the binary classification task of detecting OOD samples using the maximum softmax probability as the score (Kirchheim et al., 2022).
OOD Detection Results. As shown in Table 2, we observe that G-AUQ variants improve OOD detection performance over the vanilla baseline on 6/8 datasets, where pretrained G-AUQ obtains the best overall performance shift), but does lose some performance on GOODMoii(CCvaritt)) We note that vanilla models provided by the original benchmark generalized poorly on this particular dataset (increased training time/accuracy did not improve performance), and this behavior was reflected in our experiments. We suspect that poor generalization coupled with stochasticity may explain G-AUQ's performance here.
6/8 datasets. G-AUQ performs comparably on GOODSST2(concep)
6 FINE GRAINED ANALYSIS OF G-AUQ
Given that the previous sections extensively verified the effectiveness of G-AUQ on a variety of covariate and concept shifts across several tasks, we seek more fine-grained understanding of G-AUQ's behavior with respect to different architectures and training strategies. In particular, demonstrate that G-AUQ continues improve calibration with expressive graph transformer architectures, and that using READOUT anchoring with pretrained GNNs is an effective lightweight strategy for improving calibration of frozen GNN models.
6.1 CALIBRATION UNDER CONTROLLED SHIFTS
Recently, it was shown that modern, non-convolutional architectures (Minderer et al., 2021) are not only more performant but also more calibrated than older, convolutional architectures (Guo et al., 2017) under vision distribution shifts. Here, we study an analogous question: are more expressive GNN architectures better calibrated under distribution shift, and how does G-AUQ impact their calibration? Surprisingly,
find that more expressive architectures are not considerably better calibrated than their MPNN counterparts, and ensembles of MPNNs outperform ensembles of GTrans. Notably, G-AUQ continues to improve calibration with respect to these architectures as well.
8

---

Published as a conference paper at ICLR 2024
Experimental Setup. (I) Models. While improving the expressivity of GNNs is ac- tive area of research, positional encodings (PEs) and tuap-trranffrmer (GTran) architec- tures (Müller et al., 2023) are popular strategies due to their effectiveness and flexibility. GTrans not only help mitigate over-smoothing and over-squashing (Alon & Yahav, 2021; Topping et al., 2022) but they also better capture long-range dependencies (Dwivedi al., 2022b). Meanwhile, graph PEs help improve expressivity by differenti-- ating isomorphic nodes, and capturing structural vs. proximity information (Dwivedi et al., 2022a). Here, we ask if these enhancements translate improved calibration under distri- bution shift by comparing architectures with/without PEs and transformer vs. MPNN models. We use equivariant and stable PEs (Wang et al., 2022b), the sfate-of-the-art, "general, pow- erful, scalable" (GPS) framework with a GatedGCN backbone for the GTran, GatedGCN for the vanilla MPNN, and perform READOUT anchoring with 10 random anchors. (2) Data. In order understand calibration behavior distribution shifts become progressively more severe, we create structurally dis- torted but valid graphs by rotating MNIST images by a fixed number of degrees (Ding et al., 2021) and then creating the cor- responding super-pixe graphs (Dwivedi al., 2020; Knyazev et al., 2019; Velickovic et al., 2018). (See Appendix, Fig. 6.) Since superpixel segmentation on these rotated images will yield different superpixel k-nn graphs but leave class informa- tion unharmed, we can emulate different severities of label-preserving structural distortion shifts, We note that models are trained only using the original (09 rotation) graphs. Accuracy (see appendix) and ECE over 3 seeds are reported for the rotated graphs.
GOODCMNIST cost-cmaridte
0.5
GOODMeinn,Ps PE
size-concppt
siz-oowariste
0.5
0.25
0.0
GOODSST2
length concept
lengh covanic
0.2
0.0
Van
EZE Protr
0.0 Vania EZE Prot:
Figure 4: -ittof--iitribution cal- ibration error from applying G- AUQ in end-to-end training vs. to a pretrained model, which is a sim- ple yet effective way to use stochas- tic anchoring.
Results. In Table 3, present the OOD calibration results, with results of more variants and metrics in the supplementary Table 5 and 8. First, we observe that PEs have minimal effects on both calibration and accuracy by comparing GatedGCN with and without LPEs. This suggests that while PEs may enhance expressivity, they do not directly induce better calibration. Next, we find that while vanilla GPS is better calibrated when the distribution shift is not severe (10, 15, 25 degrees), it is less calibrated (but more performant) than GatedGCN at more severe distribution shifts (35, 40 degrees). This is in contrast known findings about vision transformers. Lastly, we see that G-AUQ continues to improve calibration across all considered architectural variants. with minimal accuracy loss. Surprisingly, however, we observe that ensembles of G-AUQ models not only effectively resolve any performance drops, they also cause MPNNs to be beiter calibrated than their GTran counterparts.
6.2 How DÒES G-AUQ PERFORM WITH PRETRAINED MODELS?
As large-scale pretrained models become the norm, it is beneficial to be able perform lightweight training that leads to safer models. Thus, we investigate if READOUT anchoring is such a viable strategy when working with pretrained GNN backbones,
only requires training stochastically centered classifier on top of a frozen backhone. (Below, we discuss results on GOODDataset, but please A.4 for results on RoMMNIST and A.12 for additional discussion.)
From Fig. 4 (and expanded in Fig. 8), we observe that across datasets, pretraining (PT) yields competitive (often superior) OOD calibration with respect to end-to-end (E2E) G-AUQ. With the exception of GOODMotif (basis) dataset, PT G-AUQ improves the OOD ECE over both vanilla and E2E G-AUQ models at comparable or improved OOD accuracy (6/8 datasets). Furthermore, PT G-AUQ also improves the ID ECE on all but the GOODMotif(sizee) (6/8), where it performs comparably to the vanilla model, and maintains the ID accuracy. Notably, as only an anchored classifier is trained, PT G-AUQ substantially reduces training time relative to E2E G-AUQ and vanilla models (see App. A.13), highlighting its strengths as a light-weight, effective strategy for improving uncertainty estimation.
Results.
CONCLUSION
We propose G-AUQ, novel training approach that adapts stochastic data centering for GNNs through newly introduced graph-specife anchoring strategies. Our extensive experiments demonstrate G- AUQ improves calibration and uncertainty estimates of GNNs under distribution shifts.

---

Published as a conference paper at ICLR 2024
8 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-ACS2-07NA27735 National Security, LLC. and is partially supported by the LLNL-LDRD Program under Project No. 22-SI-004 with IM release number LLNL- WAIL-CONF-8500 the National Science Foundation under CAREER Grant No. IIS 1845491, Army Young Investigator Award No. 19911NNF1800997, and Adobe, Amazon, Facebook, and Google faculty awards. Any opinions, findings, and conclusions recommendations expressed here those of the author(s) and do not reflect the views of funding parties. RA completed this work while at Lawrence Livermore National Security, LLC. PT thanks Ekdeep Singh Lubana and Vivek Sivaraman Narayanaswamy for useful discussions during the course of this project.
Lawrence Livermore
This work is also partially supported by
REFERENCES
Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. In Proc. Int. Conf. Learning Representations (ICLR), 2021. [Pg. 9.]
Rushil Anirudh and Jayaraman J. Thiagarajan. Out of distribution detection via neural network anchoring. In Asian Conference on Machine Learning, ACML 2022, 12-14 December 2022, Hyderabad, India, 2022. [Pg. 3.]
Beatrice Bevilacqua, Yangze Zhou, and Bruno Ribeiro. Size-invariant graph representations for graph classification extrapolations. In Proc. Int. Conf. on Machine Learning (ICML), 2021. [Pg. 22.]
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In Proc. Int. Conf. Machine Learning (ICML), 2015. [Pg. 2.]
Davide Buffeelli, Pictro Liò, and Fabio Vandin. Sizeshiftreg: regularization method for improving size-generalization in graph neural networks. In Proc. Adv. in Neural Information Processing Systems (NeurIPS), 2022. Pg. 6 and 22.]
Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, and Petar Velickovic. Principal neighbourhood aggregation for graph nets. In NeurIPS, 2020. [ [gg. 6.]
Nicki Skafte Detlefsen, Jiri Borovec, Justus Schock, Ananya Harsh, Teddy Koker, Luca Di Liello, Daniel Stancl, Changsheng Quan, Maxim Grechkin. and William Falcon. Torchmetries mea- suring reproducibility in pytorch, 2022. URL https:/ https://ttttaiitt /ttp//ggtttt..te/e/ com woon/Lightninggee torchmet.icc.. Pg. 27.]
Mucong Ding, Kezhi Kong, Jiuhai Chen, John Kirchenbauer, Micah Goldblum, David Wipf, Furong Huang, and Tom Goldstein. A closer look at distribution shifts and out-of-distribution general- ization on graphs. In NeuriPS 202/ Workshop Distribution Shifts: Connecting Methods and Applications, 2021. Pg. and 18.]
Vijay Prakash Dwivedi, Chaitanya K. Joshi, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. CoRR, 2020. Pg. 9.]
Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Graph neural networks with learnable structural and positional representations. Proc. Int. Conf. on Learning Representations (ICLR), 2022a. Pg. 4 and 9.]
Vijay Prakash Dwivedi, Ladislav Rampásek, Michael Galkin, Ali Parviz, Guy Wolf, Anh Tuan Luu, and Dominique Beaini. Long range graph benchmark. Proc. Ady. in Neural Information Processing Systems NeurIPS, Datasets and Benchmark Track, 2022b. [Pg. 9.]
Saurabh Garg. Sivaraman Balakrishnan, Zachary C. Lipton, Behnam Neyshabur, and Hanie Sedghi. Leveraging unlabeled data to predict oui-t--diitriuution performance. Proc. Int. Conf. on Learning Representations (ICLR), 2022. Pg. 28.]
10

---

Published as a conference paper at ICLR 2024
Thomas Gaudelet, Ben Day, Arian R. Jamasb, Jyothish Soman, Cristian Regep, Gertrude Liu, Jeremy B. R. Hayter, Richard Vickers, Charles Roberts, Jian Tang, David Roblin, Tom L. Blundell, Michael M. Bronstein, and Jake P. faylor-King Utilising graph machine learning within drug discovery and development. CoRR, abs/2012.05716, 2020. Pg. 1.]
Shurui Gui, Xiner Li, Limei Wang, and Shuiwang Ji. GOOD: A graph cuttof-distribution benchmark. In Proc. Adv. in Neural Information Processing Systems (NeurIPS), Benchmark Truck, 2022. Pg. 5, 6, 22, and 29.]
Devin Guillory, Vaishaal Shankar, Sayna Ebrahimi, Trevor Darrell, and Ludwig Schmidt. Predicting with confidence on unseen distributions. In ICCV, 2021. [Pg. 28.]
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. In Proc. of the Int. Conf. on Machine Learning. (ICML), 2017. [gg. 1, 8, and 27.]
Kartik Gupta, Amir Rahimi, Thalaiyasingam Ajanthan, Thomas Mensink, Cristian Sminchisescu, and Richard Hartley. Calibration of neural networks using splines. In Proc. Int. Conf. on Learning Representations (ICLR), 2021. [Pg. 1, 2, and 27.]
Arman Hasanzadeh, Ehsan Hajiramezanali, Shahin Boluki, Mingyuan Zhou, Nick Duffield, Krishna Narayanan, and Xiaoning Qian. Bayesian graph neural networks with adaptive connection sampling. In ICML, 2020. Pg. 2.]
Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and out-of-distribution examples in neural networks. In Proc. Int. Conf. on Learning Representations (ICLR), 2017. [Pg.
1, 2, and 7.]
Dan Hendrycks, Mantas Mazeika, and Thomas G. Dietterich. Deep anomaly detection with outlier exposure. Proc. Int. Conf. on Learning Representations (ICLR), 2019. Pg. 2 and 7.]
Dan Hendrycks, Nicholas Carlini, John Schulman, and Jacob Steinhardt. Unsolved problems in ML safety. CoRR, abs/2109.11/ 13916, 2021. [P.. 7.]
Dan Hendrycks, Steven Basart, Mantas Mazeika, Andy Zou, Joseph Kwon, Mohammadreza Mosta- jabi, Jacob Steinhardt, and Dawn Song. Scaling -istiitriittionn detection for real-world settings. In Proc. Int. Conf. on Machine Learning (ICML), 2022a. Pg. 7.]
Dan Hendrycks, Andy Zou, Mantas Mazeika, Leonard Tang, Bo Li, Dawn Song, and Jacob Steinhardt. Pixmix: Dreamlike pictures comprehensively improve safety measures. In Proc. Int. Conf. on Computer Vision and Pattern Recognition (CVPR), 2022b. [Pg. I and 7.]
José Miguel Hernández-Lobato and Ryan P. Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. In Proc. Int. Conf. on Machine Learning (ICML), 2015. [
Pg. 2.]
Hans Hao-Hsun Hsu, Yuesong Shen, Christian Tomani, and Daniel Cremers. What makes graph neural networks miscalibrated? In Proc. Adv. in Neural Information Processing Systems NeurIPS, 2022. [ Pg. 1, 2, and 27.]
Eyke Hüllermeier and Willem Waegeman. Aleatoric and epistemic uncertainty in machine learning: an introduction concepts and methods. Mach. Learn., 110(3), 2021. 3.]
Yiding Jiang, Dilip Krishnan, Hossein Mobahi, and Samy Bengio. Predicting the generalization gap in deep networks with margin distributions. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. [Pg. 2, 7, and 28.]
Thomas Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In ICLR, 2017. [ Pg. 6.]
Konstantin Kirchheim, Marco Filax, and Frank Ortmeier. Pytorch-ood: A library for out-of- distribution detection based on pytorch. In Workshop at the Proc. Int. Conf. on Computer Vision and Pattern Recognition CVPR, 2022. [Pg. 8.]

---

Published as a conference paper at ICLR 2024
Boris Knyazev, Graham W. Taylor, and Mohamed R. Amer. Understanding attention and generaliza- tion in graph neural networks. In Proc. Adv. in Neural Information Processing Systems (NeurIPS), 2019. [Pg. 9.]
Meelis Kull, Miquel Perelló-Nieto, Markus Kängsepp. Telmo de Menezes e Silva Filho, Hao Song. and Peter Flach. Beyond temperature scaling: Obtaining well-calibrated multi-class probabilities with dirichlet calibration. In Proc. Adv. in Neural Information Processing Systems NeurIPS, 2019.
[g.. 1, 2, and 27.]
Ananya Kumar, Percy Liang, and Tengyu Ma. Verified uncertainty calibration. In Proc. Adv. in Neural Information Processing Systems NeurlPS, 2019. [ Pg. 27.]
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Proc. Adv. in Neural Information Processing Systems (NeurIPS), 2017. [Pg. 2.]
Kimin Lee, Kibok Lee, Honglak Lee, and Jinwoo Shin. A simple unified framework for detecting -istti-iiitiitutiin samples and adversarial attacks. In Proc. Adv. in Neural Information Processing Systems NeurlPS, 2018. [Pg. 7.]
Weitang Liu, Xiaoyun Wang, John D. Owens, and Yixuan Li. Energy-based out-of-distribution detection. In Proc. Adv. in Neural Information Processing Systems NeurIPS, 2020. [ [Pg. 7.]
Matthias Minderer, Josip Djolonga, Rob Romijnders, Frances Hubis, Xiaohua Zhai, Neil Houlsby, Dustin Tran, and Mario Lucic. Revisiting the calibration of modern neural networks. In Proc. Adv. in Neural Information Processing Systems (NeurIPS), 2021. [Pg. 8.]
Christopher Morris, Nils M. Kriege, Franka Bause, Kristian Kensting, Petra Mutzel, and Marion Neumann. Tudataset: A collection of benehmark datasets for learning with graphs In ICML 2020 Workshop on Graph Representation Learning and Beyond (GRL+ 2020), 2020. URL www. www.gggaaheaaainngi... [Pg. 6.]
Luis Müller, Mikhail Galkin, Christopher Morris, and Ladislav Rampásek. Attending graph transformers. CoRR, abs/2302.04181,, 2023. Pg. 9.]
Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht. Obtaining well calibrated probabilities using bayesian binning. In Proc. Conf. on Adv. of Artificial Intelligence (AAAT), 2015. [gg. 27.]
Aviv Netanyahu, Abhishek Gupta, Max Simehowitz, Kaiqing Zhang, and Pulkit Agrawal. Learning to extrapolate: A transductive approach. In Proc. Int. Conf. on Learning Representations (ICLR),
2023. Pg. 3.]
Nathan Ng, Neha Hulkund, Kyunghyun Cho, and Marzyeh Ghassemi. Predicting out-of-domain generalization with local manifold smoothness. CoRR, abs/2207 2022. [Pg. 28.]
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D. Sculley, Sehastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. In Proc. Adv. in Neural Information Processing Systems NeurlPS, 2019. [ Pg. 1, 2, 5, 7, and 27.]
Amir Rahimi, Amirreza Shaban, Ching-An Cheng, Richard Hartley, and Byron Boots. Intra order- preserving functions for calibration of multi-class neural networks. Advances in Neural Information Processing Systems, 33:13456-13467 2020. [Pg. 27.]
Mrinank Sharma, Sebastian Farquhar, Eric Nalisnick, and Tom Rainforth. Do bayesian neural networks need to be fully stochastic? In AISTATS, 2023. [Pg. 3.]
Jayaraman J. Thiagarajan, Rushil Anirudh, Vivek Narayanaswamy, and Teer-Timo Bremer. Single model uncertainty estimation via stochastic data centering. In Proc. Adv. in Neural Information Processing Systems (NeurIPS), 2022. Pg. and 3.]
12

---

Published as a conference paper at ICLR 2024
Jake Topping. Francesco Di Giovanni, Benjamin Paul Chamberlain, Xiaowen Dong, and Michael M. Bronstein. Understanding over-squashing and bottlenecks on graphs via curvature. Proc. Int. Conf. on Learning Representations ICLR, 2022. [Pg. 9.]
Puja Trivedi, Danai Koutra, and Jayaraman J Thiagarajan. A closer look at scoring functions and generalization prediction. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023a. Pg.. 28.]
Puja Trivedi, Danai Koutra, and Jayaraman J. Thiagarajan. A closer look at model adaptation using feature distortion and simplicity bias. In Proc. Int. Conf. on Learning Representations (ICLR), 2023b. Pg. and 7.]
Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In ICIR, 2018. Pg. 9.]
Haoqi Wang, Zhizhong Li, Litong Feng, and Wayne Zhang. Vim: out-of-distribution with virtual- logit matching. In Proc. Int. Conf. on Computer Vision and Puttern Recognition (CVPR), 2022a. [
Pg. 7.]
Haorui Wang, Haoteng Yin, Muhan Zhang, and Pan Li. Equivariant and stable positional encoding for more powerful graph neural networks. In Proc. Int. Conf. on Learning Representations (ICLR), [Pg. and 9.]
2022b.
Xiao Wang, Hongrui Liu, Chuan Shi, and Cheng Yang. Be confident! towards trustworthy graph neural networks via confidence calibration. In Proc. Adv. in Neural Information Processing Systems NeurlPS, 2021. Pg. 1, 2, and 27.]
Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre-Alvise Rebuffi, Ira Ktena, Krishnamurthy Dj Dvijotham, and Ali Taylan Cemgil. A Fine-Grained Analysis on Distribution Shift. In Proc. Int. Conf. on Learning Representations (ICLR), 2022. Pg. and 2.]
Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and probabilistic perspective of generalization. In Proc. Adv. in Neural Information Processing Systems NeurlPS, 2020. [ 2.J
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In ICLR, 2019. [Pg. 6.]
Yujun Yan, Jiong Zhu, Marlena Duda, Eric Solarz, Chandra Sekhar Sripada, and Danai Koutra. Groupinn: Grouping-base based interpretable neural network for classification of limited, noisy brain data. In Proc. Int. Conf. on Knowledge Discovery & Data Mining, KDD, 2019. [ Pg. 1.]
Jianwei Yang, Jiasen Lu, Stefan Lee, Dhruv Batra, and Devi Parikh. Graph R-CNN for scene graph generation. Proc. Euro Conf. on Computer Vision (ECCV), 2018. Pg. 1.]
Gilad Yehudai, Ethan Fetaya, Eli Meirom, Gal Chechik, and Haggai Maron. From local structures to size generalization in graph neural networks. In International Conference on Machine Learning, pp. 11975-11986 PMLR, 2021. [ Pg. and 22.]
Bianca Zadrozny and Charles Elkan. Transforming classifier scores into accurate multiclass probabil- ity estimates. In Proceedings of the eighth ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 694 699, 2002. [ Pg.. 27.]
Jize Zhang. Bhavya Kailkhura, and Thomas Yong-Jin Han. Mix-n-match. Ensemble and com- positional methods for uncertainty calibration in deep learning. In Proc. Int. Conf. on Machine Learning (ICMI), 2020. [ Pg. 1, 2, and 27.]
Muhan Zhang and Yixin Chen. Link prediction based on graph neural networks. In Proc. Adv. in Neural Information Processing Systems NeurIPS, 2018. [ [Pg. 1.]
Yingxue Zhang Soumyasundar Pal, Mark Coates, and Deniz Ustebay. Bayesian graph convolutional neural networks for semi-supervised classification. In AAAI, 2019. [Pg. 2.]
Yanqiao Zhu, Yuanqi Du, Yinkai Wang, Yichen Xu, Jieyu Zhang, Qiang Liu, and Shu Wu. A survey on deep graph generation: Methods and applications. In Learning on Graphs Conference (LoG), 2022. Pg. 1.]
13

---

Published as a conference paper at ICLR 2024
A
APPENDIX
PseudoCode (Sec. A.1)
Ethies (Sec. A.2)
Reproducibility (Sec. A.3)
Details and Expanded Results for Super-pixel Graph Experiments(Sec. A.4)
Stochastic Centering on the Empirical NTK of Graph Neural Networks (Sec. A.5)
Size-Generallzatioon Dataset Statistics (Sec. A.6)
GOOD Dataset Statistics and Expanded Results (Sec. A.7)
Alternative Anchoring Strategies GOOD Datasets (Sec. A.8)
Discussion of Post-hoc Calibration Strategies (Sec. A.9)
Details of Generalization Gap Experiments (Sec. A.10)
Expanded Pretraining Results (Sec. A.12)
Runtimes (Sec. A.13)
Mean and Variance of Node Feature Anchoring Distributions (Sec. A.14)
Discussion on Anchoring Design Choices (Sec. A.15)
14

---

Published as a conference paper at ICLR 2024
A.1 PSEUDOCODE FOR G-AUQ
def init (Q80)
SOLF.HPHN1 NPMN(InpurDinn 2,H01m1)
self.MRRI2 H0NK(10)11, HD:m2)
SeLT. (HHN(H HDlr31
READOUTLayerli
ceas.classiierr self. ancher distt(Agglodaaaatss)
def make, anchor (Modeeeott)ss
class GBi:
MEAM Awgloodaaaats, dinw))
def
init {):
STDLApplideeeetts din=i)
self, make, layersi]
wwwtiinputDin,
self.AnchorDist (maane stdestd)
self.MPRIC NFWK4HDA1,.10.221
seLf,NRReL HENN( (B1x3)
def Addj
*eLf.READOUT = READOUTLayert!!
InputDin)
self.classifier LinsartHDin3, NunClass]
self.
forwardlselff
AncRop CONCAT([X C, dis-1)
self.MFReats.d Ad) Out: (BZe))
h Self.HPHN1(AncRRe
nout: (N,H01e2)
self.Niaw((h.Ad) x Out: (9, (.22222
h self.Miaa22 (h, Adj)
selr. Adj) # Out: (01e3))
GhLT.MRAaBB, Adj)
self_READOUTTh, Batch) dutz (8, 60и83)
seLt.READOUTIh, Batch)
cels.classifeer() Butz (usClass)
celf.classlfii pred
(a) Vanilla GNN
(b) G-AUQ with Node Feature Anchoring
class MW__Noddfeetture
def
inii (GNN);
SoLT.MPHN1 HPHN1 Innu(Inptttin 2,HDim1)
self self.NMerz NIM02 MPAN HDNN
HDlx3)
self, READOUT READOUTLayer1]
seff.classifae self, make anchor disttttgllldddaaatt
Linear(1BBin3, NunClass]
class GNN ReLdeeeDut
def make, anchor dist (Nodefests):
MPANL def
init (980):
VEAMMAgglisdaFaatts
IPpWttnn HD1v1)
ad-die Goussian
SoLT.HRRSS HPHN(HDÍN),
self.Anchorbist
Norsal Korsal stdnstd)
SoLT. READOUT READDUTLLaYrr(
#Corres. to GOUDLayer!
def forward(x,,dj)
self.classifier Linear(HDiad 2, Nunclass)
INputtDiie))
Aods::
(N. Inputbiee
Ancitop CONCAT(EX c, C), dim-1)
(h. Ad]1
h.tt.HNP.u(nncRRppkk.d. abut: (0,e0)
h = self (h..- Adj)
(h.f. Adj)
REABOURTTTUTT Batch)) aOUT: (8, ADis3)
SHUFFLEtO. din=0) NDie3)
selt.READOUTTh,co
CONCAT([g-C,C), dim=)))r)
pred self.classi (g)
pred self.classifaar(gnn #0ut: (8, NeeClass)
return pred
retura prod
(c) G-AUQ with Hidden Rep Anchoring
(d) G-AUQ with READOUT Anchoring
Figure 5: PseudoCode for G-AUQ. We provide simplified pseudo-code to demonstrate how anchor- ing be performed. We assume Geomereii style mini-batching. Changes with respect the vanilla GNN are shown bold. Unchanged lines are grayed out.
15

---

Published as a conference paper at ICLR 2024
Node Feature
Int. MPNN
Readout
READOUT/MPP
READOUTMLP
MLP
Anchored Inference
IG G.I
GNN.
Prediction:
GNN,
G,
G.
pby| (
Σ. ,.,D E...
X 1X
READOUT
[X,, c, l,l,s
(A X
Uncertainty:
c ~ . ())
e (
K-1 C,0 -N²
GNN.
GNN
Stochastic
(A,, (y) x,)= ~ D train
(A, X)
Calbrated Prediction:
Static
(A c3 X,) D train
Health
μ(1
Overview of G-AUQ models. Here, present conceptual overview of how G-AUQ induces partially stochastic models. This figure is complementary to 5.
16

---

Published as a conference paper at ICLR 2024
A.2 ETHICS STATEMENT
This work proposes a method to improve uncertainty estimation in graph neural networks, which has potential broader societal impacts. As graph learning models increasingly deployed in real-world applications like healthcare, finance, and transportation, it becomes crucial to ensure these models make reliable predictions and know when they may be wrong. Unreliable models can lead harmful outcomes if deployed carclessly. By improving uncertainty quantification, our work contributes towards trustworthy graph Al systems.
We also consider several additional safety-critica tasks, including generalization gap prediction for graph classification (to the best of our knowledge, we are the first to report results on this task) and OOD detection. We hope our work will encourage further study in these important areas.
However, there are some limitations. Our method requires (modest) additional computation during training and inference, which increases resource usage. Although G-AUQ, unlike post-hoc methods, does not need to be on validation dataset, evaluation of its benefits also also relies on having some out-offdiitrribtiio shifted data available. which may not always be feasible. We have seen in Table that there are tasks for which G-AUQfails to improve accuracy and/or calibration of some post-hoc methods, further emphasizing the need to perform appropriate model selection and the risks shifted validation data is not available. Finally, there are open questions around how much enhancement in uncertainty calibration translates real-world safety and performance gains.
Looking ahead, we believe improving uncertainty estimates an important direction for graph neural networks and deep learning more broadly. This will enable the development safe, reliable Al that benefits society. We hope our work inspires more research in the graph domain that focuses on uncertainty quantification and techniques that provide guarantees about model behavior, especially for safety-critical applications. Continued progress will require interdisciplinary collaboration between graph machine learning researchers and domain experts in areas where models are deployed.
models
A.3 REPRODUCIBILITY
For reproducing our experiments, have made code available at this anonymous repository. In the remainder of this appendix (specifically App. A.6, A.7), and A.10), we also provide additional details about the benchmarks and experimental setup.
17

---

Published as a conference paper at ICLR 2024
A.4 DETAILS ON SUPER-PIXES EXPERIMENTS
We provide an example of the rotated images and corresponding super-pixel graphs in Fig. 6. (Note that classes "6" and "9" may be confused under severe distribution shift, i.e. 90 degrees rotation or more. Hence, to avoid harming class information, our experiments only consider distribution shift from rotation up 40 degrees.)
Rot: 25"
Rot: 25
Rot: 40
Rot:
Rot: 280
Figure 6: Rotated Super-pixel MNIST. Rotating images prior to creating super-pixels leads to some structural distortion (Ding et al., 2021). However, we can see that the class-discriminative information is preserved, despite rotation. This allows simulating different levels of graph structure distribution shifts, while still ensuring that samples are valid.
Tables 4 and 5 provided expanded results on the rotated image super-pixel graph classification task, discussed in Sec. 6.1.
In Table 7 we focus on the calibration results on this task for GPS variants alone. Across all levels of distribution shift, the best method our strategy for applying G-AUQ a pretrained model- demonstrating that this is not just a practical choice when it is infeasible to retrain model, but can lead powerful performance by any measure. Second-ber best on all datasets is applying G-AUQ during training. further highlighting the benefits of stochastic anchoring.
In addition to the structural distribution shifts we get by rotating the images before constructing super-pixel graphs, we also simulate feature distribution shifts by adding Gaussian noise with different standard deviations the pixel value node features in the super-pixel graphs. Table 8, we report accuracy and calibration results for varying levels of distribution shift (represented by the size of the standard deviation of the Gaussian noise). Across different levels of feature distribution shift, we also that G-AUQ results in superior calibration, while maintaining competitive in many cases superior accuracy.
18

---

Published as a conference paper at ICLR 2024
Table 4: RotMNIST-AAcurant Here, we report expanded results (accuracy) the Rotated MNIST dataset, including a variant that combines G-AUQ with Deep Ens. Notably, we see that anchored ensembles outperform basic ensembles in both accuracy and calibration. (Best results for models using Deep Ens, and those not using it marked separately.)
MODEL
G-AUQ?
LPE?
Avg. Test (1)
Acc. (1)
Acc. (15)
Acc. (25)
Acc. (35)
Acc. (1)
X
0.918 10,002
20,009
+00007
GatedGCN
X
200099
+0.020
#0002
#0.005
+0.048
+0.006
#00007
+0007
=0.032
+0.088
0.852
+00039
+0,07
+0.033
GPS
100001
20.001
10.001
10.013
10.009
20000
+0.003
20003
20.012
20.014
(Pretrained)
=d.002
+0000
#0,00
0.963 10.0002
0.913
10.001
0.983
10.001
0.874 10.002
0.791
10.002
0.731
10.002
GatelGCN-DENS
0.949 =0.008
0.922 +0.008
0.907
+0.011
+0.020
0.733 +0.032
0.662 +0.046
0.965
0.943 100001
0.933 10.001
0.873 10:001
0.792 10.004
0.736 10.003
0.954 100005
0.930 10.010
0.917 10.011
0.850 10.023
0.759 10.025
0.696 10.032
GPS-DENS
0,980 1.0000
0.969 10.000
0.961 10.000
0.913 1.0 000
0.834 1.00000
0.750 100000
0.978 10.001
0.963 10,0000
0.953 100001
0.903 10.000
0.822 10.0002
0.736 1.0003
Table 5: RotMNNIST-Calirrtioon Here, we report expanded results (calibration) on the Rotated MNIST dataset, including variant that combines G-AUQ with Deep Ens. Notably, we see that anchored ensembles outperform basic ensembles in both accuracy and calibration. (Best results for models using Deep Ens. and those not using it marked separately.)
MODEL
G-AUQ
LPE?
Avg.ECE (L)
(10) (1)
(15) (J)
ECE (25) (L)
ECE (35) (4)
BCE (40) (4)
Ga
0.035 10.001
0.054 10.002
0.062
10.00
0.118 10.007
0.185 10.006
0.233 10.008
+0.002
0.053 40.002
0.001 -0.004
+0.005
0.179 +0.005
0.225 =0.005
#0.001
+0.001
0.068
+0.008
0.195
0.033
0.009
10.067
GatedGCN
X
0.036 +0.003
0.059 +0.002
0.068 -0.340
125 +0.006
0.191 +0.007
0.240 +00.088
0.022 =0.007
0.028 10.014
0.034 10.109
0. 0.062 =0.022
0. 0.109 20.019
0.141 10.019
GPS-TS
X
0.024 100001
0,041 10.001
0.049 0.001
0. 100 10.006
0,188 16.012
0.261 2:0008
X
0.026 +0.001
0.044 +0.001
0.052 -0.156
108 +0.006
0.197 +0.012
0.273 =60008
0.022 +0.001
0.037 +0.005
0.044 -0.13
0.091 +0.008
0. +0.018
0,239 +0.0188
GPS (Pretrained)
0.021 20.001
0.032 40.003
0.039 =0.116
0.083 +0.002
0.153 40.007
0.217 =0.012
X
X
0.026 +0.000
0.038 +0.001
0.042 -0.001
0.084
+0.002
0.135 +0.001
0.185 +60003
0.014 1.0000
0.018 10.005
0.021
10.005
0.096 10.012
0.069 10.002
0.114 100056
GaledGCN- DENS
X
0.024 10.001
0.038 10.001
0.043 00000
0.083 10.001
0.139 +0.004
0.181 20.002
0.017 #0.002
0.024 +0.005
0.027 =0.008
0. 0.030 =0.004
0.036 =0.012
0.059 +00003
GPS-DENS
0.016 10.001
0.026 10.002
0.030 00.00
0.006 1.00000
0.123 10.000
0.195 10.000
+-0.000
0.023 +0.002
0.027 00.003
0.055 +000044
0.103 +0.005
=00.066
Table Accuracy of GPS Variants on RotatedMNIST. We focus on the accuracy results for GPS variants on rotated MNIST dataset. Using G-AUQ (with or without pretraining) remains close in accuracy to foregoing generally within the range the standard deviation of the results.
MODEL
G-AUQ!
Avg. Test (t)
Acc. (10) (t)
Acc. (15) (t)
Acc. (25) (1)
Acc. (35) (t)
Ace. (40) (t)
10.001
0.948 10.001
0.938 10.001
0.770 1:0013
10.009
GPS
0.969 1:0.001
0.946 40.003
0.937 20.003
0.869 4-0.003
0.769 40.012
0.679 20.014
GPS (Pretrained)
0.967 4-0.002
0.945 40.004
0.934 =0.005
0.864 +-0.009
0.759 40.010
0.674 00.002
Table 7: Calibration of GPS Variants on RotatedMNIST. We focus on the calibration results for GPS variants on rotated MNIST dataset. Across the board, we see improvements from using G-AUQ, with our strategy of applying it to pretrained model doing best.
MODEL
G-AUI
Avg.ECH (1)
FCE (10) (4)
FCE (L)
BCE (25) (4)
ECH (L)
HCB (40) (4)
GPS-TS
X
0.024 +0.001
0.041 +0.001
0.049 +0.001
0.102 +0.006
0.188
=0.012
0.261 +0.008
X
0.026 +0.001
0.044 +0.001
0.052 +0.156
0.108 +0.006
0.197
-0.012
0.273 +0.008
GPS
0.022 +0.001
0.037 +0.005
0.044 10.133
0.001 +0.008
0.165
00.0
0.239 +0.018
GPS (Pretrained)
0.021 +:0
0.032 +0.003
0.039 10.116
0.083 20.002
0.153 +0 007
0.217 =0.012
19

---

Published as a conference paper at ICLR 2024
Table 8: MNIST Feature Shifts. G-AUQ improves calibration and maintains competitive improved accuracy across varying levels of feature distribution shift.
MODEL
LIMI
G-UUQ?
Calibraiin
Accuracy (1)
BCB (L)
Accurcy (1)
ICC
Accuracy (f)
BCE L3)
Accuracy (1:
DCB (1)
GaneOOCN
0.775 (0.053
0.075 200
0.53) 10:00
0.160 0.160:0008 09/07
0.356 10
0.122 (0053
0.249
0.522 10001
X
0.519
0.845
0.740650.001
0.1000-00005
0.541:30.040
01200:41.007
0.408:0.118
11242240.088
0,500800181
20

---

Published as a conference paper at ICLR 2024
A.5 STOCHASTIC CENTERING ON THE EMPIRICAL NTK OF GRAPH NEURAL NETWORKS
Using a simple grid-graph dataset and 4 layer GIN model, we compute the Fourier spectrum of the NTK. As shown in Fig. 7, find that shifts to the node features induce systematic changes to the spectrum.
Input
G-AUQ (input) G-AUQ
G-AUQ
G-AUQ
G-AUQ (READOUT)
10a
101
10-1
10
0.000
0.100
0.200
0.300
0.400
0.500 0.600 0.700
0.800
0.900
Figure 7: Stochastic Centering with the empirical GNN NTK. We find that performing constant shifts at intermediate layers introduces changes to a GNN's NTK. We include vanilla GNN NTK in black for reference. Further, note the shape of the spectrum should not be compared across subplots each subplot created with different random initialization.
21

---

Published as a conference paper at ICLR 2024
A.6 CIZERAENRRALIZIIION DATASET STATISTICS
The statistics for the size generalization experiments (see Sec. 5.1) are provided below in Table 9.
Table 9: Size Generalization Dataset Statisties: This table directly reproduced from (Buffelli et al., 2022), who in turn used statistics from (Ychudai et al.. 2021; Bevilacqua et al., 2021).
NCI
NCT109
ALL
SMALLEST 50%
LARGEST 10%
ALL
SMALLEST 50%
LARGEST 10%
CLASS A
49.95%
62.30%
19.17%
49.62%
62.04%
21.37%
CLASS B
50.04%
37.69%
80.82%
50.37%
37.95%
78.62%
# OF GRAPHS
4110
2157
412
4127
2079
421
AYG GRAPH SIZE
29
20
61
PROTEINS
DD
ALL
SMALLEST 50%
LARGEST
ALL
SMALLEST
LARGEST 10%
CLASS 59.56%
41.97%
90.17%
58.65%
35.47%
79.66%
CLASS
40.43%
58.02%
9.82%
41.34%
64.52%
20.33%
# OF GRAPHS
1113
567
112
1178
592
118
AYG GRAPH SIZE
39
15
138
284
144
746
A.7 GOOD BENCHMARK EXPERIMENTAL DETAILS
For our experiments in Sec. 5.2, we utilize the doouu-of-distributie covariate and concept splits provided by Gui al. (2022). Furthermore, we use the suggested models and architectures provided by their package. In brief, we use GIN models with virtual nodes (except for GOODMotiI) for training, and average scores over seeds. When performing stochastic anchoring at particular layer, we double the hidden representation size for that layer. Subsequent layers retain the original size of vanilla model.
When performing stochastic anchoring, we use 10 fixed anchors randomly drawn from the in- distribution validation dataset. We also train the G-AUQ for an additional 50 epochs to ensure that models are able to converge. Please see code repository for the full details.
We also include results on additional node classification benchmarks featuring distribution shift in Table 12. In Table 13, we compare models without G-AUQ to the use of G-AUQ with randomly sampled anchors at the first or second layer.
22

---

Published as a conference paper at ICLR 2024
Dunaset
Shift
ID validatio
test
OOD validation
OOD validation
validation
ID
Length
GOOD-SST2
concept
5843
5843
15142
Color
covtatttl
42000
7000
7000
7000
7000
GOOD CMNIST
concept
29400
6300
6300
14000
14000
shin
42883
140000
14000
5000
concept
Degroe
covariate
9378
1979
1979
3000
3454
8213
1979
1979
3841
3781
GODD-Cooa
concept
7278
1858
1558
3807
3997
7281
1560
1560
3706
5686
University
conteriatc concept
244
61
61
125
126
GOOD WebKB
282
60
60
106
109
GOOD-CRAS
conceqt
Table 10: Number of Graphs/Nodes per dataset.
Dataset
Model
# Model layers
Batch Size
Max Epochs
Iterations per epoch
Initial
Node Feature Dim
SOOD-SST2
GIN. Virtual
32
200/100
le-3
768
GOOD.CMNIST
GOOD CMNIST
GIN-V GirtViiuual
5
128
500
le-3
GOOD-Mett
GIN
3
32
200
le-3
GOOD-Cora
GCN
3
4096
100
10
le-3
8710
GOOD GOOD-WehKR
GCN
3
4096
100
1e-3/Se-3
1703
GOOD-CBAS
GCN
3
1000
200
36-3
Table 11: Model and hyperparameters for GOOD datasets.
23

---

Published as a conference paper at ICLR 2024
A.8 GOOD DATASET ADDITIONAL RESULTS
We also include results on additional node classification benchmarks featuring distribution shift in Table 12. In Table 13, we compare models without G-AUQto the use of G-AUQwith randomly sampled anchors at the first or second layer.
Table 12: Additional Node Classification Benchmarks. Here, we compare accuracy and calibration error of G-AUQ and "no G-AUQ" (vanilla) models on 4 node classification benchmarks across concept and covariate shifts. First, we note that across all our evaluations, without any posthoc calibration, G-AUQ is superior to the vanilla model on nearly every benchmark for better or same accuracy (8/8 benchmarks) and better calibration error (7/8), often with significant gain in calibration performance. However, due to the challenging nature of these shifts, achieving state-of- the-art calibration performance often requires the use of post-hoc calibration methods- evaluate how these posthoc methods can be elevated when combined with G-AUQ (versus the vanilla variant). When combined with popular posthoc methods, we highlight that performance improves across the board, when combined with G-AUQ (including in WehKB and CBAS- (BBAS-Concept). For example, WebKB across the 9 calibration methods considered, "G-AUQ + calibration method" improves or maintains the calibration performance of the analogous "no G-AUQ + calibration method" in 7/9 (concept) and 6/9 (covariate). In CBAS, calibration is improved or maintained as the no- no-G-UUQ version on 5/9 (concept) and 9/9 (covariate). In all cases, this is achieved with little or no compromise on classification accuracy (often improving over "no G- G-AUQ" variant). We also emphasize that, across all the 8 evaluation sets (4 datasets x 2 shift types) in Table 10, the best performance is almost always obtained with a GDUQ variant: (accuracy: 8/8) as well as best calibration (6/8) or second best (2/8).
also
Shift Concapi
Shilt Commritt
Accuracy (1)
(1)
Accuracy (t)
(J)
Daizset
Doitain
Calibration
G-G UQ
G-3 UQ
No G-G UQ
G-6 UQ
No UQ
G- UQ
No UQ
G-3 UQ
0.25822000
0.281
0.67500001
0.50819000
0.122.0112
0.11530001
0.50010000
0.026
CAGCN
0,250120005
0.28220008
0.402:0.14
0.47349.12
0.122:000L5
00092:00151
0.355400.277
0
Disiclee
0.229:3.01
0.2210022
0:4722005
0.4210005
0.206-9004
ETS
0,250120005
0
0
0.377
0.121
00084240077
0,50950.122
WALKB
Univenity
GATS
0.200230002
0.27360.01
0000000000
0.48540.00
0:122:00155
007900000
0.45589.007
0.376:00022
0:251:3.005
0:200-0011
0.319--3000
0:00140:04
0000000015
0.352400007
0.122-0.013
Orderinniian
0.258120005
0.27
1628100226
0,56110004
0.122
0.00600055
0.54BL0000
10995
Spline
0.23762012
0,257500225
0.38622004
0.122:0003
0.171:00050
0.472500011
0
VS
0.25810000
0.276 10011
0.6710.00
0.588100111
0,122
0.096
0,002100644
0.507 0LB1
0:31103001
0:55500000
0.00:00.000
0.13
0,47#0000
031820014
0.348:0.002
CAGCN
0.581
0.5
0.135-00000
0. 1284-0000
0:47:40.02
0.52 0.05
022550005
0.231=000 D.05
0.531
0.414.00007
ETS
0.581:0000
0.500
00001000000
0.1166200018
0,47:90022
0.623
0.14
Degree
0.506
0.521
0.21110001
0.508
IRM
0.56220.002
0.507
0.125
0.102200022
0,000:00000
0.522-400004
0.19
0.10
Ondertonuriant
0,51122003
0.59250000
0,213
0,47#0002
0.488200227
0.318=20042
0.196
Spline
0.571:3.00
0:50000000
00001100000
0008550.00 001
0.030:00003
0.5240.004
0.188000
VS
0.58110000
0.500 10.01
0.00610.00
0.127) 10.000
0:17:0091
0.522 10.055
0.515
0,146 EL005
1.007
0.026
100011
0.201000000
0.000.0000
0.633 10031
0.208
0.118 ELOI>
CAGEN
0.007
0:026
0.188200111
0.23620001
0.030:00004
0.634:00055
0.12Da50000
Dirichee
0,579=2.007
0.588500000
0.100000111
0.108:2.005
0.502=00077
0
0.000000.00
0.200
BTS
0.907530002
0:028-00000
0282240000
0.11:00003
0:000000011
0663
0.245
O. 0.100:-00013
Wood
0.007:3000
IRM
0.00810000
0.63110002
0.11540012
0.08
0.602:00003
0.685-10000
0.306600022
Orderinctiont
0.007:3000
0.22
0,202
Spline
0.5081:0005
0:620:10002
107754000122
0.062..005
0.501
01685-400004
1.003
0.065
2.007
0.60
01283500008
0.600200004
0.636
1.201
0.119:00005
0.33:00:1
0880-10011
0.10010011
8.151
0.700-10:015
029951007
0.108 ELDIA
CAGON
0.53.0013
0.8310011
11187100111
0,148120022
0.700:0000
0.7651005
D 0.25
0L156-
Dirichlet
0.30110.00
10.05
0.101:0011
0.071 H0.0S
0.771-0000
0:211
0
ETS
0.531003
0.827
0.146100115
0.1661:0077
0,700-0000
0,7610.577
0.28504123
Color
GATS
0.5120005
0,83100021
0.1666.000
0.177
0.700000001
0.70L=00010
0.200250.000
0.1G
IRM
0.820
0
0.142500000
135
0,72+00010
0.207=0.005
0. 0.15M:00.17
Orderincaiint
0.60000005
0.11110.00
0.173-3.009
0:700:0001
0.706-00055
0:301100011
0:191:00001
Spline
0.82sh006
0.824:0011
0.159500.000
0.36:
0 0.651=0m1
0.780:00018
0.223e50.084
0.177
0:820-3.001
0:800:0001
0110510:011
0.100:30001
0.717=00019
0:800-00000
0,212400009
0.152:0.011
24

---

Published as a conference paper at ICLR 2024
Table 13: Layerwise Anchoring for Node Classification Datasets with Intermediate Representa- tion Distributions. Here, we provide preliminary results for performing layerwise anchoring when performing node classification. We fit gaussian distribution over the representations (similar to node feature anchoring) and then sample anchors from this distribution. We fit gaussian distribution over the representations (similar to node feature anchoring) and then sample anchors from this distribution. We see that these alternative strategies do provide benefits in some cases, but overall, our original input node feature anchoring strategy is more performant.
S0 Groupe
Grosia
DCE
Acuneyytt)
DCE
Darael
Dasin
Calvadis
AB
Rediel
Racker
NTG
Reckerl
Roakerl
N:G
Rickel
3adia?
NvG.
Realen
Redles?
Desiii
0171:300
020-4
335
8521
0286LE
0,30
26
115
0211
042
Sytac
0,52
1100000
082 100g
6357-1005
(05 80
013-4:07 17
Dikiki
0.00
Coca
Degas
ETS
61425541
0.581000
13/500m
039955
0.75000
0.3010015
13
2
255
0382 812
850740011
0,020000
63
1820001
Oslctiraion
0.00100
122618
0:300005
0.100
Dasckit
0.77
0355:180
0.80100
0.20
1580-182
125-48:
000
15574183
038
(38744
180
0.33
0370mm
1511-990
113 82
0.23mm0
0.00et 10
Descii
831423000
034:000
0317:1912
0/
0320000
0.28
0/22-1
6.20-1100
1120-1
3:709905
0101-1003
2211-3001
Sytac
821720000
02511000
64
0.8770090
0481LERG
8:201000
625111.11
0.67210011
637
025
Table 14: Layerwise Anchoring for Node Classification Datasets with Random Shuffling. Here, we provide preliminary results for performing layerwise anchoring when performing node clas- sification. We use random shuffling (similar to the proposed hidden layer strategy) to create the interemediate representations. We see that these alternative strategies do provide benefits
Abomsc (1)
Actuncy:1) (1)
ICE
Doanc
Demuh
Ne.Ga UV
Blcch
Ne
Beahh
NoGa
Rech
Bach
Diskiis
030000
EdITEER
CIRAS
Citar
FTS
0:40-0
1323
251
K34
0.243000
63210415
B49TLE807
Unletratie
0.39100
0.3010063
Diskia
821500001
04222005
0121000
136110
Caza
Dogos
0381119
00010000
B4KLEES
1311100
Sylue
0.6mmm
0.5
113
Y5
0.00
0.2
647mmmm
0.23mm0
0363 m
33
Weol
0.30
0100mmm
00004e 80
64
0. 1/11e0n
006600
8oll md
026mmm:
Sttec
0001000
0.77U00
1301a002
10210006
1000000
15
(12001 80
03029116
030
6380-1077
0.3mmmn
018005
FTS
022mm00
0 (0leno
02755200 m
01944007
013745 60
13
251
02321001
0.3611315
0011000
4 123000
030lemm
0130-1:0
0/21-1
//2eemxx
0d35et 313
Sptec
02021181
149611
0UBLERS
8:221000
012210001
(197.1813
6.17100111
25

---

Published as a conference paper at ICLR 2024
Table 15: Alternative Anchoring Strategies. Here, we consider an alternative anchoring formulation for graph classification. Namely, instead of shuffling features across the batch (denoted Batch in the table), we perform READOUT anchoring by fitting normal distribution over the hidden representations. We then randomly sample from this distribution create anchors. Conceptually, this is similar to the node feature anchoring strategy. One potential direction of future work that is permitted by this formulation to optimize the parameters of this distribution given signal from an appropriate auxiliary task or loss. For example, we could perform alternating optimization where the GNN trained to minimize the loss, and the mean and variance of the anchoring distribution optimized to minimize the expected calibration error separate calibration dataset. While a rigorous formulation left to future work, we emphasize that the potential for improving the anchoring distribution, and thus controlling corresponding hypothesis diversity, is in fact unique benefit of G-AUQ.
BiNk
200 81000
048
0.80
0221
1100
414
1000
5017
510 55155
0113
530
1500
130 $4100
0311 1010
LO8
26

---

Published as a conference paper at ICLR 2024
A.9 POST-HOC CALIBRATION STRATEGIES
Several post hoc strategies have been developed for calibrating the predictions of a model. These have the advantage of flexibility, as they operate only on the ouiputs of a model and do not require that any changes be made to the model itself. Some methods include:
Temperature scaling (TS) (Guo et al., 2017) simply scales the logits by a temperature parameter T to smooth the predictions. The scaling parameter T can be tuned on validation set.
Ensemble temperature scaling (ETS) (Zhang et al., 2020) learns an ensemble of lemperature-scaled predictions with uncalibrated predictions (T' 1) and uniform proba- bilistic outputs (T oc).
Vector scaling (VS) Guo et al. (2017) scales the entire output vector of class probabilities, rather than just the logits.
Multi-class isotonic regression (IRM) (Zhang et al., 2020) is multiclass generalization of the famous isotonic regression method (Zadrozny Elkan, 2002)): it ensembles predictions and labels, then learns a monotonically increasing function map transformed predictions to labels.
rnder-invariant calibration (Rahimi al., 2020) uses a neural network to learn an intra- order-preserving calibration function that can preserve model's top-k predictions.
Spline calibration instead uses splines to fit the calibration function (Gupta et al., 2021).
Dirichlet calibration (Kull et al., 2019) models the distribution of outputs using a Dirichlet distribution, using simple log-transformation of the uncalibrated probabilities which then passed to a regularized fully connected neural network layer with softmax activation.
For node classification, some graph-specific post-hoc calibration methods have been proposed. CaGCN (Wang et al., 2021) uses graph structure and an additional GCN to produce node-wise temperatures. GATS (Hsu et al., 2022) extends this idea by using graph attention to model the influence of neighbors' temperatures when learning node-wise temperatures. We use the post hoc calibration baselines provided by Hsu et al. in our experiments.
All the above methods, and others, may applied to the output of any model including one using G-AUQ. As we have shown, applying such post hoc methods to the outputs of the calibrated models may improve uncertainty estimates even more. Notably, 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; 7; Ovadia et al., 2019). While poorly calibrated Cls are over/under confident in other their predictions, calibrated CIs more trustworthy and can also improve performance safety-critical tasks which implicitly require reliable prediction probabilities (see Sec. 5). We report the top-1 label expected calibration error (ECE) (Kumar et al., 2019; Detlefsen al., 2022). Formally, let p; be the top-1 probability, e; be the predicted confidence, b; a uniformly sized bin in [0, 1]. Then,
ECE =
by -p)ll
27

---

Published as a conference paper at ICLR 2024
A.10 DETAILS ON GENERALIZATION GAP PREDICTION
Accurate estimation of the expected generalization error on unlabeled datasets allows models with unacceptable performance 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-Ievel scores, S(x;) 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: Prrorr := AcCtageet XX E. I(S({x;; F) > T)ll where T is tuned by minimizing GEP eπror on the validation dataset. We use the confidences obtained by the different baselines as sample-level scores, S(x;) corresponding to the model's 's expectation that a sample is correct. The MAE between the estimated error and true error reported on both in- and out-of -distribution test splits provided by the GOOD benchmark.
A.11 RESULTS ON GENERALIZATION ERROR PREDICTION
GEP Experimental Setup. GEPs (Garg et al., 2022; Ng et al., 2022; Jiang al., 2019; Trivedi et al., 2023a; Guillory et al., 2021) aggregate sample-level scores capturing a model's uncertainty about maximum softmax
the correctness of prediction into dataset-level error estimates. Here, probability for scores and a thresholding mechanism as the GEP. (See Appendix A.10 for more details.) We consider READOUT anchoring with both pretrained and end-to-end training, and report the mean absolute eπror between the predicted and true target dataset accuracy on the OOD test split.
GEP Results. As shown in Table 16, both pretrained and end-to-end G-AUQ outperform the vanilla model on 7/8 datasets. Notably, we see that pretrained G-AUQ is particularly effective as it obtains the best performance across 6/8 datasets. This not only highlights its utility strategy for improving uncertainty estimates without sacrificing accuracy, but also emphasizes that importance of structure, lieu of full stochasticity, when estimating GNN uncertainties.
flexible, light- weight
Table 16: GOOD-Datasets, Generalization Error Prediction Performance. The MAE between the predicted and true test error the OOD test split is reported. G-AUQ variants outperform BOODDooo((Basis.Covarr
being the exception). Pretrained G- AUQ is particularly effective at this task as it achieves the best performance overall 6/8 datasets. Promisingly, we see that regular G-AUQ improves performance over the vanilla model on 6/8 datasets (even if it is not the best overall). We further observe that performing generalization error prediction is more challenging under covariate shift than concept shift on the GOODCMNIST. GOODMotii((asis)) and GOODMotif(Size) datasets. On these datasets, the MAE is almost twice as large than their respective concept shift counterparts. across methods. GOODSST2 is the exception, where concept shift is in fact more challenging. To the best our knowledge. we are the first to investigate generalization eтror prediction on GNN-based tasks under distribution shift. Understanding this behavior further an interesting direction of future work.
vanilla models on 7/8 datasets
CMNIST (Coloe)
MecifLPE (Blds)
ModifL (Size)
SS12
Method
Cencepti.)
(1)
Concepai))
Conceptt.)
Covariate(L)
Cancept(l)
Covsriate
0.089
0.045
0.570
0.507
0.044
0.190
0.072
0022
0.007
0.528
0.066
G-SUQ
0.018
0.012
0.016
0.114 0.004
28

---

Published as a conference paper at ICLR 2024
A.12 ADDITIONAL STUDY ON PRETRAINED G-AUQ
For the datasets and data shifts on which we reported out-of-distribution calibration error of pretrained vs. in-training G-AUQ earlier in Fig. 4, now report additional results for in-distribution and out-of distribution accuracy as well as calibration error. We also include results for the additional GOODMotif-base benchmark for completeness, noting that the methods provided by the original benchmark (Gui et al., 2022) generalized poorly to this split (which may be related to why G- AUQ methods offer little improvement over the vanilla model.) Fig. 8 shows these extended results. By these additional metrics, we again see the competitiveness of applying G-AUQ to pretrained model versus using it in end-to-end training.
GOODSST2
GOODCNNIST
GOODMSIUPE
GOODWARPPE
lenghhecneept
1,0
100
1.0
1.5
0.75
0.75
50.00
0.50
0.50
80.2
0.26
0.00
0.00
0.0
0.
0.00
0.00
0.04
0.22
0.10
10.10
0.84
0:44
0.00
0.02
0.01
0.02
602
0.
0.00
0.00
0.00
0.00
0.6
1.0
0.4
6.76
0.c
0.50
0.4
0.2
0.2
0.26
0.0
0.0
0.0
60
6.0
0.00
0.0
0.00
06
ue
0.3
000 0006
0.075
0.10
0.030
0.2
0.10
0.2
02
0.026
0.06
0.0
Van
22EP
0.0
Von ExEPrrt
0.000
Von. ExEProt
60 Van. 225 Pros
60 Van 020 Prof:
Van E2E Prot:
0.0
Van.
E2EP Pxy. 0.00
Von. ExE
Figure Evaluating Pretrained G-AUQ. Here, we report the performance of pretrained G- AUQ models end-to-end and vanilla models with respect in-distribution and out-of-distribution accuracy as well as expected calibration error. With the exception of the GOODMotif (basis) dataset, pretrained G-AUQ improves the OOD ECE over both the vanilla model and end-to-end G-AUQ comparable or improved OOD accuracy on 7/8 datasets. Furthermore, pretrained G-AUQ also improves the ID ECE on all but the GOODMotif (size) datasets (6/8), where performs comparably to the vanilla model, and maintains the ID accuracy. (We note that all methods are comparably better calibrated on the GOODMotif ID) data than GOODCMIST/GOODSS because there may exist simple shortcuts available in the GOODMotif dataset that can be used on the ID test set effectively.) Overall, these results clearly demonstrate that pretrained G-AUQ does offer some performance advantages over end-to-end G-AUQ and does so at reduced training times (see Table. A.13). For example, on GOODCMNIST (covariate shift), pretrained G-AUQ is not only 50% faster than end-to-end G-AUQ, it also improves OOD accuracy and OOD ECE over both the vanilla and end-to-end G-AUQmodels.
IDD data; we suspect this
29

---

Published as a conference paper at ICLR 2024
A.13 RUNTIME TABLE
Table 17: Runtimes. We include the runtimes both training per epoch (in seconds) and performing calibration. Reducing stochasticity can help reduce computation (L1 L3).. Cost can also be reduced by using a pretrained model.
GOODCMNIST
GOODSST2
GOODMetifLPE
TTriining
Training (S)
Inference
Training (S)
Inference
Vanilia
18.5
25.8
10.8
18.5
3.8
4.5
Temp. Scaling
23,5
10.8
13.4
DEns (Ens Size=3)
18.456 Ens Size
59.4
10.795 x Ens Size
29.0
3.8 Ens Size
11.8
G-AUQ (L1. anchors)
22.1
181.5
15,9
17.1
15.5
G-AUQ (1.2, 10)
22.4
148.6
12.7
15.5
5.8
11.8
G-AUQ (HiddenRep.
18.5
28.0
13.8
19.6
G-AUQ (Pretr HiddenRep, 10)
8.6
27.8
6.8
16.0
2.5
6.4
30

---

Published as a conference paper at ICLR 2024
A.14 MEAN AND VARIANCE OF NODE FEATURE GAUSSIANS
Table 18:
Mean and Variance of Node Feature Anchoring Gaussians. We report the mean and variance of the Gaussian distributions fitted to the input node features. Because the input node features vary in size, we report aggregate statistics over the mean and variance corresponding to each dimension. For example, Min(Mu) indicates that are reports the minimum mean over the d-dim set of means.
Dataset
Domain
Shift
Min (Mu)
Max (Mu)
Mcan (Mu)
Sid. (Mu)
Min (Std)
Max (Std)
Mcan (Sid)
Std (Sid)
GOODSST2
length
concept
-4.563
0.69
40.011
0.278
0,163
0,803
0.242
0.049
GOODSST2
length
covariato
-4.902
0.684
-0.01
0.3
0.175
0.838
0.255
0.05
GOODCMNIST
color
concept
0.117
0.133
0.127
0.008
0.092
0.097
0.095
0.003
GODDCMNIST
0.102
G00DMoifT.PE
size
covariate
0.003
0.021
0.011
0.008
0.835
1.728
1.248
0.377
GO
00.06
+0002
0.003
0.542
1.114
0.783
0.242
concept
GOODMoifT.PE.
basis
covariate
0.007
40.002
-0.044
0.002
0.808
1.13
1:251
0.469
GODDWehKB
university
concept
0.95
0.049
0.099
0.001
0.095
GOODWebKB
university
covariatc
0.934
0.104
0.001
0.164
0.098
GOODCora
degree
concept
0.507
0.007
0.017
0.001
0.5
0.061
0.051
GOODCora
degree
GOODCBAS GOODCBAS
color
covariate
0.394
0.591
0.471
0.083
0.142
0.492
0.403
0.174
color
concept
0.23
0.569
0.4
144
0,168
0.495
0.39
0.152
GOODCMNIST Color, Concept: Anchoring Distribution vs. Input Features
Dim: 0
Dim: 1
Dim: 2
0.0
0.3
Figure 9: GOODCMNIST, Concept, Anchoring Distribution. We plot the mean and variance of the fitted anchoring distribution vs. the true feature distribution for each input dimension. We observe there is mismatch between the empircal distribution and the fitted Gaussian. However, did not find this mismatch to harm the effectiveness of G-AUQ.
GOODCMNIST Color, Covariate: Anchoring Distribution Input Features
Dim: 0
Dim:
Dim:
-0.2
0.0
0.2
0.4
-0.2
0.0
0.2
0.4
-0.2
0.0
0.2
0.4
Figure 10: GOODCMNIST, Covariate, Anchoring Distribution. We plot the mean and variance of the fitted anchoring distribution the true feature distribution for each input dimension. We observe there is mismatch between the empircal distribution and the fitted Gaussian. However, we not find this mismatch harm the effectiveness of G-AUQ.
31

---

Published as a conference paper at ICLR 2024
Table 19: Number of Parmeters per Model.We provide the number of parameters in the vanilla and modified parameter as follows. Note, that the change in parameters is architecture and input dimension dependent. For example, GOODCMNIST, and GOODSST2 use GIN MPNN layers. Therefore, when changing the layer dimension, are changing the dimension of its internal MLP. not an error that intermediate layer G-AUQhave the same number of parameters, this due to the architecture: these layers are the same size in the vanilla model. Likewise, GOODCora's 's input features have dimension is 8701, so doubling the input layer's dimension appears to add signficant number of parameters. We do not believe this
Dutaset
GOODCMNIST
GOODMotif
GOODSST2
GOODCORA
GOODWebKB
GOODCBAS
Baseline
2001310
911403
1732201
2816770
695105 1206005
185104
G-AUQ(NEA)
2003110
913803
2193001
5429770
186304
G-AUQ(LI)
2360110
1633203
2091001
G-AUQ(L2)
23601 10
1633203
2091001
G-AUQ(L3)
2360110
G-AUQ(L4)
2360110
G-AUQ(Readuut)
2004310
912303
1732501
32

---

Published as a conference paper at ICLR 2024
A.15 EXPANDED DISCUSSION ON ANCHORING DESIGN CHOICES
Below, we expand upon some of the design choices for the proposed anchoring strategies.
When performing node featuring anchoring, how does fitting Gaussian distribution to the input node features help manage the combinatorial stochasticity induced by message passing?
Without loss of generality, consider node classification setting, where every sample is assigned unique anchor. Then, due to message passing, after hops, given node's representation will have aggregated information from its / hop neighborhood. However, since each node in this neighborhood has a unique anchor, we see that any given node's representation is not only stochastic due to its own anchor but also that of its neighbors. For example, if any of its neighbors are assigned different anchor, then the given node's 's representation will change, even if its own anchor did not. Since this behavior holds true for all nodes and each of their respective neighborhoods, we loosely refer to this phenomenon having combinatorial complexity, as effectively marginalizing out the anchoring distribution would require handling any and all changes to all (-hop neighbors. In contrast, when performing anchored image classification,, the representation of sample is only dependent unique, corresponding anchor, and is not influenced by the anchors of other samples. To this end, using the fitted Gaussian distribution helps manage this complexity, since changes to the anchors of node's 's I-hop neighborhood are simpler to model as they require only learning to marginalize out a Gaussian distribution (instead of the training distribution). Indeed, for example, if were to assume simplified model where message passing only summed node neighbors, the anchoring distribution would remain Gaussian after rounds of message passing since the sum of Gaussian is still Gaussian (the exact parameters of the distribution would depend on the normalization used however).
its
33


Reviews:
[{"review_id": "A2FK11Or34", "reviewer": "Reviewer_b8ZS", "summary": "The paper proposes a training framework for GNNs that is designed to improve the intrinsic uncertainty estimates. The adopted strategy is to adapt the principle of stochastic data centering to graph data. This involves introducing novel graph anchoring techniques. The paper demonstrates that the methodology can support partially stochastic GNNs. Experimental results in the paper suggest that the partial stochasticity is sufficient; it also has the advantage of providing a mechanism for incorporating pre-trained models. The paper reports experiments investigating the impact of covariate, concept and graph size shifts and demonstrates that the proposed technique leads to better calibrated GNNs, both for node and graph classification. Additional experiments illustrate how the approach performs for out-of-distribution detection and generalization gap estimation.", "questions": "Q1. Why does the proposed method appear to work relatively poorly (at least in terms of often making ECE worse and sometimes considerably worse) when used in conjunction with Dirichlet calibration in the node classification setting?\nQ2. Figure 2 is challenging to understand. Insufficient detail is provided. What is the meaning of L1, L2, L3, N/A? What does it mean “generally performs well across datasets and architectures” – which bars am I comparing to draw this conclusion? The accuracy and ECE behaviour seems quite inconsistent across different datasets and architectures – sometimes increasing, sometimes decreasing, sometimes going up then down. How do I determine when something “performs well”?\nQ3. “We emphasize that this distribution is only used for anchoring and does not assume that the dataset’s node features are normally distributed.” This sentence seems to imply that it does not matter if the anchoring distribution matches the data distribution or not. But then why fit the Gaussian distribution to the training data? The fitting procedure implies that there is a need to have an anchor distribution that matches (at least to some degree) the feature distribution. So how close does it need to be to a Gaussian distribution?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1. The paper introduces a novel approach for improving the intrinsic uncertainty estimates of GNNs by translating the stochastic centering strategy to the graph domain. This is non-trivial, both for node and graph classification.\n \nS2. The paper reports on experiments exploring (i) node classification under distribution shift (concept and covariate shifts); (ii) calibration under distribution shift for graph classification; (iii) how the approach impacts the calibration of more expressive models such as graph transformers. The experiments are thorough and examine multiple interesting questions. \n\nS3. The experiments support the interesting observation that the network does not need to be fully stochastic in order to provide improved uncertainty estimates. This paves the way to combine the p", "weaknesses": "W1.\tSome of the methods employed to translate stochastic centering to the graph domain appear somewhat heuristic, or are at least the text describing the methodology does not provide sufficient detail to perceive the design principles. For example, the node feature anchoring fits a Gaussian distribution, but there is no explanation as to why a Gaussian is selected and no discussion as to whether the mismatch between the fitted anchor distribution and the feature distribution has a negative impact or could be a concern. The text states that the introduction helps to “manage the combinatorial stochasticity induced by message passing” but it does not elaborate on this to explain why or how. For the graph anchoring, there is a random shuffling of the node features over the entire batch. There is no discussion of this design choice – it doesn’t seem obvious to me that this is the only thing one could choose to do (or the optimal). \n\nW2.\tDistilling the methodological contributions, we see that they involve: (i) node anchoring via fitting a Gaussian distribution and drawing an anchor from this fitted distribution; (ii) hidden layer anchoring by randomly shuffling the node features after the r-th layer. After these steps to construct appropriate anchors in the graph domain, there is effectively a standard application of the stochastic centering approach. The technical methodological contribution is thus not particularly substantial. On the other hand, the experiments are thorough and provide a good balancing contribution. \n\nW3.\tSome of the results are not presented in a particularly helpful manner and are not described or discussed in much detail. For example, the observations for node classification essentially boil down to “our method works better”. The table contains interesting elements such as the proposed method failing to improve (WebKB, CBAS – Concept) or substantially increasing (Cora) the ECE when combined with Dirichlet calibration. But there is no discussion of this. In general there is not a significant attempt to draw detailed conclusions from the obtained results – similar comments apply to Section 5.3 where again the conclusions are “both pretrained and end-to-end G-∆UQ outperform the vanilla model on 7/8 datasets” and “G-∆UQ variants improve OOD detection performance over the vanilla baseline on 6/8 datasets”. Insights beyond “works better” make a paper much stronger and more insightful. \n\nThere is room for improvement in some of the figures and the explanations of how they are being interpreted. Figure 2 is a particular case in point – L1, L2, L3, N/A are not clearly defined. The text states that “READOUT anchoring generally performs well” but does not explain how we should interpret the figure to come to this conclusion. It’s not obvious what “performs well” means – what is an acceptable deterioration in performance. The behaviour over datasets and architectures differs considerably and should be discussed. \n\nIn terms of assessing the variability of performance, the paper reports standard deviations over a few trials, but does not make any attempt to assess the statistical significance of the results or to specify confidence intervals on the reported means.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bbGAJl7DF7", "reviewer": "Reviewer_bgJx", "summary": "The paper proposed a new training framework inspired by the stochastic anchored training in computer vision. In the framework two (partial) stochastic anchoring techniques (node feature anchoring and hidden layer anchoring) are designed for sake of the better uncertainty estimation and calibration in GNNs. The experiments on node classification and graph classificaiton validate the effectiveness of the framework.", "questions": "See weakness", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The idea that introducing the anchoring training in GNNs is new and interesting.\n\n2. The paper conducted comprehensive evalutions on the calibration of GNNs under different settings. Specifically, the evaluation on calibration of GNNs under distribution shift is new and significant. The experimental results show that on most tasks the proposed method can achieve lower calibration error.", "weaknesses": "1.Since the paper focused on calibration and uncertainty estimation of GNNs. The concepts of calibration and uncertainty estimation of GNNs should be introduced in the paper.\n\n2.The paper doesn't provide sufficient discussion or theories to justify the methods provided in the paper. For instance,  in node feature anchoring why authors sample the anchors from the Gaussian Distribution? How to determine the value of $u$ and $\\sigma$?\n\n3.The paper claimed that the framework can improve the uncertainty estimation. However, how the method can improve the uncertainty estimation is still unclear.\n\n4.The organization of the experiment part is messy. More setup details should be clarified.\n\n5.Some typos in the paper and the title of Table 2 is missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XteOB9wzsw", "reviewer": "Reviewer_Cjvm", "summary": "This paper extends the ∆-UQ model to graph learning to enhance intrinsic GNN uncertainty estimates. It proposes several training techniques to introduce partial stochasticity for node and graph-level classification tasks. Extensive experimental results validate the effectiveness of the proposed method. Interestingly, the graph-level training method can be applied to pre-trained graph models.", "questions": "1. In Section 2 (Preliminaries), the notation $\\\\mathbf{X}$ is somewhat unclear. It seems to represent some distributions, but there is also the concatenation $[\\\\mathbf{X} - \\\\mathbf{C}, \\\\mathbf{C}]$, which is described as \"channel-wise concatenating two images\". Please clarify the notations regarding distributions, samples and training (test) sets.\n\n2. On page 3 (Section 3.1 Node Feature Anchoring), where are $\\\\mathbf{A}_c$ and $Y_c$ defined?\n\n3. On page 4, \"the anchor/query node feature pair\" is defined as $[\\\\mathbf{X}_i − \\\\mathbf{C}|| \\\\mathbf{X}_i]$, where $\\\\mathbf{C}^{N \\\\times d} \\\\sim \\\\mathcal{N} (\\\\mu, \\\\sigma^2)$. How are the shapes of $\\\\mathbf{X}_i$ and $\\\\mathbf{C}$ aligned?\n\n4. For the node classification task, anchoring is sampled from a learned Gaussian. The notations used seem unclear to the reviewer. Could you elaborate on the learning procedures, the loss function used, and the dimensions of the Gaussians? During inference, stochasticity arises from sampling from such Gaussians. Is it possible to display qualitative results of the learned Gaussians, including details such as variance and mean?\n\n\n5. a) For the graph classification task, there seems to be no extra learnable component to create the anchors, but the subsequent MPNN layers are modified to accept features of $d_r \\\\times 2$ dimensions. How does this subtle change affect the number of trainable parameters?\\\nb) The stochasticity appears to arise from the random shuffle of the node features across the entire batch, and the subsequent selection of anchoring $\\\\mathbf{c}^{1 \\\\times d}$. Does creating anchors from a Gaussian distribution work in this case?\\\nc) Conversely, is it feasible to use intermediate representations to construct anchors for the node classification task? Technically, it seems like these could substitute for the learned Gaussian.\\\nd) Furthermore, when randomly shuffling the node features across the entire batch (potentially involving cross-graph shuffling), are there observable impacts attributed to the size of the mini-batch?\n\n6. Can this method be generalized to graph/node/edge level regression tasks?\n\n7. The design of shuffling across batches in this paper appears conceptually somewhat similar to the mixup method [1] in the feature space. A comparison between them would be interesting.\n\n8. In Fig. 6, Tables 4/5, the effects of applying G-∆UQ on top of the pretrained model (GPS) seem very marginal.\n\n[1] mixup: Beyond Empirical Risk Minimization, ICLR 2018.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Extensive experimental results on various datasets and setups demonstrate the effectiveness of the proposed method.\n\n2. The introduced method is fairly simple and easy to understand. It seems that it can be applied to many existing models in a plug-and-play manner, making it flexible and extendable.\n\n3. The paper also explores the potential of applying this method to pre-trained models, which is an interesting aspect.", "weaknesses": "1. Some notations and definitions in the paper are unclear, such as the symbols used for samples and distributions, and the shape of tensors during concatenation.\n\n2. The experiments conducted focus solely on classification tasks. Consideration for regression tasks seems to be missing.\n\n3. There are limited performance gains observed with the use of pre-trained models.\n\nThese points are further detailed below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)