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

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

Speakers: Zichao Yue, Zhenhui Deng, Zhiru Zhang

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

Graph Neural Networks (GNNs) have emerged as a cornerstone in modern machine learning, demonstrating remarkable success in diverse applications ranging from fraud detection to circuit functional reasoning. However, their widespread adoption, especially on large-scale graphs, is often hampered by significant scalability challenges. The "neighbor explosion problem," inherent to traditional Message Passing GNNs (MPGNNs), leads to an exponential increase in memory and computational demands as the network deepens, making training on massive graphs prohibitively expensive. This talk, presented by Zichao Yue alongside collaborators Zhenhui Deng and Professor Zhiru Zhang, delves into Pre-Propagation GNNs (PPGNNs), a promising alternative approach designed to circumvent these scalability bottlenecks.

The core of this research lies in a comprehensive characterization and optimization effort for PPGNNs. While PPGNNs offer theoretical advantages in efficiency and accuracy by decoupling feature propagation from transformation, their practical performance in real-world systems has been underexplored. Yue and his team address this gap by systematically evaluating PPGNNs across various graph datasets, identifying critical system-level challenges like data loading bottlenecks and the unique "input expansion problem." Their work culminates in the development of targeted optimizations that not only unlock the theoretical benefits of PPGNNs but also push them to the forefront of scalable graph learning, offering an efficient and economically viable solution for training GNNs on graphs with hundreds of millions of nodes.

Background

▶ Watch: Introduction to GNNs and message passing framework (0:00)

Graph Neural Networks, particularly those based on the message passing framework, operate by iteratively aggregating features from neighboring nodes and transforming them into new representations. A single GNN layer involves collecting features from one-hop neighbors, aggregating them, and then applying a transformation. Stacking multiple layers extends this process, allowing features to propagate across multi-hop neighborhoods. While powerful, this propagation mechanism is the root of the neighbor explosion problem. For instance, training a three-layer GNN on a mini-batch of labeled nodes necessitates including all three-hop neighbors, which can quickly encompass a substantial portion, if not the entirety, of a large graph.

The mainstream strategy to mitigate this scalability challenge in MPGNNs is graph sampling. Methods like neighbor sampling aim to generate smaller, manageable subgraphs for training, trading off accuracy for efficiency. While state-of-the-art samplers like the one proposed at NeurIPS 2023 continue to evolve, they face inherent limitations. Firstly, they don't eliminate the scalability problem entirely; large subgraphs are often still required to maintain acceptable accuracy on massive original graphs. Secondly, and critically for certain applications, sampling methods can break the functional integrity of computation graphs (e.g., logic circuits). This structural disruption leads to significant accuracy degradation in downstream tasks that rely on the graph's precise topology.

In response to these persistent challenges, Pre-Propagation GNNs (PPGNNs) have emerged as a novel paradigm. The fundamental intuition behind PPGNNs is to decouple the feature propagation step from the learnable transformation step. By adopting a simple, non-learnable aggregation scheme for propagation, this computationally intensive phase can be moved into a one-time pre-processing step. This strategic shift aims to eliminate the neighbor explosion problem, preserve graph functionality, and ultimately provide a more scalable and accurate solution for graph learning, especially on complex or very large graphs where traditional MPGNNs struggle.

Key Findings

▶ Watch: Why graph sampling fails for large computation graphs (4:00)

The comprehensive characterization of PPGNNs by Yue and his team yielded several critical findings that challenge existing assumptions and pave the way for practical, scalable graph learning solutions:

1. Competitive Accuracy on Common Datasets: Contrary to some assumptions, PPGNNs (including models like SGC, SIGN, and HOGA) achieve comparable accuracy to traditional MPGNNs (such as GraphSAGE and GAT) on a range of commonly used graph benchmark datasets. This empirically validates their applicability beyond specialized computation graphs.

1. Vanilla PPGNNs Lack Real-world Efficiency: Despite theoretical advantages of eliminating neighbor explosion and redundant sampling, vanilla implementations of PPGNNs often exhibit longer epoch times compared to highly optimized MPGNNs. This indicates that their theoretical efficiency does not automatically translate to practical system performance without tailored optimizations.

