通过虚拟节点学习实现快速且分布式的等变图神经网络
机器学习
2025-06-25 v1 人工智能
摘要
等变图神经网络（Equivariant GNNs）在 diverse scientific applications 中取得了显著的成功。然而，现有方法在扩展到大型几何图时面临严峻的效率挑战，并且在输入图被稀疏化以实现可计算性时会出现显著的性能下降。为解决这些限制，我们引入了 FastEGNN 和 DistEGNN 两个针对大规模几何图的等变 GNN 的新型改进。FastEGNN 采用一种小且有序的虚拟节点集合有效地近似实际节点的大型无序图。具体而言，我们为不同的虚拟节点实现不同的消息传递和聚合机制，以确保其互不相同，并最小化虚拟坐标与真实坐标之间的最大均值差异（MMD），以实现全局分布性。这种设计使 FastEGNN 能够在保持高精度的同时高效处理大规模稀疏图。对于极大规模的几何图，我们提出了 DistEGNN，这是一种分布式扩展，其中虚拟节点充当不同设备中子图之间的全局桥梁，在保持一致性的同时大幅降低内存和计算开销。我们在四个具有挑战性的领域进行了全面评估：N 体系统（100 个节点）、蛋白质动力学（800 个节点）、Water-3D（8,000 个节点）以及我们新的 Fluid113K 基准（113,000 个节点）。结果表明我们的模型在效率和性能方面均表现出色，为大规模等变图学习开辟了新的能力。代码可在 https://github.com/GLAD-RUC/DistEGNN 查阅。
引用
@article{arxiv.2506.19482,
title = {Fast and Distributed Equivariant Graph Neural Networks by Virtual Node Learning},
author = {Yuelin Zhang and Jiacheng Cen and Jiaqi Han and Wenbing Huang},
journal= {arXiv preprint arXiv:2506.19482},
year = {2025}
}