通过波形松弛提升GNN训练精度,解决跨分区信息滞后问题。
WaveGAS: Waveform Relaxation for Scaling Graph Neural Networks
- 引入波形松弛思想,在GAS中多次前向传播优化历史嵌入
- 相比原GAS,节点嵌入精度显著提升,甚至超越全图训练
- 适合大规模图神经网络训练,尤其资源受限场景
随着现实世界图数据规模持续增长,为克服训练图神经网络(GNN)时的资源限制,已有诸多技术被提出。其中,GNNAutoScale(GAS)通过图划分实现有限显存下的训练,并存储历史嵌入向量,从邻接分区的一跳邻居中检索,以保留跨分区关键信息。然而,这些来自上一轮迭代的历史嵌入相较于GAS估算的嵌入存在滞后,导致训练算法出现近似误差。且误差随层数累积,影响最终节点嵌入质量。为此,本文提出两项改进:其一,WaveGAS受波形松弛启发,在反向传播前进行多轮前向传播,逐步优化历史嵌入与梯度的估计;其二,引入梯度追踪机制,存储并利用更准确的历史梯度。实验表明,WaveGAS在提升GAS性能的同时,实现了更高精度,甚至优于全图训练方法,归功于其对节点嵌入的鲁棒估计。
原文摘要 · Abstract (English)
With the ever-growing size of real-world graphs, numerous techniques to overcome resource limitations when training Graph Neural Networks (GNNs) have been developed. One such approach, GNNAutoScale (GAS), uses graph partitioning to enable training under constrained GPU memory. GAS also stores historical embedding vectors, which are retrieved from one-hop neighbors in other partitions, ensuring critical information is captured across partition boundaries. The historical embeddings which come from the previous training iteration are stale compared to the GAS estimated embeddings, resulting in approximation errors of the training algorithm. Furthermore, these errors accumulate over multiple layers, leading to suboptimal node embeddings. To address this shortcoming, we propose two enhancements: first, WaveGAS, inspired by waveform relaxation, performs multiple forward passes within GAS before the backward pass, refining the approximation of historical embeddings and gradients to improve accuracy; second, a gradient-tracking method that stores and utilizes more accurate historical gradients during training. Empirical results show that WaveGAS enhances GAS and achieves better accuracy, even outperforming methods that train on full graphs, thanks to its robust estimation of node embeddings.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。