Paper: Distribution Shift Resilient GNN via Mixture of Aligned Experts (ICLR-2024)

Reviewer comment (Experiments):
It is not clear how well the 5 designed distribution shifts can well cover real-world graph distribution shifts. The authors listed 5 distribution shifts in the paper (add edge, drop edge, feature noise, subgraph, drop nodes), but in fact there may be more graph distribution shifts than that. For example, there may be a systematic change in link preference (i.e. nodes tend to link with a different type of neighbors), adding malicious nodes (e.g. malicious users in a trading system). Maybe the authors can justify how well the 5 designed shifts can cover real-world distribution shifts.

Author response:
### **Comment & Question 1: Generality of the distribution shifts covered by the transform functions.**

Good question! We provide the response from two angles:


**(a) For general domain**: In our experiments, we mainly use the five stochastic transform functions, which are universal graph augmentations as listed in Zhao et al., (2021) [3]. In our code implementation, we have also included additional transform functions as shown in Appendix B. We believe these transform functions, while not exhaustive, still cover a wide range of distribution shifts observing from our experimental results.

Nevertheless, we agree that the real graph distribution shifts can go beyond any possible combinations of the predefined transform functions. In that case, the assumption may not hold, meaning that GraphMETRO may not capture and precisely mitigate the unknown distribution shift. This scenario could always possibly exist due to the lack of information about the testing distribution or its domain knowledge. We include it as a limitation in Appendix 5, while we further discuss how we could alleviate the problem with additional information.

**(b) For specific domains where additional knowledge is available**: In fact, knowing the tendency of the distribution shifts, such as increasing malicious users in a trading system, would be very helpful in constructing the transform functions that can cover the target distribution shifts well. We believe that such knowledge can come from two sources: 
- **Domain knowledge**, e.g., on molecular datasets, the transform function could be adding additional carbon structure to a molecule (while preserving its functional groups). Or, in a particular social network, transform functions can be defined from known user behaviors. 
- **Leveraging a few samples from target distribution**. This is in fact in line with the reviewer’s Comment #4 regarding graph transfer learning. Specifically, with the guide from a few target samples, we can select more relevant transform functions by, e.g., measuring the distance of the extrapolated datasets under a certain transform function with the target samples in the embedding space.

 We hope this response can alleviate your concern about our applicability.

Reviewer comment (Experiments):
It is not clear how the proposed GraphMETRO can handle imbalanced distribution shifts within the same graph. For example, it may happen that some subgraphs in the graph get denser, while other subgraphs get sparser (e.g. some topics gain interest, while others lose). How will GraphMETRO respond to this kind of shifts?

Author response:
### **Comment & Question 2: The mechanism of GaphMETRO in handling heterogeneous shifts within one graph.**

Good question! For simplicity, suppose we have two transform functions, i.e., adding edges and dropping edges. Given a node classification task and the objective in Eq. (3), the model is trained on the extrapolated datasets based on the transform functions. After that, given an unseen graph with imbalanced distribution shifts, the gating model outputs scores to identify nodes likely to experience increasing or decreasing degrees, while others might adhere to the original distribution. Then, for each node, each expert takes its multihop subgraph, and outputs its referential invariant representation w.r.t. the correposing transform function. These expert outputs and gating model results are then aggregated to form the final representation. Assuming accurate predictions by the gating function, nodes in denser/sparser subgraphs are represented by expert models corresponding to adding/dropping edges. As each expert is trained to create invariant representations, the final node representations remain unaffected by their individual distribution shifts.

Reviewer comment (Theory):
It is not clear how graph pre-training can address the OOD generalization issue. Pre-training trains the model to observe a wide range of graphs and should be helpful in improving generalization.

Author response:
### **Comment 3: How to leverage graph pretraining.**

If we understand correctly, the reviewer was asking how to leverage a pretrained model to further aid the training of GraphMETRO. Please let us know if otherwise. 

This is in fact an interesting point! We believe a model pretrained on a wide variety of data can be very helpful to initialize the gating model, which is required to output the mixture of a node or graph (after it is finetuned on the extrapolated dataset). By enhancing the gating model's predictive capability regarding mixtures, GraphMETRO's final representation should become more resilient. This becomes particularly advantageous when dealing with graphs not previously encountered in the extrapolated dataset. Moreover, it is possible that the pretrained model will also benefit the expert models, while one minor concern would be that the expert model may tend to be similar instead of dedicating to generating invariant representation w.r.t. their corresponding transform function. We added the above discussion to Appendix F as a future work. Thanks again for this comment!

