WaveGAS: Waveform Relaxation for Scaling Graph Neural Networks

2025-02-27

AI Summary

To address the problem of accumulated multi-layer errors and degraded node embedding quality in GNNAutoScale (GAS), caused by reusing stale historical embeddings across partitions, this paper proposes WaveGNN: an iterative forward-optimization and gradient-tracking framework inspired by waveform relaxation. Its core innovation is the first application of waveform relaxation to distributed GNN training—employing multiple lightweight forward passes to dynamically refine historical embeddings while jointly tracking and reusing gradients across iterations to suppress error propagation. WaveGNN is fully compatible with graph partitioning and embedding caching, enabling near-full-graph accuracy under GPU memory constraints. Extensive experiments on multiple large-scale graph benchmarks demonstrate that WaveGNN significantly outperforms GAS and other state-of-the-art baselines, achieving higher model accuracy and improved embedding robustness.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Search and Optimization: Distributed Search
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies

Abstract

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.

Problem

• Reduces approximation errors in GNN training
• Improves accuracy of historical embeddings
• Enhances node embedding estimation across partitions

Innovation

• WaveGAS refines historical embeddings
• Multiple forward passes improve accuracy
• Gradient-tracking enhances training precision