[论文解读] LMC: Fast Training of GNNs via Subgraph Sampling with Provable Convergence

Zhihao Shi, Xize Liang+1 more|arXiv (Cornell University)|Feb 2, 2023

Advanced Graph Neural Networks被引用 5

一句话总结

LMC 是一种用于训练图神经网络（GNNs）的新型子图采样方法，通过在反向传播中使用消息传递框架来补偿采样过程中丢弃的消息，从而实现可证明的收敛性。该方法通过高效地结合最新和历史消息，准确估计小批量梯度，实现了比当前最先进方法更快的收敛速度和更优的性能。

ABSTRACT

The message passing-based graph neural networks (GNNs) have achieved great success in many real-world applications. However, training GNNs on large-scale graphs suffers from the well-known neighbor explosion problem, i.e., the exponentially increasing dependencies of nodes with the number of message passing layers. Subgraph-wise sampling methods -- a promising class of mini-batch training techniques -- discard messages outside the mini-batches in backward passes to avoid the neighbor explosion problem at the expense of gradient estimation accuracy. This poses significant challenges to their convergence analysis and convergence speeds, which seriously limits their reliable real-world applications. To address this challenge, we propose a novel subgraph-wise sampling method with a convergence guarantee, namely Local Message Compensation (LMC). To the best of our knowledge, LMC is the {\it first} subgraph-wise sampling method with provable convergence. The key idea of LMC is to retrieve the discarded messages in backward passes based on a message passing formulation of backward passes. By efficient and effective compensations for the discarded messages in both forward and backward passes, LMC computes accurate mini-batch gradients and thus accelerates convergence. We further show that LMC converges to first-order stationary points of GNNs. Experiments on large-scale benchmark tasks demonstrate that LMC significantly outperforms state-of-the-art subgraph-wise sampling methods in terms of efficiency.

研究动机与目标

• 为解决现有子图采样方法在 GNN 训练中因梯度估计不准确而导致的收敛不稳定和收敛缓慢问题。
• 通过实现具有可证明收敛保证的高效小批量训练，克服大规模图上 GNN 的邻居爆炸问题。
• 设计一种在反向传播中补偿被丢弃消息的方法，且不会带来指数级计算成本。
• 仅通过局部邻域交互实现对 GNN 一阶驻点的收敛，避免递归扩展邻域。

提出的方法

• 通过将梯度误差分解为偏差（来自被丢弃的消息）和方差（来自随机采样），提出无偏的小批量梯度。
• 使用消息传递框架重新表述反向传播，以检索并补偿在子图采样过程中被丢弃的消息。
• 提出一种补偿机制，结合不完整的最新消息和历史消息，以高效纠正梯度偏差。
• 使用可学习的缩放因子 $\beta_i = \text{score}(i)\alpha$ 来平衡历史消息的陈旧性与准确性，其中得分函数基于局部到全局度数比。
• 通过消息补偿维持精确的梯度估计，确保收敛到 GNN 的一阶驻点。
• 通过直接在采样子图上运行各种 GNN 架构，将该方法应用于多种模型，实现广泛适用性。

实验结果

研究问题

• RQ1 尽管在小批量之外丢弃了消息，子图采样方法是否仍能在 GNN 训练中实现可证明的收敛？
• RQ2 如何在不引入高计算成本的前提下，有效补偿反向传播中因丢弃消息而引入的偏差？
• RQ3 在梯度补偿中同时使用最新消息和历史消息，对收敛速度和模型准确率有何影响？
• RQ4 在小批量设置下，该方法与全批量训练相比在性能上表现如何？
• RQ5 该方法是否能在大规模图学习任务中保持高效率，同时实现收敛保证？

主要发现

• LMC 收敛到 GNN 的一阶驻点，使其成为首个具有可证明收敛性的子图采样方法。
• 在 Ogbn-arxiv、FLICKR、REDDIT 和 PPI 等大规模基准测试中，LMC 相较于 GAS 最多减少 30% 的每轮训练时间，相较于 FM 最多减少 40%。
• 在小批量设置下（如 1 或 40），LMC 的预测性能与全批量 GD 相当，而基线方法则显著表现不佳。
• LMC 在前向和反向传播中均使用了 100% 的采样子图消息，与全批量的内存使用量一致，避免了消息丢失。
• 消融实验表明，当 $\alpha = 0.4$ 且 $\text{score}(i) = 2x - x^2$ 时，在小批量设置下性能最佳，有效降低了消息陈旧性的影响。
• 在批量大小为 1 时，LMC 在 Ogbn-arxiv 上达到 71.65% 的准确率，优于所有基线方法。