LMC: Fast Training of GNNs via Subgraph Sampling with Provable Convergence

2023-02-02

International Conference on Learning Representations

AI Summary

Neighbor explosion in message-passing GNN training on large-scale graphs severely burdens computation, and existing subgraph sampling methods—while alleviating this burden—introduce substantial gradient bias, undermining convergence guarantees. Method: We propose the first subgraph sampling framework with rigorous theoretical convergence guarantees. By modeling the backward propagation mechanism of message passing, we design a local message compensation strategy that precisely recovers the contributions of discarded neighbors during both forward and backward passes. Our approach integrates gradient bias correction with non-convex optimization analysis. Contribution/Results: We formally prove that the algorithm converges to a first-order stationary point. Experiments on multiple large-scale benchmark tasks demonstrate that our method significantly outperforms existing sampling schemes—achieving faster convergence, more accurate gradient estimation, and superior generalization—while maintaining both theoretical rigor and practical efficacy.

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.

Problem

• Addresses neighbor explosion in GNN training
• Ensures convergence in subgraph-wise sampling methods
• Improves gradient accuracy and training efficiency

Innovation

• Subgraph sampling for GNN training
• Local Message Compensation (LMC) technique
• Convergence guarantee with message retrieval