Paper: Accurate and Scalable Estimation of Epistemic Uncertainty for Graph Neural Networks (ICLR-2024)

Reviewer comment (Writing):
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.

Author response:
Thanks for the comments. We’ve updated the main paper with better notation. Please see sections 2 and 3 for the updated notations, highlighted in blue. We summarize our anchoring strategies with the updated notation below.  

Let a graph classification GNN consisting of  $( 1 \dots \ell)$ message passing layers $(\text{MPNN})$, a graph-level readout function ($\text{READOUT}$), and classifier head ($\text{MLP}$) is defined as follows: $\mathbf{X}^{\ell} = \text{MPNN}^{\ell} \left(\mathbf{X}^{\ell - 1}, \mathbf{A} \right)$, $\mathcal{G}=  \text{READOUT} \left(\mathbf{X}^{\ell} \right)$, and $\hat{Y} =\text{MLP}\left( \mathcal{G} \right)$ where $\mathbf{X}^{\ell} \in \mathbb{R}^{N  \times d_{\ell}}$ is the intermediate node representation at layer $\ell$, $\mathcal{G} \in\mathbb{R}^{1 \times d_{\ell}}$ is the graph representation, and $\hat{Y} \in \{0,1\}^q$ is the predicted label. When performing node classification, the $\text{READOUT}$ layer is not included, and instead node-level predictions are outputted: $\hat{Y}=\text{MLP}\left( \mathbf{X}^{\ell}\right)$. Then, we can define the anchoring strategies as follows. *(Note, we show an abrdiged version here, please see the text for full details.)*

- **Node feature anchoring:** $\mathbf{X}^{1} = \text{MPNN}^{1}(\mathbf{X}^0 - \mathbf{C} || \mathbf{C}, \mathbf{A})$, where the dimension of the $\text{MPNN}^1$ has been changed to accommodate $d \times 2$ dimensional input.

- **Hidden layer anchoring:** Let $r > 2$ be the layer at which anchoring is applied. Then, $\mathbf{X}^{r-1} = \text{MPNN}^{r-1}(\mathbf{X}^{r-2}, \mathbf{A})$ and $\mathbf{C}=\text{SHUFFLE}(\mathbf{X}^{r-1}, dim=0)$ , where  $\mathbf{X}^{r} = \text{MPNN}^{r}(\mathbf{X}^{r-1} - \mathbf{C} || \mathbf{C}), \mathbf{A})$ is the anchored representation. Here, $\text{MPNN}^r$ ‘s input dimension is changed to $d_{r-1} \times 2$

- **READOUT anchoring:** $\mathbf{G} = \text{READOUT}(\text{MPNN}^{1 \ldots \ell}(\mathbf{X}^0, A),BatchVec)$, where $\mathbf{C}=\text{SHUFFLE}(\mathbf{G}, dim=0)$ and the anchored representation is defined as  $\mathbf{Y} = \text{MLP}(\mathbf{G} - \mathbf{C} || \mathbf{C})$.

Reviewer comment (Evaluation):
There are limited performance gains observed with the use of pre-trained models.

Author response:
Thanks for the feedback. We’ve expanded our discussion in the paper to make the benefits of the pre-trained G-ΔUQ clearer, and discuss key takeaways here.

- **Comments about Fig 6.** With the exception of the GOODMotif (basis) dataset, *pretrained G-ΔUQ improves the OOD ECE over both the vanilla model and end-to-end G-ΔUQ at comparable or improved OOD accuracy (6/8 datasets)*. Furthermore, pretrained G-ΔUQ also improves the *ID ECE on all but the GOODMotif (size) datasets (6/8)*, where it performs comparably to the vanilla model, and maintains the ID accuracy.  This demonstrates that pretrained G-ΔUQ does offer some performance advantages over end-to-end G-ΔUQ, though we agree the absolute benefits may appear a bit small.
 
    However, this assessment did not take into account the computational benefits and the flexibility of being able to use a pre-trained backbone. Indeed, only an anchored classifier is trained for the pretraining setting; this substantially reduces training time relative to end-to-end G-ΔUQ, *where training time is ~40% less than vanilla or node-feature G-ΔUQ training (We’ve added new runtime results in App., Table 17 to illustrate this.)* Moreover, as larger, pretrained models become widely available, pretrained G-ΔUQ offers an effective, lightweight solution for improving uncertainty estimates (as evidenced by the improvements over the Vanilla model in Fig 6.). We highlight that this is a unique benefit of our stochastic centering formulation.

