Slicing Input Features to Accelerate Deep Learning: A Case Study with Graph Neural Networks

TL;DR

This paper introduces SliceGCN, a feature-slicing method for large-scale graph neural network training that improves scalability and stability by reducing memory usage and communication, especially on large datasets.

Contribution

SliceGCN is a novel distributed GNN training approach that slices node features across GPUs, maintaining accuracy and reducing communication compared to existing methods.

Findings

SliceGCN improves efficiency on large datasets.

Feature fusion and slice encoding enhance training stability.

SliceGCN has a potentially parameter-efficient design.

Abstract

As graphs grow larger, full-batch GNN training becomes hard for single GPU memory. Therefore, to enhance the scalability of GNN training, some studies have proposed sampling-based mini-batch training and distributed graph learning. However, these methods still have drawbacks, such as performance degradation and heavy communication. This paper introduces SliceGCN, a feature-sliced distributed large-scale graph learning method. SliceGCN slices the node features, with each computing device, i.e., GPU, handling partial features. After each GPU processes its share, partial representations are obtained and concatenated to form complete representations, enabling a single GPU's memory to handle the entire graph structure. This aims to avoid the accuracy loss typically associated with mini-batch training (due to incomplete graph structures) and to reduce inter-GPU communication during message passing (the forward propagation process of GNNs). To study and mitigate potential accuracy reductions due to slicing features, this paper proposes feature fusion and slice encoding. Experiments were conducted on six node classification datasets, yielding some interesting analytical results. These results indicate that while SliceGCN does not enhance efficiency on smaller datasets, it does improve efficiency on larger datasets. Additionally, we found that SliceGCN and its variants have better convergence, feature fusion and slice encoding can make training more stable, reduce accuracy fluctuations, and this study also discovered that the design of SliceGCN has a potentially parameter-efficient nature.