Reviewer comment (Novelty):
It may be good to discuss some previous works in graph transfer learning. Graph transfer learning addresses the problem that the source graph has a different distribution with the target graph, but the target graph should be given beforehand and is thus less difficult than the problem in this paper. Nevertheless, it may be good to discuss them and clarify the differences, e.g. (Zhang et al. 2019), (Wu et al. 2020).

Author response:
**Comment 4: Related works on graph transfer learning.**

Yes! We agree that graph transfer learning is a relevant topic. Thank you for providing the additional related work! We added the discussion on these works in our revision, and we repeat it here for your convenience:

```It is also worth mentioning that graph domain adaptation (Zhang et al., 2019; Wu et al., 2020), different from the problem studied in this work, commonly relies on limited labeled samples from target datasets for improved transferability. For instance, to generate domain adaptive network embedding, DANE (Zhang et al., 2019) uses shared weight graph convolutional networks and adversarial learning regularization, while UDA-GCN (Wu et al., 2020) employs attention mechanisms to merge global and local consistencies.```

Moreover, we believe this is actually relevant to `Comment & Question 1` regarding the generality of the distribution shifts covered by the transform functions. In our future works, we added the following discussion:

```...Leveraging a few samples from target distribution (i.e., domain adaptation). Specifically, we can leverage the samples from the target distribution to inform the selection or construction of transform functions, which can better guarantee the distribution shifts are covered by the transform functions.  For example, we can select more relevant transform functions by, e.g., measuring the distance of the extrapolated datasets under a certain transform function with the target samples in the embedding space. We believe this would be an interesting future direction.```

Reviewer comment (Novelty):
DANE: Domain adaptive network embedding. Zhang et al. IJCAI 2019.

Author response:
**Comment 4: Related works on graph transfer learning.**

Yes! We agree that graph transfer learning is a relevant topic. Thank you for providing the additional related work! We added the discussion on these works in our revision, and we repeat it here for your convenience:

```It is also worth mentioning that graph domain adaptation (Zhang et al., 2019; Wu et al., 2020), different from the problem studied in this work, commonly relies on limited labeled samples from target datasets for improved transferability. For instance, to generate domain adaptive network embedding, DANE (Zhang et al., 2019) uses shared weight graph convolutional networks and adversarial learning regularization, while UDA-GCN (Wu et al., 2020) employs attention mechanisms to merge global and local consistencies.```

Moreover, we believe this is actually relevant to `Comment & Question 1` regarding the generality of the distribution shifts covered by the transform functions. In our future works, we added the following discussion:

```...Leveraging a few samples from target distribution (i.e., domain adaptation). Specifically, we can leverage the samples from the target distribution to inform the selection or construction of transform functions, which can better guarantee the distribution shifts are covered by the transform functions.  For example, we can select more relevant transform functions by, e.g., measuring the distance of the extrapolated datasets under a certain transform function with the target samples in the embedding space. We believe this would be an interesting future direction.```

Reviewer comment (Theory):
The proposed approach assumes that the distribution shift on a graph dataset happens due to some underlying transformations. While this assumption looks plausible at first, it seems quite difficult to identify all distribution shifts with the assumption since the transformations are treated independently when combined together.

Author response:
**(1) How does GraphMETRO identify all distribution shifts from transform functions if they are treated independently when combined together?** 

If we understand correctly, by “treated independently”, the reviewer is referring to the first term in our objective $\text{BCE}(\phi(\tau^{(k)}(\mathcal{G})), Y (\tau^{(k)}))$, where we formulate predicting the distribution shifts types of a jointly transformed graph as a binary multiclass classification problem. We believe the difficulty of this task comes from both the property of transform functions and the expressiveness of the gating model. 

