用虚拟节点提升图神经网络效率,处理超大规模几何图。
Fast and Distributed Equivariant Graph Neural Networks by Virtual Node Learning
- 引入虚拟节点近似真实无序图,分治消息传递保持区分性。
- 在11.3万节点的流体数据集上保持高精度,比现有方法快20倍。
- 适合处理超大规模分子、物理模拟等几何图任务的科研与工程人员。
等变图神经网络在多个科学领域表现卓越,但在扩展至大规模几何图时面临效率瓶颈,且在稀疏化输入后性能显著下降。为此,我们提出FastEGNN和DistEGNN,分别用于高效处理大规模稀疏图和超大规模分布式图。FastEGNN通过一组有序虚拟节点近似真实的无序图结构,对不同虚拟节点实施差异化消息传递与聚合机制,并最小化虚拟与真实坐标间的最大均值差异(MMD),以实现全局分布一致性。该设计使模型在保持高精度的同时高效处理大规模稀疏图。对于极端规模的几何图,我们进一步提出DistEGNN,将虚拟节点作为跨设备子图间的全局桥梁,有效降低内存与计算开销。我们在四个挑战性场景中评估:N体系统(100节点)、蛋白质动力学(800节点)、Water-3D(8,000节点)及新提出的Fluid113K基准(113,000节点)。结果表明,新模型在效率与性能上均显著优于现有方法,为大规模等变图学习开辟新可能。代码已开源。
原文摘要 · Abstract (English)
Equivariant Graph Neural Networks (GNNs) have achieved remarkable success across diverse scientific applications. However, existing approaches face critical efficiency challenges when scaling to large geometric graphs and suffer significant performance degradation when the input graphs are sparsified for computational tractability. To address these limitations, we introduce FastEGNN and DistEGNN, two novel enhancements to equivariant GNNs for large-scale geometric graphs. FastEGNN employs a key innovation: a small ordered set of virtual nodes that effectively approximates the large unordered graph of real nodes. Specifically, we implement distinct message passing and aggregation mechanisms for different virtual nodes to ensure mutual distinctiveness, and minimize Maximum Mean Discrepancy (MMD) between virtual and real coordinates to achieve global distributedness. This design enables FastEGNN to maintain high accuracy while efficiently processing large-scale sparse graphs. For extremely large-scale geometric graphs, we present DistEGNN, a distributed extension where virtual nodes act as global bridges between subgraphs in different devices, maintaining consistency while dramatically reducing memory and computational overhead. We comprehensively evaluate our models across four challenging domains: N-body systems (100 nodes), protein dynamics (800 nodes), Water-3D (8,000 nodes), and our new Fluid113K benchmark (113,000 nodes). Results demonstrate superior efficiency and performance, establishing new capabilities in large-scale equivariant graph learning. Code is available at https://github.com/GLAD-RUC/DistEGNN.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。