We are grateful for your detailed suggestions! We provide responses below:

Comment 1,3: A thorough review on related works and highlights on the novelty

Thanks for the comment! We added more related works, including those in [2]:

• Added Related work: Recent methods for improving OOD generalization in GNNs include OOD-GNN, which uses nonlinear decorrelation to eliminate spurious correlations [3]; OOD-GMixup, which creates virtual OOD samples by perturbing graph rationale space [5]; and OODGAT, which models interactions between in-distribution and OOD nodes [6]. G-Mixup generates synthetic graphs by interpolating graphons [1], while AIA uses adversarial augmentation to create new environments while preserving stable features [7]. GREA employs environment replacement for better graph rationalization [4], and DPS constructs diverse populations from source domains to train an equi-predictive GNN [8].
• Novelty of GraphMETRO with comparison: GraphMETRO differs fundamentally by decomposing complex distributional shifts into interpretable shift components, enabling it to mitigate a combinatorial number of potential shift variations. While previous works [1,3,4,5] consider heterogeneity, GraphMETRO offers a more fine-grained level of heterogeneity by modeling a continuous space of environments, handling infinite potential shift scenarios. Its concept of referential invariant representation and expert model alignment aid generating robust representations.

Please let us know if this clears the concern! Due to the character limit, we are happy to provide more detailed comparisons with individual works.

Comment 2: Detailed theoretical analysis why GraphMETRO addresses the problem and outperforms baselines.

Thank you for the opportunity to clarify the theory. Here is part of the updated Appendix A.

---

Theoretical Guarantees

Theorem 1 (Shift-Invariance) For any graph $\mathcal{G}$ and shift component $\tau_i$, the encoder $h$ satisfies: $h(\tau_i(\mathcal{G})) = h(\mathcal{G})$ ...

Theorem 2 (Composition of Shifts) For any graph $\mathcal{G}$ and combination of $k$ shift components $\tau^{(k)} = \tau_{i_1} \circ \tau_{i_2} \circ ... \circ \tau_{i_k}$, the encoder $h$ approximately satisfies: $h(\tau^{(k)}(\mathcal{G})) \approx h(\mathcal{G})$ ...

Theorem 3 (Generalization Bound) Let $\mathcal{L}(\cdot, \cdot)$ be the cross-entropy loss. For any distribution $\mathcal{D}_t$ resulting from a combination of shift components in $\tau^{(k)}$, the generalization error of GraphMETRO satisfies:

$$\mathbb{E} _{(\mathcal{G}, y) \sim \mathcal{D} _t} [\mathcal{L}(f(\mathcal{G}), y)] \leq \mathbb{E} _{(\mathcal{G}, y) \sim \mathcal{D} _s} [\mathbb{E} _{\tau^{(k)}} [\mathcal{L}(f(\tau^{(k)}(\mathcal{G})), y)]] + \epsilon$$

where $\epsilon$ is a small constant depending on the complexity of the model and the number of samples.

Proof:

1. Our training objective minimizes:

$$\mathcal{L} _\text{train} = \mathbb{E} _{(\mathcal{G}, y) \sim \mathcal{D} _s} [\mathbb{E} _{\tau^{(k)}} [\mathcal{L}(f(\tau^{(k)}(\mathcal{G})), y)]]$$

2. Recall that $f = \mu \circ h$, where $\mu$ is implemented as a single linear layer with a softmax output, and $\mathcal{L}$ is the cross-entropy loss. This combination is strictly convex with respect to the inputs to $\mu$ if they are not perfectly collinear. Therefore, by Jensen's inequality:

$$\mathbb{E} _{\tau^{(k)}} [\mathcal{L}(f(\tau^{(k)}(\mathcal{G})), y)] > \mathcal{L}(\mu(\mathbb{E} _{\tau^{(k)}}[h(\tau^{(k)}(\mathcal{G}))]), y)$$

3. This implies:

$$\mathcal{L} _\text{train} > \mathbb{E} _{ (\mathcal{G}, y) \sim \mathcal{D} _s} [\mathcal{L}(\mu(\mathbb{E} _{\tau^{(k)}}[h(\tau^{(k)}(\mathcal{G}))]), y)],$$

hence, minimizing $\mathcal{L}_\text{train}$ also minimizes the right-hand side of the inequality.