- **Comments about GPS:**  Tables 5 and 6 contain several axes (Architecture, G-ΔUQ vs. Vanilla, LPE vs w/o LPE, Ensemble vs. Single Model, Pretrained vs. End-to-End) that makes it difficult to immediately parse the benefits of Pretraining with respect to GPS. Thus, *we have added a new table (Tables 6/7 in the updated appendix) where we pull out the comparable rows and see that pretrained G-ΔUQ clearly improves the calibration at comparable accuracy relative to the Vanilla GPS model, especially as the distribution shift becomes larger.* We’ve updated the paper text to better discuss these axes.

Reviewer comment (Writing):
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.

Author response:
Thanks for the comments. We’ve updated the main paper with better notation. Please see sections 2 and 3 for the updated notations, highlighted in blue. We summarize our anchoring strategies with the updated notation below.  

Let a graph classification GNN consisting of  $( 1 \dots \ell)$ message passing layers $(\text{MPNN})$, a graph-level readout function ($\text{READOUT}$), and classifier head ($\text{MLP}$) is defined as follows: $\mathbf{X}^{\ell} = \text{MPNN}^{\ell} \left(\mathbf{X}^{\ell - 1}, \mathbf{A} \right)$, $\mathcal{G}=  \text{READOUT} \left(\mathbf{X}^{\ell} \right)$, and $\hat{Y} =\text{MLP}\left( \mathcal{G} \right)$ where $\mathbf{X}^{\ell} \in \mathbb{R}^{N  \times d_{\ell}}$ is the intermediate node representation at layer $\ell$, $\mathcal{G} \in\mathbb{R}^{1 \times d_{\ell}}$ is the graph representation, and $\hat{Y} \in \{0,1\}^q$ is the predicted label. When performing node classification, the $\text{READOUT}$ layer is not included, and instead node-level predictions are outputted: $\hat{Y}=\text{MLP}\left( \mathbf{X}^{\ell}\right)$. Then, we can define the anchoring strategies as follows. *(Note, we show an abrdiged version here, please see the text for full details.)*

- **Node feature anchoring:** $\mathbf{X}^{1} = \text{MPNN}^{1}(\mathbf{X}^0 - \mathbf{C} || \mathbf{C}, \mathbf{A})$, where the dimension of the $\text{MPNN}^1$ has been changed to accommodate $d \times 2$ dimensional input.

- **Hidden layer anchoring:** Let $r > 2$ be the layer at which anchoring is applied. Then, $\mathbf{X}^{r-1} = \text{MPNN}^{r-1}(\mathbf{X}^{r-2}, \mathbf{A})$ and $\mathbf{C}=\text{SHUFFLE}(\mathbf{X}^{r-1}, dim=0)$ , where  $\mathbf{X}^{r} = \text{MPNN}^{r}(\mathbf{X}^{r-1} - \mathbf{C} || \mathbf{C}), \mathbf{A})$ is the anchored representation. Here, $\text{MPNN}^r$ ‘s input dimension is changed to $d_{r-1} \times 2$

- **READOUT anchoring:** $\mathbf{G} = \text{READOUT}(\text{MPNN}^{1 \ldots \ell}(\mathbf{X}^0, A),BatchVec)$, where $\mathbf{C}=\text{SHUFFLE}(\mathbf{G}, dim=0)$ and the anchored representation is defined as  $\mathbf{Y} = \text{MLP}(\mathbf{G} - \mathbf{C} || \mathbf{C})$.

Reviewer comment (Reproducibility):
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?

Author response:
Thanks for the comments! We *updated the paper with Pytorch-Style pseudo-code (Figures 1,2 and Appendix A.1) to more intuitively explain the learning procedures, and clarified the notations to make the text easier to follow (Sec 3).*

In brief, G-∆UQ does not change the loss function when training. Thus, we use the CrossEntropyLoss to train both anchored and vanilla models. Anchors are uniquely re-sampled for training samples at each epoch (each sample gets a new anchor).  At inference, we average predictions over $K$ random anchors, and use the prediction variance over the $K$-predictions to estimate uncertainty (as shown in Sec 2.). When performing node feature anchoring, the dimension of the Gaussian is determined by the dataset’s node feature dimension. We’ve updated Table 9 in the appendix to include this information.

