[论文解读] Training Graph Neural Networks with 1000 Layers

Guohao Li, Matthias Müller+2 more|arXiv (Cornell University)|Jun 14, 2021

Neural Networks and Applications被引用 8

一句话总结

该论文提出可逆图神经网络（RevGNN），通过可逆连接、参数共享和平衡模型，将GNN的内存复杂度从O(L)降低至O(1)，从而实现在单张GPU上训练深度达1001层的GNN。该方法实现了最先进性能，其中RevGNN-Deep在ogbn-proteins数据集上达到87.74 ± 0.13 AUC，将模型深度记录提升了整整一个数量级。

ABSTRACT

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.

研究动机与目标

• 为解决现实应用中深度图神经网络（GNN）高GPU内存消耗的问题。
• 在不增加内存占用的前提下，实现超深GNN（超过1000层）的训练，突破标准架构的O(L)内存复杂度限制。
• 在保持或提升大规模图基准性能的同时，提高参数与内存效率。
• 探索可逆连接、参数共享和平衡模型作为替代小批量采样或子图划分的可扩展GNN训练方法。

提出的方法

• 在GNN中使用可逆残差连接，消除对中间激活值存储的需求，将内存复杂度从O(L)降低至与深度L无关的O(1)。
• 采用组卷积和参数共享以减少参数量，实现在极低内存开销下构建宽或深的网络。
• 提出一种深度图平衡GNN（DEQ-GNN），通过隐式微分和根求解方法，求解无限深度网络的平衡点。
• 应用隐式微分技术，实现对平衡点的反向传播而无需存储中间状态，维持常数内存使用。
• 将可逆GNN与小批量采样结合，进一步降低与节点数量相关的内存消耗。
• 采用可学习的消息聚合机制和GraphNorm，提升图属性预测任务中的训练稳定性和性能。

实验结果

研究问题

• RQ1 GNN中的可逆连接是否能实现与网络深度无关的常数内存复杂度，从而支持极深模型的训练？
• RQ2 与标准GNN及其他参数高效架构相比，可逆GNN在大规模图基准上的内存消耗、参数量和性能表现如何？
• RQ3 尽管训练时间更长，使用可逆连接训练的超参数化GNN是否能超越现有SOTA模型？
• RQ4 归一化技术（如GraphNorm）与BatchNorm相比，对深度可逆GNN性能有何影响？
• RQ5 平衡模型（DEQ-GNN）是否能在保持O(1)内存和参数复杂度的同时，实现与深度可逆GNN相当的性能？

主要发现

• RevGNN-Deep（1001层，每层80个通道）成功在单张消费级GPU上训练，创下GNN模型深度新纪录。
• RevGNN-Deep在ogbn-proteins数据集上达到87.74 ± 0.13的ROC-AUC，优于现有方法。
• RevGNN-Wide（448层，每层224个通道）在相同基准上达到88.24 ± 0.15 AUC，展现出高深度与高宽度下的卓越性能。
• 可逆GNN在所有实验中均显著优于基线模型，且内存消耗与深度无关。
• 在RevGNN中使用GraphNorm显著提升了性能，在ogbg-molhiv上达到78.62 AUC，优于使用BatchNorm的77.82 AUC。
• RevGNN在ogbn-products上优于SGC和SIGN，测试准确率达到82.16%，显著低于SGC的74.87%和SIGN的77.60%，且内存使用量大幅降低。