PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Graph Neural Networks with 1000 Layers

Guohao Li, Matthias Müller, Bernard Ghanem, Vladlen Koltun

arXiv:2106.07476v3cs.LGcs.AIcs.SI

TL;DR

Deep GNNs face severe GPU-memory constraints on large graphs, limiting the use of deeper and wider over-parameterized models. The paper studies reversible, weight-tied, and equilibrium architectures to reduce memory or parameter costs, finding that reversible GNNs achieve strong state-of-the-art results. The approach also incurs increased training time, especially on very large datasets.

  • Problem

    Large graphs make activation memory a bottleneck for deep GNN training, while reducing parameters can limit the benefits of deeper and wider models.

  • Method

    The paper investigates grouped reversible connections, weight tying, and equilibrium models to reduce memory complexity or parameter cost in deep GNNs.

  • Results

    Reversible GNNs significantly outperform current state-of-the-art methods on several datasets, including 88.24% ROC-AUC for RevGNN-Wide on ogbn-proteins.

  • Takeaways & Limitations

    Reversible connections enable over-parameterized GNNs that are an order of magnitude deeper while consuming less memory than current models.

  • Takeaways & Limitations

    The methods increase training time by less than 40%, which can add several days on very large datasets.

Abstract

from arXiv · show

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 \pm 0.13$ and $88.24 \pm 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. Please visit our project website https://www.deepgcns.org/arch/gnn1000 for more information.

1. Introduction

Large graphs make GPU memory a bottleneck for deep GNNs, while reducing parameters can undermine the benefits of larger models. This work investigates reversible, grouped, weight-tied, and equilibrium architectures to improve memory or parameter efficiency while preserving strong performance.

  • Large real-world graphs contain millions or billions of nodes and edges, making limited GPU memory a practical bottleneck for GNN training.
  • Existing mini-batch methods reduce memory by sampling nodes or partitioning graphs, but add hyperparameters and can break important graph structures when sampling is too aggressive.
  • Grouped reversible graph connections reduce activation memory complexity from O(L) to O(1), making memory consumption independent of network depth.
  • Weight-tied GNNs share weights across layers, while graph equilibrium GNNs solve an infinite-depth network with root finding and implicit differentiation.
  • The broader analysis finds reversible architectures most powerful for state-of-the-art performance, while weight tying and equilibrium models offer parameter-efficiency trade-offs.
  • RevGNN-Deep uses 1001 layers and 80 channels, achieving 87.74% ROC-AUC with 2.86 GB of GPU memory, while RevGNN-Wide achieves 88.24% ROC-AUC.

2. Related Work

Related work addresses deep-GNN training through architectural, normalization, propagation, and sampling strategies. These methods seek to improve depth or scalability, but memory remains a central constraint for large graphs.

  • Early GNNs learned graph representations through stable states or spectral graph convolutions based on graph-Laplacian eigendecomposition.
  • Many state-of-the-art GNNs remain shallow because over-smoothing and vanishing gradients make deep training difficult.
  • Normalization and regularization methods include DropEdge, DropConnect, PairNorm, weight normalization, differentiable group normalization, and GraphNorm.
  • Large graphs with millions of nodes and edges make full-batch training infeasible under typical memory limits.
  • GraphSAGE samples neighborhoods but incurs exponential memory growth with depth, while VR-GCN stores historical activations with O(LND) memory.
  • FastGCN, Cluster-GCN, and GraphSAINT reduce memory through layer-wise or subgraph sampling, with trade-offs involving sparse connections or graph-structure preservation.

3. Building Deep GNNs