We wish to clarify that the Gaussians obtained from the node features are fixed at initialization and are not updated (or learned) during training. For example, when performing node classification, let $\mathbf{X}^{N \times d}$ be the set of training node features. We compute the mean, $\mu$, and standard deviation, $\sigma$ for each dimension $d$ independently, and then initialize a $d$ dimensional Normal with mean $= [\mu_1 \dots \mu _d]$ and scale $= [\sigma_1 \dots \sigma_d]$.  We can then straight-forwardly sample from this distribution when training to obtain anchors.

For graph classification, this procedure is slightly modified since each sample (graph) has its own set of node features. Therefore, we first create the set of node features in the dataset by aggregating node features from individual graph samples. In detail, let $\mathcal{D} = {\mathcal{G}_1 \times \mathcal{G}_t}$ be the training dataset, where each $\mathcal{G}_i$ has $\mathbf{X}^{N \times d}$ node features. The aggregated set of node features would be of dimension of $(\sum^t N_i \times d)$. Given this set, the same procedure as above is used to calculate the normal.

Since the dimension of the node feature anchoring Gaussians can vary per dataset, we provide summary statistics computed over the per-dimension means and standard deviations, and have included them in the appendix (Table 18).

Reviewer comment (Reproducibility):
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?

Author response:
You are correct that there are no extra learnable components for creating the anchors. Only the dimension of the layer at which anchoring is applied is modified, namely by doubling its input dimension to accommodate the anchor and residual sample. For example, if anchoring is applied to a 4-layer GNN at the 2nd MPNN layer, we modify MPNN_2 ’s input dimension, but $\text{MPNN}_{1,3,4}$ and $\text{MPNN}_2$ output dimension remain the same.

The exact change to the number of parameters depends on the underlying architecture, the selected intermediate layer, and input dimension . For example, changing the dimension of a GIN layer will require changing the dimension of its MLP, while READOUT anchoring only changes the dimension of the classifier (a single linear layer in our experiments).  *We include this information in the appendix (Table 19). Note that while the number of parameters increases, the task has also increased a bit in difficulty, so we do not consider this increase as affecting the model’s capacity.*

Reviewer comment (Experiments):
In the graph classification task, 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?

Author response:
It is indeed conceptually possible to create a Gaussian distribution over the intermediate representations that can be used to perform hidden representation anchoring on graph classification tasks. *We’ve added some preliminary results to Tables 15 in the appendix evaluating this formulation.* We find that the performance is generally comparable with respect to our original shuffling-based formulation and indeed, there are some improvements. Performing automatic anchoring strategy selection is an interesting direction of future research.

Reviewer comment (Experiments):
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.

Author response:
It is indeed feasible to use intermediate representations to construct anchors when performing node classification in a manner similar to Sec. 3.2. *Below and in the updated appendix, we provide some preliminary results demonstrating this is possible and that G-∆UQ is able to provide partially stochastic models for node classification.* We note that we focused on using node featuring anchoring for node classifications tasks as this formulation is most closely related to the original stochastic centering formulation [1].
In brief, [1], which focuses on image classification tasks, uses input-space transformations (image-image subtraction, and channel wise concatenation) to create anchored representations and fully stochastic models, where each anchored image constitutes a single sample. Similarly, with node classification, each node can be considered a “sample” since the loss is computed with respect to its predicted label. Then, by using node feature anchoring, we can also perform input-space transformations to create per-sample anchored representations and fully stochastic GNNs.

[1] Single Model Uncertainty Estimation via Stochastic Data Centering, Thiagarajan et al. NeurIPS'22

Reviewer comment (Experiments):
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?

Author response:
We used the same batch size as indicated by the original benchmark code-bases (GOODBenchmark,GPS, SizeShiftReg) to ensure fair comparison with a well-tuned vanilla baseline. *We do not anticipate negative impacts caused by cross-graph shuffling or observable impacts attributed to the size of mini-batch.* Indeed, in our present formulation, cross-graph shuffling is already occurring as we use PytorchGeometric style mini-batching (which stacks the node-features of different graphs). Reducing the batch-size would lead to less cross-graph shuffling. However, there is also increased risk of training instability due to noisier gradient updates, and cross-graph shuffling is not in fact harmful to G-ΔUQ’s performance.