4. Consider any target distribution $\mathcal{D}_t$ resulting from a combination of shift components in $\tau^{(k)}$. By definition, we can express $\mathcal{D}_t$ as a set of $\tau^{(k)}(\mathcal{G})$, where $\mathcal{G} \sim \mathcal{D} _s, \tau^{(k)} \sim P(\tau^{(k)})$. Here $P(\tau^{(k)})$ is some distribution over possible combinations of shift components.

5. Therefore:

$$\mathbb{E} _{(\mathcal{G}, y) \sim \mathcal{D} _t} [\mathcal{L}(f(\mathcal{G}), y)] = \mathbb{E} _{(\mathcal{G}, y) \sim \mathcal{D} _s} [\mathbb{E} _{\tau^{(k)} \sim P(\tau^{(k)})} [\mathcal{L}(f(\tau^{(k)}(\mathcal{G})), y)]] \leq \mathbb{E} _{(\mathcal{G}, y) \sim \mathcal{D} _s} [\mathbb{E} _{\tau^{(k)}} [\mathcal{L}(f(\tau^{(k)}(\mathcal{G})), y)]] = \mathcal{L} _\text{train}$$

The inequality here holds because our training objective considers a wider range of transformations than those in the actual target distribution.

6. Equations in steps (3) and (5) show that minimizing $\mathcal{L} _\text{train}$ implies finding a model with low true risk and one more invariant to $\tau^{(k)}(\mathcal{G})$, since step (3) shows the loss is lower if $\forall \tau \in \text{supp}(\tau^{(k)})$, $h(\tau(\mathcal{G})) = \mathbb{E} _{\tau^{(k)}}[h(\tau^{(k)}(\mathcal{G}))]$.

7. The gap between the true risk and the empirical risk can be bounded by a constant $\epsilon$. Therefore, we get:

$$\mathbb{E} _{(\mathcal{G}, y) \sim \mathcal{D} _t} [\mathcal{L}(f(\mathcal{G}), y)] \leq \mathbb{E} _{(\mathcal{G}, y) \sim \mathcal{D} _s} [\mathbb{E} _{\tau^{(k)}} [\mathcal{L}(f(\tau^{(k)}(\mathcal{G})), y)]] + \epsilon$$

This gives GraphMETRO advantages over baselines, especially with heterogeneous environments, which generally assume the environment space $\mathbb{E}$ in $P(Y | f(\mathcal{G}), e\in\mathbb{E}) \approx P(Y | f(\mathcal{G}))$ involves relatively simple components.

Summary

We thank you for your approval on our motivation, presentation, and effectiveness. We appreciate it if you could reconsider your evaluation if some concerns on novelty and theoretical analysis are addressed. Thanks!

Reference: Please see the PDF.

We appreciate your efforts and insightful comments! To address your concerns, we provide detailed responses below.

---

Comment 1: Does each expert encode undesired information from the reference model?

Great catch! Previously we conducted experiments to explore the impact of aligning each expert with the reference model. We found that results without alignment were poor, with final representations appearing mixed which we observed by t-SNE visualization.

To formally validate this observation, here are the results (repeated three times for $\lambda=0$):

In our paper, we had attributed this effect to two reasons:

• Each expert model may develop its own unique representation space, leading to information loss when aggregated with other expert outputs.
• Aggregating independent representations creates a mixed space with high variance, which is difficult for predictor heads like MLPs to process and capture interactions among diverse representations.

Therefore, training without alignment poses challenges. However, the reviewer raised an interesting point about the risks of learning undesirable biases from the reference model. We believe the risk can be mitigated by preventing the reference model from overfitting on training data, thereby reducing bias during the training. We have added this ablation result to the Appendix and included this discussion in our methodology section.

---

Comment 2: Can the expert model consistently produce invariant features using GNN or MLP architectures?

Great question! The short answer is yes. Please see Figure 4(a) and Section 4.3, where we compute an invariant matrix showing that the representation produced by the expert closely matches that of the reference model. This indicates that the expert model consistently produces (approximately) invariant features with respect to their shift components.

Please let us know if this answers your question; we are happy to provide further validation if needed.

---

Comment 3: How does the choice of the parameter $k$ impact the model’s generalizability?

Thanks for asking! In Appendix E: Study on the Choice of Transform Functions, we have explored varying numbers of transform functions, i.e., $k$.

Due to character limits, we can't include the full content here. In summary, when $k$ is too large, performance may decline due to noise or "confusion" for the gating model. Please see the detailed analysis in Appendix E.