- Firstly, some transform functions are **inherently disentangled**, e.g., adding nodes feature noise and random subgraph extraction. In this case, there will be certain distinction between any pair from these three data distributions, i.e., (graphs with node noise, random subgraph graphs, random subgraphs with node noise), which the gating model can easily tell. 
- While some transform functions can be **essentially similar**, e.g., drop path and drop edges, this won’t affect the performance of our method as long as each expert outputs the corresponding invariant representation. 
- Lastly, indeed, there could be more **complex combinations of the transform functions**, which poses challenges to the gating model’s expressiveness in identifying the combinations. However, this challenge may be minor in the practice. Specifically, we observe fairly high accuracy performances of the gating model, which are above 85% and 73% averagely on extrapolated datasets with one transformation and multiple transformations, respectively.

We added the above discussion to Appendix F to enable a more comprehensive view towards our methodology. We hope this can alleviate your concern on our gating model’s performance in identifying the distribution shift types.

Reviewer comment (Experiments):
The entire framework requires a set of predefined transformation classes to learn the model, and all necessary transformations need to be identified beforehand. However, it is unclear whether the set of transformations used in the experiments is enough to cover complex distributions causing the distribution shift.

Author response:
**(2) How does the predefined transform functions cover complex distributions causing the distribution shift?**

This is also a great question! We believe there are two angles for this question. 

**(a) For general domain**: In our experiments, we mainly use the five stochastic transform functions, which are universal graph augmentations as listed in Zhao et al., (2021) [3]. In our code implementation, we have also included additional transform functions as shown in Appendix B. We believe these transform functions, while not exhaustive, still cover a wide range of distribution shifts observing from our experimental results.

Nevertheless, we agree that the real graph distribution shifts can go beyond any possible combinations of the predefined transform functions. In that case, the assumption may not hold, meaning that GraphMETRO may not capture and precisely mitigate the unknown distribution shift. This scenario could always possibly exist due to the lack of information about the testing distribution or its domain knowledge. We include it as a limitation in Appendix 5, while we further discuss how we could alleviate the problem with additional information.

**(b) For specific domains where additional knowledge is available**: In fact, knowing the tendency of the distribution shifts, such as increasing malicious users in a trading system, would be very helpful in constructing the transform functions that can cover the target distribution shifts well. We believe that such knowledge can come from two sources: 
- **Domain knowledge**, e.g., on molecular datasets, the transform function could be adding additional carbon structure to a molecule (while preserving its functional groups). Or, in a particular social network, transform functions can be defined from known user behaviors. 
- **Leveraging a few samples from target distribution**. Specifically, with the guide from a few target samples, we can select more relevant transform functions by, e.g., measuring the distance of the extrapolated datasets under a certain transform function with the target samples in the embedding space.

Reviewer comment (Experiments):
Most of the transformation requires a set of hyperparameters, e.g., dropout probability. However, each mixture only models a single instance of the transformation but not the entire class of transformations. Given that the hyperparameters are selected via the validation set, there is no evidence that the same configurations can work for the test set since it may be caused by the different hyperparameters of the same transformation type.

Author response:
**(3) What is the complexity of the transform functions and how does it affect generalization?** 

Interesting question! In fact, our implementation and framework could easily avoid selecting hyperparameters on the transform functions. Specifically, we can make multiple transform functions of the same type with different ranges of hyperparameters. Specifically, GraphMETRO allows three edge dropping transform functions, $\tau_1^{\alpha_1}, \tau_2^{\alpha_2}, \tau_3^{\alpha_3}$, where $\alpha_i$ (i=1, 2, 3) are three different ranges of edge dropping probabilities, e.g., [0.1,0.3], [0.3, 0.6], [0.6,0.9], representing different transform extents. Thus, given an input from the validation dataset, the gating model will highlight the transform function which simultaneously selects the corresponding hyperparameter that matches the distribution of the validation set. Interestingly, this idea is in the same spirit as how DARTS [3]  proposes to perform architecture search by formulating the task in a differentiable manner. 

In our previous experiments, we did try this scheme where we replaced a single edge dropping transform function with the ratio range [0.3, 0.5] to three transform functions as mentioned above. While we didn’t see a significantly different performance in that case, we believe this would be a flexible solution which avoids the need to conduct hyperparameter selection. 

We include the above discussion in our open discussion and future works (Appendix F). We hope this response can alleviate your concern about the applicability of our method.