Reviewer comment (Experiments):
Can this method be generalized to graph/node/edge level regression tasks?

Author response:
Applying G-ΔUQ to regression tasks is a natural extension and our model formulation does conceptually permit it. However, since regression tasks require different definitions of uncertainty to account for real-valued outputs, we suspect that our inference procedure and calibration strategy will need to be changed. This is definitely part of our future work; thanks for the suggestion! Lastly, we briefly note that while we focus on graph and node classification tasks in the paper due to the readily available distribution shift benchmarks, G-ΔUQ can in fact be easily extended to link prediction by using a similar formulation to node feature anchoring. The main difference would be that both the “head” and “tail” nodes of the predicted link would be anchored, prior to being fed into the binary link-predictor head.

Reviewer comment (Experiments):
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.

Author response:
Mixup, which trains models on convex combinations of pairs of samples and labels, regularize neural networks to prefer simple linear behavior, which can help generalization, robustness and create smoother uncertainty estimates. While G-ΔUQ also conceptually combines pairs of samples, our overarching goals and formulation are fundamentally different. Importantly, mixup uses label information, whereas G-ΔUQ is only an input transformation that does not use the label.

Our objective is to provide reliable uncertainty estimates that can then be used in various uncertainty based tasks. *To this end, we use anchoring as a mechanism for sampling different functional hypotheses from a single model.* Sampling different hypotheses allows us to obtain reliable epistemic uncertainty estimates that can then be used to improve confidence estimates on a variety of downstream uncertainty based tasks (Calibration, OOD Detection, Generalization Gap Prediction). *In contrast, mixup is a data augmentation strategy, so it does not offer a mechanism for sampling different hypotheses at inference time, and can only regularize the behavior of a single functional hypothesis.* We note that it is conceptually possible to combine G-ΔUQ and mixup to support both the sampling of different functions for obtaining epistemic uncertainty estimates, while also ensuring that these hypotheses exhibit more linear behavior. We thank the reviewer for their suggestion and will conduct this experiment.

[1] mixup: Beyond Empirical Risk Minimization, Zhang et al. ICLR 2018.

Reviewer comment (Evaluation):
In Fig. 6, Tables 4/5, the effects of applying G-∆UQ on top of the pretrained model (GPS) seem very marginal.

Author response:
Thanks for the feedback. We’ve expanded our discussion in the paper to make the benefits of the pre-trained G-ΔUQ clearer, and discuss key takeaways here.

- **Comments about Fig 6.** With the exception of the GOODMotif (basis) dataset, *pretrained G-ΔUQ improves the OOD ECE over both the vanilla model and end-to-end G-ΔUQ at comparable or improved OOD accuracy (6/8 datasets)*. Furthermore, pretrained G-ΔUQ also improves the *ID ECE on all but the GOODMotif (size) datasets (6/8)*, where it performs comparably to the vanilla model, and maintains the ID accuracy.  This demonstrates that pretrained G-ΔUQ does offer some performance advantages over end-to-end G-ΔUQ, though we agree the absolute benefits may appear a bit small.
 
    However, this assessment did not take into account the computational benefits and the flexibility of being able to use a pre-trained backbone. Indeed, only an anchored classifier is trained for the pretraining setting; this substantially reduces training time relative to end-to-end G-ΔUQ, *where training time is ~40% less than vanilla or node-feature G-ΔUQ training (We’ve added new runtime results in App., Table 17 to illustrate this.)* Moreover, as larger, pretrained models become widely available, pretrained G-ΔUQ offers an effective, lightweight solution for improving uncertainty estimates (as evidenced by the improvements over the Vanilla model in Fig 6.). We highlight that this is a unique benefit of our stochastic centering formulation.