We also would like to mention that our framework involves a relatively small number of extra hyperparameters, including $k$ and $\lambda$ in the training objective, compared to training a standard GNN. We hope this study provides more insight into the mechanism of our framework.

---

Summary

We hope to address your concerns with:

1. Our additional ablation study and discussion on expert alignment.
2. Reference to the study on whether expert models consistently produce invariant features.
3. Reference to the study on the parameter $k$.

We also appreciate your support and positive feedback on our presentation, novelty, and effectiveness. We hope our response justifies the properties of GraphMETRO training, leading to improved soundness and interpretability. We would greatly appreciate your reconsideration of the empirical soundness of our work given our response.

Thank you again!

We thank the reviewer for the insightful comments. Please find our responses to each point below.

---

Comment 1: Real-world graph shifts may be more varied than defined shifts

Great question! We discussed this in Appendix G: Open Discussion and Future Works, focusing on The applicability of GraphMETRO under complex scenarios and potential mitigations:

• Specific Domains: Real-world graphs can vary significantly, but domain knowledge and observations help address this issue. Insights into distribution shifts, such as the rise of malicious users in trading systems, aid in creating transform functions that cover target shifts. This knowledge can come from:

  • Domain Knowledge: In molecular datasets, transform functions might add carbon structures while preserving functional groups. In social networks, they can be based on known user behaviors.
  • Domain Adaptation: Using a few samples from the target distribution can guide the selection or creation of transform functions. This ensures coverage of distribution shifts. For example, selecting transform functions by measuring the distance between extrapolated datasets and target samples in the embedding space is a promising future direction.

• General Domains: Our experiments utilize five stochastic transform functions, universal graph augmentations as listed in Zhao et al. (2021). While not exhaustive, these functions effectively cover a wide range of distribution shifts, supported by our results. We provide a theoretical guarantee (see our response to reviewer ymW6) that bounds the generalization error on testing data under the assumption of shared support between testing and training transformations.

• Further Considerations: Although our approach is robust, real-world graph distribution shifts can be highly complex, and even mild assumptions may not always hold. Such scenarios occur due to limited knowledge of the testing distribution, which is a common challenge for out-of-distribution generalization methods. We see this as an opportunity for future research to proactively capture and anticipate distribution shift types.

We appreciate your insights! We hope this discussion addresses your concerns.

---

Comment 2: Adding latest OOD methods such as OOD-GNN, OODGAT-ATT, and OOD-GMixup

Thanks for bringing up these relevant works! We included two more baselines: OODGAT [1] (for node classification only) and G-Mixup [4] (for graph classification only), though we couldn't find existing code for OOD-GNN [2] and OOD-GMixup [3]. Please let us know if otherwise.

The results for OODGAT and G-Mixup are repeated three times. For G-Mixup, we used the same GNN architectures and common hyperparameters (e.g., lr) specified in Appendix B. For OODGAT, due to its unique architecture design, we used their original implementation including hyperparameters (e.g., $a=0.9, b=0.01$).

We found that OODGAT performs well on the Twitch dataset. However, G-Mixup struggled to generate effective synthetic node features for semantic graphs, particularly when node dimensions are large (768 on both Twitter and SST2).

We have added more discussions about these three works to the related work. Please see our response to reviewer ymW6 if you are interested!

---

Comment 3: Experimental support of the alignment design of different expert models

Great catch! We previously conducted experiments to explore the impact of aligning each expert with the reference model. We found that results without alignment were poor, with final representations appearing mixed, as observed through t-SNE visualization.

To formally validate this observation, here are the results (repeated three times for $\lambda=0$):

In our paper, we attributed this effect to two reasons:

• Each expert model may develop its own unique representation space, leading to information loss when aggregated with other expert outputs.
• Aggregating independent representations creates a mixed space with high variance, which is difficult for predictor heads like MLPs to process and capture interactions among diverse representations.

Therefore, training without alignment poses challenges. We hope this validates our alignment design.

---

Summary

We are grateful for your positive evaluation of our presentation, novelty, and effectiveness. We further provide the following responses:

1. An in-depth discussion about applicability.
2. Additional experiments on baselines.
3. An additional ablation study on the alignment impact.

We hope our responses address your concerns. We would like to highlight that our main contribution is framing the graph OOD problem into a simple yet novel formulation that is more tractable, as well as proposing the MoE training framework with theoretical guarantees.

Overall, we believe our work makes good contributions to the field of graph distribution learning, and we would appreciate your reconsideration. Thank you for your efforts again!

