GraphFM: Improving Large-Scale GNN Training via Feature Momentum
This addresses scalability issues in GNN training for researchers and practitioners, though it appears incremental as it builds on existing momentum techniques.
The paper tackles the challenge of training large-scale graph neural networks (GNNs) for node classification by proposing feature momentum (FM) to incorporate historical embeddings, which alleviates the neighborhood explosion problem and achieves promising performance on multiple datasets.
Training of graph neural networks (GNNs) for large-scale node classification is challenging. A key difficulty lies in obtaining accurate hidden node representations while avoiding the neighborhood explosion problem. Here, we propose a new technique, named feature momentum (FM), that uses a momentum step to incorporate historical embeddings when updating feature representations. We develop two specific algorithms, known as GraphFM-IB and GraphFM-OB, that consider in-batch and out-of-batch data, respectively. GraphFM-IB applies FM to in-batch sampled data, while GraphFM-OB applies FM to out-of-batch data that are 1-hop neighborhood of in-batch data. We provide a convergence analysis for GraphFM-IB and some theoretical insight for GraphFM-OB. Empirically, we observe that GraphFM-IB can effectively alleviate the neighborhood explosion problem of existing methods. In addition, GraphFM-OB achieves promising performance on multiple large-scale graph datasets.