Accelerating GNN Training through Locality-aware Dropout and Merge

TL;DR

LiGNN introduces a hardware-based approach that enhances data locality in GNN training by using locality-aware dropout and memory access merging, resulting in significant speedups and reduced DRAM accesses without sacrificing accuracy.

Contribution

LiGNN proposes a novel locality-aware dropout and memory merging technique to improve data locality and accelerate GNN training on hardware platforms.

Findings

Achieves 1.48~3.02x speedup over state-of-the-art methods.

Reduces DRAM accesses by 34%~55%.

Lowers DRAM row activations by 59%~82%.

Abstract

Graph Neural Networks (GNNs) have demonstrated significant success in graph learning and are widely adopted across various critical domains. However, the irregular connectivity between vertices leads to inefficient neighbor aggregation, resulting in substantial irregular and coarse-grained DRAM accesses. This lack of data locality presents significant challenges for execution platforms, ultimately degrading performance. While previous accelerator designs have leveraged on-chip memory and data access scheduling strategies to address this issue, they still inevitably access features at irregular addresses from DRAM. In this work, we propose LiGNN, a hardware-based solution that improves data locality by applying dropout and merge techniques during neighbor aggregation to accelerate GNN training. Unlike conventional algorithm-level dropout methods that primarily aim to improve accuracy while overlooking hardware costs, LiGNN introduces a locality-aware feature dropout mechanism. This approach selectively drops node features with data locality awareness, effectively reducing irregular DRAM accesses without compromising model accuracy. Moreover, by leveraging detailed knowledge of memory layout and organization-including critical alignment constraints-LiGNN strategically merges memory accesses during neighbor aggregation at the DRAM row level, guided by GNN-level semantics. This optimization significantly improves data locality with minimal additional cost. Under the commonly adopted 0.5 dropout rate, LiGNN outperforms state-of-the-art methods, delivering a 1.48~3.02x speedup, reducing DRAM accesses by 34%~55%, and lowering DRAM row activations by 59%~82%, all while maintaining model accuracy.