Reviewer comment (Experiments):
This work only considers the distribution shift in the graph structure but not in the labels.

Author response:
Great point! We think the issue of label distributional shift, while important, is orthogonal and complementary to the focus of our current study. To elaborate, label distributional shifts exert analogous impacts across various modalities, such as graphs or images. Moreover, existing methods [1,2] designed to tackle label distributional shifts can be seamlessly integrated into our proposed framework. Such integration would necessitate minimal adjustments, potentially involving modifications to the loss function or the training pipeline. We added this as a future work in Appendix F.

Reviewer comment (Evaluation):
Which of the results are statistically significant in Table 1? For graph classification tasks, the proposed model seems marginally better than the others.

Author response:
Thanks for the question! We compute the p-value of our method against the best baselines method as follows:

| | WebKB | Twitch | Twitter | SST2 | 
|:--|:--|:--|:--|:--|
|p-value|< 0.001| 0.023 | 0.042| 0.081|

Given the cut-off threshold as 0.05, we believe the performances of GraphMETRO are statistically significant on WekGB, Twitch, and Twitter datasets, while on the SST2, we see relatively weak evidence. We added the p-value results to our revision and hope our response can alleviate your concern on our improvements.

Reviewer comment (Presentation):
Figure 2 is not easy to digest since there is only a single label along the radial axis. Could you provide the exact numbers?

Author response:
**Question 3: Numerical results on Figure 2**

Thanks for pointing it out! We included all of the numerical results of Figure 2 in Appendix E, while showing the results on DBLP below.

| |i.i.d. (0)|noisy feature (1)|add edge (2)|drop edge (3)|drop node (4)|random subgraph (5)|
|:--|:--|:--|:--|:--|:--|:--|
|ERM | 85.71 | 84.48 | 71.08 | 79.69 | 83.41 | 76.9|
|ERM-Aug | 85.66 | 85.29 | 74.85 | 82.34 | 84.44 | 72.81|
|GraphMETRO | 85.92 | 85.78 | 76.61 | 82.95 | 84.98 | 81.32|

| |(4, 5)|(3, 5)|(2, 5)|(1, 5)|(2, 4)|(1, 4)|     (2, 3)|(1, 3)|
|:--|:--|:--|:--|:--|:--|:--|:--|:--|
|ERM | 70.4 | 77.63 | 81.99 | 79.69 | 70.55 | 71.52 | 77.73 | 79.59|
|ERM-Aug | 74.16 | 81.04 | 83.65 | 68.62 | 74.01 | 68.27 | 81.13 | 84.49|
|GraphMETRO | 76.18 | 81.71 | 84.26 | 80.31 | 75.1 | 71.05 | 81.85 | 87.14|

Across all of the synthetic environments, GraphMETRO averagely outperforms ERM and ERM-Aug by 3.20% and 2.45%, respectively.

Reviewer comment (Evaluation):
The test accuracy for the synthetic dataset for some transformations is relatively lower than the other transformations. Why some transformations is harder to identify than others for some datasets?

Author response:
### **Question 4: Why does the test accuracy vary across different transformations?**

Great question! Here we summarize three possible reasons:

- **Information Preservation in Transformations:** Certain transformations retain more informative features than others. For instance, in the REDDIT-BINARY (graph classification task), the random subgraph transformation may retain more graph label-related information compared to dropping edges, as the latter tends to lose more global information. This discrepancy in testing performance, where dropping edges outperforms random subgraph extraction, could be due to the preservation of crucial information. However, conclusions may vary across datasets or tasks depending on how information influences final predictions. For CiteSeer (a node classification task), a random subgraph might preserve more local node information, potentially explaining why its testing performance surpasses dropping edges in this specific task.

- **Complexity of transformation:** Certain transformations inherently generate more diverse graphs than others. If the model lacks the expressive capacity to capture such diversity, it may lead to a decline in testing performance.

- **Model Sensitivity:** Certain transformations may be easier for a model to learn due to compatibility with specific model architectures. This extends beyond transformation complexity and emphasizes how different model architectures may prefer learning particular distributions from one of the extrapolated datasets, which can also contribute to the difference in the testing performance.

We included the above discussion to Appendix F: Open Discussions. We hope this response can answer your question and improve the soundness of our work.

