Efficient GNN Training Through Structure-Aware Randomized Mini-Batching

TL;DR

This work tackles the scalability bottleneck of training Graph Neural Networks on large graphs by addressing a conflict: random mini-batching aids convergence but hurts on-chip cache locality, while structure-based batching improves per-epoch speed but can impair accuracy. The authors propose COMM-RAND, a structure-aware randomized mini-batching framework that introduces controllable community bias in root-node partitioning and neighborhood sampling to balance per-epoch efficiency with convergence. Across four benchmarks and multiple GNN models, COMM-RAND achieves up to 2.76x faster per-epoch performance and around 1.8x faster total training on average, with final validation accuracy within approximately 2 percentage points of strong baselines; hyperparameter tuning can further improve accuracy under the same budget. The approach also leverages software caches and scales with cache capacity, while incurring modest pre-processing overhead, making it practical for real-world large-scale graph learning and adaptable to various hardware and model choices.

Abstract

Graph Neural Networks (GNNs) enable learning on realworld graphs and mini-batch training has emerged as the de facto standard for training GNNs because it can scale to very large graphs and improve convergence. Current mini-batch construction policies largely ignore efficiency considerations of GNN training. Specifically, existing mini-batching techniques employ randomization schemes to improve accuracy and convergence. However, these randomization schemes are often agnostic to the structural properties of the graph (for eg. community structure), resulting in highly irregular memory access patterns during GNN training that make suboptimal use of on-chip GPU caches. On the other hand, while deterministic mini-batching based solely on graph structure delivers fast runtime performance, the lack of randomness compromises both the final model accuracy and training convergence speed. In this paper, we present Community-structure-aware Randomized Mini-batching (COMM-RAND), a novel methodology that bridges the gap between the above extremes. COMM-RAND allows practitioners to explore the space between pure randomness and pure graph structural awareness during mini-batch construction, leading to significantly more efficient GNN training with similar accuracy. We evaluated COMM-RAND across four popular graph learning benchmarks. COMM-RAND cuts down GNN training time by up to 2.76x (1.8x on average) while achieving an accuracy that is within 1.79% points (0.42% on average) compared to popular random mini-batching approaches.

Figures (10)

• Figure 1: Community-based Graph Reordering:Assigning community members consecutive IDs improves locality
• Figure 2: Impact of entirely community-based mini-batching:Removing randomization leads to (a) accuracy loss and/or (b) a net training slowdown due to delayed convergence
• Figure 3: Biased Root Partitioning:Communities in the training set are shown with different colors. For the sake of simplicity, we do not show community-wide randomization for the community-aware schemes in the bottom row. In practice, communities are also shuffled as whole blocks in addition to randomizing the contents within community boundaries.
• Figure 4: Biased Neighborhood Selection:Biased neighborhood selection for the graph in Figure \ref{['fig:reordering-example']}. Intra-community edges are highlighted with blue and inter-community edges are highlighted with red. For a 2-layer GNN, uniform random neighborhood sampling creates a larger sub-graph compared to a biased scheme where intra-community neighbors are selected with a higher probability.
• Figure 5: Impact of COMM-RAND across different GNN training metrics (rows) and input graphs (columns):All results (besides accuracy) are normalized to the baseline ( RAND-ROOTS & $p=0.5$). The y-axis for val. acc. does not start from 0.