---

[1] Song et al. Learning on Graphs with Out-of-Distribution Nodes. KDD 2022.

[2] Li et al. OOD-GNN: Out-of-Distribution Generalized Graph Neural Network.

[3] Lu et al. Graph Out-of-Distribution Generalization with Controllable Data Augmentation.

[4] Han et al. G-Mixup: Graph Data Augmentation for Graph Classification. ICML 2022.

We are grateful for your positive feedback and detailed suggestions! We provide responses below to address your remaining concerns.

Comment 1: The considered distribution shifts may not fully represent real-world shifts.

Thanks for the comment! We discussed this in Appendix G: Open Discussion and Future Works, focusing on The applicability of GraphMETRO under complex scenarios and potential mitigations:

• Specific Domains: Real-world graphs can vary significantly, but domain knowledge and observations help address this issue. Insights into distribution shifts, such as the rise of malicious users in trading systems, aid in creating transform functions that cover target shifts. This knowledge can come from:

  • Domain Knowledge: In molecular datasets, transform functions might add carbon structures while preserving functional groups. In social networks, they can be based on known user behaviors.
  • Domain Adaptation: Using a few samples from the target distribution can guide the selection or creation of transform functions. This ensures coverage of distribution shifts. For example, selecting transform functions by measuring the distance between extrapolated datasets and target samples in the embedding space is a promising future direction.

• General Domains: Our experiments utilize five stochastic transform functions, universal graph augmentations as listed in Zhao et al. (2021). While not exhaustive, these functions effectively cover a wide range of distribution shifts, supported by our results. We provide a theoretical guarantee (see our response to reviewer ymW6) that bounds the generalization error on testing data under the assumption of shared support between testing and training transformations.

• Further Considerations: Although our approach is robust, real-world graph distribution shifts can be highly complex, and even mild assumptions may not always hold. Such scenarios occur due to limited knowledge of the testing distribution, which is a common challenge for out-of-distribution generalization methods. We see this as an opportunity for future research to proactively capture and anticipate distribution shift types.

We appreciate your insights! We hope this discussion addresses your concerns.

---

Comment 2: Notations

Thanks! We corrected the mentioned notations and rigorously checked other notations, including those used in our response to reviewer ymW6 on the detailed theoretical analysis.

---

Question 1: How to obtain the ground truth labels of the types of distribution shift on the test set?

Good question! In Figure 3: Accuracy on synthetic distribution shifts, the testing environments on synthetic datasets are also constructed by augmentation. The difference is that testing environments could be augmented by multiple transform functions, resulting in testing graphs that are unseen during training. For real-world graphs, the ground truth distribution shifts are not available; therefore, we conducted qualitative analysis in Section 4.4.

---

Question 2: How did you train the baseline generalization/robustness models? How did you ensure a fair comparison between GraphMETRO and the baseline methods?

Since we are using real datasets from the GOOD benchmark, which provides implementation of the OOD baseline methods and base GNN models, we aligned our settings (Appendix C, Table 3) with them.

Specifically, We adopted the same encoder and classifier from their implementation. Results of the baseline methods, except for the Twitter dataset, are officially reported by the GOOD benchmark. The results on the Twitter dataset are not available in the official leaderboard but were also produced using their code base.

Therefore, we ensured the dataset settings (splits, node features) and architecture components (encoder, classifier, pooling layer, etc.) were the same to ensure a fair comparison.

---

Summary

We want to thank you for your insightful questions and approval of our motivation, novelty, and effectiveness.

• For your applicability concern, we conducted more discussion to make it clear.
• For notations, we hope we have resolved this issue with more rigorous checks.

We sincerely hope all of your concerns are addressed. We appreciate your support and we are happy to engage further if there are any other points we missed! Thank you!

We appreciate your comments! We believe the major concern might be due to a bit of misunderstanding on the referential invariant representations and our objective. Please feel free to correct us if otherwise! Here we provide justification from a theoretical perspective for simplicity.

---

Comment 1, 2: GraphMETRO's Invariant Representations

1) Definition of invariant representations

We define an invariant representation with respect to a stochastic transform function $\tau$ as:

$\xi^*(\mathcal{G}) = \xi^*(\tau(\mathcal{G}))$

This aligns with the concept of stability across different environments, where each $\tau$ creates a different environment.

2) Why each expert is invariant w.r.t. a shift component