- **Comments about GPS:**  Tables 5 and 6 contain several axes (Architecture, G-ΔUQ vs. Vanilla, LPE vs w/o LPE, Ensemble vs. Single Model, Pretrained vs. End-to-End) that makes it difficult to immediately parse the benefits of Pretraining with respect to GPS. Thus, *we have added a new table (Tables 6/7 in the updated appendix) where we pull out the comparable rows and see that pretrained G-ΔUQ clearly improves the calibration at comparable accuracy relative to the Vanilla GPS model, especially as the distribution shift becomes larger.* We’ve updated the paper text to better discuss these axes.

Reviewer comment (Writing):
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.

Author response:
Thanks for the feedback. We have *updated the introduction* so uncertainty estimation and calibration are presented earlier to the reader. (We note that Appendix A.9 also contains information on different post-hoc calibration methods and defines the expected calibration error.)

Reviewer comment (Theory):
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$?

Author response:
Thanks for the questions. We have (i) added *additional discussion in Sec 3*. to explain the goals of the anchoring distribution/strategy to provide context for our design decisions, (ii) improved the description of the proposed strategies, and (iii) added pytorch-style pseudocode *to improve clarity (Fig 1, Fig 2, Fig 5).*

Conceptually, anchoring provides a mechanism for sampling multiple hypotheses from a single model such that their variability can be aggregated to provide reliable epistemic uncertainty estimates. *Our introduced anchoring strategies are specifically designed to (i) handle discrete, variable-sized inputs and (ii) induce diverse hypotheses with controlled stochasticity.* In particular, node feature anchoring is proposed to provide a fully stochastic formulation that is compatible with graph data and message passing neural networks (MPNNs), while Hidden Layer and Readout Anchoring not only support graph data/MPNNs, but also induce partially stochastic models that are able to provide functionally diverse hypotheses. Empirically, we find that balancing functional diversity and stochasticity leads to better performance.

**Hidden Layer and Readout Anchoring:** While the success of model ensembles has been attributed to their ability to leverage diverse hypotheses, it was recently demonstrated that partial stochasticity can perform comparably to fully stochastic networks at significantly less cost [1]. This suggests that in addition to the *amount of diversity*, the *effective or functional diversity* of the sampled hypotheses is also critical for performance. Since hidden layer and readout anchoring are performed after multiple rounds of message passing, the corresponding representations capture both node-feature and topological information. *By defining the residual transformation and anchoring distribution with respect to these intermediate representations, we are able to sample hypotheses from a latent space containing higher-level, disentangled semantics.* Thus, even though the layers prior to anchoring are deterministic, the sampled hypotheses are functionally diverse (e.g., they sample decision rules that rely upon different higher-order features). We empirically find, through extensive experiments across datasets, shifts, tasks and architectures, that this improves performance. Note that hidden layer and readout anchoring are, by design, able to handle discrete and variable-sized input. However, to support fully stochastic networks, we proposed *node feature anchoring.*

[1] Do Bayesian Neural Networks Need To Be Fully Stochastic? Sharma et al. AISTATS'23

*Please see Fig 2 for pseudocode explaining this process.* In brief, when performing node classification, let $\mathbf{X}^{N \times d}$ be the set of training node features. We compute the mean, $\mu$, and standard deviation, $\sigma$ for each dimension $d$ independently, and then initialize a $d$ dimensional Normal with mean $= [\mu_1 \dots \mu _d]$ and scale $= [\sigma_1 \dots \sigma_d]$.  We can then straight-forwardly sample from this distribution to obtain anchors.

For graph classification, this procedure is slightly modified since each sample (graph) has its own set of node features. Therefore, we first create the set of node features in the dataset by aggregating node features from individual graph samples. In detail, let $\mathcal{D} = {\mathcal{G}_1 \dots \mathcal{G}_t}$ be the training dataset,  where each $\mathcal{G}_i$ has $\mathbf{X}^{N \times d}$ node features. The aggregated set of node features would be of dimension of $((\sum^t N_i) \times d))$. Given this set, the same procedure as above is used to calculate the normal.

*Since the dimension of the node feature anchoring Gaussians can vary per dataset, we provide summary statistics over the per dimension means and standard deviations, and have included them in appendix A.14.*

