Faster Convergence in Mini-batch Graph Neural Networks Training with Pseudo Full Neighborhood Compensation

Download PDF

Abstract

Graph Neural Networks (GNNs) have achieved remarkable success in various graph-related tasks. However, training GNNs on large-scale graphs is hindered by the neighbor explosion problem, rendering full-batch training computationally infeasible. Mini-batch training with neighbor sampling is a widely adopted solution, but it introduces gradient estimation errors that slow convergence and reduce model accuracy. In this work, we identify two primary sources of these errors: (1) missing gradient contributions from unsampled target nodes, and (2) inaccuracies in messages computed from sampled nodes. While existing methods largely focus on mitigating the second source, they often overlook the first, resulting in incomplete gradient estimation. To address this gap, we propose the Pseudo Full Neighborhood Compensation (PFNC) framework, which leverages historical information to simultaneously compensate for both errors. PFNC is designed to integrate seamlessly with any neighbor sampling technique and significantly lowers memory demands by maintaining only a partial cache of historical embeddings and gradients. Theoretical analysis demonstrates that PFNC provides a closer approximation to the ideal gradient, enhancing convergence. Extensive experiments across multiple benchmark datasets confirm that PFNC accelerates convergence and improves generalization across diverse neighbor sampling strategies.