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

TL;DR

GraphMETRO introduces a novel Graph Neural Network architecture employing a Mixture-of-Experts approach to effectively model and mitigate complex, natural distributional shifts in graph data, achieving state-of-the-art results.

Contribution

It proposes a new MoE-based GNN architecture with a novel alignment objective to handle diverse distributional shifts in graph data.

Findings

Achieves 67% improvement on WebKB dataset.

Outperforms previous methods on four real-world datasets.

Effectively models complex distributional shifts.

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.