Training Graph Neural Networks with 1000 Layers
Figure 1. ROC-AUC score vs. GPU memory consumption on the ogbn-proteins dataset. Reversible models consistently achieve the same or better performance as the baseline using only a fraction of the memory. Weight-tied and equilibrium models offer a good performance to parameter efficiency trade-off.
Deep graph neural networks (GNNs) have achieved excellent results on various tasks on increasingly large graph datasets with millions of nodes and edges. However, memory complexity has become a major obstacle when training deep GNNs for practical applications due to the immense number of nodes, edges, and intermediate activations. To improve the scalability of GNNs, prior works propose smart graph sampling or partitioning strategies to train GNNs with a smaller set of nodes or sub graphs. In this work, we study reversible connections, group convolutions, weight tying, and equilibrium models to advance the memory and parameter efficiency of GNNs. We find that reversible connections in combination with deep network architectures enable the training of overparameterized GNNs that significantly outperform existing methods on multiple datasets. Our models RevGNN-Deep (1001 layers with 80 channels each) and RevGNN-Wide (448 layers with 224 channels each) were both trained on a single commodity GPU and achieve an ROC-AUC of 87.74 ± 0.13 and 88.24 ± 0.15 on the ogbn-proteins dataset. To the best of our knowledge, RevGNN-Deep is the deepest GNN in the literature by one order of magnitude.
Figure 2. GPU memory consumption vs. number of layers for ResGNN and RevGNN, our adaptation with reversible connections.
Figure 3. GPU memory consumption vs. number of layers for WT-ResGNN and WT-RevGNN
Figure 4. GPU memory vs. number of layers/iterations for ResGNN and DEQ-GNN.
Figure 5. GPU memory consumption vs. number of layers for ResGNN and RevGNN with full-batch and mini-batch training on the ogbn-products.
Table 1. Comparison of complexities. L is the number of layers, D is the number of hidden channels, N is of the number of nodes, B is the batch size of nodes and R is the number of sampled neighbors of each node. K is the maximum Broyden iterations.