提出消息不变性机制,让图神经网络训练快10倍且不丢精度。
Accurate and Scalable Graph Neural Networks via Message Invariance
- 用消息不变性把耗时的跨批次消息传递转为快速的组内传递
- 在百万级节点、百亿边的大图上提速数个数量级,精度损失极小
- 适合大规模图数据的快速训练,尤其对显存受限场景友好
基于消息传递的图神经网络在诸多实际应用中表现优异。对于一个采样的目标节点小批量,消息传递过程分为组内消息传递(MP-IB)和组外到组内消息传递(MP-OB)。然而,MP-OB依赖高阶外部邻居,导致计算量随层数呈指数增长。由于邻居爆炸问题,大部分节点和边需驻留于GPU显存,使许多GNN难以应用于大规模图。为此,我们提出一种高效且准确的小批量方法——拓扑补偿(TOP),仅通过组内消息传递即能获得完整消息传递的输出,无需昂贵的组外传递。其核心是引入消息不变性概念,定义消息不变变换,将耗时的MP-OB转化为快速的MP-IB,保证修改后的MP-IB输出与完整消息传递一致。实验表明,相较于现有方法,TOP在超大规模图(百万节点、百亿边)上速度提升达数量级,精度下降微乎其微。
原文摘要 · Abstract (English)
Message passing-based graph neural networks (GNNs) have achieved great success in many real-world applications. For a sampled mini-batch of target nodes, the message passing process is divided into two parts: message passing between nodes within the batch (MP-IB) and message passing from nodes outside the batch to those within it (MP-OB). However, MP-OB recursively relies on higher-order out-of-batch neighbors, leading to an exponentially growing computational cost with respect to the number of layers. Due to the neighbor explosion, the whole message passing stores most nodes and edges on the GPU such that many GNNs are infeasible to large-scale graphs. To address this challenge, we propose an accurate and fast mini-batch approach for large graph transductive learning, namely topological compensation (TOP), which obtains the outputs of the whole message passing solely through MP-IB, without the costly MP-OB. The major pillar of TOP is a novel concept of message invariance, which defines message-invariant transformations to convert costly MP-OB into fast MP-IB. This ensures that the modified MP-IB has the same output as the whole message passing. Experiments demonstrate that TOP is significantly faster than existing mini-batch methods by order of magnitude on vast graphs (millions of nodes and billions of edges) with limited accuracy degradation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。