Reviewer comment (Novelty):
The motivation behind the proposed method is not adequately substantiated. The primary basis given is that "previous research has concentrated on addressing specific types of distribution shifts." However, this overlooks a plethora of prior works in the field. Contrary to the suggestion that graph shift heterogeneity is under-explored, numerous studies have delved into learning the "environment generators" for GNNs to detect graph shifts, as exemplified by [https://arxiv.org/abs/2202.02466]. Other works have focused on learning shift-specific transformations, such as [https://arxiv.org/abs/2211.02843]. Consequently, there exists a wide spectrum of approaches to tackle graph shift heterogeneity. The choice of approach in this paper, especially the emphasis on MOE, requires a more detailed and robust justification to elucidate its relevance and significance.

Author response:
Thank you for this comment! We believe there might be a bit of misunderstanding due to our different definitions of *“graph shift heterogeneity”*. We firstly discuss the related works mentioned and then justify our statement:

**(1) Related works**

Please see Section 2 where we discussed the paper mentioned by the reviewer, i.e., EERM (Wu et al., 2022a). We repeat part of it here for you convenience:

```The prevailing invariant learning approaches assume that there exist an underlying graph structure (i.e., subgraph) (Wu et al., 2022c; Li et al., 2022b;a) or representation (Arjovsky et al., 2019; Wu et al., 2022a; Chen et al., 2022; Bevilacqua et al., 2021; Zhang et al., 2022) that is invariant to different environments and / or causally related to the label of a given instance. However, these approaches focus on group patterns without explicitly considering nuanced (instance-wise) distribution shifts, making their applicability limited.```

Moreover, we apologize for missing the recent interesting work by **Sui et al. [1]** which officially came out two days before the ICLR abstract deadline. We added it to our revision, thank you! 

Specifically, Sui et al. [1] proposed a graph data augmentation strategy that alleviates covariate shift by generating diverse and invariant causal features. However, the trainable augmenter they used may not distill diverse augmentations or construct unseen perturbations. Moreover, Sui et al. [1] test its method only on graph classification tasks, while GraphMETRO can be applied to both node and graph classification tasks. Besides, we have discussed graph augmentation and attention-based methods in our related works, and we added more recent works on graph OOD [2,3,4], and we hope our response clears your concern on the related work discussion.


**(2) The definition of graph shift heterogeneity**

In this work, we refer to **“heterogeneous shifts”** as multiple and different levels of shifts which vary across different instances (nodes or graphs), as illustrated in the example in the abstract. While we agree that the existing invariant learning approches can accommodate multiple distribution shifts, it could be hard for them to tackle nuanced distribution shifts for individual instances (nodes or graphs) since the distribution shifts are inferred from variance across multiple data environments. If GraphMETRO's approach were described via environments, we would have a combinatorial number of such environments in training (the product of all different subsets of nodes and all their possible distinct shifts). Thankfully, GraphMETRO avoids this combinatorial explosion by considering **a mixture of transformations as a proxy** for the target distribution shifts rather than invariance to whole-graph environment shifts. This is the type of heterogeneity we are interested in our paper. 

**(3) Regarding our original motivation statement**

While the statement pointed out by the reviewer serves as our primary motivation, we would like to note that we did not claim all of the previous works fall into this category. And we have provided detailed discussion about three lines of research in the related work section.

However, we agree that we could make this statement border to cover the previous invariant learning methods. To improve the clarity, we change the statement from "previous works mostly focus on addressing specific types of distribution shifts" to "**previous works mostly focus on addressing specific types of distribution shifts or inferring distribution shifts from data environments…**”. We also modified our introduction correspondingly, we hope this will better position our work.

**(4) The choice of our MoE design**

- The choice of our approach comes as a consequence of our motivation to model the graph/instance shift heterogeneity. As mentioned, mitigating multiple and nuanced distribution shifts simply goes beyond certain distribution shift types or environment construction as seen in the previous methods. 
- Thus, GraphMETRO takes a different path, i.e., predicting a mixture of transformations as the proxy of the target distribution shifts. This enables the prediction of multiple different distribution shifts and the flexibility to model fine-grained heterogeneity since the mixture can be varied across different instances. We then tackled the proxy to mitigate the target distribution shifts. Intuitively, this solution provides a **“middle ground”** to deem graph generalization as an equivalent mixture, which, we believe, is a more tractable solution.

We updated our paper to make the above point more clear. We genuinely hope our answer can justify the motivation and solve your concern.

Reviewer comment (Theory):
The presented assumption seems overly broad and lacks specificity. Additionally, the architectural design appears to be somewhat arbitrary. Consequently, it's challenging to discern the functionality, its underlying mechanism, and its improvements over existing methods.

Author response:
Thanks for this comment! In Section 3.1, we have discussed the assumptions in detail with specifications. We repeat part of it here for your convenience:

```Assumption 1 essentially states that the distribution shifts (whatever they are) can be decomposed into several mixture components of stochastic graph transformations. For example, on a social network dataset, each mixture component can represent different patterns of user behavior or network dynamics shifts. Specifically, one mixture component might correspond to increased user activity, while another could signify a particular trend of interaction within a certain group of users. Such a mixture pattern is common and well-studied in the real-world network datasets (Newman, 2003; Leskovec et al., 2005; 2007; Peel et al., 2017).```

We hope the analysis provides detailed illustration to the assumption. We are happy to add more discussion if you think anything  is still unclear!

---

### **Comment 2 & 3: Implementation of our method**

- **Model architectures**: 
Please see Tables 2 and 3 in Appendix A for detailed model architecture information. Specifically, we use the same encoders and classifiers from GOOD benchmarkfor real-world datasets to ensure fair comparisons. We employ backbones based on the best ERM performance for synthetic datasets. Moreover, GraphMETRO is model-agnostic, which consistently improves performance across varied model architectures.

- **Stochastic transformation**: Please see Appendix B where we include the introduction of the stochastic transform functions. 

- **Optimization process**: Please see Table 2 and 3 (Appendix A), where we included hyperparameters on each dataset. In the second paragraph after Eq. (3), we also described our training pipeline: `“...we set apart the other loss terms from backpropagating to it to avoid interference with the training of the gating model… We optimize the objective via stochastic gradient descent”`. 

- **For a more intricate optimization strategy:** This is a great catch! In fact, we did try to pretrain the gating model for several epochs as warm up before training the whole model in an end-to-end fashion via the objective Eq. in (3). However, we didn’t notice a statistically significant difference in their performance, which can be due to that expert models take more time to convergence (since they need to align with the base model during the training) compared to the gating model.

Due to the space limitation, we had to include most of the implementation details in the appendix while we added pointers in the main paper. However, we will try to make it more detailed and feel free to let us know if anything is missing! 

---

### **Comment 2 (Cont.): Understand the functionality, underlying mechanism, and performance gain**

Thanks for the comment. For clarity, we provide the following pointers to the paper:

- In Section 4.3, we provide a study to reveal the underlying **mechanism** of GraphMETRO, i.e., each expert excels in generating invariant representations concerning a stochastic transform function, which provide a solid foundation in generating referential invariant representations w.r.t. the specific transformations and further guarantee the generalization.

- In Appendix C, we study the **impact of the MoE architecture** on model performance, which shows that the model performance may decline if the expressiveness of the expert model decreases. 

- In Appendix D, we study the **impact of the stochastic transform function** on model performance, where we also provide a detailed discussion of the modeling mechanism.

We believe the above studies and discussion in our paper provides an in-depth view, highlighting the roles of our objective, architecture, and stochastic transform function.  Please let us know if any of these perspectives is still unclear, and we can further improve our experimental study.

Reviewer comment (Novelty):
The novelty is not well justified and comparison with recent methosd is not sufficient.

Author response:
### **Comment 1: Regarding the novelty of GraphMETRO.**

Thanks for the comment! We believe our novelty comes from the proposal of an equivalent mixture for graph OOD and the construction of our training framework, as detailed below:

- **An equivalent mixture for graph OOD**: The key challenge we faced to mitigate multiple and nuanced distribution shifts is the intrinsic complexity and heterogeneity of graph distribution shifts, which simply goes beyond certain distribution shift types [6,7,8,9] or environment construction as seen in the previous methods [1,2,3]. GraphMETRO takes a different path, i.e., predicting a mixture of transformations as the proxy of the target distribution shifts, where the mixture can be varied across different instances, and then tackled the proxy to mitigate the target distribution shifts. We believe the high-level idea is succinct, nevertheless, novel, in the sense that it provides a “middle ground” to deem graph generalization as an equivalent mixture that is more tractable.

- **Training framework**: With the guide of our formulation, the training framework is still non-trivial due to two problems, i.e., “how to provide supervision for predicting the mixture” and “how to ensure the experts corresponding to mixture components are compatible when working as a whole”. Specifically, GraphMETRO solves the first problem by conducting graph extrapolation. This is somewhat similar to the spirit of graph pretraining in the sense that we inject heterogeneity to promote the expressiveness of the gating model in recognizing the mixture components. For the second problem, we introduce the concept of Referential Invariant Representation, along with the novel objective in Eq. (3) to enforce the invariance and compatibility. It is worth mentioning that the model performance is much worse than the reported numbers (e.g., 2.7\% lower on Twitch dataset) without the compatibility constraint, indicating the proposed referential invariance concept is indispensable. 

We added more justification in the introduction (updated in the revision). We genuinely hope our responses can solve your concerns about the novelty of our work.

Reviewer comment (Novelty):
Some of recent papers on out-of-distribution learning on graphs are not discussed [3-5].

Author response:
### **Comment 2: Comparison with recent methods [3,4,5].**

In compacting the paper to fit in the page limit we mistakenly did not include these relevant references, we apologize. We added discussion in the revised version. Here we summarize these works and point out their key differences with our method: 

- In particular, **Yang et al. [3]** explore molecule representation learning in out-of-distribution (OOD) scenarios. They achieve this by directing the molecule encoder to utilize stable and environment-invariant substructures relevant to the labels without the need for environmental labels. 
- Similarly, **Sui et al. [4]** introduces causal attention modules to identify key invariant subgraph features that can be described as causing the graph label. The type of OOD task that Sui et al. [4] considers assumes the graph label is caused by a subgraph, which is quite different from ours. Moreover, both Yang et al. [3] and Sui et al. [4] consider tasks where the graph label is caused by a subgraph. 
- **Ma et al.[5]** is an interesting theoretical work which studies GNN generalization and examines their fairness, showing that the test subgroup's distance from the training set impacts GNN performance. Ma et al.[5], as far as we could assess, does not propose any specific architecture to solve the type of OOD tasks we consider in our work.

Overall, the goal of GraphMETRO is to be invariant to a mixture of selected stochastic transform functions (and the mixture can vary across different instances), which is a more flexible and general solution. 
We added a discussion of these works to our revision, hopefully providing a more comprehensive comparison and literature overview. We hope our responses can solve your concerns about the related work.

Reviewer comment (Novelty):
The authors argued that "previous works mostly focus on addressing specific types of distribution shifts", which seems inproper and incorrect. E.g., the typical works for graph OOD learning EERM [1] and DIR [2] do not assume the type of distribution shifts in their problem formulation.

Author response:
### **Comment 3 & Question 1: Related works on invariant learning and clarification on our statement.**

Thanks for pointing it out! Below we clarify the statement and clear potential misunderstanding:

**(1) How does GraphMETRO compare with invariant learning methods like DIR and EERM?** 

Please see our related work section where we discussed these two papers, i.e., DIR (Wu et al., 2022c) and EERM (Wu et al., 2022a). We repeat part of it here for you convenience:

```The prevailing invariant learning approaches assume that there exist an underlying graph structure (i.e., subgraph) (Wu et al., 2022c; Li et al., 2022b;a) or representation (Arjovsky et al., 2019; Wu et al., 2022a; Chen et al., 2022; Bevilacqua et al., 2021; Zhang et al., 2022) that is invariant to different environments and / or causally related to the label of a given instance. However, these approaches focus on environmental patterns without explicitly considering nuanced (instance-wise) distribution shifts, making their applicability limited.```

Besides, we also provide a more in-depth comparison in our point (3) below to highlight our key technical originality.

**(2) Regarding our statement about previous works:** 

While the statement serves as our primary motivation, we would like to note that we did not claim all of the previous works fall into this category, and we have provided detailed discussion about three lines of research in the related work section.

To improve the clarity, we change the statement from "previous works mostly focus on addressing specific types of distribution shifts" to **"previous works mostly focus on addressing specific types of distribution shifts or inferring distribution shifts from data environments** (which is highly limited when confronted with nuanced distribution shifts)”. We also modified our introduction correspondingly. Thanks for letting us know our statement could be misinterpreted.

Reviewer comment (Novelty):
How does the model compare with existing invariant learning-based models for graph OOD generalization, e.g., EERM [1] and DIR [2]? What is the key technical originality of this work?

Author response:
### **Comment 3 & Question 1: Related works on invariant learning and clarification on our statement.**

Thanks for pointing it out! Below we clarify the statement and clear potential misunderstanding:

**(1) How does GraphMETRO compare with invariant learning methods like DIR and EERM?** 

Please see our related work section where we discussed these two papers, i.e., DIR (Wu et al., 2022c) and EERM (Wu et al., 2022a). We repeat part of it here for you convenience:

```The prevailing invariant learning approaches assume that there exist an underlying graph structure (i.e., subgraph) (Wu et al., 2022c; Li et al., 2022b;a) or representation (Arjovsky et al., 2019; Wu et al., 2022a; Chen et al., 2022; Bevilacqua et al., 2021; Zhang et al., 2022) that is invariant to different environments and / or causally related to the label of a given instance. However, these approaches focus on environmental patterns without explicitly considering nuanced (instance-wise) distribution shifts, making their applicability limited.```

Besides, we also provide a more in-depth comparison in our point (3) below to highlight our key technical originality.

**(3) Why do we say our method could be more broad than the existing invariant learning approaches?** 

- Invariant subgraph learning approaches, e.g., [1,2], consider variance of constructed data environments, which are designed very differently compared to our work. 
While they can accommodate multiple distribution shifts (as in multiple environments), these focus on patterns within each environment and ignore the variety across instances (e.g., shifts at the resolution of nodes), which may not be well-captured by the environment assignments. 
- GraphMETRO considers that specific parts of the test graph may have different shifts. Particularly, our goal is to make the generalization to unknown testing distribution more adaptive and broad, as opposed to limiting the distribution shifts to being invariant to specific types of subgraphs. 

In other words, if GraphMETRO's approach were described via environments, we would have a combinatorial number of such environments in training (the product of all different subsets of nodes and all their possible distinct shifts). Thankfully, GraphMETRO avoids this combinatorial explosion by considering a mixture of transformations as a proxy for the target distribution shifts rather than invariance to environment shifts.

**(4) Key technical originality compared to invariance learning (going more deeply)**

Another interesting view to see the innovation of GraphMETRO is that it breaks the typical invariant learning formulation, which assumes the data is manipulated by the environment variables (and then can be “decoded” into multiple environments). Instead, GraphMETRO sees the distribution shifts on an instance as a mixture, which is represented by the score vector output by the gating function over the basis of the transform functions. In other words, GraphMETRO can produce infinite environments as the elements in the score vector are continuous. One can see that once we limit the output domain of the gating function into, e.g., binary {0, 1}, GraphMETRO can also produce a limited number of environments (if we categorize the instances based on the score vector), which covers the environment construction in invariant learning. Moreover, as mentioned, we propose the concept of referential invariant representation with a base model $\xi_0$, which is also different from previous works on invariant learning. We added the above discussion to Appendix F to improve the depth of our analysis.

Reviewer comment (Experiments):
Can the proposed model handle multiple different types of distribution shifts that simultanenously exist in data?

Author response:
### **Question 3: Can GraphMETRO handle multiple different types of distribution shifts that simultaneously exist in data?**

Yes! The distribution shift types corresponding to the gating outputs with high scores will be tackled during training. That is, if the gating output highlights multiple mixture components, their corresponding distribution shift types will be handled jointly.

Reviewer comment (Experiments):
Can the proposed model tackle distribution shifts and out-of-distribution generalization on molecular graphs?

Author response:
### **Question 4: Can GraphMETRO  tackle distribution shifts on molecular graphs?**

That is a great idea! Yes, GraphMETRO can be applied to molecular datasets if one designs transform functions to cover typical molecular variants. For instance, a transform function may add carbon structures to the molecules. These domain-specific transform functions are outside the scope of our work, however, we believe these would be interesting future work directions!