GraphMETRO: Mitigating Complex Graph Distribution Shifts via Mixture of Aligned Experts
Authors: Shirley Wu
Kaidi Cao
Bruno Ribeiro
James Zou
Jure Leskovec
Venue: NeurIPS 2024 poster (2024)
Keywords: Graph Neural Network
Distribution Shifts
Generalization
Mixture-of-expert model
Primary area: graph_neural_networks
TL;DR: GraphMETRO utilizes a mixture-of-experts architecture to effectively handle complex distribution shifts in graph data, achieving state-of-the-art results on benchmark datasets.

Abstract: Graph data are inherently complex and heterogeneous, leading to a high natural diversity of distributional shifts. However, it remains unclear how to build machine learning architectures that generalize to the complex distributional shifts naturally occurring in the real world. Here, we develop GraphMETRO, a Graph Neural Network architecture that models natural diversity and captures complex distributional shifts. GraphMETRO employs a Mixture-of-Experts (MoE) architecture with a gating model and multiple expert models, where each expert model targets a specific distributional shift to produce a referential representation w.r.t. a reference model, and the gating model identifies shift components. Additionally, we design a novel objective that aligns the representations from different expert models to ensure reliable optimization. GraphMETRO achieves state-of-the-art results on four datasets from the GOOD benchmark, which is comprised of complex and natural real-world distribution shifts, improving by 67% and 4.2% on the WebKB and Twitch datasets. Code and data are available at https://github.com/Wuyxin/GraphMETRO.

GraphMETRO: Mitigating Complex Graph Distribution Shifts via Mixture of Aligned Experts
Shirley Wu
Kaidi Cao Stanford University
Stanford University
shirwu@cs, stanford..d...o.
kaidicao0cs stanford.edu
Bruno Ribeiro"
James Zou*
Jure Leskovec*
Purdue University
Stanford University
Stanford University jureecs stanford edu
ribeirob@purdue. edu
j Sammes2@0ssstaafoorooo
Abstract
Graph data are inherently complex and heterogeneous, leading to a high natural diversity of distributional shifts. However, it remains unclear how to build machine learning architectures that generalize to the complex distributional shifts naturally occurring in the real world. Here, develop GraphMETTRO, a Graph Neural Net- work architecture that models natural diversity and captures complex distributional shifts. GraphMETRO employs Mixture-of-E offeEsstts (MoE) architecture with gat- ing model and multiple expert models, where each expert model targets specific distributional shift produce referential representation w.r.f. a reference model, and the gating model identifies shift components. Additionally, we design a novel objective that aligns the representations from different expert models to ensure reli- able optimization. GraphMETRO achieves state-of-the- art results on four datasets from the GOOD benchmark, which is comprised of complex and natural wolld distribution shifts, improving by 67% and 4.2% on the WebKB and Twitch datasets. Code and data are available https: : // gi thub. com/ http///yxi//Graaahilet /GraphMETRO....
1 Introduction
The intricate nature of real-world graph data introduces wide variety of distribution shifts and heterogeneous graph variations [47, 30, 43, 26). For in- stance, in a social graph, some user nodes can experience reduced activity and profile alterations, while other user nodes may see increased interactions. More broadly, such shifts go beyond the group-wise pattern and further con- tribute to the heterogeneous nature of graph data. In Figure 1, we provide a real-world example on a webpage network dataset, where, besides the general distribution shift from source to target distribution, two webpage nodes U1 and U2 in the target domain exhibit varying de- grees of change in their content features. These inherent shifts and complexity accurately characterize the dynam- ics of real-world graph data, e.g., social networks [2, 19] and ecommerce graphs [76).
M1
data
An example on WebKB [51, 20]. illustrate (1) The distribution shift from source to target (the thick arrow in the upper right) and (2) Instance-wise het- erogeneity in the target distribution (the thin arrows pointing to i1 and u2).
Figure 1:
Above the diverse graph variants, Graph Neural Networks (GNNs) [21, 25, 10] have become prevailing method for downstream graph tasks. Standard evaluation often
Equal Senior Authorship
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Ds
Graph size
Avg. node degree ?
Encoder ht
Nigned
space
$1
Gatling model
0.06
Aga
Feature noise 1
9
G. 4.
$3
c
Dt
g
0.08
50
Ttafomesce mostll
(a) High-level concept of GraphMETRO
GraphMETRO Architecture
Figure 2: Overview of GraphMETRO on graph classification tasks. (a) High-level Concept: As
simple example, the distribution shift from a target graph G Dj to source distribution D, is decomposed along three shift dimensions: graph size (51). node degree (E2), and feature noise (Ex).
Note that the shift components can be customized and expanded based downstream tasks.
(b) Architecture: Given an input graph, the gating model jl decomposes the instance-specific distribution shift into the contributions from the shift components. Then, each expert model & (i > 0)
is tasked with generating referential invariant representations (cf. Section for the definition) w. w.r.f. an assigned shift component. Eo is reference model used for aligning the representation spaces of the expert models. The final representation is aggregated from the experts output and is referentially invariant to any distribution shifts, which is then input to the classifier.
adopts random data splits for training and testing GNNs. However, it overlooks the complex distributional shifts naturally occurring in the real world. Moreover, compelling evidence shows that GNNs are extremely vulnerable to graph data shifts [79, 26, 20]. Thus, our goal is to build GNN models that generalize better to real-world data splits and graph dynamics described carlier.
Previous research on GNN generalization has mainly focused on two lines: (I) Data- Data-auumentttion training procedures that learn environment-robue predictors by augmenting the training data with the environment changes. For example, works have looked at distribution shifts related to graph size [49, 14]. node features [26, 8, 27], and node degree or local structure [65, 39), assuming that the target data adhere to a designated shift type. (2) Learning
wetroommnttinvariian representations or predictors either through inductive biases learned by the model [69, 67], through regularization [4, 32, 75], or a combination of both [72, 11, 80].
However, the real-world distribution shifts and graph dynamics are unknown. Specifically, the distribution shift could be any fusion of multiple shift dimensions, each characterized by unique statistical properties [26, 20, 50], which is rarely covered by single-dimension synthetic augmentation or fixed combinations of shift dimensions used in data augmentation approaches. Moreover, as seen in Figure 1, graph data may involve instance- wise heterogeneity, lacking stable properties from which invariant predictors can be learned [47, 30]. Here, the standard strategy of learning invariant predictors or representations must contend with combinatorially large number of distribution shift variations. Thus, previous works may not be well-equipped to address this challenging task effectively.
Here we propose a novel and general framework, GraphMETRO. The key to approach is to decompose any unknown shift into multiple shift components and learn predictors that can adapt to the graph heterogeneity observed in the target data. Figure 2a shows an example of our method on graph-level tasks, where the shift from the target graph data G E D: to the source distribution D, decomposed into two strong shift components controlling feature noise and graph size, while the shift component controlling average node degree is identified as irrelevant. Specifically, the shift components are constructed such that each possesses unique statistical characteristics. Moreover, the contribution of each shift component to the shift is determined by an influence function that encodes the given graph G and source distribution D.. This design enables breaking down the generalization problem into (1) inference on strong shift components and their contributions as surrogates for any distributional or heterogeneous shifts, and (2) mitigation toward the surrogate shifts, where the individual shift components are interpretable and more tractable.

---

For the first subproblem, we design a hierarchical architecture composed of gating model and multiple expert models, inspired by the mixture-of-experts (MoE) architecture [24]. As shown in Figure 2b, the gating model takes any given node graph data and identifies strong shift components that govern the localized distribution shift, while each expert model corresponds to an individual shift component. Second, to further mitigate surrogate distribution shifts, we train the expert models to generate referentially invariant representations with respect to their corresponding shift components, which are then aggregated as the final representation vector. Moreover, the expert outputs must align properly in a common representation space to prevent extreme divergence in the aggregated design novel objective to ensure smooth training process. Finally, during the evaluation process, we integrate outputs from both the gating and expert models for final representations.
representation. Consequently,
This process effectively generates invariant representations across complex distributional shifts. To highlight, our method achieves the best performance on four node- and graph-level tasks from the GOOD benchmark [20], which involves a diverse set of natural distribution shifts such as user language shifts in gamer networks and university domain shifts in university webpage net- works. GraphMETRO achieves a 67% relative improvement over the state-of-the-art on the WebKB dataset [51]. On synthetic datasets, our method outperforms Empirical Risk Minimization (ERM) by 4.6% on average. To the best of our knowledge, GraphMETRO is the first to explicitly target complex distribution shifts that resemble real-workd settings.
The key benefits of GraphMETRO are as follows:
A novel paradigm:
GraphMETRO provides new approach to aid GNN generalization by decomposing and mitigating complex distributional shifts via a mixture-of-experr architecture.
Superior performance: It outperforms tft----r-t---ttt methods on real-world datasets with natural splits and shifts, demonstrating promising generalization ability.
Enhanced interpretability: GraphMETRO offers insights into the shift types of graph data by identifying and interpreting strong shift components.
2 Related Works
Invariant learning for graph OOD. The prevailing invariant learning approaches assume that there exist an underlying graph structure (i.e., subgraph) [69, 36, 34, 71, 60, 83, 37] or representation [I, 67, 6, 3, 81, 68, 11, 8, 42] that is invariant to different environments and/or causally related to the label of a given instance. For example, DIR [69] constructs interventional distributions and distills causal subgraph patterns to make generalizable predictions for graph-level tasks. However, this line of research focuses on group patterns without explicitly considering instance heterogeneity. Therefore, the standard invariant learning approaches are not well-equipped to mitigate the complex distribution shifts our context. See Appendix A for an in-depth comparison.
Data augmentation for graph OOD, GNNs demonstrate robustness to data perturbations when incorporating augmented views of graph data [7]. Previous works have explored augmentation with respect to graph sizes [85, 4, 84], local structures [40, 38], feature metrics [15], and graphons[33]. For example, DOD-G-Mixup [22] creates virtual OOD samples by perturbing the graph rationale space. Recently, Jin et al. [23] proposed adapting testing graphs to transformed graphs with patterns similar to the training graphs. Other approaches conduct augmentation implicitly via attention mechanisms [45, 66). For example. GSAT [45] injects stochasticity into attention weights to block label- information. Nevertheless, this line of research may not effectively solve the challenging problem, since unseen distribution shifts may not be covered by the distribution of augmented graphs. Moreover, it may lead to degradation of in in-distribution performance due to GNNs' limited expressiveness in encoding a broad distribution.
Instance heterogeneity for graph OOD, Recent methods [41, 59, 61, 77, 35, 74] have explic- itly considered instance heterogeneity for improving OOD generalization in GNNs. For example, OOD-GNN [35] mitigates instance- wise heterogeneity by eliminating spurious correlations between irrelevant and relevant graph representations through nonlinear decorrelation and sample reweighting. Yao et al. [74] focus on explicitly model domain correlations and spurious features and adapt to each test instance's 's unique distribution shifts. While these methods explicitly consider instance heterogeneity in graph OOD problems, they often focus on specific types of distribution shifts or rely on the assumption that target data adhere to certain designated shift types.
3

---

