GraphMETRO: Mitigating Complex Graph Distribution Shifts via Mixture of Aligned Experts

Authors: shirley  wu
kaidi  cao
bruno  ribeiro
james  zou
jure  leskovec
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 22:22:05

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.

Main paper: 
 introduction: 
 The intricate nature of real-world graph data introduces a wide variety of distribution shifts and [51, 20] . It illustrates (1) The distribution shift from source to target (the thick arrow in the upper right) and (2) Instance-wise heterogeneity in the target distribution (the thin arrows pointing to u 1 and u 2 ). heterogeneous graph variations [47, 30, 43, 26] . For instance, 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 contribute 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 u 1 and u 2 in the target domain exhibit varying degrees of change in their content features. These inherent shifts and complexity accurately characterize the dynamics of real-world graph data, e.g., social networks [2, 19] and ecommerce graphs [76] .𝒖 𝟏 𝒖 𝟐 𝒖 𝟐 𝒖 𝟏Above the diverse graph variants, Graph Neural Networks (GNNs) [21, 25, 10] have become a prevailing method for downstream graph tasks. Standard evaluation often As a simple example, the distribution shift from a target graph G ∈ D t to a source distribution D s is decomposed along three shift dimensions: graph size (ξ 1 ), node degree (ξ 2 ), and feature noise (ξ 3 ). Note that the shift components can be customized and expanded based on downstream tasks. (b) Architecture: Given an input graph, the gating model µ decomposes the instance-specific distribution shift into the contributions from the shift components. Then, each expert model ξ i (i > 0) is tasked with generating referential invariant representations (cf. Section 3 for the definition) w.r.t. an assigned shift component. ξ 0 is a 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 earlier.Previous research on GNN generalization has mainly focused on two lines: (1) Data-augmentation training procedures that learn environment-robust 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 environment-invariant 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 a 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 our 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 ∈ D t to the source distribution D s is 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 s . 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 a 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 or 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 representation. Consequently, we design a novel objective to ensure a smooth training process. Finally, during the evaluation process, we integrate outputs from both the gating and expert models for final representations. 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 networks. 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-world settings.The key benefits of GraphMETRO are as follows:• A novel paradigm: GraphMETRO provides a new approach to aid GNN generalization by decomposing and mitigating complex distributional shifts via a mixture-of-experts architecture. • Superior performance: It outperforms state-of-the-art 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. 
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 [1, 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 in 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, OOD-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-irrelevant 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-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 explicitly 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 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.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-experts 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] , particularly 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 a 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 to 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 a more general solution applicable to both.3 Method Problem formulation. For simplicity, we consider a graph classification task and later extend it to node-level tasks. Let D s be the source distribution and D t be an unknown target distribution. We are interested in the natural graph distribution shifts. Our goal is to learn a model f θ with high generalization ability. The standard approach is Empirical Risk Minimization (ERM), i.e.,EQUATIONwhere L denotes the loss function and y is the label of the graph G. However, the assumptions underlying ERM can be easily violated, making θ * suboptimal. Moreover, since the distribution shift is unknown and cannot provide supervision for model training, the direct optimization of Eq 1 is intractable.The gating model ϕ and expert models {ξ i } K i=0 allow for adaptive handling of heterogeneous shifts:EQUATIONwhere w i (G) = ϕ(G) i are instance-dependent weights. 
shift components: 
 Based on the common mixture pattern studied in real-world networks [29, 30, 50] , we propose the following informal assumption:Assumption 1 (An equivalent mixture for distribution shifts) Let the distribution shift between the source D s and target D t distributions be the result of an unknown intervention in the graph formation mechanism. We assume that the resulting shift in D t can be modeled by up to k out of K classes of stochastic transformations applied to each instance in the source distribution D s (k ≤ 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 on the constructed environments2 .Graph extrapolation as shift components. To construct the shift components, we employ a 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 noise, and random edge removal [54] . The i-th class, governed by the i-th shift component, defines a stochastic transformation τ i that transforms an input source graph G into an output graph τ i (G), where i = 1, . . . , K. For instance, τ i can be defined to randomly remove edges with an edge-dropping probability in the range of [0.3, 0.5]. Note that the extrapolation aims to construct the basis of shifts rather than directly conducting data augmentation, as explained in Eq 3 later.In light of the shift components, we formulate the generalization problem as two separate phases:• Surrogate estimation: Identify a mixture of shift components as 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 to resolve the surrogate shift.Overview. Inspired by the mixture-of-experts (MoE) architecture [24] , the core idea of GraphMETRO is to build a hierarchical architecture composed of a gating model and multiple expert models, where the gating model predicts the influence of the shift components on a 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 a 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, a 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 i , the weight on the i-th component, strives to be sensitive to the stochastic transformation τ i but insensitive to the application of other stochastic transformations τ j , j ̸ = 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 as ξ i : G → R v , where v is the hidden dimension, and we use z i = ξ i (G) to denote the output representation. Each expert model essentially produces invariant representations [48] with respect to the distribution shift controlled by its assigned shift component. However, independently optimizing each expert without properly aligning the expert's output space is incompatible 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 a 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 1 (Referential Invariant Representation) Let G be an input graph and let τ be an arbitrary stochastic transform function, with domain and co-domain in the space of graphs. Let ξ 0 be a model that encodes a graph into a representation. A referential invariant representation w.r.t. the given τ is denoted as ξ * (G), where ξ * is a function that maps the original data G to a highdimensional representation ξ * (G) such that ξ 0 (G) ≈ ξ * (τ (G)) holds for every G ∈ supp(D s ), where supp(D s ) denotes the support of D s . We refer to ξ 0 as the reference model.Thus, the representation space of the reference model serves as an intermediate to align different experts, while each expert ξ i has its own ability to produce referential invariant representations w.r.t. a stochastic transform function τ i , 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 invariant representations for individual shift components. This ensures model expressiveness while increasing memory usage due to multiple encoders. To alleviate this concern, we provide an alternative approach. Specifically, we can construct a shared module, e.g., a GNN encoder, among the expert models, coupled with a specialized module, e.g., an MLP, for each expert. We discuss the impact of architecture choices on model performance in the experiment section.The MoE workflow. Given a node or graph instance, the gating model assigns weights w ∈ R K+1 over the expert models, indicating the mixture of shift components on the instance. The output weights, being conditional on the input 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) = Aggregate({(ϕ(G) i , ξ i (G)) | i = 0, 1, . . . , K})where h is the encoder of f . The aggregation function can be a weighted sum over the expert outputs or a selection function that selects the expert output with maximum weight, e.g.,EQUATIONAssuming the distribution shift on an instance is controlled by any single shift component, we haveh(τ i (G)) = ξ i (τ i (G)) ≈ ξ 0 (G) = h(G) for i = 0, . . . , K, where ξ i (τ i (G)) ≈ ξ 0 (G)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 or control strengths.For clarity, we define τ (k) as a joint stochastic transform function composed of any k or fewer transform functions out of the K transform functions. We refer to the scenario where h produces referentially invariant representations w.r.t. τ (k) as τ (k) -invariance. To extend k to higher orders (k > 1), we design the objective in Section 3.3, which enforces h to satisfy τ (k) -invariance, ensuring model generalization when multiple shifts exist. After that, a classifier µ takes the aggregated representation from Eq 2 for prediction tasks. Thus, we have f = µ • h as the mixture-of-experts model.As shown in Figure 2b , we consider three trainable modules, i.e., the gating model ϕ, the expert models {ξ i } K i=0 , and the classifier µ. We propose the following objective:min θ L f = min θ (L 1 + L 2 ), whereEQUATION• L 1 : Y (τ (k) ) ∈ {0, 1} K+1 is the ground truth vector, and its i-th element is 1 if and only if τ i composes τ (k) . BCE is the Binary Cross Entropy. This term indicates that the gating model ϕ is optimized to accurately predict a mixture of shift components. • L 2 : CE is the Cross Entropy function. d(•, •) is a 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 , z 2 ) = 1 n ∥z 1 -z 2 ∥ F = 1 n n i=1 (z 1i -z 2i ) 2, and we use λ = 1 for all the experiments. The second loss term optimizes the expert models and the classifier, and we prevent it from backpropagating to the gating model to avoid interference. Specifically, L 2 aims to improve the encoder's performance in predicting graph classes and achieves referential alignment with the reference model ξ 0 via the distance function. Note that, when k > 1, L 2 also enforces h to be invariant to multiple shifts via the τ (k) -invariance condition.We optimize our model via stochastic gradient descent, where τ (k) is sampled at each gradient step. Overall, GraphMETRO yields a MoE 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 task-specific classifier that utilizes robust and invariant representations for class prediction. Node classification tasks. While we introduce our method following a graph-level task setting, GraphMETRO is readily adaptable for node-level tasks. Instead of generating graph representations, GraphMETRO is capable of producing node-level invariant representations. Additionally, we apply stochastic transform functions to the subgraph containing a target node and identify its shift components, which is consistent with the objective in Equation 3.Interpretability. The gating model of GraphMETRO predicts the shift components on the node or 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.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 O(K 2 |D s |), where |D s | is the size of the source dataset.The training objective enforces alignment between expert outputs and a reference model:EQUATION 
experiments: 
 We perform systematic experiments on both real-world (Section 4.1) and synthetic datasets (Section 4.2) to validate the generalizability of GraphMETRO under complex distribution shifts.We perform experiments on real-world datasets, which introduce complex and natural distribution shifts. In these scenarios, the test distribution may not precisely 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 a natural challenge of applying GNNs trained on some university data to other unseen data. Twitch is a binary classification task that predicts whether a user streams mature content, with nodes split mainly by user language domains. Twitter and GraphSST2 are real-world grammar tree graph datasets, where graphs from different domains differ in sentence length and language style, posing a direct challenge of generalizing to different language lengths, styles, and contexts. 3Baselines. 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 robustness/generalization 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 use an individual GNN encoder for each expert in the experiments. Additionally, we include the results of using a shared module among experts in Appendix D.1 due to space limitations. For evaluation metrics, 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.Results. In Table 1 , we 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-specific information during training.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 distribution shifts. This adaptability is crucial, as real-world graph data often exhibit unpredictable shifts that can 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 on 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 .Following the experiments on real-world datasets, we perform experiments on synthetic datasets to further inspect and validate the effectiveness of our approach. 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 GraphMETRO. 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 B for detailed settings and hyperparameters.Results. Figure 3 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 IMDB-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 a 2.9% and 2.0% boost on Reddit-BINARY and DBLP, respectively. This could be attributed to the increased model expressiveness of the MoE architecture or weak distribution shifts that can exist in the randomly split testing datasets.A key insight from GraphMETRO is that each expert excels in generating invariant representations specifically for a shift component. To delve into the modeling mechanism, we denote I ∈ R K×K as an invariance matrix. This matrix quantifies the sensitivity of expert ξ i to the j-th shift component. Specifically, for i ∈ [K] and j ∈ [K], we haveI ij = E G∼Ds E τj [d(ξ i (τ j (G)), ξ 0 (G))]Ideally, for a given shift component, the representation produced by the corresponding expert should be most similar to the representation produced by the reference model. That is, the diagonal entries I ii should be smaller than the off-diagonal entries I ij for j ̸ = i 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.With the trained MoE model, we aim to understand the distribution shifts in the target distribution.Here we conduct case studies on the WebKB and Twitch datasets. Specifically, we first validate the gating models' ability to identify mixtures, which is a 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 human-understandable interpretations and compute the average mixture across G ∈ D t as 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 on the Twitch dataset. The results align with dataset structures, i.e., WebKB's natural shifts across different university domains and Twitch's language-based shifts. While quantitatively validating these observations in complex graph distributions remains a challenge, we aim to explore these complexities in greater depth in future work, which can potentially offer insights into real-world graph dynamics. 
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, a mixture-of-aligned-experts 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 as mixtures, represented by the gating function'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 a reference model is a 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.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: Our data and code are available at https://anonymous.4open.science/ status/GraphMETRO Guidelines:• The answer NA means that paper does not include experiments requiring code. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or 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 or unintentional) misuse of the technology. • If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe 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)? Answer: [NA] Justification: [NA]  
 
 Appendix: We provide a 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.Consider a graph classification task with the following causal model:C E G Ywhere G is the input graph, Y is the label, E is an unobserved environment variable, and C is an unobserved confounder.Existing approaches primarily focus on learning environment-invariant predictors f (G) such that:EQUATIONHowever, 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.GraphMETRO addresses these limitations through its novel architecture and training objective:1. Decomposition of shifts: Instead of learning a single invariant predictor, GraphMETRO decomposes the complex shift into K shift components:EQUATIONwhere each E i represents a specific type of graph transformation.We now provide theoretical guarantees for GraphMETRO's performance:Theorem 1 (Shift-Invariance) For any graph G and shift component τ i , the encoder h satisfies:EQUATIONProof 1 Given the gating model's ability to identify shift components and the expert models' invariance properties:EQUATIONEQUATION)where the second equality holds because the gating model identifies τ i , 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 2 (Composition of Shifts) For any graph G and combination of k shift componentsEQUATIONProof 2 We prove this theorem by induction on 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.LetEQUATIONEQUATIONThe approximation in the second line comes from the fact that the gating model ϕ may not perfectly identify the shift component τ i k when applied after τ (k-1) . However, our training objective L 2 explicitly minimizes:EQUATIONThis ensures that even for compositions of shifts, the output of h remains close to the reference model ξ 0 , which is invariant to all shifts.Therefore, by induction, the theorem holds for any k ≥ 1.Theorem 3 (Generalization Bound) Let L(•, •) be the cross-entropy loss. For any distribution D t resulting from a combination of shift components in τ (k) , the generalization error of GraphMETRO satisfies:EQUATION)where ϵ is a small constant depending on the complexity of the model and the number of samples.Proof 3 1) Our training objective minimizes:EQUATION2) Recall that f = µ • h, where µ is implemented as a single linear layer with a softmax output, and L is the cross-entropy loss. This combination is stictly convex with respect to the inputs to µ if they are not perfectly collinear. Therefore, by Jensen's inequality:EQUATION3) This implies:EQUATION) hence, minimizing L train also minimizes the left hand side of the inequality. 4) Now, consider any target distribution D t resulting from a combination of shift components in τ (k) . By definition, we can express D t as:EQUATION)where P (τ (k) ) is some distribution over the possible combinations of shift components.Therefore:EQUATION)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 L train implies both finding a model with low true risk and a model that is more invariant to τ (k) (G), since Equation (21) shows the loss is lower if∀τ ∈ supp(τ (k) ), h(τ (G)) = E τ (k) [h(τ (k) (G))].6) The gap between the true risk and the empirical risk can be bounded by a constant ϵ that depends on the complexity of the model and the number of samples, according to standard statistical learning theory. Therefore, we get:EQUATIONThese theoretical results demonstrate that GraphMETRO can effectively handle complex, heterogeneous 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 a 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.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., τ (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.Model architecture and optimization. We summarize the model architecture and hyperparameters for synthetic experiments (Section 4.2) in Table 2 . We use the Adam optimizer with weight decay set to 0. The encoder (backbone) architecture, including the number of layers and hidden dimensions, is selected based on validation performance from the ERM model and then fixed for each encoder during GraphMETRO training. For the real-world datasets, we use the same encoder and classifier from the implementation of the GOOD benchmark 4 . 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. 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.We built a library of 11 stochastic transform functions on top of PyG5 , and used 5 of them in our synthetic experiments for demonstration purposes. Each function allows for one or more hyperparameters to control the degree of the transformation, such as the probability parameter in a Bernoulli distribution for dropping edges. A certain amount of randomness is retained in each stochastic transform function, ensuring diversity in the generated graphs. loss of useful information. The predictor heads, such as MLPs, struggle to process these mixed representations. The alignment mechanism is thus crucial for maintaining a coherent representation space, allowing the model to capture interactions more effectively and improving overall performance. We investigate how the choice and number of stochastic transform functions impact the performance of GraphMETRO , ranging from 2 to 7 functions. These functions are applied in the following sequential order:[noisy_node_feat, add_edge, drop_edge, drop_node, random_subgraph, drop_path, node_feat_shift]We use the first n 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 n 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 to a decline in 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 to the target distribution shifts, and (2) the gating model's expressiveness being insufficient to handle a larger number of transformations, leading to noisier predictions.Tables 6 and 7 present the numerical results on synthetic datasets corresponding to Figure 3 , enabling a 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 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 be more similar, but the 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 a pre-trained model from a 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, as the score vector is continuous. When restricting the gating function to binary outputs, GraphMETRO can simulate finite environments, akin to the environment construction in invariant learning. Additionally, the concept of referential invariant representation using the base model ξ 0 sets GraphMETRO apart from previous invariant learning approaches.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 our results. However, real-world distribution shifts may go beyond the predefined transforms, and in such cases, GraphMETRO might struggle to capture and mitigate unknown shifts. This is a limitation when the test distribution or domain knowledge is insufficient. • Specific domains: In certain domains, additional knowledge can help infer distribution shifts, such as an increase in malicious users in a trading system. This knowledge can guide the construction of transform functions to better cover the target distribution shifts. Specifically, two sources of knowledge can be used: i) Domain knowledge, e.g., in 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 distribution (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 for future work.Label distribution shifts. In this work, we focus on distribution shifts in graph structures and features. Extending GraphMETRO to handle label distribution shifts would be a 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 as modifying the loss function or training pipeline.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to 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 a 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 one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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.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.We 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 on expert knowledge.https://github.com/divelab/GOOD/tree/GOODv1https://github.com/pyg-team/pytorch_geometric 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.


2. 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.


3. 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. 


__Reference__

1. Yao et al., Improving Domain Generalization with Domain Relations, ICLR 2024.
2. Hu et al., Open Graph Benchmark: Datasets for Machine Learning on Graphs, NeurIPS 2020.
3. 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.
questions: 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: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: (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.

(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 $=$.
questions: 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.

(2)   How did you train the baseline generalization/robustness models? How did you ensure a fair comparison between GraphMETRO and the baseline methods?
limitations: limitations: The authors have discussed the limitations in the appendix.
paper_summary: paper_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.
strength_weakness: strength_weakness: •	The 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.

•	The 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.

•	The alignment design of different expert models lacks experimental support, and it is recommended to supplement with corresponding ablation experiments.
questions: questions: •	Why 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?

•	Should the expert model and Gating be trained separately or together, and why was it designed this way?
limitations: limitations: please refer to the weakness
paper_summary: paper_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.
strength_weakness: strength_weakness: 1.	The 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.

2.	Although 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.

3.	The 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.

[1] Li, Haoyang, et al. "Disentangled Graph Self-supervised Learning for Out-of-Distribution Generalization." Forty-first International Conference on Machine Learning. 2024
questions: questions: Based on the weakness section, the questions are as follows:

1.	Does each expert encode undesired information from the reference model?
2.	Can the expert model consistently produce invariant features using GNN or MLP architectures?
3.	How does the choice of the parameter k impact the model’s generalizability?
limitations: limitations: The authors adequately addressed the limitations and potential negative societal impact of the work.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. There are no detailed theoretical analyses to demonstrate why the model can well address the problem and perform better than the baselines.
3. 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.
questions: questions: Could you provide detailed theoretical analyses to demonstrate why the model can excellently address the problem and perform better than the baselines?
limitations: limitations: Yes

Comments: title: Thanks!
Thank you
About L_2 and DRO
comment: We thank Reviewer uY6m for the positive consent and for providing a clear summary of our work's applicability!
Thank you for your reply! Your response has resolved my concerns about the generalizability of the method. Since you can use domain knowledge and observations to define the types of distribution shifts of a specific application, and also due to the theoretical guarantee, it seems that GraphMETRO might be useful for various potential applications. I would like to raise my score to 7.
We thank reviewer ymW6 for the approval! We will make sure all references show full names appropriately.
Thank the authors for the detailed responses. I will increase my score.

A minor suggestion for the revised paper is to keep the newly added 8 references in the General Response PDF in the same format as the references in the original paper. I noticed that the formatting of references doesn't seem to be standardized, with abbreviations and full names being mixed up.
The authors addressed my concern with additional experiments and an ablation study, so I remain my positive score 6.
Thank you for your feedback and for raising the score!

To sum up, GraphMETRO's objective is for the experts to output referentially invariant representations, and further for the final MoE model to output invariant representations.

We have made explicit clarifications in the current revision, which should avoid any confusion regarding these distinctions. Thank you again for your constructive comments!
The author's response has effectively addressed my questions regarding how and why GraphMetro works, therefore I raise my score to 5. However, I still suggest that the authors make additional efforts to improve the clarity of the term "invariant" (e.g., in lines 75-76). This term may be misleading, as it refers to generating invariant representations with respect to referential representation, rather than representations that remain stable across different environments.
Thanks for the quick reply! We're glad to know that our response has clarified the method!

Yes, $k$ in $\mathcal{L}_2$ includes $k=0$. This mainly optimizes the reference model on the original training data with CE loss (since the second term in $\mathcal{L}_2$ is 0), making it an "in-distribution" expert.

**GraphMETRO and DRO**: This is a thoughtful point! Yes both GraphMETRO and DRO enforce model robustness. One of the essential differences is perhaps on how they model shifts (a continuous mixture of shift components v.s. defined perturbation sets). Due to the infinite number of mixtures, standard DRO’s worst-case minimization may not be directly applicable, leading us to seek more tractable solutions in GraphMETRO.
Thank you for the authors' response. Your reply has given me a clearer understanding of how the proposed method works. The proposed _GraphMetro_ seems to be a variant of Distribution Robust Optimization (DRO), where different distribution shifts are not too far from the referential distribution to enforce model robustness. I have one more question: Does $k$ include $0$ in $\mathcal{L}\_2$?