The paper develops deep and parameter-efficient GNN architectures that reduce activation or parameter costs, enabling substantially deeper models under fixed GPU memory. Reversible connections provide depth-independent activation memory, while weight tying and equilibrium models improve parameter efficiency.

  • Over-parameterized GNNs: Deep residual GNN activations require O(LND) memory, so memory grows linearly with network depth.L is the number of layers, N the number of vertices, and D the vertex-feature size.
  • Over-parameterized GNNs: Grouped reversible GNN blocks reconstruct inputs during backpropagation and reduce activation memory to O(ND), independent of depth.Only the output features of the last block need to remain in GPU memory.
  • Parameter-Efficient GNNs: Weight-tied GNNs share weights across layers, keeping parameter count constant with depth; reversible weight tying also reduces activation memory to O(ND).Weight-tied residual GNNs retain O(LND) activation complexity, whereas weight-tied reversible GNNs reconstruct features during backpropagation.
  • Over-parameterized GNNs: RevGNN-80 exceeds 1000 layers at constant memory and reaches 87.06% ROC-AUC, while RevGNN-224 reaches 87.41% ROC-AUC at 448 layers.The reversible models use saved memory to increase depth or width relative to ResGNN-64, which reaches 85.94% ROC-AUC at 112 layers and cannot go deeper on commodity hardware.
  • Parameter-Efficient GNNs: Graph equilibrium GNNs represent infinite-depth weight-tied networks and compute equilibrium states with root finding and implicit differentiation.The approach avoids storing intermediate states and has the memory and parameter cost of a single layer; its performance depends on the iteration threshold.

4. Over-parameterized Deep GNNs

The section evaluates reversible GNNs across datasets, operators, depths, and training schemes, showing that they combine strong performance with substantially lower memory use. RevGNN variants achieve state-of-the-art results while supporting very deep or wide models.

  • 4.1. State-of-the-art Results: RevGNN-Deep uses 1001 layers and RevGNN-Wide uses 448 layers, achieving new state-of-the-art results on ogbn-proteins.RevGNN-Deep uses 10.5% of GPU training memory, while RevGNN-Wide uses about 29%; training takes 13.5 and 17.1 days, respectively, for 2000 epochs.
  • 4.2. Application to Different GNN Operators: RevGNNs achieve new state-of-the-art results on ogbn-arxiv across GCN, GAT, and self-knowledge-distillation variants.RevGCN-Deep reaches 73.01% accuracy using 1.84GB, RevGAT-Wide reaches 74.05% versus 73.97% for UniMP v2, and RevGAT-SelfKD reaches 74.26%.
  • 4.2. Application to Different GNN Operators: RevGCN reduces memory by more than 75% while improving accuracy by 0.76% against a comparable ResGCN.The comparison uses 28 layers, with RevGCN using 180 channels and ResGCN using 128 channels.
  • 4.3. Full-batch vs. Mini-batch Training: Mini-batch training reduces RevGNN memory consumption to 44% and raises ogbn-products accuracy from 78.77% to 82.16%.The study combines reversible architectures with simple random-clustering mini-batches and evaluates with full-batch inference.
  • 4.5. More Ablation Studies: RevGNN reaches 82.16 ± 0.15% accuracy on ogbn-products, exceeding SGC at 74.87 ± 0.25% and SIGN at 77.60 ± 0.13%.The RevGNN uses 56 layers and 160 hidden channels with random-clustering mini-batch training.
  • 4.5. More Ablation Studies: On ogbg-ppa, RevGNN and ResGNN both achieve around 77% accuracy, but RevGNN uses only 16% of ResGNN's GPU memory.Both models use 28 layers and approximately 2.3M parameters.

5. Conclusion and Discussion

The work reduces GNN memory complexity with respect to depth from O(L) to O(1), enabling deeper and wider over-parameterized networks. This scalability comes with an additional training-time cost that can become substantial on very large datasets.

  • O(L) to O(1) memory complexity with respect to depth enables GNNs that are an order of magnitude deeper than current state-of-the-art models.Reversible GNNs remove memory costs associated with depth, allowing width to increase for additional representational power.
  • Over-parameterized reversible networks significantly outperform current models on standard benchmarks while consuming less memory.
  • Less than 40% additional training time can still extend training by several days on very large datasets.The paper identifies reducing GNN training time as an avenue for future research.

A. Grouping of RevGNN

Grouping reduces the parameter count of reversible GNNs, although the current implementation can increase GPU memory usage as group size grows. Performance remains similar for some deeper configurations.

  • More than 30% fewer parameters result when group size increases from 2 to 4.The experiments use 224 hidden channels on ogbn-proteins.
  • 112-layer networks achieve the same performance, while group size 4 uses around 67% of the parameters of group size 2.
  • GPU memory rises from 7.30 GB to 11.05 GB as group size increases from 2 to 4 in the current implementation.The authors conjecture that inefficient implementation causes this increase.