In contrast, our method introduces a novel paradigm that decomposes any unknown shift into multiple shift components and learns predictors that can adapt to the graph heterogeneity observed in the target data. By leveraging a mixture-of-expeer architecture, our approach can handle complex distribution shifts without assuming specific shift types or relying solely on group patterns.
Mixture-of-expert models. The applications of mixture-of-expert models (MoE) [24, 57] have largely focused on their efficiency and scalability [13, 12, 53, 9], partieularly in image and language domains. For image domain generalization, Li et al. [31] focus on neural architecture design and integrate expert models with vision transformers to capture correlations in the training dataset that may benefit generalization, where an expert is responsible for group of similar visual attributes. Puigcerver et al. [52] observed improved robustness by adopting MoE models in the image domain. In the graph domain, differently motivated from our work, Wang et al. [64] consider experts as information aggregation models with varying hop sizes to capture different ranges of message passing, aiming to improve model expressiveness on large-scale data.
GraphMETRO is the first to design a mixture-of-expert model specifically tailored address complex distribution shifts in graphs, coupled with a novel objective for producing invariant representations. While previous methods mostly focus on either node- or graph-level tasks, GraphMETRO is more general solution applicable to both.
3 Method
Problem formulation. For simplicity, we consider a graph classification task and later extend to node-level tasks. Let D. be the source distribution and D, be an unknown target distribution. We are interested in the natural graph distribution shifts. Our goal is to learn model so with high generalization ability. The standard approach is Empirical Risk Minimization (ERM), ie..
0 = arg min E(g. C (fe (9), 3),
where L denotes the loss function and y is the label of the graph G. However, the assumptions underlying ERM can be easily violated, making 6* suboptimal. Moreover, since the distribution shift is unknown and cannot provide supervision for model training, the direct optimization of Eq is intractable.
3.1 Shift Components
Based on the common mixture pattern studied in real-world networks [29, 30, 50], we propose the following informal assumption:
Assumption (An equivalent mixture for distribution shifts) Let the distribution shift between
the source DD and target Dj distributions be the result of an unknown intervention in the graph formation mechanism. We assume that the resulting shift in Dj can be modeled by up out of K classes of stochastic transformations applied to each instance in the source distribution DD (k S K).
Assumption 1 essentially states that any distribution shift can be decomposed into k shift components
of stochastic graph transformations. The assumption simplifies the generalization problem by enabling the modeling of individual shift components that constitute the shift and their respective contributions to the overall distribution shift. While this assumption is generally applicable, as observed in the experiments, we include a discussion on scenarios that fall outside the scope of this assumption in Appendix F. Previous works [28, 69, 67] implicitly infer such shift components from the data environments constructed based on the source distribution. However, distilling diverse shift components from the source data is challenging due to the complexity of the graph distribution shifts and largely depends the constructed environments?.
Graph extrapolation as shift components. To construct the shift components, we employ data extrapolation technique based on the source data. In particular, we introduce K independent classes of transform functions, including multihop subgraph sampling, the addition of Gaussian feature
In other words, if the distribution shifts were described via environment assignments, one would have a combinatorial number of such environments, i.e., the product of all different subsets of nodes and all their possible distinct shifts.

---

noise, and random edge removal [54]. The i-th class, governed by the i-th shift component, defines a stochastic transformation 7; that transforms an imput source graph & into an output graph (9)) where i 1,..., K. For instance, can be defined to randomly remove edges with an edge-dropping probability the range of [0.3, 0.5]. Note that the extrapolation aims to construct the basis of shifts rather than directly conducting data augmentation, explained in Eq 3 later.
3.2 Mixture of Aligned Experts
In light of the shift components, we formulate the generalization problem as two separate phases:
Surrogate estimation: Identify mixture of shift components the surrogate for the target shifts, where the mixture can vary across different node or graph instances to capture heterogeneity.
Mitigation and aggregation: Mitigate individual shift components, followed by aggregating the representations output by each expert resolve the surrogate shift.
Overview. Inspired by the mixture-of-experts (MoE) architecture [24], the core idea of GraphMETRO build hierarchical architecture composed of gating model and multiple expert models, where the gating model predicts the influence of the shift components on given instance. For the expert models, we design each to handle an individual shift component. The experts produce referential representations invariant to their designated shift component, with the representations aligned in a common representation space. Finally, our architecture combines the expert outputs into a final representation, which our training objective ensures is invariant to the stochastic transformations within the mixture distribution. We detail each module as follows:
Gating model. We introduce a GNN á as the gating model, which takes any graph as input and outputs weight vector w on the shift components. The weight vector suggests the most probable shift components from which the input graph originates. For example, in Figure 2b, given an unseen graph with decreased graph size and node feature noise, trained gating model should assign large weights to the corresponding shift components and small values to the irrelevant ones. Note that ø should be such that w;, the weight on the i-th component, strives be sensitive the stochastic transformation т; but insensitive to the application of other stochastic transformations Tj. j A i. This way, determining whether the i-th component is present should not depend on other components.
Expert models. We build K expert models, each corresponding to a shift component. Formally, we denote an expert model & 9 + R" where v is the hidden dimension, and we use z; (g) to denote the output representation. Each expert model essentially produces invariant representa- tions [48] with respect the distribution shift controlled by its assigned shift component. However, independently optimizing each expert without properly aligning the expert's 's output space is incom- patible with model training. Specifically, an expert model may learn its own unique representation space, which may cause information loss when its output is aggregated with other expert outputs. Moreover, aggregating independent representations results in mixed representation space with high variance, which makes it difficult for the predictor head, such as multi-layer perceptrons (MLPs), to capture the interactions and dependencies among these diverse representations and output rational predictions. Thus, aligning the representation spaces of experts is necessary to ensure compatibility and facilitate stable model training. To align the experts' output spaces properly, we introduce the concept of referential invariant representation:
Definition (Referential Invariant Representation) Let G be input graph and let T be ar- bitrary stochastic transform function, with domain and co-domain in the spuce of gruphs. Let So be model that encodes graph into representation. A referential invariant representation w.r.t. the given T is denoted as (G), where g* is a function that maps the original data & to a high- dimensional representation E" (G) such that Eo (g) se E* (T()) holds for every G E supp(D.), where supp(D.) denotes the support of D.. We refer to Eo as the reference model.
Thus, the representation space of the reference model serves as an intermediate to align different experts, while each expert & has its own ability produce referential invariant representations w.r.t. a stochastic transform function T; i= 1...., K. We include the reference model as a special "in-distribution" expert model on the source data.
Architecture design for the expert models. Further, we propose two architecture designs for the expert models. A straightforward way is to construct (K + 1) GNN encoders to generate referential
5

---