[1] motivated stochastic centering through the lens of neural tangent kernel (NTK), where they showed that constant shifts to the input $(\mathbf{X}-\mathbf{C})$ induce analogous shifts to the effective NTK, i.e. different hypotheses. They then demonstrated that a single (anchored) model can replicate this behavior if the input bias (e.g. the anchor, $\mathbf{C}$)  is treated as a random variable during training (see Fig. 1) and marginalized out during inference. While [1] used the training samples as anchors and computed $(\mathbf{X}-\mathbf{C})$ for image classification, this is not applicable to graph datasets Subtracting two adjacency matrices will introduce connectivity artifacts and subtracting input node features is also not advisable because features are not IID and may belong to different sized graphs. *By using the proposed Gaussian anchoring distribution, we are able to sample anchors for arbitrarily sized graphs, without introducing connectivity artifacts or drastically deviating from the original training landscape (evidenced by similar ID accuracies).*

**Controlling Stochasticity.** Without loss of generality, consider a node classification setting, where every sample is assigned a unique anchor. Then, due to message passing, after $\ell$ hops, a given node’s representation will have aggregated information from its $\ell$ hop neighborhood. *As a result, any given node’s representation is stochastic not only due to its own anchor but also that of its neighbors. Since this behavior holds true for all nodes and each of their respective neighborhoods, we loosely refer to this phenomenon as having combinatorial complexity*, as effectively marginalizing out the anchoring distribution would require handling any and all changes to all $\ell$-hop neighbors. In contrast, when performing anchored image classification, the representation of a sample is only dependent on its own anchor, and is not influenced by the other samples. To this end, using the fitted Gaussian distribution helps manage this complexity, *since changes to the anchors of a node’s $\ell$-hop neighborhood are simpler to model as they require only learning to marginalize out a Gaussian distribution (instead of the training distribution).* 

[1] Single Model Uncertainty Estimation via Stochastic Data Centering, Thiagarajan et al. NeurIPS'22

Reviewer comment (Experiments):
The organization of the experiment part is messy. More setup details should be clarified.

Author response:
Thanks for the feedback. We will provide more set-up details in the expanded appendix and will open-source a code-base for replicating our experiments. (While a working anonymous github repo link has been provided; we will clean this before officially releasing). Further, since we are focused on evaluating the benefits of G-ΔUQ on various distribution shift benchmarks and tasks, we adhere to published experimental protocols to ensure fair comparison.

Reviewer comment (Writing):
Some typos in the paper and the title of Table 2 is missing.

Author response:
Thanks for the catch! We’ve updated the paper after some additional editing.

Reviewer comment (Theory):
Some 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).

Author response:
> ***(W1) The text describing the methodology does not provide sufficient detail to perceive the design principles.***

We’ve expanded and moved some of our experiments to the appendix so that we can discuss this in more detail. In particular, we’ve *added some additional discussion to Sec 3 and Sec 2* that conceptually explains the goals of the anchoring and strategy for distribution design to provide context for our design decisions. *We’ve also updated the appendix to include Pytorch-style pseudo code (App.1) to more clearly explain the proposed variants.* 

> ***(W1.1/Q3.1) Is it problematic if the fitted anchor distribution and the feature distribution do not match?***

The anchoring distribution is used to sample different functional hypotheses from the model, but does not directly carry information that can be used to learn features. Indeed, the anchored representation $[\mathbf{X}-\mathbf{C} ||\mathbf{C}]$ contains both the residual $(\mathbf{X}-\mathbf{C})$ and the anchor ($\mathbf{C}$) so information is not in fact lost during stochastically centering, as the original sample can be recovered as ($\mathbf{X}- \mathbf{C} + \mathbf{C} = \mathbf{X}$). Therefore, **mismatch between the fitted anchor distribution and the feature distribution is not a major concern**. The bigger concern is if the anchoring distribution leads to an untenable amount of stochasticity where the model is in fact unable to effectively marginalize out the anchoring distribution when training. We elaborate on this below, and note this is partially why we proposed the Gaussian feature distribution. *Lastly, to provide some empirical evidence that potential mismatch is not a concern, we’ve added additional results to the appendix which plot the true distribution and fitted distribution. (App. A.14).* Despite the mismatch, G-ΔUQ still outperforms the vanilla model (Fig. 9, Fig. 10).

Reviewer comment (Evaluation):
Some 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.

Author response:
> ***(W3). Some of the results are not presented in a particularly helpful manner and are not described or discussed in much detail.***

