[논문 리뷰] Graph Mixture of Experts: Learning on Large-Scale Graphs with Explicit Diversity Modeling

Haotao Wang, Ziyu Jiang외 6명|arXiv (Cornell University)|2023. 04. 06.

Advanced Graph Neural Networks인용 수 7

한 줄 요약

이 논문은 다양한 그래프 구조를 위한 모델 용량을 향상시키기 위해 수신장이 다른 전문화된 집합 전문가들로 노드를 동적으로 라우팅하는 그래프 믹스처 오브 은사(GMoe)를 제안한다. GMoE는 희소 전문가 활성화를 통해 추론 효율성을 유지하면서 ogbg-molhiv에서 최대 1.81% 향상된 ROC-AUC와 ogbg-molbbbp에서 1.40% 향상된 성능을 기록한다.

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 레이어마다 다양한 인덕티브 바이어스를 가진 다수의 학습 가능한 집합 전문가를 도입하며, 이는 다양한 힙 크기를 가진 수신장 범위를 제공한다.
• 각 노드는 국소적 이웃 구조와 특징에 기반해 라우터 네트워크를 통해 가장 적합한 전문가(또는 전문가들)를 동적으로 선택한다.
• 전문가들은 유사한 이웃 패턴을 가진 노드 그룹에 특화되도록 훈련되어, 효율적이고 타겟된 정보 집합을 가능하게 한다.
• 전문가 과도 사용을 방지하고 전문가 앙상블 전반의 균형 잡힌 활용을 보장하기 위해 로드 밸런싱 손실을 적용한다.
• GCN 및 GIN과 같은 기존 GNN 백본과의 통합이 가능하여 즉시 성능 향상을 얻을 수 있다.
• 동적 라우팅은 희소 활성화를 허용하여, 전문가 수가 많아져도 추론 효율성을 유지한다.

실험 결과

연구 질문

• RQ1 Mixture-of-Experts 프레임워크는 계산 효율성을 유지하면서도 다양한 그래프 구조를 효과적으로 모델링할 수 있는가?
• RQ2 이웃 패턴에 기반한 동적 전문가 라우팅은 이질적인 그래프 간의 GNN 일반화 능력을 어떻게 향상시키는가?
• RQ3 다양한 힙 크기를 가진 전문가를 통합할 경우 GNN의 장거리 종속성 학습 능력은 어느 정도 향상되는가?
• RQ4 대규모 그래프 벤치마크에서 GMoE는 표준 GNN 및 MoE 기준 모델에 비해 정확도와 효율성 측면에서 어떻게 비교되는가?

주요 결과

• GMoE는 non-MoE 기준 모델 대비 ogbg-molhiv에서 ROC-AUC를 1.81% 향상시키고, ogbg-molbbbp에서 1.40% 향상시켰다.
• GMoE는 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는 높은 추론 효율성을 유지하여 트리리온 규모 응용에 적합하다.