1. Data Loading Dominates PPGNN Training Time: Detailed profiling revealed that data loading constitutes the dominant bottleneck in vanilla PPGNN training. This highlights a critical, previously underexplored area for performance improvement.

1. Unique Input Expansion Problem: PPGNNs introduce a distinct scalability challenge: the input expansion problem. During pre-processing, node features expand by a factor of K (number of operators) multiplied by R (number of hops). This expansion can cause medium-sized graphs to exceed GPU memory capacity and large graphs to exceed host main memory, posing a significant hurdle for training.

1. System-Level Optimizations are Essential: The team successfully developed and implemented several system-level optimizations (efficient batch assembly, data prefetching, chunk reshuffling, and direct storage access). These optimizations are crucial for realizing the theoretical benefits of PPGNNs in real systems.

1. Significant Throughput Improvements: With the proposed optimizations, PPGNNs demonstrate dramatic speedups. On medium-sized graphs, they achieve up to 136 times faster training throughput and an average of 7.4 times faster compared to MPGNNs. For large graphs, speedups reach up to 77 times higher throughput and an average of 9.9 times faster. Overall, optimized PPGNNs are 8.6 times faster on average.

1. Pareto Frontier Advancement: The optimizations effectively push PPGNNs to the Pareto frontier on the efficiency-accuracy trade-off diagram, making them a compelling choice for scalable graph learning.

1. Economic Solution for Very Large Graphs: The integration of storage-based training with direct storage access provides an acceptable throughput even for graphs exceeding host memory, offering an economic and scalable solution for extremely large-scale GNN training.

Technical Deep Dive

▶ Watch: Theoretical training efficiency and accuracy benefits of PPGNNs (8:00)

The core innovation of Pre-Propagation GNNs (PPGNNs) lies in the decoupling of feature propagation from transformation. In contrast to the iterative, coupled aggregation and transformation steps of traditional Message Passing GNNs (MPGNNs), PPGNNs move the propagation into a distinct, one-time pre-processing phase. This is achieved by adopting a simple, often non-learnable aggregation scheme (e.g., using adjacency matrix operators) during propagation. This means there's "nothing to learn" during this step, transforming it into a fixed computation that can be performed once and stored.

The general architecture of PPGNNs involves two main phases. First, in the pre-processing phase, several operators are derived from the graph's adjacency matrix. These operators define how features propagate across different hops. For instance, an operator might represent one-hop neighbors, another two-hop neighbors, and so on. Node features are then propagated according to these operators, generating hop-wise features. The crucial aspect here is that this is a one-time overhead; the pre-processed features, which now include information from various hops, can be stored and reused for subsequent training runs.

In the subsequent training phase, these pre-processed node features are loaded. These features are expanded along two additional dimensions: an operator dimension (K, representing the number of different propagation operators used) and a hop dimension (R, representing the maximum number of hops considered). This results in high-dimensional tensors that encapsulate rich multi-hop neighborhood information. These tensors are then fed into a usually dense, learnable model (e.g., a multi-head attention-based structure as seen in HOGA, which treats hop-wise features as sequences, or simpler linear models like SGC and SIGN) for the actual GNN training.

The theoretical advantages of this architecture are twofold. Firstly, training efficiency is significantly boosted because PPGNNs eliminate the neighbor explosion problem. This allows for straightforward mini-batch training without complex and redundant node sampling processes, leading to substantially less data movement and computation—up to two orders of magnitude reduction compared to traditional MPGNNs, as per the analysis. Secondly, PPGNNs offer superior accuracy, particularly for computation graphs. By avoiding graph sampling, PPGNNs preserve the inherent functionality and topology of these graphs, preventing the accuracy degradation common with MPGNNs. For example, HOGA demonstrated over 10% accuracy improvement on circuit functional reasoning tasks compared to MPGNNs.

However, translating these theoretical advantages into practical system performance presented unique challenges. Detailed profiling revealed that data loading became the dominant bottleneck, consuming a disproportionate amount of training time. Furthermore, the pre-processing step introduced the input expansion problem: the original node features expand by a factor of K * R. For a medium-sized graph, this expansion can push the feature set beyond GPU memory capacity, while for large graphs, it can exceed even the host's main memory.

To address these challenges, the researchers proposed a suite of system-level optimizations:

1. Efficient Batch Assembly: The primary motivation was the observation that host-side kernel launching dominated epoch time. This technique focuses on host-side kernel fusion for the batch assembly process, effectively reducing overhead by combining multiple small kernel calls into fewer, larger ones.

1. Data Prefetching: To overlap data loading with computation, a GPU-side double buffer is employed. This allows the GPU to process one batch of data while the next batch is being loaded, achieving a pipelined execution and hiding data transfer latencies.

1. Chunk Reshuffling: This technique introduces a form of insufficient reshuffle, a common practice shown to have negligible impact on task accuracy. Its key benefit is offloading the data batch assembly task from the host CPU to the GPU. Instead of loading individual node features, data is loaded in larger "chunks," which not only speeds up loading but also paves the way for storage-based training by enabling more efficient access patterns.

1. Direct Storage Access: Combined with chunk reshuffling, this technique facilitates storage-based training. For graphs too large to fit into host memory, data chunks can be directly accessed from persistent storage (e.g., SSDs). This provides an economic and scalable solution, allowing GNN training on graphs that would otherwise be intractable due to memory constraints.

These optimizations collectively target the identified bottlenecks, transforming PPGNNs from a theoretically promising approach into a practically efficient and scalable solution for graph learning.

Experimental Setup & Results

▶ Watch: Two crucial questions about PPGNNs in real systems (9:00)

To comprehensively characterize PPGNNs, the researchers conducted experiments across six diverse graph datasets. Three were classified as medium-sized, with node counts ranging from 1 to 2 million nodes, primarily used for exploring the accuracy-efficiency trade-off. The other three were large graphs, each containing over 100 million nodes, specifically chosen for investigating scalability challenges and solutions.

For model comparison, the study included representative PPGNN architectures such as SGC, SIGN, and HOGA (their prior work, which uses a multi-head attention-based structure). These were benchmarked against traditional MPGNNs, specifically GraphSAGE and GAT. It's important to note that the MPGNNs chosen were representative of the family, rather than necessarily the absolute state-of-the-art models, to facilitate a fair comparison between the two GNN paradigms.

The experiments were conducted on a single server, multi-GPU setting. Key metrics included accuracy on downstream tasks, convergence rate (number of epochs to reach convergence), epoch time, and overall training throughput.

Initial accuracy comparisons revealed that PPGNNs achieved comparable accuracy to MPGNNs on the selected common graph datasets. This was a crucial finding, as it demonstrated that PPGNNs are not solely specialized for computation graphs but can perform competitively across a broader range of applications. Furthermore, PPGNNs showed an on-par or faster convergence rate, requiring a similar or fewer number of epochs to reach convergence.

However, the initial assessment of training efficiency presented a challenge. Vanilla PPGNN implementations, without specific optimizations, exhibited even longer epoch times compared to MPGNNs that leveraged existing system-level optimizations. This counter-intuitive result prompted a deeper investigation, which, as detailed earlier, pinpointed data loading as the primary bottleneck, consuming a significant portion of the epoch time. The input expansion problem was also quantified, where node features expanded by a factor of K (operators) R* (hops), often exceeding available GPU or even host memory.

The impact of the proposed system-level optimizations was dramatic. On the three medium-sized graphs, the optimized PPGNNs achieved an impressive 136 times faster training throughput in the best case, with an average speedup of 7.4 times faster compared to MPGNNs. For the three large graphs, the optimizations resulted in up to 77 times higher training throughput and an average of 9.9 times faster. Across all evaluated graphs, the optimized PPGNNs were, on average, 8.6 times faster than traditional MPGNNs.

A particularly significant result was the performance of storage-based training. Even when utilizing this approach for graphs exceeding host memory, the optimized PPGNNs achieved an acceptable throughput. This demonstrated the economic viability of training GNNs on extremely large graphs, offering a solution that was previously intractable or prohibitively expensive due to memory constraints. These results collectively illustrate that the proposed optimizations successfully push PPGNNs to the Pareto frontier in the efficiency-accuracy trade-off space, making them a highly competitive and scalable solution for graph learning.

Practical Implications

▶ Watch: Vanilla PPGNNs surprisingly slower due to data loading (11:00)

The research on characterizing and optimizing Pre-Propagation GNNs carries significant practical implications for practitioners, infrastructure teams, model builders, and deployers working with large-scale graph data.

