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

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.

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. 2. The model design is easy to understand. The paper provides a detailed explanation of the proposed model. 3. 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. 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

Could you provide detailed theoretical analyses to demonstrate why the model can excellently address the problem and perform better than the baselines?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

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.

Strengths

1. This paper proposes a framework that is able to decompose any distribution shift into multiple shift components. 2. Based on the MoE structure, the authors design an architecture with a gating model and multiple expert models to identify strong distribution shift components. 3. The proposed method achieves state-of-the-art results on real-world GOOD benchmarks and provides interpretability of distribution shifts on unknown datasets

Weaknesses

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

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?

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors adequately addressed the limitations and potential negative societal impact of the work.

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.

Strengths

• This paper proposes a novel method to addresses distribution shifts in real-world graph data from a MOE perspective. • This paper is well-written and it provides a detailed experimental analysis. • This empirical validation demonstrates the effectiveness and applicability of the proposed method.

Weaknesses

• 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

• 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?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

please refer to the weakness

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.

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. (2) The model design is novel and reasonable. It also has good interpretability. (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. (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

(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?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have discussed the limitations in the appendix.

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.

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. 2. The proposed framework can be applied for both node and graph-level classification. 3. 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. 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

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?

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

N/A