Thanks for the comments. We’ve moved the generalization error prediction experiments to the appendix so that we can expand upon the presented results and *summarize the interesting insights below. We have updated the paper with this expanded discussion. Changes are shown in blue.*

Reviewer comment (Evaluation):
There 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.

Author response:
>  ***(Q4) Figure 2 is challenging to understand. Insufficient detail is provided.***

Thanks for bringing this up. We will clarify the text around the figure – this figure corresponds to an ablation study that illustrates the impact of “partial stochasticity” on accuracy and calibration – i.e., only a part of the network is stochastic while the rest is deterministic. G-ΔUQ is uniquely capable of inducing stochasticity in only a part of the network. *As such, L1, L2, and L3 correspond to performing anchoring (and thereby making the remainder of the network stochastic) at the first, second or third/READOUT layer. The bar in gray corresponds to the vanilla (no G-ΔUQ) model for comparison.*
Since we are doing intermediate layer anchoring here, the best performance is highly dependent on the network architecture and the task. We see this reflected in the plots, where different layers are optimal across different networks and tasks. *However, a consistent trend is that across all datasets and architectures the best accuracy and ECE is always one of the G-ΔUQ variants.*
Additionally, we note that we are the first to empirically study the impact of partial stochasticity on accuracy and calibration, and our results suggest that it can give big boosts in performance, and enable the use of pre-trained models (for e.g. L3).

Reviewer comment (Evaluation):
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?

Author response:
- ***(W3/ Q1) Why does the proposed method appear to work relatively poorly in terms of ECE when used with Dirichlet calibration in the node classification setting?***

Table 12 contains many different axes, so while it appears that our proposed method performs worse with Dirichlet, this is not a shortcoming of our method. *Instead, it is a shortcoming of Dirichlet as a post hoc calibration strategy. Indeed, with the exception of (WebKB, covariate shift), applying Dirichlet non-trivially decreases the accuracy and calibration on 7 ⁄ 8 datasets with respect to both the vanilla and G-ΔUQ models (where no post-hoc strategy has been applied).* This clearly indicates that Dirichlet is not a robust post hoc calibration strategy under distribution shifts. One explanation for this behavior is that Dirichlet fits a distribution to the uncalibrated probabilities on the in-distribution validation data using a regularized MLP, and the resulting distribution does not align with OOD data. Indeed, the failure of Dirichlet in this setting highlights the importance of evaluating calibration and other uncertainty based tasks on the realistic setting of OOD data. Furthermore, it supports the development of methods like G-ΔUQ which produce intrinsically better uncertainty estimates, even when additional validation/calibration datasets are not available.

Reviewer comment (Presentation):
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”?

Author response:
>  ***(Q4) Figure 2 is challenging to understand. Insufficient detail is provided.***

Thanks for bringing this up. We will clarify the text around the figure – this figure corresponds to an ablation study that illustrates the impact of “partial stochasticity” on accuracy and calibration – i.e., only a part of the network is stochastic while the rest is deterministic. G-ΔUQ is uniquely capable of inducing stochasticity in only a part of the network. *As such, L1, L2, and L3 correspond to performing anchoring (and thereby making the remainder of the network stochastic) at the first, second or third/READOUT layer. The bar in gray corresponds to the vanilla (no G-ΔUQ) model for comparison.*
Since we are doing intermediate layer anchoring here, the best performance is highly dependent on the network architecture and the task. We see this reflected in the plots, where different layers are optimal across different networks and tasks. *However, a consistent trend is that across all datasets and architectures the best accuracy and ECE is always one of the G-ΔUQ variants.*
Additionally, we note that we are the first to empirically study the impact of partial stochasticity on accuracy and calibration, and our results suggest that it can give big boosts in performance, and enable the use of pre-trained models (for e.g. L3).

Reviewer comment (Theory):
“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?

Author response:
> ***(W1.4/Q3.4) 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?***

You are correct that we desire that “the anchored distribution matches (at least to some degree) the feature distribution.” As discussed above, this is because the anchoring distribution can be seen as a random variable which the model must learn to (partially) marginalize for convergence. By keeping the anchoring distribution closer to input features (even imperfectly), it is a bit *easier on the model to train, since the model does not have to handle vastly different input ranges* or drastically diverge from the training manifold. *We’ve added summary statistics of the fitted distribution to the appendix in Table 18.*