B.1. Datasets and Frameworks

Experiments use three Open Graph Benchmark datasets under standard OGB evaluation protocols and implement the models with established deep-learning and graph-learning frameworks.

  • Datasets: Experiments cover ogbn-proteins, ogbn-arxiv, and ogbn-products using the standard OGB 1.2.4 data splits and evaluation protocol.
  • Frameworks: The implementation uses PyTorch 1.6.0 and PyTorch Geometric 1.6.1, except GAT experiments, which use DGL 0.5.3.
  • Frameworks: The reversible and deep equilibrium modules are implemented using MemCNN and DEQ, respectively.

B.2. Hyperparameters and Experimental Settings

The experiments combine dataset-specific training settings, mini-batch partitioning where needed, multi-view inference, and peak GPU-memory measurement. Ablations compare reversible, weight-tied, and equilibrium designs under controlled configurations.

  • Ogbn-proteins: Ogbn-proteins mini-batch training uses random partitioning, one sampled subgraph per SGD step, max aggregation, 0.1 dropout, and 2000 training epochs.
  • Ogbn-proteins: A 112-layer RevGNN achieves 87.02 ROC-AUC on ogbn-proteins among the compared 112-layer models.Reversible connections reduce memory consumption and enable wider RevGNNs.
  • Multi-view inference: Table 7 reports layers, channels, views, and partitions for multi-view inference ablations, with more views and fewer partitions preferred.
  • Ogbn-proteins: Multi-view inference improves ROC-AUC by about 0.4% from 1 to 10 views for both RevGNN-Deep and RevGNN-Wide.The method uses different graph partitions across inference passes to expose vertices to different neighbor sets.
  • Other datasets: Ogbn-arxiv uses an undirected graph with self-loops and full-batch training and testing, while ogbn-products compares full-batch and randomly partitioned mini-batch training.
  • GPU memory measurement: Peak GPU memory is measured during the first training epoch because intermediate computations and network parameters contribute beyond node-feature storage.The reported value is obtained with torch.cuda.max_memory_allocated() and can be below NVIDIA-SMI’s actual peak.

B.4. Correlation of Model Predictions

RevGNN, weight-tied RevGNN, and DEQ-GNN produce strongly correlated predictions despite differing depths, iterations, and parameterization.

  • 0.8571 Pearson correlation links RevGNN with 1000 layers and WT-RevGNN-224 with 7 layers.
  • 0.8565 Pearson correlation links RevGNN with 1000 layers and DEQ-GNN-224 with 56 iterations.
  • 0.8948 Pearson correlation is reported between WT-RevGNN-224 with 7 layers and DEQ-GNN-224 with 56 iterations.

C.1. Gradient Checkpointing and Model Parallelism √

Gradient checkpointing uses more memory than RevGNN in the reported comparison, while model parallelism is identified as a possible route to wider models.

  • Gradient Checkpointing: 2.5X more memory is consumed by ResGNN with gradient checkpointing than RevGNN for 112-layer, 64-channel models.The reported memory usage is 5.22 G versus 2.09 G, with similar performance and training time on ogbn-proteins.
  • Gradient Checkpointing: Similar performance and training time are reported for ResGNN with gradient checkpointing and RevGNN on ogbn-proteins.
  • Model Parallelism: Model parallelism is orthogonal to the proposed approach and may enable RevGNN models with greater width across multiple GPUs.

C.2. Going Deeper and Datasets

The benefit of increasing GNN depth depends on the dataset: it is strong on ogbn-proteins but smaller on ogbn-products and ogbn-arxiv, leaving the relevant dataset factors unresolved.

  • Less improvement from very deep models is observed on ogbn-products and ogbn-arxiv.
  • It remains unclear which datasets benefit most from depth and overparameterization.The paper identifies dataset size, graph modality, and graph learning task as factors for future investigation.
  • Pre-training overparameterized GNNs on larger protein datasets is proposed for applications including drug discovery and protein structure prediction.
Loading 2106.07476v3…