Graph learning trains on fewer nodes then rebuilds full graph efficiently
Decoupled Prune-and-Reconstruct: Training on a Compact Core with Double-Layer Sparsification
Machine LearningArtificial Intelligence
Summary
Training graph neural networks on very large graphs uses too much memory, limiting their usefulness. The authors present a method called DPR that first prunes the graph to a smaller set of important nodes to fit memory constraints, then uses this smaller set to reconstruct the full graph for accurate inference. This approach reduces memory use significantly during training while recovering most of the accuracy lost by pruning. It also runs inference much faster than some existing methods and scales to graphs with millions of nodes.
What this means in practice
- •For machine learning engineers: Train graph neural networks on very large graphs using less VRAM by pruning then reconstructing nodes to maintain accuracy.
- •For real-time system developers: Implement faster graph inference workflows by using DPR to deliver all-node inference with reduced memory and improved speed.
Authors
Siqi Gu, Gaozhen Ma, Liheng Huang, Weiqi He, Jiawen Luo, Jiaming Wu, Lixin Zheng
Abstract
On million-node graphs, the memory wall arrives before the compute wall: a GNN that never fits in VRAM cannot deliver the all-node inference its accuracy promises. The standard remedy, pruning, has always been an accuracy tax - paid once and never refunded. DPR refuses this framing. Its central insight is that the mask that shrinks training can also scaffold recovery: one compact node set bounds both computations of graph learning, serving as the encoder's training graph and, at inference, as the anchor from which a sparse structure-aware decoder rebuilds every pruned node (double-layer sparsification). The order is inverted - prune for efficiency, then reconstruct for accuracy - and the tax is repaid: the deployed 30%-budget model trains in 1,450 MiB of peak memory on one 6 GB GPU and recovers 96.9% of its dense-reconstruction reference, beating a same-GPU rerun of the edge-sparsification baseline DSpar by +1.36 pp at 50.1% lower peak training memory; against the same-protocol full-batch no-prune encoder the deployed point pays a measured 2.64 pp accuracy tax for a 2.8x smaller training footprint. Pruning without reconstruction leaves the encoder at 0.378 within the same 40-epoch window (Fig. 4), and all-node serving runs 46x faster than sampling-based GraphSAGE, giving up approx. 5.4 pp of test accuracy. The framework scales to a 2.45M-node production graph at 6,285 MiB of peak training memory; extended variants (long-range context, external text features, a CAN-bus application) are reported as extensions and do not compete with the deployed method, and the 41.8%-budget M4 point, scored on a coarsened super-node graph with its own split, is reported separately rather than compared with the original-node baselines. DPR thus turns aggressive node pruning from an accuracy tax into a safe training accelerator for VRAM-limited all-node inference.