[论文解读] 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

一句话总结

该论文提出图混和专家（GMoE），一种GNN架构，通过为具有不同感受野的节点动态路由至专业聚合专家，增强模型对多样化图结构的建模能力。GMoE在ogbg-molhiv上提升ROC-AUC达1.81%，在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层引入多个可学习的聚合专家，每个专家具有不同的归纳偏置，包括不同跳数以适应不同的感受野范围。
• 每个节点通过路由器网络动态选择最合适的专家（或专家组合），依据其局部邻域结构与特征。
• 专家被训练以专门处理具有相似邻域模式的节点子集，实现高效且有针对性的信息聚合。
• 引入负载均衡损失，防止专家过度使用，确保专家集合中各专家的利用更加均衡。
• 该模型可无缝集成至现有GNN主干网络（如GCN与GIN），实现即插即用的性能提升。
• 动态路由机制支持稀疏激活，即使专家数量庞大，也能保持高效的推理速度。

实验结果

研究问题

• RQ1 Mixture-of-Experts框架是否能在保持计算效率的前提下，有效建模GNN中的多样化图结构？
• RQ2 基于邻域模式的动态专家路由如何提升GNN在异构图上的泛化能力？
• RQ3 引入具有不同跳数的专家在多大程度上能增强GNN对长距离依赖关系的建模能力？
• RQ4 在大规模图基准测试中，GMoE与标准GNN及MoE基线模型相比，在准确率与效率方面表现如何？

主要发现

• 与非-MoE基线相比，GMoE在ogbg-molhiv上提升ROC-AUC 1.81%，在ogbg-molbbbp上提升1.40%。
• 在ogbn-proteins上实现ROC-AUC 0.95%的提升，在ogbl-ddi上Hits@20提升0.89%。
• 经过充分训练后，GMoE超越单专家GCN与GIN，表明初始性能滞后可通过专家专业化得到克服。
• 采用λ=0.1的负载均衡策略相比λ=0，实现ROC-AUC提升2.58%，凸显其在性能优化中的关键作用。
• 将GMoE集成至Neural Fingerprints后，准确率达到82.72% ± 0.53%，相比原始模型提升0.4%，且计算量保持不变。
• 尽管模型容量更大，GMoE仍通过稀疏专家激活保持高推理效率，适用于万亿规模图的应用场景。