Firstly, PPGNNs are established as a viable and competitive alternative to traditional Message Passing GNNs, not just on specialized computation graphs but also on commonly used graph datasets. This broadens the applicability of PPGNNs, offering a powerful tool for scenarios where graph sampling might be problematic or where the sheer scale of the graph makes MPGNN training infeasible. For model builders, this means exploring PPGNN architectures like SGC, SIGN, or HOGA could yield comparable or even superior accuracy, particularly when graph integrity is paramount (e.g., in logic circuit analysis).

For infrastructure teams and deployers, the work highlights a critical insight: the theoretical efficiency gains of PPGNNs do not automatically translate to practical system performance. Without tailored optimizations, vanilla PPGNN implementations can be slower than optimized MPGNNs. This underscores the necessity of a holistic approach that considers the entire system stack. The identified bottlenecks, primarily data loading and the input expansion problem, serve as crucial focal points for system design. Infrastructure teams should prioritize efficient data pipelines, potentially leveraging techniques like host-side kernel fusion, GPU-side double buffering, and chunk-based data access to maximize PPGNN throughput.

The proposed optimizations offer concrete strategies for achieving significant speedups. The average 8.6x speedup and peak 136x speedup over MPGNNs translate directly into reduced training times, lower computational costs, and faster iteration cycles for model development. This can be a game-changer for companies dealing with rapidly growing graph datasets.

Furthermore, the successful implementation of storage-based training with acceptable throughput provides an economic solution for training GNNs on graphs that are simply too large to fit into GPU or even host main memory. This opens up new possibilities for analyzing truly massive networks, democratizing access to GNN capabilities for datasets previously considered intractable due to hardware limitations. Practitioners no longer need to invest in exorbitantly large memory systems to handle such graphs; instead, optimized storage access can provide a cost-effective alternative.

Finally, while PPGNNs offer compelling advantages, it's important to consider trade-offs. The pre-processing phase, while a one-time overhead, can be computationally intensive for extremely large graphs and may require significant storage for the expanded features. Model builders should also be mindful that while PPGNNs achieve competitive accuracy, they might not always surpass the absolute state-of-the-art MPGNNs on every specific metric or dataset, especially if those MPGNNs incorporate highly specialized inductive biases. The choice between PPGNNs and MPGNNs will therefore depend on the specific application, graph characteristics, available hardware, and the relative importance of accuracy, training time, and memory footprint. The open-sourced code from this research provides a valuable starting point for practitioners to experiment and adapt these optimizations to their unique environments.

Key Takeaways

• PPGNNs Decouple Propagation: Pre-Propagation GNNs (PPGNNs) tackle GNN scalability by decoupling feature propagation from transformation, moving the propagation into a one-time pre-processing step.
• Vanilla PPGNNs Face Bottlenecks: Despite theoretical advantages, vanilla PPGNN implementations are often bottlenecked by data loading and suffer from a unique "input expansion problem" (features expand by K*R), leading to slower training than optimized MPGNNs.
• System Optimizations are Crucial: Tailored system-level optimizations, including efficient batch assembly, data prefetching, chunk reshuffling, and direct storage access, are essential to unlock PPGNNs' practical efficiency.
• Significant Speedups Achieved: Optimized PPGNNs demonstrate dramatic training throughput improvements, achieving up to 136 times faster and an average of 8.6 times faster training compared to traditional Message Passing GNNs (MPGNNs).
• Competitive Accuracy & Functionality Preservation: PPGNNs achieve comparable accuracy to MPGNNs on common graph datasets and excel on computation graphs by preserving their functional integrity, which graph sampling methods often disrupt.
• Economic Solution for Large Graphs: The integration of storage-based training provides an economic and scalable solution for training GNNs on graphs that exceed host memory capacity, making very large-scale graph learning more accessible.

About the Speaker(s)

Zichao Yue is a researcher who presented this work, a collaborative effort with Zhenhui Deng and his advisor, Professor Zhiru Zhang. Their research focuses on addressing the scalability challenges of graph neural networks, particularly through the lens of pre-propagation GNN architectures and their system-level optimization. Professor Zhiru Zhang is affiliated with the academic institution where this research was conducted, guiding the team's efforts in developing efficient and scalable solutions for machine learning on graphs.