invariant representations for individual shift components. This ensures model expressiveness while increasing memory usage due multiple encoders. To alleviate this concern, we provide an alternative approach. Specifically, we can construct a shared module, e.g., GNN encoder, among the expert models, coupled with specialized module, e-g,, an MI.,, for each expert. We discuss the impact of architecture choices model performance in the experiment section.
The MoE workflow. Given a node or graph instance, the gating model assigns weights w IR g+++ over the expert models, indicating the mixture of shift components on the instance. The output weights, being conditional on the imput instance, enable the depiction of heterogeneous distribution shifts that vary across instances. After that, we obtain the output representations from the expert models, which eliminate the effect of the corresponding shift component. Then, the final representation is computed via aggregating the representations based on the weight vector, i.e.,
h(g) = Aggregatel (.)))))))))))) i 0, 1, K})
where h the encoder of f. The aggregation function can be weighted sum over the expert outputs or selection function that selects the expert output with maximum weight, e.g.,
(G) Softmax(w) [2o...., _k]
Assuming the distribution shift on an instance controlled by any single shift component, we have h(T((G) = (G)) Eo( h()) for O..... K, where (G) (o( holds according to Definition 1. This indicates that h automatically produces referentially invariant representations while allowing heterogeneity across different instances, e.g. different shift types control strengths. For clarity, we define 7(k) as joint stochastic transform function composed of any k fewer transform functions out of the K transform functions. We refer to the scenario where h produces referentially invariant representations w.r.t. 7(²) as -nvarrinne To extend higher orders (k 1), we design the objective Section 3.3, which enforces model generalization when multiple shifts exist. After that, a classifier p takes the aggregated representation from Eq for prediction tasks. Thus, we have f f o as the mixture-of-experrs model.
satisfy ( ensuring
3.3 Training Objective
As shown in Figure 2b, consider three trainable modules, i.e., the gating model ø the expert models { and the classifier μ. We propose the following objective:
min CJ = min L1 + L2), where
#
C1 B(G.y) E. (()( (((()))
C2 B(G.)) (o) (E(N(M(T(*)(T) v)) +X d(k(( (o)))
L1: (k)) E {0, 1] is the ground truth vector, and its i-th element is if and only if r. composes (kk) BCE is the Binary Cross Entropy. This term indicates that the gating model ø is o optimized to accurately predict a mixture of shift components.
C2: CE is the Cross Entropy function. d(-,) -) is distance function between two representations, and is a parameter controlling the strength of the distance penalty. In the experiments, we use the Frobenius norm as the distance function, i.e,, d(z1, Z2) = 1 ". | - Zz l|. for all the experiments. The second loss the classifier, and we prevent it from backpropagating to the gating model to avoid interference. Specifically, La aims improve the encoder's performance in predicting graph classes and achieves referential alignment with the reference model Eo via the distance function. Note that, when k > 1, L2 also enforces h to be invariant to multiple shifts via the (k) -invariine condition.
, 22:)² 2 optimizes the expert models and
and we
We optimize our model via stochastic gradient descent, where (()) is sampled at each gradient step. Overall, GraphMETRO yields a MoB model, comprising a gating model with high predictive accuracy, expert models that are aligned and can generate invariant representations in a shared representation space, and a spsc-fp classifier that utilizes robust and invariant representations for class prediction.

---

Table 1: Test results on the real-world datasets.
We compute the p-value between the results of GraphMETRO and the state-of- offte-of-the-ee wi-the-arrt methods. The results of GraphMETRO is repeated five times.
Node classification
Graph classification
Require domain
information
WebKB
Twitch
Twitter
SST2
ERM
14.29 3.24
48.95 3.19
56.44 0.45
80.52 f: 1.13
No
DANN
15.08 + 0.37
48.98
3.22
55.38 + 2.29
80.53 + 1.40
No
IRM
13.49 0.75
47.21
+ 0.98
55.09 + 2.17
80.75 + 1.17
Yes
VREx
14.29 = 3.24
48.99 3.20
55.98 + 1.92
80.20 + 1.39
Yes
GroupDRO
17.20 + 0.76
0.76
47.20 + 0.44
56.65 + 1.72
81.67 + 0.45
Yes
Deep Coral
13.76 + 1.30
49.64 + 2.44
55.16 + 0.23
78.94 + 1.22
Yes
SRGNN
13.23 2.93
47.30
1.43
NA
NA
Yes
BERM
24.61 = 4.86
51.34 1.41
NA
NA
OODGAT
14.41 = 1.10
49.38 0.87 NA
NA
NA
DIR
55.68
2.21
81.55 1.06
No
G-Mixup
NA
53.32
2.75
77.43 + 1.97
GSAT
NA
NA
56.40
1.76
81.49 + 0.76
No
CIGA
NA
NA
55.70
1.39
80.44 + 1.24
No
GraphMETRO
41.11 + 7.47
53.50 2.42
57.24
2.56 0.042
81.87
0.22
No
p-value
0.001
0.023
0.081
3.4 Discussion and Analysis
Node classification tasks. While we introduce our method following a graph-level task setting, GraphMETRO readily adaptable for node-level tasks. Instead of generating graph representations. GraphMETRO is capable of producing node-level invariant representations. Additionally, we ap- ply stochastic transform functions to the subgraph containing a target node and identify its shift components, which is consistent with the objective in Equation 3.
The gating model of GraphMETRO predicts the shift components the node graph instance, which provides interpretations and insights into the distribution shifts in unknown datasets. In contrast, existing research on GNN generalization [69, 45, 6, 67] often lacks proper identification and analysis of distribution shifts prevalent in real-world datasets. This creates a gap between human understanding of graph distribution shifts and the actual graph dynamics. To bridge this gap. we offer an in-depth study of the experiments to demonstrate GraphMETRO insights into the complexity of real graph distributions.
Interpretability.
Computational cost. The forward process of f requires O(K) encoder passes, using the weighted sum aggregation from (K + 1) expert outputs. Since the extrapolation process increases the dataset size by a factor of (K + 1), the training computation complexity is ((K2)D,)), where [D,] the size of the source dataset.
4 Experiments
We perform systematic experiments on both real-world (Section 4.1) and synthetic datasets (Sec- tion 4.2) validate the generalizability of GraphMETRO under complex distribution shifts.
4.1 Applying GraphMETRO to Real-world Datasets
We perform experiments on real-world datasets, which introduce complex and natural distribution shifts. In these scenarios, the test distribution may not preciscly align with the mixture mechanism encountered during training.
Datasets. We use four classification datasets, i.e. WebKB [51], Twitch [55]. Twitter [78], and GraphSST2 [78, 58], using the dataset splits from the GOOD benchmark [20], which exhibit various real-world covariate shifts. Specifically, WebKB is a 5-class prediction task that predicts the classes of university webpages, with nodes split based on different university domains, demonstrating natural challenge of applying GNNs trained on some university data to other unseen data. Twitch is binary classification task that predicts whether user streams mature content, with nodes split mainly by user language domains. Twitter and GraphSST2 real-worke grammar tree graph datasets,
7

---

nobry frature DS1
moddm sph
ealy Irature
modom Cagrsphi1
Robty feature
modaa
aod odge (41
Fode
and edge 141
node (2)
aud edge 141
node 12)
add odge 14
note 12)
zrop edge 120
dros edga 15)
dras edge 15)
drep voge 131
Aug Gp
(4.5)
0.8
0.0
(2. 31
12, 31
51
12.
12. 51
12
(2. 5)
12.
0..5
(1,
11.
Figure 3: Accuracy synthetic distribution shifts. The first row shows the testing accuracy on single shift components. We label the distribution by the clockwise order. The second row shows the testing accuracy on distribution shifts with multiple shift components, where each testing distribution a composition of two different transformations. For example, (1, 5) denotes a testing distribution where each graph is controlled by random subgraph (1) and noisy feature (5) shift components. We include the numerical values in Appendix E.
where graphs from different domains differ in sentence length and language style, posing direct challenge of generalizing to different language lengths, styles, and contexts."
Baselines. We use ERM and domain generalization baselines, including DANN [17], IRM [1], VREx [28]. GroupDRO [56], and Deep Coral [62]. Moreover, we compare GraphMETRO with robustnessigeneralizatinn techniques for GNNs, including DIR [69], OODGAT [59], GSAT [45], and CIGA [6] for graph classification tasks, and SR-GCN [85], EERM [67], and G-Mixup [22] for node classification tasks.
Training and evaluation. We Additionally, we include the results of using a shared module among experts in Appendix D. due to space limitations. For evaluation metries, we use ROC-AUC on Twitch and classification accuracy on the other datasets following [20]. See Appendix B for details about the architectures and optimizer.
individual GNN encoder for each expert in the experiments.
Results. In Table
observe that GraphMETRO consistently outperforms the baseline models across all datasets. It achieves notable improvements of 67.0% and 4.2% relative to EERM on the WebKB and Twitch datasets, respectively. When applied to graph classification tasks, Graph- METRO shows significant improvements, as the baseline methods exhibit similar performance levels. Importantly, GraphMETRO can be applied to both node- and graph-level tasks, whereas many graph-specific methods designed for generalization are limited to one of these tasks. Additionally, GraphMETRO does not require any domain-speciffe information during training.
Main Conclusion. The observation that GraphMETRO is the best-performing method demonstrates its significance for real-world applications, as it excels in handling unseen and wide-ranging distribu- tion shifts. This adaptability is crucial, as real-world graph data often exhibit unpredictable shifts that affect model performance. Thus, GraphMETRO versatility ensures its reliability across diverse domains, safeguarding performance in complex real-world scenarios. In Appendix D.2 and D.3, we provide two studies the impact of the alignment term controlled by À and the stochastic transform function choices on the model performance, analyzing the sensitivity and success of GraphMETRO.
4.2 Inspect GraphMETRO on Synthetic Datasets
Following the experiments on weal-world datasets, we perform experiments on synthetic datasets to further inspect and validate the effectiveness of our approach.
*ee specifically exclude datasets with synthetic shifts from the GOOD benchmark. We leave the applications to molecular datasets in the GOOD benchmark for future work, as it requires designing shift components based expert knowledge.
8

---

6
0.5
6x
18.5%
14.0%
14.8%
13.7%
8x
21.0%
27.7%
-0.3
6
29.8%
6
Noly feature
nodes
Add edges adges
Randorn subgraph
Noisy feature
Add edges
Random subgraph
Drop node
Drop edges
(b) Mixture of distribution shifts on WebKB (left)
(a) Invariance matrix Twitter dataset
and Twitch (right) identified by GraphMETRO.
Figure 4: (a) Invariance matrix on the Twitter dataset. Lighter colors indicate a higher invariance of representations produced by each expert. Small values on the diagonal elements of the invariance matrix indicate that each expert excels at generating invariant representations w.r.f. the specific shift component. (b) Mixture of distribution shifts identified by GraphMETRO. Higher values indicate strong shift component in the testing distribution.
Datasets. We use graph datasets from citation and social networks. For node classification tasks, we use DBLP [16] and CiteSeer [73]. For graph classification tasks, we use REDDIT-BINARY and IMDB-MULTI [46]. See Appendix B for dataset processing and details of the transform functions.
Training and evaluation. We adopt the same encoder architecture for Empirical Risk Minimization (ERM), ERM with data augmentation (ERM-Aug), and the expert models of GraROMETRO. For ERM-Aug training, we augment the training datasets using the same transform functions we used to construct the testing environments. Finally, we select the model based on the in-distribution validation accuracy and report the testing accuracy on each environment from five trials. See Appendix for detailed settings and hyperparameters.
Results. Figure illustrates our model's performance across single (the first row) and multiple (the second row) shift components. In most test distributions, GraphMETRO exhibits significant improvements or performs on par with two other methods. Notably. on the MUDB-MULTI dataset with noisy node features, GraphMETRO outperforms ERM-Aug by 5.9%, and it enhances performance on DBLP by 4.4% when dealing with random subgraph sampling. In some instances, GraphMETRO even demonstrates improved results on in-distribution datasets, such as 2.9% and 2.0% boost on Reddit- BINARY and DBLP, respectively. This could be attributed the increased model expressiveness of the MoE architecture or weak distribution shifts that can exist in the randomly split testing datasets.
4.3 Invariance Matrix for Inspecting GraphMETRO
key insight from GraphMETRO is that each expert excels in generating invariant representations specifically for shift component. To delve into the modeling mechanism, denote I I K x as an invariance matrix. This matrix quantifies the sensitivity of expert & to the j-th shift component. Specifically, for i [K] and E [K]. we have
lg Eg-D, B,, (G) (G)
Ideally, for a given shift component, the representation produced by the corresponding expert should be most similar the representation produced by the reference model. That is, the diagonal entries Is should be smaller than the off- off-diagonal entries I, for j * and i = 1...., K. In Figure 4a, we visualize the normalized invariance matrix computed for the Twitter dataset, revealing a pattern that aligns with the analysis. This demonstrates that GraphMETRO effectively adapts to various distribution shifts, indicating that our approach generates consistent invariant representations for each of the shift components.
4.4 Distribution Shift Discovery
With the trained MoE model, we aim to understand the distribution shifts in the target distribution. Here conduct case studies on the WebKB and Twitch datasets. Specifically, we first validate the

---

gating models' ability to identify mixtures, which is multitask binary classification with (K + 1) classes. The gating models achieve high accuracies of 92.4% on WebKB and 93.8% on the Twitch dataset. As mixtures output by gating models identify significant shift components on an instance, we leverage them as uuman-understandable interpretations and compute the average mixture across g E D, the global mixture on the target distribution. The results in Figure 4b show that the shift component, increased edges, dominates on the WebKB dataset, while the shift components controlling, e.g., node features and decreasing nodes, show large effects the Twitch dataset. The results align with dataset structures, i.e., WebKB's natural shifts across different university domains and Twitch's s language- based shifts. While quantitatively validating these observations in complex graph distributions remains a challenge, we aim to explore these complexities greater depth in future work, which can potentially offer insights into real-world graph dynamics.
5 Conclusion and Future Work
This work mitigates the challenge of improving the generalization of Graph Neural Networks (GNNs) to real-world data splits and dynamic graph distributions. To tackle these shifts, we introduce GraphMETRO, conture-of-allgnee--ee
architecture, which models graph distribution shifts as mixtures of shift components, each controlling shifts in unique directions with varying complexity.
GraphMETRO distinguishes itself from traditional invariant learning methods, which often rely on environment variables to partition data. Instead, our method treats distribution shifts mixtures, represented by the gating function's 's score vector, allowing for infinite environments due to the continuous nature of the score. When restricted to binary outputs, GraphMETRO can simulate finite environments, making it flexible and versatile. Furthermore, the introduction of referential invariant representation via reference model is key innovation of our approach.
Experimental results demonstrate that GraphMETRO consistently outperforms baseline methods on real-world datasets, achieving significant improvements. Additional synthetic studies and case analyses further validate the method's effectiveness and adaptability across diverse scenarios.
In future work, we aim to explore the broader applicability of GraphMETRO, including potential extensions to address label distributional shifts. Detailed discussions on these directions are provided in Appendix F.
10

---

References
[1] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization, arXiv preprint arXiv:1907.0288 2019.
[2] Tanya Y. Torger-Wolf- and Jared Saia. A framework for amalysis of dynamic social networks. In SIGKDD, 2006.
[3] Beatrice Bevilacqua, Yangze Zhou, and Bruno Ribeiro. Size-invariant graph representations for graph classification extrapolations. In JCML, 2021.
[4] Davide Buffelli, Pietro Lió, and Fabio Vandin. Sizeshiftreg: regularization method for improving sizeggeneralization in graph neural networks. In NeurIPS, 2022.
[5] Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Aréchiga, and Tengyu Ma. Learning imbalanced datasets with aubel-distribution--ae
margin loss. In NeurIPS, 2019.
[6] Yonggiang Chen, Yonggang Zhang, Yatao Bian, Han Yang, Kaili Ma, Binghui Xie, Tongliang Liu, Bo Han, and James Cheng. Learning causally invariant representations for out-of- distribution generalization graphs In NeurlPS, 2022.
[7] Kaize Ding, Zhe Xu, Hanghang Tong, and Huan Liu. Data augmentation for deep graph learning: A survey. SIGKDD, 2022.
[8] Mucong Ding, Kezhi Kong, Jiuhai Chen, John Kirchenbauer, Micah Goldblum, David Wipf, Furong Huang, and Tom Goldstein. A closer look at distribution shifts and duttrf-distiibuttin generalization on graphs. In NeurlPS DistShift, 2021.
[9] Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten P. Bosma, Zongwei Zhou, Tao Wang, Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen S. Meier-Hellsterm, Toju Duke, Lucas Dixon, Kun Zhang, Quoc V. Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. Glam: Efficient scaling of language models with mixture-of-experts. In ICML, 2022.
[10] Vijay Prakash Dwivedi, Chaitanya K. Joshi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. J. Mach. Learn. Res., 2023.
[11] Shaohua Fan, Xiao Wang, Chuan Shi, Peng Cui, and Bai Wang. Generalizing graph neural networks on duitr--disttibution graphs. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
[12] William Fedus, Jeff Dean, and Barret Zoph. A review of sparse expert models in deep learning. abs/2209001667, 2022.
[13] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. 2022.
[14] W. Feng, J. Zhang, Y. Dong, Y. Han, H. Luan, Q. Xu, Q. Yang, E. Kharlamov, and J. Tang. Graph random neural networks for sumi-super learning on graphs. Advances in Neural Information Processing Systems, 33:22092-223 22103, 2020.
[15] Wenzheng Feng, Jie Zhang, Yuxiao Dong, Yu Han, Huanbo Luan, Qian Xu, Qiang Yang, Evgeny Kharlamov, and Jie Tang. Graph random neural networks for remi-supervised learning on graphs. In NewrIPS, 2020.
[16] Xinyu Fu, Jiani Zhang, Ziqiao Meng. and Irwin King. MAGNN: metapath aggregated graph neural network for heterogeneous graph embedding. In WWW, 2020.
[17] Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversari training of neural networks. The journal of machine learning research, 2016.
[18] Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E, Dahl. Neural message passing for quantum chemistry. In ICML, 2017.

---

[19] Derek Greene, Dónal Doyle and Padraig Cunningham. Tracking the evolution of communities in dynamic social networks. In ASONAM, 2010.
[20] Shurui Gui, Xiner Li, Limei Wang, and Shuiwang Ji. GOOD: A graph out-of-distribution benchmark. In NeurIPS, 2022.
[21] William L. Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In NeurIPS, 2017.
[22] Xiaotian Han, Zhimeng Jiang, Ninghao Liu, and Xia Hu. G-mixup: Graph data augmentation for graph classification. In International Conference on Machine Learning. ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, 2022.
[23] Wei Jin, Tong Zhao, Jiayuan Ding, Yozen Liu, Jiliang Tang, and Neil Shah. Empowering graph representation learning with test-time graph transformation. In ICLR, 2023.
[24] Michael I. Jordan and Robert A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural Comput., 1994.
[25] Thomas N. Kipf and Max Welling, Semi-supervised classification with graph convolutional networks. ICLR, 2017.
[26] Boris Knyazev, Graham W. Taylor, and Mohamed R. Amer. Understanding attention and generalization in graph neural networks. In NeurlPS, 2019.
[27] K. Kong, G. Li, M. Ding, Z. Wu, C. Zhu, B. Ghanem, G. Taylor, and T. Goldstein. Robust optimization as data augmentation for large-scale graphs. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 60-69, 2022.
[28] David Krueger, Ethan Caballero, Joem-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville, Out-of-distribttion generalization via risk extrapo- lation (REx). In ICML, 2021.
[29] Jure Leskovec, Jon M. Kleinberg, and Christos Faloutsos. Graphs over time: densification laws, shrinking diameters and possible explanations. In SIGKDD. ACM, 2005.
[30] Jure Leskovec, Jon M. Kleinberg, and Christos Faloutsos. Graph evolution: Densification and shrinking diameters. ACM Trans. Knowl. Discov. Data, 2007.
[31] Bo Li, Yifei Shen, Jingkang Yang, Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of-experts are domain generalizable learners. In ICLR, 2023.
[32] H. Li, Z. Zhang, X. Wang, and Zhu. Disentangled graph contrastive learning with indepen- dence promotion. IEEE Transactions on Knowledge and Data Engineering, 2022.
[33] Haoyang Li, Xin Wang, Ziwei Zhang, and Wenwu Zhu. Out-of-distribution generalization on graphs: A survey. CoRR, 2022.
[34] Haoyang Li, Ziwei Zhang, Xin Wang, and Wenwu Zhu. Learning invariant graph representations for -ut-of-distribution generalization. In NeurIPS, 2022.
[35] Haoyang Li, Xin Wang. Ziwei Zhang, and Wenwu Zhu. OOD-GNN: out-of-distribution generalized graph neural network. IEEE Trans. Knowl. Data Eng., 2023.
[36] Sihang Li, Xiang Wang. An Zhang, Yingxin Wu, Xiangnan He, and Tat-Seng Chua. Let invariant rationale discovery inspire graph contrastive learning. In ICMI. 2022.
[37] Wanyu Lin, Hao Lan, and Baochun Li. Generative causal explanations for graph neural networks. In ICML, 2021.
[38] Gang Liu, Tong Zhao, Jiaxin Xu, Tengfei Luo, and Meng Jiang. Graph rationalization with environment-based augmentations. In KDD '22: The 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Washington, DC, USA, August 14 18, 2022, 2022.
12

---

[39] S. Liu, R. Ying. H. Dong, L. Li, T. Xu, Y. Rong, P. Zhao, J. Huang, and D. Wu. Local augmentation for graph neural networks. In International Conference on Machine Learning. pages 14042. PMLR, 2022.
[40] Songtao Liu, Rex Ying, Hanze Dong, Lanqing Li, Tingyang Xu, Yu Rong, Peilin Zhao, Junzhou Huang, and Dinghao Wu. Local augmentation for graph neural networks. In ICMI., 2022.
[41] Bin Lu, Xiaoying Gan, Ze Zhao, Shiyu Liang. Luoyi Fu, Xinbing Wang, and Chenghu Zhou. Graph out-offdistribution generalization with controllable data augmentation. CoRR, 2023.
[42] Jiaqi Ma, Junwei Deng, and Qiaozhu Mei. Subgroup generalization and faimess of graph neural networks. In NeurlPS, 2021.
[43] Julian J. McAuley and Jure Leskovec. Learning to discover social circles in ego networks. In NeurlPS, 2012.
[44] Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, and Sanjiv Kumar. Long-tail learning via logit adjustment. ICLR, 2021.
[45] Siqi Miao, Mia Liu, and Pan Li. Interpretable and generalizable graph learning via stochastic attention mechanism. ICMI., 2022.
[46] Christopher Morris, Nils M. Kriege, Franka Bause, Kristian Kersting, Petra Mutzel, and Marion Neumann. Tudataset: A collection of benchmark datasets for learning with graphs. In ICML 2020 Workshop Graph Representation Learning and Beyond (GRL+ 2020), 2020. URL www. graphlearning.......... io.
[47] M. E. J. Newman. Mixing patterns in networks. Phys. Rev. E, 67:026126, Feb 2003.
[48] Sinno Jialin Pan, Ivor W. Tsang, James T. Kwok, and Qiang Yang. Domain adaptation via transfer component analysis. IEEE Trans. Neural Networks, 2011.
[49] H. Park, S. Lee, S. Kim, J. Park, J. Jeong, K.-M. Kim, J.-W. Ha, and H. J. Kim. Metropolis- hastings data augmentation for graph neural networks. Advances in Neural Information Pro- cessing Systems, 34, 2021.
[50] Leto Peel, Jean- Charles Delvenne, and Renaud Lambiotte. Multiscale mixing patterns in networks. 2017.
[51] Hongbin Pei, Bingzhe Wei, Kevin Chen-Chuan Chang, Yu Lei, and Bo Yang. Geom-gen: Geometric graph convolutional networks. ICLR, 2020.
[52] Joan Puigcerver, Rodolphe Jenatton, Carlos Riquelme, Pranjal Awasthi, and Srinadh Bhojana- palli. On the adversarial robustness of mixture of experts. In NeurlPS, 2022.
[53] Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, An- dré Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. In NeurlPS, 2021.
[54] Yu Rong, Wenbing Huang, Tingyang Xu, and Junzhou Huang. Dropedge: Towards deep graph convolutional networks on node classification. In /CLR, 2020.
[55] Benedek Rozemberezki and Rik Sarkar. Characteristic functions on graphs: Birds of feather, from statistical descriptors to parametric models. In CIKM, 2020.
[56] Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint urXiv: /e////9910081..11 2019.
[57] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of- experts layer. In ICLR, 2017.
13

---

[58] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Y. Ng, and Christopher Potts. Recursive deep models for semantic compositionality over sentiment treebank. In EMNLP, 2013.
[59] Yu Song and Donglin Wang. Learning graphs with out-of-distribution nodes. In KDD '22: The 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Washington, DC, USA, August 14 - 18, 2022, 2022.
[60] Yongduo Sui, Xiang Wang. Jiancan Wu, Min Lin, Xiangnan He, and Tat-Seng Chua. Causal attention for interpretable and generalizable graph classification. In SIGKDD, 2022.
[61] Yongduo Sui, Qitian Wu, Jiancan Wu, Qing Cui, Longfei Li, Jun Zhou, Xiang Wang, and Xiangnan He. Unleashing the power of graph data augmentation covariate distribution shift. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10- 16, 2023, 2023.
[62] Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In ECCV, 2016.
[63] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. ICLR, 2018.
[64] Haotao Wang, Ziyu Jiang, Yan Han, and Zhangyang Wang. Graph mixture of experts: Learning on large-scale graphs with explicit diversity modeling. 2023.
[65] L. Wu, H. Lin, Y. Huang, and S. Z. Li. Knowledge distillation improves graph structure augmentation for graph neural networks. In Neural Information Processing Systems, 2022.
[66] Man Wu, Shirui Pan, Chuan Zhou, Xiaojun Chang, and Xingquan Zhu. Unsupervised domain adaptive graph convolutional networks. WWW, 2020.
[67] Qitian Wu, Hengrui Zhang, Junchi Yan, and David Wipf. Handling distribution shifts on graphs: An invariance perspective. In ICLR, 2022.
[68] Y. Wu, A. Bojchevski, and H. Huang. Adversarial weight perturbation improves generalization in graph neural networks. In Association for the Advancement of Artificial Intelligence, 2023.
[69] Ying-Xin Wu, Xiang Wang, An Zhang. Xiangnan He, and Tat seng Chua. Discovering invariant rationales for graph neural networks. In ICLR, 2022.
[70] Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanic Jegelka. How powerful are graph neural networks? In ICLR, 2019.
[71] Nianzu Yang, Kaipeng Zeng, Qitian Wu, Xiaosong Jia, and Junchi Yan. Learning substructure invariance for out-of-distribution molecular representations. In NeurIPS, 2022.
[72] Y. Yang, Z. Feng, M. Song, and X. Wang. Factorizable graph convolutional networks. Advances in Neural Information Processing Systems, 33:20286-20296, 2020.
[73] Zhilin Yang, William W. Cohen, and Ruslan Salakhutdinov. Revisiting semi-supervised learning with graph embeddings. In ICML, 2016.
[74] Huaxiu Yao, Xinyu Yang, Xinyi Pan, Shengchao Liu, Pang Wei Koh, and Chelsea Finn. Improving domain generalization with domain relations. In ICLR, 2024.
[75] Gilad Yehudai, Ethan Fetaya, Eli A. Meirom, Gal Chechik, and Haggai Maron. From local structures to size generalization in graph neural networks. In ICML, 2021.
[76] Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William I.. Hamilton, and Jure Leskovec. Graph convolutional neural networks for web-scale recommender systems. In KDD. ACM, 2018
[77] Junchi Yu, Jian Liang, and Ran He. Finding diverse and predictable subgraphs for graph domain generalization. CoRR, 2022.
14

---

[78] Hao Yuan, Haiyang Yu, Shurui Gui, and Shuiwang Ji. Explainability in graph neural networks: taxonomic survey, IEEE Trans. Pattern Anal. Mach. Intell., 2023.
[79] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In ICLR, 2017.
[80] S. Zhang, K. Kuang. J. Qiu, J. Yu, Z. Zhao, H. Yang, Z. Zhang, and F. Wu. Stable prediction graphs with agnostic distribution shift. arXiv preprint arXiv:2110003865, 2021.
[81] Zeyang Zhang, Xin Wang, Ziwei Zhang, Haoyang Li, Zhou Qin, and Wenwu Zhu. Dynamic graph neural networks under spatio-temporal distribution shift. In NeurIPS, 2022.
[82] Tong Zhao, Yozen Liu, Leonardo Neves, Oliver J. Woodford, Meng Jiang, and Neil Shah. Data augmentation for graph neural networks. AAAI, 2021.
[83] Y. Zhou, G. Kutyniok, and B. Ribeiro. Ood link prediction generalization capabilities of message-passing gnns larger test graphs. Advances in Neural Information Processing Systems, 2022.
[84] Yangze Zhou, Gitta Kutyniok, and Bruno Ribeiro. OOD link prediction generalization capabili- ties of message- paasagg-passing gnns in larger test graphs. In NeurIPS, 2022.
[85] Qi Zhu, Natalia Ponomareva, Jiawei Han, and Bryan Perozzi. Shift-robust gnns: Overcoming the limitations of localized graph training data. 2021.
15

---

A Theoretical Analysis
We provide theoretical justification for why GraphMETRO can effectively address complex graph distribution shifts and outperform existing approaches. Our analysis focuses on three key aspects: (1) the limitations of existing methods, (2) how GraphMETRO overcomes these limitations, and (3) the theoretical guarantees of our approach.
A.1 Limitations of Existing Approaches
Consider graph classification task with the following causal model:
where is the input graph, is the label, E is an unobserved environment variable, and C' is an unobserved confounder.
Existing approaches primarily focus learning environment-invariaat predictors f(g) such that:
P(Y))(G), E P(Y)S(G), Ve E
However, these methods face significant challenges when:
The environment space E is vast and complex.
The distribution shifts are heterogeneous across instances.
The shifts involve multiple interacting components.
A.2 GraphMETRO's Approach
GraphMETRO addresses these limitations through its novel architecture and training objective:
1. Decomposition of shifts: Instead of learning single invariant predictor, GraphMETRO decom- poses the complex shift into K shift components:
= (E,, Eg. Ex)
(5)
where each E; represents specific type of graph transformation.
2. Mixture-of-Experte The gating model and expert models { allow for adaptive handling of heterogeneous shifts:
K Σ 2=0
hG))
w. (G)
(G)
(6)
where (G) o(G); instance-depen dependent weights.
Referential alignment: The training objective enforces alignment between expert outputs and reference model:
C2 EGGG) E,00 [CE(u(h((1()) + +
16

---

A.3 Theoretical Guarantees
now provide theoretical guarantees for GraphMETRO's performance:
Theorem ( For any graph G and shift component Tú, the encoder h satisfies:
h(r;(G)) = h())
Proof Given the gating model's ability to identify shift components and the expert models' invari- properties:
K
((T((9)) Σ
wy (T) ()(
(9)
j=0
w;(T:(()) ((((()))
(10)
-w E9(
(11)
h(g)
(12)
where the second equality holds because the gating model identifies and the third equality follows from the definition of referential invariant representation.
This theorem guarantees that GraphMETRO can handle individual shift components. We can extend this to combinations of shifts:
Theorem (Composition of Shifts) For any graph 9 and combination of k. shift components ((6)
Ti O Tig o o Tig the encoder h approximately satisfies:
h(z(() h(g)
(13)
Proof We prove this theorem by induction k, the number of shift components.
Base case (k = 1): This is directly given by Theorem 1.
Inductive step: Assume the theorem holds for k 1 shift components. We need to prove it holds for k
shift components.
Let (b) Tie ((--1) where -(k-1) Tr, Ti, Tik-1
((T()(()) N(T,, (+(+-2)
(14)
h(T (k- (G)) (by Theorem 1)
(15)
h()) (by induction hypothesis)
(16)
The approximation in the second line comes from the fact that the gating model ó may not perfectly identify the shift component Tie when applied after (1k-1), However, our training objective La explicitly minimizes:
E,ao]d((( (((xs)((((-)((()))))
(17)
This ensures that for compositions of shifts, the output of h remains close the reference model Eo, which invariant to all shifts.
Therefore, by induction, the theorem holds for any l 2 1.
Theorem 3 (Generalization Bound) Let )(:,)) be the cross-entropy loss. For any distribution DD resulting from a combination of shift components in 7(0), the generalization error of GraphMETRO satisfies:
(CC))) 3)] < E(G,y) [E((( (G)), v)]l +
(18)
where is small constant depending the complexity of the model and the number of samples.
17

---

Proof 3 1) Our training objective minimizes:
Lirin = (.G.)),,,(LC (C(TT///(((((0))))0 (G),, ())
(19)
2) Recall that u h, where is implemented single linear layer with softmax output, and the cross-entropy loss. This combination is stictly convex with respect to the inputs to μ if they perfectly collinear. Therefore, by Jensen's inequality:
E, (C((((((((() w)1 > C(u(EE (G((()
(20)
3) This implies:
Livan > B)G (E( ( (M((( v)),
(21)
hence, minimizing Lmain also minimizes the left hand side of the inequality.
4) Now, consider any target distribution Dj resulting from combination of shift components in 7(k)
By definition, we can express Dj as:
D. D. (T
D., ~ (()(8))}
(22)
where P(T(k)) is some distribution over the possible combinations of shift components.
5) Therefore:
((C,)))(((((())))) = E(G.o) Pro v)ll
S B, (B, [C(f(T (G)).
(23)
Cunin
The inequality in the second line holds because our training objective considers a wider range of transformations than those in the actual target distribution.
6) Equations (21) and (23) show that minimizing Lirain implies both finding model with low true risk and a model that is more invariant to - 7(k)(G), since Equation (21) shows the loss is lower if ddT E su)p(r(")), h(T(G)) = E, ( (()(((())))))
6) The gap between the risk and the empirical risk can be bounded by constant that depends on the complexity of the model and the number of samples, according to standard statistical learning theory. Therefore, we get:
(C()))),((((()))
(GG,)),, (CU +
(24)
These theoretical results demonstrate that GraphMETRO can effectively handle complex, heteroge- neous graph distribution shifts by:
Decomposing shifts into manageable components.
Adaptively combining expert models to handle instance-specific shifts.
Ensuring invariance to individual and combined shift components.
Providing tractable upper bound on the generalization error for shifted distributions.
Compared to existing approaches that struggle with vast environment spaces or heterogeneous shifts, GraphMETRO's adaptive mixture-of-experts architecture and alignment-based training objective provide a more flexible and scalable solution for real-world graph distribution shifts.
B Experimental Details
Experimental settings on synthetic datasets. We randomly split each dataset into training (80%), validation (20%), and testing (20%) subsets. We consider transformations for k: = 2, i.e., 7(2), which includes both single transformations and compositions of two different transformation functions. For the compositions, we exclude trivial combinations (e.g.. adding and dropping edges) and combinations that may result in an empty graph (e.g., random subgraph sampling and node dropping). These transformations are applied to the testing datasets to create multiple variants for testing environments.
18

---

Model architecture and optimization. We summarize the model architecture and hyperparameters for synthetic experiments (Section 4.2) in Table We use the Adam optimizer with weight decay set to O. The encoder (backbone) architecture, including the number of layers and hidden dimensions, selected based on validation performance from the ERM model and then fixed for each encoder during GraphMETRO training.
Node classification
Graph classification
DBLP
CiteSeer
IMDB-MULI REDDIT-BINARY
Backbone
Graph Attention Networks (GAT) [63]
Activation
PeLU
Dropout
0.0
Number of layers
3
3
2
2
Hidden dimension
64
32
128
128
Global pool
NA
NA
global add pool
glohal add pool
Epoch Batch
200
100
100
NA
NA
32
32
ERM Learning
Ic-3
le-3
lc-4
lc-3
GraphMETRO Learning
1c-3
1e-3
1c-4
lc-3
Table 2: Architecture and hyperparameters synthetic experiments.
For the real-world datasets, we the same encoder and classifier from the implementation of the GOOD benchmark*. The results for the baseline methods, except for Twitter (recently added to the benchmark), are reported by the GOOD benchmark. We summarize the architecture and hyperparameters used for real-world experiments below.
Node classification
Graph classification
WebKB
Twitch
Twitter
SST2
Graph Convolutional Network
Graph Isomorphism Network [70] w/ Virtual node [18]
Backbone
[25]
Activation
ReLU
Dropout
0.5
Number of layers
3
Hidden dimension
300
Global pool
NA
NA
global mean pool
global mean pool
Epoch Batch size
100
100
200
200
NA
NA
32
32
ERM Learning rate
le-3
le-3
le-3
le-3
GraphMETRO Learning rate
le-2
le-2
le-3
le-3
Table 3: Architecture and hyperparameters on real-world datasets.
For all datasets, we conduct a grid search for GraphMETRO learning rates due to the difference in architecture compared to traditional GNN models. GraphMETRO uses multiple GNN encoders, serving as expert modules.
C Stochastic Transform Functions
We built library of 11 stochastic transform functions on top of PyG*, and used 5 of them in our synthetic experiments for demonstration purposes. Each function allows for one or more hyperparameters control the degree of the transformation, such as the probability parameter in Bernoulli distribution for dropping edges. A certain amount of randomness is retained in each stochastic transform function, ensuring diversity in the generated graphs.
'https///iithub http:::///ituuu..oo///ggg www.ccoo//oo./eecccc.occ http://yyettee//rrooroogggotiiit
19

---

stochsttic.trrnssorn..
"mask_edge_feaa" feat'; MaskEdgeFeat (p, Hill_value,
moisy__dge__eeae : NoisyEdgeFeat
"edge_feat_shiitt' EdgeFeatShift
"mask_node_feat"
MaskNodeFeat (p, fill_value),
p
NoisyNodeFFett (p),
snode_feat_shifte NodeFeatShift (p),
"add_edge"" AddEdge(p).
"dgop DropEdge
cdrop_node". DropNode
idrop_pathl DropPath(p),
'random_ subbrrppp RandomSubgraph(4:
}
We observed that different sets or numbers of transform functions can impact model performance. Specifically, we use stochastic transform functions as the foundation for the decomposed target distribution shifts. Ideally, these functions should be diverse and cover different potential aspects of distribution shifts. However, using large number of transform functions increases the demand on the gating model's expressiveness, as it must distinguish between different transformed graphs. Additionally, more transform functions increase computational cost due to the larger number of experts. An ablation study Appendix D.3 further validates this analysis.
In practice, the stochastic transform functions proved effective on real-world datasets, suggesting their ability represent various distribution shifts. Exploring common base transform functions to better capture real-world distribution shifts would be an interesting direction for future research.
D
Ablation Studies
D.I
Design Choices of Expert Models
WebKB
Twitch
Twitter
SST2
GraphMETRO (original)
41.11
53.50
57.24
81.87
GraphMETRO (w/o LI)
23.22
50.58
56.14
78.98
GraphMETRO (Shared)
31.14
52.69
57.15
81.68
Table 4: Experiment results comparing different design choices for expert models. Results are averaged over five runs.
In the main paper, we discussed the trade-off between model expressiveness and memory utilization in expert model design. Here, we investigate configuration where multiple experts share GNN encoder but use individual MLPs customize their output representations. Table 4 presents the comparative results.
Our experiments show performance decrease when sharing the GNN encoder, which attribute to limitations in the expressiveness of the customized modules. This may hinder alignment with the reference model and reduce the experts' ability to remain invariant specific shift components. The concept of "being invariant to all shifts" using a shared module seems insufficient in this case. Nevertheless, this configuration still outperforms baseline models from Table 1, thanks the gating model's ss ability to selectively use relevant experts and the objective function's 's ability to generate invariant representations.
D.2 Alignment Design
When the alignment term is removed (A = 0), performance drops significantly, especially for WebKB, where accuracy falls from 41.11 to 18.79. This suggests that without alignment, the expert models develop distinct representation spaces, which, when aggregated, lead to higher variance and
20

---

WebKB
Twitch
Twitter
SST2
GraphMETRO (original) GraphMETRO (A 0)
41.11
53.50 50.88
57.24 56.97
81.87
18.79
81.15
Table 5; Validating GraphMETRO design to align expert models with the reference model.
loss of useful information. The predictor heads, such as MLPs, struggle to process these mixed representations. The alignment mechanism is thus crucial for maintaining coherent representation space, allowing the model to capture interactions more effectively and improving overall performance.
D.3
Choice of Transform Functions
0.60
0.40
0.35
0.55
0.30
0.50
3
6
#Stochastic transform functions
#Stochastic transform functions
(a) WebKB
(b) Twitter
Figure 5: Impact of transform function choices model performance. Each number of transform functions corresponds to specific set of transformations.
We investigate how the choice and number of stochastic transform functions impact the performance of GraphMETRO ranging from 2 to functions. These functions are applied in the following sequential order:
[noisy_node_fedt add _dg_, drop_edge, trop_node, sandon_subgraph, drop_path, mode_feat_shifft)
We use the first functions and their paired combinations (excluding trivial combinations like adding and dropping edges) during the training of GraphMETRO Due to computational constraints, we do not explore all possible combinations of the distinct functions but focus on specific sets of transformations.
Figure 5 shows the results on the WebKB and Twitter datasets. A consistent trend emerges: increasing the number of stochastic transform functions generally leads decline performance. For example, performance on WebKB drops from 42.4% to 31.9%. This decline can be attributed to: (1) some stochastic functions introducing noise unrelated the target distribution shifts, and (2) the gating model's ss expressiveness being insufficient to handle a larger number of transformations, leading to noisier predictions.
E
Numerical results of Accuracy on Synthetic Distribution Shifts
Tables 6 and 7 present the numerical results on synthetic datasets corresponding to Figure 3, enabling more detailed interpretation of the results. Additionally, we compute the average performance across different extrapolated testing datasets, showing an overall improvement.
F
Open Discussion and Future Works
Performance of the gating model. The performance of GraphMETRO depends in part on how effec- tively the gating model can identify distribution shifts from the transform functions. Some functions,
21

---

DBLP
CiteSeer
ERM
ERM-Aug
GraphMETRO
ERM
ERM-Aug
GraphMETRO
i.d. (0)
85.66
75.80
76.00
78.01
random subgruph (1)
84.48
85.29
85.78
75.47
75.82
77.01
drop node (2)
71.08
74.85
76.61
62.21
63.89
66.22
drop edge (3)
79.69
82.34
82.95
71.48
73.24
77.00
add odge (4)
83.41
84.44
84.98
74.29
74.87
77.26
noisy features (S)
76.90
72.81
85.28
88.43
(1, 77.63
81.04
81.71
70.37
71.42
74.97
(2. 81.99
83,65
84.26
73,60
74.06
76.11
(1, 4)
79.69
68.62
80.31
84,47
86.36
88.56
(2, 4)
70.55
74.01
75.10
62.13
63.53
65.73
5)
71.52
68.27
66.89
67.32
77.73
81.13
81.85
70.19
72.21
76.77
(3. 5)
79.59
84.49
87.14
78.24
73.29
89.18
(4, 5)
70.40
74.16
76.18
61.64
63.53
66.42
Average
77.88
78.63
81.08
72.29
72.41
76.36
Table 6: Numerical results synthetic node classification datasets
IMDB-MULTI
REDDIT-BINARY
ERM
ERM-Aug
GraphMETRO
ERM
ERM-Aug
GraphMETRO
Li.. (0)
50.17
49.28
49.16
72.93
73.02
75.94
random subgraph (1)
34.30
39.94
45.86
62.59
69.03
71.22
drup (2)
48.73
70.01
72.27
72.26
drop edge (3)
49. 66
48.99
48.83
59.13
70.55
72.51
add edge (4)
49.64
48.14
48.90
65.
67.28
69.34
noisy features (5)
50.17
49.28
49.16
68.66
68 68.50
66.79
(2, 3)
34.55
40.32
45.11
58.72
64.06
66.50
(1. 4)
34.32
40.28
46.01
59,40
65.29
34.57
40.17
46.79
61.34
66.02
66.71
(I.
49.31
48.36
48.68
65.89
66.88
68.09
(2, 5)
50.51
48.78
48.79
68.72
69.77
68.76
(3, 5)
49.38
47.72
48.35
55.36
65.21
64,87
(1,
48.36
61.08
62.57
34.62
39.88
46.15
62.99
68.68
68.34
Average 44.31
45.58
47.82
63.71
67.56
68.51
Table 7: Numerical results on synthetic graph classification datasets
like adding node feature noise and extracting random subgraphs, are inherently disentangled, making it easy for the gating model to differentiate between these distributions. Other functions, such as dropping paths and dropping edges, may more similar, but method remains robust as long as each expert produces the corresponding invariant representation. More complex combinations of transforms pose a greater challenge for the gating model's expressiveness. To address this, initializing the gating model with pre-trained model from diverse dataset may enhance its ability to predict mixtures, improving performance on unseen graphs
Comparison with invariant learning methods. GraphMETRO differs from traditional invariant learning, where environments are constructed using environment variables. Instead, GraphMETRO views distribution shifts on an instance as a mixture, represented by the score vector from the gating function. This approach enables the creation of infinite environments, the score vector continuous. When restricting the gating function to binary outputs, GraphMETRO can simulate finite environments, akin the environment construction invariant learning. Additionally, concept of referential invariant representation using the base model &o sets GraphMETRO apart from previous invariant learning approaches.
Applicability of GraphMETRO, A key question is how well the predefined transform functions capture complex distribution shifts.
General domain: In our experiments, we primarily use five universal graph augmentations (as listed in [82]). Our code also includes additional transforms (Appendix C). While these transforms are not exhaustive, they cover a wide range of shifts observed in results. However, real-world
22

---

distribution shifts may beyond the predefined transforms, and in such cases, GraphMETRO might struggle capture and mitigate unknown shifts. This is limitation when the test distribution domain knowledge is insufficient.
Specific domains: In certain domains, additional knowledge can help infer distribution shifts, such increase in malicious users in trading system. This knowledge guide the construction of transform functions to better cover the target distribution shifts. Specifically, two sources of knowledge can be used: Domain knowledge, e.g., molecular datasets, transform functions could add carbon structures to molecules while preserving functional groups, or in social networks, known user behaviors can guide transformations. ii) Leveraging samples from the target dis- tribution (i.e., domain adaptation), where samples from the target can inform the selection of relevant transforms. For example, by measuring the distance between the extrapolated datasets under specific transforms and the target samples in the embedding space, more relevant transform functions can be selected. This presents an interesting direction future work.
Label distribution shifts. In this work, we focus on distribution shifts in graph structures and features. Exiending GraphMETRO to handle label distribution shifts would be complementary and interesting direction. Label shifts affect various modalities, including graphs and images, and existing methods [44, 5] designed for label shifts could be integrated into our framework with minimal adjustments, such modifying the loss function or training pipeline.
23

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justifiiaton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We highlighted our contributions and the scope which is the generalization of Graph Neural Networks under complex distributions and natural graph variations.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
24

---

Justification:
We have discussed our limitations such as selecting the transform functions in details. please refer to the Appendix.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We have made the discussion about our training pipeline and how the other factors such hyperparameters can impact the final results.
Guidelines:
25

---

The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our data and code are available at wope///ooooouuuuup.......... science/ status /raphMER
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
squides/CodeSubmisse
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ inips.cc/publin http////bwlc//guiessss.ebeeeeeee http:///wgee/eeeeeeeeeceeee.eeee http://iiiigiiieee////esuuiiiiii SonodeSSumissiionf
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
26

---

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have made the discussion about our training pipeline and how the other factors such as hyperparameters can impact the final results.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
Table presents the computed p-value to indicate the statistical significance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula,
call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
We have provided the information on our GPUs used for training.
Guidelines:
The answer NA means that the paper does not include experiments.
27

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https::/ httpp////ccppccehthccceeeeeeeee
Answer: [Yes]
Justification:
We strictly followed the NeurlPS Code of Ethics
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
[NA]
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
28

---

Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We have cited all the sources we used to conduct the experiments.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
29

---

Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "peMfVJEUaT", "reviewer": "Reviewer_U8um", "summary": "The study investigates how to use mixture-of-experts (MoE) to model complex distribution shifts. It explores leveraging various expert modules, each tailored to different distribution shifts, to generate embeddings that enhance OOD generalization on graphs.", "questions": "1. How does GraphMetro perform on other datasets, such as OGBG-Molbace, OGBG-Molbbbp, or DrugOOD, which exhibit typical and single distribution shifts like scaffold shift?", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. This work proposes an novel idea to make use of representations from various distribution shifts for OOD generalization, which is different from previous invariant learning and distribution robustness optimization (DRO) literatures.\n\n2. The proposed framework can be applied for both node and graph-level classification.\n\n3. GraphMetro can achieve significant performance improvement in some datasets with complex distribution shifts, such as WebKB.", "weaknesses": "1. The terminology \"invariant\" in this work appears to be inconsistent with its usage in the invariant learning literature. In invariant learning, \"invariant\" refers to the features that perform stably across different environments. However, in this work, some claims regarding \"invariant\" seem misleading. For instance, the statement: \"we train the expert models to generate invariant representations with respect to their corresponding shift components\" (lines 69-70) is unclear to me. It is not clear how the expert model can generate invariant representations given the training objective in Equation 3. Different experts are optimized using empirical risk minimization (ERM) with specific data augmentations, and there seems to be no regularization to enforce the model to learn invariant features.\n\n\n2. In my opinion, GraphMetro actually utilizes spurious (style) features in different distribution shifts learned by the expert models, rather than invariant features as discussed in the paper. When the test dataset exhibits specific distribution shifts that align with the data augmentations in the training set, the gating model can select the representations learned by a subset of experts and make use of them.  A similar work that leverages spurious features for OOD generalization is [1], where the distribution shifts are reflected in different environments. Please correct me if I am wrong, further discussions regarding this are also welcomed.\n\n\n3. One of my concern of this method is that the distribution shifts are correlated with random data augmentation, which may restrict its efficacy. For instance, consider scaffold shift or assay shift, which are prevalent in molecular[2] or drug-discovery[3] datasets. It is unclear whether GraphMetro would be effective for such datasets, where there is only one type of distribution shift, which may not align with a specific type of random data augmentation. \n\n\n__Reference__\n\n1. Yao et al., Improving Domain Generalization with Domain Relations, ICLR 2024.\n2. Hu et al., Open Graph Benchmark: Datasets for Machine Learning on Graphs, NeurIPS 2020.\n3. Ji et al., DrugOOD: Out-of-Distribution Dataset Curator and Benchmark for AI-Aided Drug Discovery – a Focus on Affinity Prediction Problems with Noise Annotations, AAAI 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n## **Comment 1, 2: GraphMETRO's Invariant Representations**\n\n### **1) Definition of invariant representations**\n\nWe define an invariant representation with respect to a stochastic transform function $\\tau$ as:\n\n$\\xi^*(\\mathcal{G}) = \\xi^*(\\tau(\\mathcal{G}))$\n\nThis aligns with the concept of stability across different environments, where each $\\tau$ creates a different environment.\n\n### **2) Why each expert is invariant w.r.t. a shift component**\nEach expert $\\xi_i$ produces **referential invariant representations (c.f. definition 1)** with respect to its shift component $\\tau_i$:\n\n$\\xi_0(\\mathcal{G}) = \\xi_i(\\tau_i(\\mathcal{G}))$\n\nThis 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`.\n\n**Even with a single expert, the model's representation remains invariant,** as guaranteed by our Shift-Invariance Theorem, explained in point 4 later.\n\n### **3) Expert model invariance w.r.t. a shift component given the objective**\n\nThe distance term in Equation 3:\n\n$d(h(\\tau^{(k)}(\\mathcal{G})), \\xi_0(\\mathcal{G}))$\n\nencourages each expert to align its output with the reference model $\\xi_0$. This, combined with the gating model, leads to invariant representations by MoE.\n\n### **4) How the MoE model (gating + experts) generates invariant representations**\n\n**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})$\n\n**Proof:** Given the gating model's ability to identify shift components and the expert models' invariance properties:\n\n$$\nh(\\tau_i(\\mathcal{G})) = \\sum_{j=0}^K w_j(\\tau_i(\\mathcal{G})) \\cdot \\xi_j(\\tau_i(\\mathcal{G}))\n= w_i(\\tau_i(\\mathcal{G})) \\cdot \\xi_i(\\tau_i(\\mathcal{G}))\n= w_i(\\mathcal{G}) \\cdot \\xi_0(\\mathcal{G})\n= h(\\mathcal{G})\n$$\n\nThe second equality holds because the gating model identifies $\\tau_i$, and the third follows from the definition of referential invariant representation.\n\n**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.\n\n**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.\n\nWe 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.\n\nWe are grateful for the opportunity to address these points and present them better in our revision.\n\n**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!\n\n---\n\n## **Comment 2: Applying GraphMetro on molecular datasets**\n\nThanks for the insightful comments. We provide answers from several angles:\n\n- **Adaptability to Molecular Graphs**:\n   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.\n\n- **Incorporating Domain Knowledge**: \n   As discussed in `Appendix G`, we can use domain knowledge to construct relevant transform functions. For molecular datasets, this could involve:\n   - Adding or modifying carbon structures while preserving functional groups.\n   - Altering specific molecular substructures relevant in scaffold shifts.\n   - Simulating changes in molecular properties that occur in assay shifts.\n\n- **Single Shift Scenario**:\n   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`)\n\n- **Transfer Learning Approach**:\n   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.\n\nWe believe these adaptations would be interesting future directions to apply GraphMETRO on molecular datasets. Thank you for highlighting it!\n\n## **Summary**\n\nWe 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!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yuQiO7rDJr", "reviewer": "Reviewer_uY6m", "summary": "This paper proposes a novel model, GraphMETRO, that applies mixture of experts (MoE) to facilitate generalizable graph neural network (GNN). It uses different expert networks to mitigate different types of distribution shifts, such as adding edges, dropping nodes, noisy features, etc. When facing a distribution shift, the paper assumes that the distribution shift can be decomposed as a combination of different types of distribution shifts. Each expert network can generate invariant embeddings under a specific type of distribution shift. A gating network is in charge of deciding the types of distribution shifts. Therefore, GraphMETRO can automatically identify the types of distribution shifts and generate invariant embeddings.", "questions": "(1)   I am wondering how you obtain the ground truth labels of the types of distribution shift on the test set. I understand that for the training set, you can do data augmentation by adding different types of distribution shifts, so you have the ground truth, but for the test set, I am wondering about it.\n\n(2)   How did you train the baseline generalization/robustness models? How did you ensure a fair comparison between GraphMETRO and the baseline methods?", "limitations": "The authors have discussed the limitations in the appendix.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The studied problem is interesting and worth studying, and the motivation is meaningful. MoE has been used in domain transfer and domain generalization in vision and NLP tasks, but using it in GNN to facilitate distribution shift is novel and worth studying.\n\n(2) The model design is novel and reasonable. It also has good interpretability.\n\n(3) Extensive experiments verify the effectiveness of GraphMETRO.", "weaknesses": "(1)   The actual distribution shift can be very complex, but the paper only considers a few basic types of distribution shifts. They might not be able to represent some distribution shifts in the real world.\n\n(2)   Some writing notations are not strict enough and can be improved. For example, in Definition 1 (referential invariant representation), in line 196, it says $\\epsilon_0(G) = \\epsilon^*(\\tau(G))$, maybe it should be $\\approx$ rather than $=$. Also, in line 218, it might be better to use $\\approx$ instead of $=$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your positive feedback and detailed suggestions! We provide responses below to address your remaining concerns. \n\n## **Comment 1: The considered distribution shifts may not fully represent real-world shifts.**\n\nThanks 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:\n- **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:\n  - **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.\n  - **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.\n\n- **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.\n\n- **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.**\n\nWe appreciate your insights! We hope this discussion addresses your concerns.\n\n---\n\n### **Comment 2: Notations**\nThanks! We corrected the mentioned notations and rigorously checked other notations, including those used in our response to `reviewer ymW6` on the detailed theoretical analysis. \n\n---\n\n### **Question 1: How to obtain the ground truth labels of the types of distribution shift on the test set?**\n\nGood 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.\n\n---\n\n### **Question 2: How did you train the baseline generalization/robustness models? How did you ensure a fair comparison between GraphMETRO and the baseline methods?**\n\nSince 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. \n\nSpecifically, 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. \n\nTherefore, we ensured the **dataset settings (splits, node features)** and **architecture components (encoder, classifier, pooling layer, etc.)** were the same to ensure a fair comparison. \n\n---\n\n## **Summary**\n\nWe want to thank you for your insightful questions and approval of our motivation, novelty, and effectiveness. \n\n- For your applicability concern, we conducted more discussion to make it clear. \n- For notations, we hope we have resolved this issue with more rigorous checks.\n\nWe 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!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vHnAUhhmTH", "reviewer": "Reviewer_LHfk", "summary": "This paper proposed GraphMETRO, a novel Graph Neural Network (GNN) architecture that try to models complex and natural distribution shifts in real-world graph data. The key innovation is a mixture-of-experts (MoE) approach, where the model decomposes the distribution shift into multiple shift components, each handled by a dedicated expert model. The proposed  GraphMETRO is expected to identify the shift type for a given input graph, and  generate representations that are invariant to the assigned shift components. The experiments show that GraphMETRO can capture heterogeneous shifts across instances and outperform baseline methods on diverse real-world datasets exhibiting different distribution shifts.", "questions": "•\tWhy distributional difference can be decomposed into the proposed five shifts? How to deal with complex distribution shifts in the real world that are out of the scope of these five shifts?\n\n•\tShould the expert model and Gating be trained separately or together, and why was it designed this way?", "limitations": "please refer to the weakness", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "•\tThis paper proposes a novel method to addresses distribution shifts in real-world graph data from a MOE perspective.\n\n•\tThis paper is well-written and it provides a detailed experimental analysis. \n\n•\tThis empirical validation demonstrates the effectiveness and applicability of the proposed method.", "weaknesses": "•\tThe paper proposes five types of graph distribution shifts (adding edges, deleting edges, feature noise, subgraphs, deleting nodes), but real-world graph distribution shifts may be more varied, such as systematic changes in link preferences and the addition of malicious nodes.\n\n•\tThe comparison methods in the experiment lack the latest OOD methods such as OOD-GNN, OOD-GAT-ATT, and OOD-GMixup and suggested to supplement the experimental results on more datasets.\n\n•\tThe alignment design of different expert models lacks experimental support, and it is recommended to supplement with corresponding ablation experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful comments. Please find our responses to each point below.\n\n---\n\n## **Comment 1: Real-world graph shifts may be more varied than defined shifts**\n\nGreat question! We discussed this in `Appendix G: Open Discussion and Future Works`, focusing on `The applicability of GraphMETRO` under complex scenarios and potential mitigations:\n- **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:\n  - **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.\n  - **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.\n\n- **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.\n\n- **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.**\n\n\nWe appreciate your insights! We hope this discussion addresses your concerns.\n\n---\n\n## **Comment 2: Adding latest OOD methods such as OOD-GNN, OODGAT-ATT, and OOD-GMixup**\n\nThanks 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.\n\n|| WebKB     | Twitch| Twitter    | SST2       |\n|:--|:--|:--|:--|:--|\n| OODGAT| 14.41±1.10| 49.38±0.87 | -| -|\n| G-Mixup| -| -| 53.32±2.75 | 77.43±1.97 |\n| GraphMETRO| 41.11±7.47| 53.50±2.42 | 57.24±2.56 | 81.87±0.22 |\n\nThe 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$).\n\nWe 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).\n\nWe have added more discussions about these three works to the related work. Please see our response to `reviewer ymW6` if you are interested!\n\n---\n\n## **Comment 3: Experimental support of the alignment design of different expert models**\n\nGreat 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.\n\nTo formally validate this observation, here are the results (repeated three times for $\\lambda=0$):\n\n|                    | WebKB | Twitch | Twitter | SST2  |\n|--------------------|-------|--------|---------|-------|\n| $\\lambda=0$        | 18.79 | 50.88  | 56.97   | 81.15 |\n| $\\lambda=1$ (original results) | 41.11 | 53.50  | 57.24   | 81.87 |\n\nIn our paper, we attributed this effect to two reasons:\n\n- Each expert model may develop its own unique representation space, leading to information loss when aggregated with other expert outputs.\n- 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.\n\n**Therefore, training without alignment poses challenges.** We hope this validates our alignment design.\n\n---\n\n## **Summary**\n\nWe are grateful for your positive evaluation of our presentation, novelty, and effectiveness. We further provide the following responses:\n\n1. An in-depth discussion about applicability.\n2. Additional experiments on baselines.\n3. An additional ablation study on the alignment impact.\n\nWe 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.\n\nOverall, 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!\n\n---\n\n[1] Song et al. Learning on Graphs with Out-of-Distribution Nodes. KDD 2022.\n\n[2] Li et al. OOD-GNN: Out-of-Distribution Generalized Graph Neural Network.\n\n[3] Lu et al. Graph Out-of-Distribution Generalization with Controllable Data Augmentation.\n\n[4] Han et al. G-Mixup: Graph Data Augmentation for Graph Classification. ICML 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "msOATzGASt", "reviewer": "Reviewer_t31K", "summary": "This paper introduces a novel Graph Neural Network (GNN) architecture designed to handle complex distribution shifts in graph data. GraphMETRO uses a Mixture-of-Experts (MoE) model, with multiple expert models each targeting specific distribution shifts. Additionally, this paper uses a gating model to identify the most portable shifts. The approach aims to produce invariant representations, aligning expert outputs for smooth optimization. GraphMETRO demonstrates state-of-the-art results on various datasets, outperforming existing methods in generalizing to real-world distribution shifts.", "questions": "Based on the weakness section, the questions are as follows:\n\n1.\tDoes each expert encode undesired information from the reference model?\n2.\tCan the expert model consistently produce invariant features using GNN or MLP architectures?\n3.\tHow does the choice of the parameter k impact the model’s generalizability?", "limitations": "The authors adequately addressed the limitations and potential negative societal impact of the work.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1.\tThis paper proposes a framework that is able to decompose any distribution shift into multiple shift components.\n\n2.\tBased on the MoE structure, the authors design an architecture with a gating model and multiple expert models to identify strong distribution shift components.\n\n3.\tThe proposed method achieves state-of-the-art results on real-world GOOD benchmarks and provides interpretability of distribution shifts on unknown datasets", "weaknesses": "1.\tThe alignment of expert representations with the reference model using the Frobenius norm might lead to the incorporation of undesired information (e.g., node degree or feature noise). The impact of this on the generalizability of the proposed framework remains unclear.\n\n2.\tAlthough each expert model is intended to correspond to an individual shift component, the outputs of a GNN or MLP may include extraneous information beyond the distribution shift [1]. This raises questions about whether the expert model can consistently produce invariant features.\n\n3.\tThe number of experts is determined by the parameter k, which influences the identification of individual shift components. The choice of k could impact the generalizability of the framework, yet the effects of different k values on the framework’s performance are not clearly elucidated.\n\n[1] Li, Haoyang, et al. \"Disentangled Graph Self-supervised Learning for Out-of-Distribution Generalization.\" Forty-first International Conference on Machine Learning. 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your efforts and insightful comments! To address your concerns, we provide detailed responses below.\n\n---\n\n## **Comment 1: Does each expert encode undesired information from the reference model?**\n\nGreat 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.\n\nTo formally validate this observation, here are the results (repeated three times for $\\lambda=0$):\n\n|                    | WebKB | Twitch | Twitter | SST2  |\n|--------------------|-------|--------|---------|-------|\n| $\\lambda=0$        | 18.79 | 50.88  | 56.97   | 81.15 |\n| $\\lambda=1$ (original results) | 41.11 | 53.50  | 57.24   | 81.87 |\n\nIn our paper, we had attributed this effect to two reasons:\n\n- Each expert model may develop its own unique representation space, leading to information loss when aggregated with other expert outputs.\n- 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.\n\n**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.\n\n---\n\n## **Comment 2: Can the expert model consistently produce invariant features using GNN or MLP architectures?**\n\nGreat 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.\n\nPlease let us know if this answers your question; we are happy to provide further validation if needed.\n\n---\n\n## **Comment 3: How does the choice of the parameter $k$ impact the model’s generalizability?**\n\nThanks for asking! In `Appendix E: Study on the Choice of Transform Functions`, we have explored varying numbers of transform functions, i.e., $k$. \n\nDue 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.\n\nWe 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.\n\n---\n\n## **Summary**\n\nWe hope to address your concerns with:\n\n1. Our additional ablation study and discussion on expert alignment.\n2. Reference to the study on whether expert models consistently produce invariant features.\n3. Reference to the study on the parameter $k$.\n\nWe 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.\n\nThank you again!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9PWhPbRme1", "reviewer": "Reviewer_ymW6", "summary": "The manuscript presents a model for handling complex distribution shifts in graph data. The proposed GraphMETRO model employs a mixture-of-experts (MoE) architecture, with a gating model to identify distributional shifts and multiple expert models to generate shift-invariant representations. This method aims to generalize graph neural networks (GNNs) to non-synthetic distribution shifts occurring naturally in real-world graph data, achieving state-of-the-art results on several datasets.", "questions": "Could you provide detailed theoretical analyses to demonstrate why the model can excellently address the problem and perform better than the baselines?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper studies an interesting research problem that is handling complex distribution shifts in graph data. This research problem has been very hot recently.\n2. The model design is easy to understand. The paper provides a detailed explanation of the proposed model.\n3. The experiment demonstrates the effectiveness of the model. The performance improvement on some comparisons seems to be significant.", "weaknesses": "1. Although the motivation is clear to me, the novelty is one of my concerns since explicitly considering instance heterogeneity for graph OOD problems is not new. I think the authors could address my concern by presenting more differences with related works clearly.\n2. There are no detailed theoretical analyses to demonstrate why the model can well address the problem and perform better than the baselines.\n3. The reviews of the related works are not enough. Please refer to the paper Out-Of-Distribution Generalization on Graphs: A Survey for a more comprehensive discussion of the related works on graph OOD problems. For example, the discussions on invariant learning in section 2 are too brief. The main claim “the standard invariant learning approaches are not well-equipped to mitigate the complex distribution shifts” is a little arbitrary considering the graph OOD literature.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your detailed suggestions! We provide responses below:\n\n## **Comment 1,3: A thorough review on related works and highlights on the novelty**\n\nThanks for the comment! We added more related works, including those in [2]:\n\n- **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].\n- **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.\n\nPlease let us know if this clears the concern! Due to the character limit, we are happy to provide more detailed comparisons with individual works.\n\n## **Comment 2: Detailed theoretical analysis why GraphMETRO addresses the problem and outperforms baselines.**\n\nThank you for the opportunity to clarify the theory. Here is part of the updated Appendix A.\n\n---\n\n### **Theoretical Guarantees**\n**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})$ ...\n\n**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})$ ...\n\n**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:\n\n$$\n\\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\n$$\n\nwhere $\\epsilon$ is a small constant depending on the complexity of the model and the number of samples.\n\n**Proof:**\n\n1) Our training objective minimizes:\n\n$$\n\\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)]]\n$$\n\n2) 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:\n\n$$\n\\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)\n$$\n\n3) This implies:\n\n$$\n\\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)],\n$$\n\nhence, minimizing $\\mathcal{L}_\\text{train}$ also minimizes the right-hand side of the inequality.\n\n4) 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.\n\n5) Therefore:\n\n$$\n\\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)]]\n\\leq \\mathbb{E} _{(\\mathcal{G}, y) \\sim \\mathcal{D} _s} [\\mathbb{E} _{\\tau^{(k)}} [\\mathcal{L}(f(\\tau^{(k)}(\\mathcal{G})), y)]]\n= \\mathcal{L} _\\text{train}\n$$\n\nThe inequality here holds because our training objective considers a wider range of transformations than those in the actual target distribution.\n\n6) 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}))]$.\n\n7) The gap between the true risk and the empirical risk can be bounded by a constant $\\epsilon$. Therefore, we get:\n\n$$\n\\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\n$$\n\nThis 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.\n\n## **Summary**\n\nWe 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!\n\nReference: Please see the PDF.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# **General response**

We truly appreciate the reviewers' efforts and valuable suggestions in reviewing our paper. We are glad that most reviewers reached a positive consensus on our work's motivation, presentation, novelty, and experimental effectiveness. Since we received a decent number of reviews, we provide a summary on the reviewers’ major feedback and our corresponding actions: 

| | Reviewer U8um| Reviewer uY6m| Reviewer LHfk| Reviewer t31K| Reviewer ymW6| [Action] Referring to existing content | [Action] Modification | 
|:---|:---|:---|:---|:---|:---|:---|:---|
|**Motivation**| NA | “problem is interesting…motivation is meaningful”| NA | NA | “This paper studies an interesting research problem” | `NA`|`NA`|
|**Presentation**|“Presentation: good” | “Presentation: good” | “well-written and provides a detailed experimental analysis.” | “Presentation: good”| “The model design is easy to understand” |`NA`|`NA`|
|**Novelty**|”This work proposes a novel idea”|“The model design is novel and reasonable… good interpretability.”|”proposes a novel method to addresses distribution shifts” | “a framework that is able to decompose any distribution shift into multiple shift components.” | “....instance heterogeneity for graph OOD problems is not new…the authors could address my concern by presenting more differences with related works” | `NA` | RE Reviewer ymW6: `We justify our novelty with extensive comparison with related works` |
|**Empirical Improvements**|”can achieve significant performance improvement…such as WebKB.”| “Extensive experiments verify the effectiveness” | “This empirical validation demonstrates the effectiveness and applicability” |”achieves state-of-the-art results on real-world GOOD benchmarks and provides interpretability…”| “The performance improvement on some comparisons seems to be significant.” | `NA`|`NA`|
|**Applicability**| ”...can be applied for both node and graph-level classification”; Question about applying GraphMETRO on molecular graphs| “The actual distribution shift can be very complex” | “The paper proposes five types of graph distribution shifts…but real-world graph distribution shifts may be more varied” | NA | NA | `We point the reviewers to Appendix G, where we had discussed the applicability of GraphMETRO including on molecular datasets.` | `NA` |
| **Clarification**| (1) Definition of “invariance” and (2) Objective (Equation 3) | Suggest to improve notations | NA | NA | NA | `NA`  |RE Reviewer U8um: `We provide in-depth discussions on "invariance" and our objective`; RE Reviewer uY6m: `We improved our notations`;  |
| **Extended Evaluation**  | NA |NA | Suggest to add the latest OOD methods and more datasets  | NA| NA| `NA` | `We include two more baselines: (1) OODGAT suggested by the reviewer, and (2) G-Mixup (ICML 2022)`| 
| **Property of GraphMETRO**  | NA | NA | (1) Impact (ablation) of aligning expert representations | (1) Impact of aligning expert representations, (2) Whether the expert model can consistently produce invariant features, (3) The effects of different k values on the framework’s performance|Request on detailed theoretical analyses| RE Reviewer LHfk and t31K: `We point the reviewers to` (2) `Section 4.3, where we validated the invariance of the embeddings produced by the expert models`, and (3) `Appendix E where we study on the effect of the $k$ value` | RE Reviewer ymW6: `We provide more detailed theoretical analysis` RE Reviewer LHfk and t31K: (1) `We conduct ablation study to measure the impact of alignment`;  | 

------


# **Summary**

- We address comments on the applicability and properties of GraphMETRO by referring to our existing content. 
- For the other concerns, we mainly made modifications to the experiments and comparisons with recent graph OOD methods. 

We thank the reviewers for their suggestions, which have made our work more solid. We hope our responses clarify any confusion and alleviate the concerns. 

We would be thrilled if you could let us know whether your concerns have been addressed or if you have any follow-up questions!

— Best,

  Authors of Paper 5517

Meta-review: This paper proposes GraphMETRO, a graph neural network architecture designed to handle complex distribution shifts in graph data. It uses a mixture-of-experts approach with multiple expert models targeting specific distribution shifts and a gating model to identify shifts. The goal is to produce shift-invariant representations that generalize well to out-of-distribution data.

The reviewers unanimously agree to accept this paper provided that the additional experiments and theoretical analyses are included in a camera-ready version of the paper. Some reviewers further have suggestions to improve notations and the overall writing. 
Overall, the authors have addressed the major concerns raised by reviewers and the paper appears to make a solid contribution to graph OOD generalization. The mixture-of-experts approach is novel and the empirical results are strong.
Decision: Accept (poster)