[論文レビュー] Graph Mixture of Experts: Learning on Large-Scale Graphs with Explicit Diversity Modeling

Haotao Wang, Ziyu Jiang+6 more|arXiv (Cornell University)|Apr 6, 2023

Advanced Graph Neural Networks被引用数 7

ひとこと要約

本稿では、異なる受容 field を持つ専門的な集約エキスパートにノードを動的にルーティングすることで、多様なグラフ構造に対するモデル容量を向上させる GNN アーキテクチャである Graph Mixture of Experts (GMoE) を提案する。GMoE は、スパースなエキスパート活性化により推論効率を維持したまま、ogbg-molhiv で最大 1.81%、ogbg-molbbbp で最大 1.40% の ROC-AUC 向上を達成する。

ABSTRACT

Graph neural networks (GNNs) have found extensive applications in learning from graph data. However, real-world graphs often possess diverse structures and comprise nodes and edges of varying types. To bolster the generalization capacity of GNNs, it has become customary to augment training graph structures through techniques like graph augmentations and large-scale pre-training on a wider array of graphs. Balancing this diversity while avoiding increased computational costs and the notorious trainability issues of GNNs is crucial. This study introduces the concept of Mixture-of-Experts (MoE) to GNNs, with the aim of augmenting their capacity to adapt to a diverse range of training graph structures, without incurring explosive computational overhead. The proposed Graph Mixture of Experts (GMoE) model empowers individual nodes in the graph to dynamically and adaptively select more general information aggregation experts. These experts are trained to capture distinct subgroups of graph structures and to incorporate information with varying hop sizes, where those with larger hop sizes specialize in gathering information over longer distances. The effectiveness of GMoE is validated through a series of experiments on a diverse set of tasks, including graph, node, and link prediction, using the OGB benchmark. Notably, it enhances ROC-AUC by $1.81\%$ in ogbg-molhiv and by $1.40\%$ in ogbg-molbbbp, when compared to the non-MoE baselines. Our code is publicly available at https://github.com/VITA-Group/Graph-Mixture-of-Experts.

研究の動機と目的

• 実世界の GNN 応用において、均一な集約メカニズムが一般化性を制限するという課題に対処すること。
• 計算コストが著しく増大するか、学習性が損なわれることなく、大規模グラフにおける GNN の性能を向上させること。
• ノードの周辺構造に基づいて、動的かつ適応的なエキスパート選択を可能にし、専門性とインダクティブバイアスの多様性を促進すること。
• 多様なグラフ学習タスクにおいて、標準的な GNN や MoE ベースラインと比較して優れた精度-効率トレードオフを達成すること。

提案手法

• GMoE は、各 GNN 層ごとに、異なるインダクティブバイアス（特に異なるホップサイズを有する）を持つ複数のトレーニング可能な集約エキスパートを導入する。
• 各ノードは、その局所的近傍構造と特徴に基づいて、ルーター ネットワークを介して最も適切なエキスパート（複数可）を動的に選択する。
• エキスパートは、類似した近傍パターンを持つノードのサブグループに特化するように訓練され、効率的かつ的確な情報集約を可能にする。
• エキスパートの過剰使用を防ぎ、エキスパート エンsemble 全体での利用をバランスさせるために、負荷バランス損失が適用される。
• GCN や GIN などの既存の GNN バックボーンとシームレスに統合可能であり、即座に性能向上が得られる。
• 動的ルーティングにより、スパースな活性化が実現され、多数のエキスパートを有しても推論効率が維持される。

実験結果

リサーチクエスチョン

• RQ1 Mixture-of-Experts フレームワークは、計算効率を維持したまま、GNN における多様なグラフ構造を効果的にモデル化できるか？
• RQ2 近傍構造に基づく動的エキスパートルーティングは、異種グラフ間での GNN の一般化性をどのように向上させるか？
• RQ3 異なるホップサイズを持つエキスパートを組み込むことで、GNN における長距離依存性学習はどの程度向上するか？
• RQ4 大規模グラフベンチマークにおいて、GMoE は標準的な GNN や MoE ベースラインと比較して、精度と効率の両面で優れているか？

主な発見

• GMoE は、非-MoE ベースラインと比較して、ogbg-molhiv で 1.81%、ogbg-molbbbp で 1.40% の ROC-AUC 向上を達成した。
• ogbn-proteins では ROC-AUC が 0.95% 向上し、ogbl-ddi では Hits@20 が 0.89% 向上した。
• 十分な訓練を経た後、単一エキスパートの GCN や GIN を上回る性能を示し、エキスパートの特化により初期の性能遅れが克服されることを示した。
• λ=0.1 の負荷バランスにより、λ=0 と比較して ROC-AUC が 2.58% 向上し、これが性能向上において極めて重要な役割を果たしていることが示された。
• Neural Fingerprints に GMoE を統合することで、計算量を変更せずに精度が 82.72% ± 0.53% に向上し、元のモデルを 0.4% 上回った。
• 容量が大きくなっても、スパースなエキスパート活性化のおかげで、GMoE は高い推論効率を維持しており、トリンティリオンスケールの応用に適している。