arXiv:2504.13266cs.LG2025-04被引 2

提出优化方案,让大规模图神经网络训练快15倍以上

Graph Learning at Scale: Characterizing and Optimizing Pre-Propagation GNNs

  • 预传播GNN通过预处理解耦特征传播与训练,理论解决邻居爆炸问题
  • 发现数据加载是训练效率瓶颈,输入扩展导致可扩展性挑战
  • 新方案使训练吞吐量提升15倍,速度比采样方法快两数量级

图神经网络(GNN)广泛用于图结构中的节点嵌入学习,通常采用消息传递机制。然而,该方法在层数增加时引发邻居爆炸问题,导致计算和内存需求呈指数增长。图采样已成为主流的规模化方法,但未能彻底解决问题。预传播GNN(PP-GNNs)是一类新型模型,通过预处理将特征传播与训练解耦,在理论上解决了邻居爆炸问题。然而其实际优势与系统级优化仍不明确。本文全面分析了PP-GNNs,对比其与基于采样的方法在训练效率、可扩展性和准确性方面的表现。尽管精度相当,我们发现数据加载是训练效率的主要瓶颈,输入扩展是关键的可扩展性挑战。为此,提出优化的数据加载方案与定制化训练方法,在大型图基准上平均使PP-GNN训练吞吐量提升15倍,相较采样基线最高提速达两个数量级。代码已开源:https://github.com/cornell-zhang/preprop-gnn。

原文摘要 · Abstract (English)

Graph neural networks (GNNs) are widely used for learning node embeddings in graphs, typically adopting a message-passing scheme. This approach, however, leads to the neighbor explosion problem, with exponentially growing computational and memory demands as layers increase. Graph sampling has become the predominant method for scaling GNNs to large graphs, mitigating but not fully solving the issue. Pre-propagation GNNs (PP-GNNs) represent a new class of models that decouple feature propagation from training through pre-processing, addressing neighbor explosion in theory. Yet, their practical advantages and system-level optimizations remain underexplored. This paper provides a comprehensive characterization of PP-GNNs, comparing them with graph-sampling-based methods in training efficiency, scalability, and accuracy. While PP-GNNs achieve comparable accuracy, we identify data loading as the key bottleneck for training efficiency and input expansion as a major scalability challenge. To address these issues, we propose optimized data loading schemes and tailored training methods that improve PP-GNN training throughput by an average of 15$\times$ over the PP-GNN baselines, with speedup of up to 2 orders of magnitude compared to sampling-based GNNs on large graph benchmarks. Our implementation is publicly available at https://github.com/cornell-zhang/preprop-gnn.

图神经网络训练优化大规模预处理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。