Each expert $\xi_i$ produces referential invariant representations (c.f. definition 1) with respect to its shift component $\tau_i$:

$\xi_0(\mathcal{G}) = \xi_i(\tau_i(\mathcal{G}))$

This ensures stability under its transformation, aligned with the reference model $\xi_0$. The necessity of referential invariant representation is justified in our response to reviewer t31K.

Even with a single expert, the model's representation remains invariant, as guaranteed by our Shift-Invariance Theorem, explained in point 4 later.

3) Expert model invariance w.r.t. a shift component given the objective

The distance term in Equation 3:

$d(h(\tau^{(k)}(\mathcal{G})), \xi_0(\mathcal{G}))$

encourages each expert to align its output with the reference model $\xi_0$. This, combined with the gating model, leads to invariant representations by MoE.

4) How the MoE model (gating + experts) generates invariant representations

Theorem 1 (Shift-Invariance) For any graph $\mathcal{G}$ and shift component $\tau_i$, the encoder $h$ satisfies: $h(\tau_i(\mathcal{G})) = h(\mathcal{G})$

Proof: Given the gating model's ability to identify shift components and the expert models' invariance properties:

$$h(\tau_i(\mathcal{G})) = \sum_{j=0}^K w_j(\tau_i(\mathcal{G})) \cdot \xi_j(\tau_i(\mathcal{G})) = w_i(\tau_i(\mathcal{G})) \cdot \xi_i(\tau_i(\mathcal{G})) = w_i(\mathcal{G}) \cdot \xi_0(\mathcal{G}) = h(\mathcal{G})$$

The second equality holds because the gating model identifies $\tau_i$, and the third follows from the definition of referential invariant representation.

The distance term in Equation 3 regularizes, enforcing invariance across shifts. This extends to combinations of shifts through our training objective and expert output alignment.

Our approach differs from simple ERM with data augmentation by adaptively combining expert models to handle instance-specific shifts and ensuring invariance to individual and combined shift components.

We sincerely appreciate the reviewer's astute observation and thank you for bringing this to our attention. We strive to make this clearer in the revision. We clearly distinguished between the standard definition of invariant representation and our concept of referential invariant representation, including lines 199 and 243, which are the only two instances in our methodology section. We also recognize that our statement (in an attempt to avoid confusion with new concepts) at lines 69-70 can be more precise, and we have revised it to accurately reflect the concept of referential invariant representation.

We are grateful for the opportunity to address these points and present them better in our revision.

Related work: This is interesting work! We see the intuitive similarity that both works move from specializing in a single environment to achieving invariance across multiple environments. The core difference is perhaps that they explicitly model domain correlations and spurious features, while we focus on experts that contribute to global invariance. We added it to related work as well. Thanks for bringing this relevant work to us!

---

Comment 2: Applying GraphMetro on molecular datasets

Thanks for the insightful comments. We provide answers from several angles:

• Adaptability to Molecular Graphs: GraphMETRO is designed to be adaptable with user-defined transform functions. For molecular graphs, we can incorporate domain-specific transform functions mimicking scaffold or assay shifts rather than relying solely on generic graph augmentations.

• Incorporating Domain Knowledge: As discussed in Appendix G, we can use domain knowledge to construct relevant transform functions. For molecular datasets, this could involve:

  • Adding or modifying carbon structures while preserving functional groups.
  • Altering specific molecular substructures relevant in scaffold shifts.
  • Simulating changes in molecular properties that occur in assay shifts.

• Single Shift Scenario: GraphMETRO also works with a single shift where there will be an in-distribution expert (aka, reference model) and an OOD expert for one type of shift (e.g., scaffold shift). In a more flexible design, the gating model can be adapted/learned to focus on more relevant components since GraphMETRO doesn't require perfect alignment to be effective (c.f., Theorem 3 for reviewer ymW6)

• Transfer Learning Approach: In Appendix G, we discussed a transfer learning setting leveraging a few samples from the target distribution to inform the selection or construction of transform functions with manual addition, ensuring better coverage of distribution shifts in molecular datasets.

We believe these adaptations would be interesting future directions to apply GraphMETRO on molecular datasets. Thank you for highlighting it!

Summary

We are grateful for your positive evaluation on our novelty, applicability on both node and graph tasks, and effectiveness. We further provide clarifications on the invariant learning terminology and applicability on molecular data. We would very much appreciate it if you could reconsider your evaluation if some concerns are addressed. Thank you very much!