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

2025-04-17

AI Summary

Graph Neural Networks (GNNs) suffer from the “neighbor explosion” problem, causing computational and memory overhead to grow exponentially with network depth. While Pre-propagation GNNs (PP-GNNs) theoretically circumvent this issue, their practical efficiency bottlenecks and system-level optimization opportunities remain poorly understood. This work presents the first systematic analysis identifying data loading and input feature expansion as the primary performance bottlenecks in PP-GNN training. To address them, we propose: (i) memory-aware data loading, (ii) hierarchical feature caching, (iii) optimized batch scheduling, and (iv) a PP-GNN-specific training paradigm. Evaluated on multiple large-scale graph benchmarks, our approach achieves an average 15× improvement in training throughput over baseline PP-GNN implementations and up to 100× speedup relative to sampling-based GNNs—while preserving model accuracy. Our contributions provide a reproducible, system-level optimization framework for efficient PP-GNN deployment.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Search and Optimization: Learning to Search
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Social Networks and Social Media: Influence propagation, information diffusion, and the prediction on networks
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

Abstract

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$ imes$ 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.

Problem

• Addresses neighbor explosion in GNNs via pre-propagation
• Optimizes data loading for PP-GNN training efficiency
• Improves scalability of PP-GNNs on large graphs

Innovation

• Decouples feature propagation via pre-processing
• Optimizes data loading for efficiency
• Tailors training methods for scalability