提出两级分解的AdaGrad方法,加速图神经网络分布式训练。
Two-level domain-decomposition AdaGrad method for scalable training of graph neural networks

- 将图分块后交替在全局与子图上优化,降低通信开销。
- 相比基线方法,在相同计算量下性能提升最高22%。
- 适合大规模图数据的分布式训练场景,尤其关注效率提升。
图神经网络(GNN)是处理图结构数据的强大框架,但其高效训练在分布式环境下仍具挑战,主要源于消息传递机制导致节点耦合,引发高计算成本、内存占用和通信开销。为此,我们提出一种新型域分解(DD)变体——DD-AG2m,它是结合二阶曲率信息与动量的AdaGrad方法(AG2m)的改进版本。该方法在原始(全局)图与划分后的子图之间交替进行AG2m优化。为进一步以低成本融入全局信息,我们进一步提出两级变体(2DD-AG2m),在每个子域内随机采样节点构建粗粒度图,于其上执行全局优化步骤。数值实验涵盖图分类、节点级回归与时空预测任务,结果表明,所提方法在达到相同预测性能时,计算成本降低4-8倍;在固定计算成本下,预测性能相比基线AG2m最高提升22%。
原文摘要 · Abstract (English)
Graph neural networks (GNNs) have emerged as a powerful framework for learning from graph-structured data. However, their efficient training remains challenging, particularly in distributed computing environments. This challenge arises from the use of message passing, which couples all graph nodes, leading to expensive optimization steps, high memory requirements, and substantial communication overhead. To alleviate these limitations, we propose a novel domain-decomposition (DD) variant of AG2m, an AdaGrad method enhanced with second-order curvature information and momentum, denoted by DD-AG2m. The proposed DD-AG2m alternates between AG2m optimization on the original (global) graph and AG2m optimization on the partitioned graphs. To incorporate global information at reduced cost, we further introduce a two-level variant (2DD-AG2m) that performs global optimization steps on a coarse graph obtained by randomly subsampling nodes within each subdomain. Numerical experiments spanning graph classification, node-level regression, and spatiotemporal forecasting tasks demonstrate that the proposed DD methods reduce the computational cost required to achieve the same predictive performance by a factor of 4-8. Moreover, for the fixed computational cost, they improve the predictive performance of GNNs by up to 22% compared with the baseline AG2m.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。