[논문 리뷰] Training Graph Neural Networks with 1000 Layers

Guohao Li, Matthias Müller외 2명|arXiv (Cornell University)|2021. 06. 14.

Neural Networks and Applications인용 수 8

한 줄 요약

이 논문은 가역 연결, 가중치 묶음, 평형 모델을 통해 메모리 복잡도를 O(L)에서 O(1)로 감소시켜 단일 GPU에서 최대 1001층의 GNN을 훈련할 수 있도록 하는 가역 그래프 신경망(RevGNN)을 제안한다. 이 방법은 상태 기반 성능을 달성하며, 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 메모리 소비 문제를 해결하기 위해.
• 표준 아키텍처의 O(L) 메모리 복잡도를 초월하여 깊이가 1000층이 넘는 GNN을 메모리 사용량을 늘리지 않고 훈련시킬 수 있도록 하기 위해.
• 대규모 그래프 벤치마크에서 성능을 유지하거나 향상시키면서 파라미터 및 메모리 효율성을 향상시키기 위해.
• 소형 배치 샘플링이나 부분 그래프 분할과 같은 기존 방법의 대안으로 가역 연결, 가중치 묶음, 평형 모델을 탐색하기 위해.

제안 방법

• 깊이 L에 대해 O(L)에서 O(1)로 메모리 복잡도를 감소시키기 위해 GNN에 가역 잔여 연결을 사용하여 중간 활성값을 저장할 필요 없이 처리한다.
• 군집 컨볼루션과 가중치 묶음을 활용하여 파라미터 수를 줄여 메모리 오버헤드를 최소화하면서도 넓거나 깊은 네트워크를 가능하게 한다.
• 무한 깊이의 네트워크의 평형점을 은닉 미분과 근의 탐색 방법을 사용해 풀어내는 깊이 있는 그래프 평형 GNN(DEQ-GNN)을 도입한다.
• 중간 상태를 저장하지 않고도 평형점에서 역전파를 수행하기 위해 은닉 미분을 적용하여 일정한 메모리 사용량을 유지한다.
• 가역 GNN과 소형 배치 샘플링을 조합하여 노드 수에 비례한 메모리 소비를 추가로 줄인다.
• 학습 가능한 메시지 집계와 GraphNorm을 사용하여 그래프 성질 예측 작업에서의 훈련 안정성과 성능을 향상시킨다.

실험 결과

연구 질문

• RQ1 GNN에서 가역 연결이 네트워크 깊이에 관계없이 일정한 메모리 복잡도를 달성할 수 있는가? 이는 극도로 깊은 모델의 훈련을 가능하게 하는가?
• RQ2 가장 깊은 GNN과 다른 파라미터 효율적인 아키텍처와 비교했을 때, 가역 GNN은 대규모 그래프 벤치마크에서 메모리, 파라미터 수, 성능 측면에서 어떻게 다른가?
• RQ3 가역 연결을 사용해 훈련된 과도한 파라미터를 가진 GNN은 더 긴 훈련 시간에도 불구하고 기존 최고 성능 모델을 능가할 수 있는가?
• RQ4 GraphNorm과 BatchNorm과 같은 정규화 기법이 깊은 가역 GNN의 성능에 미치는 영향은 어떠한가?
• RQ5 평형 모델(DEQ-GNN)은 O(1)의 메모리 및 파라미터 복잡도를 유지하면서도 깊은 가역 GNN과 유사한 성능을 달성할 수 있는가?

주요 결과

• 80개 채널을 가진 1001층의 RevGNN-Deep는 일반 소비자용 GPU에서 성공적으로 훈련되어 GNN 분야에서 새로운 깊이 기록을 수립했다.
• RevGNN-Deep는 ogbn-proteins 데이터셋에서 87.74 ± 0.13 AUC를 기록하여 기존 방법들을 능가했다.
• 224개 채널을 가진 448층의 RevGNN-Wide는 동일한 벤치마크에서 88.24 ± 0.15 AUC를 기록하여 높은 깊이와 넓이를 동시에 확보한 초우수한 성능을 보였다.
• 가역 GNN은 깊이에 관계없이 기준 모델보다 메모리의 일부분만 사용하면서도 일관되게 뛰어난 성능을 보였다. 메모리 비용은 깊이와 무관했다.
• RevGNN에서 GraphNorm을 사용함으로써 ogbg-molhiv에서 성능이 향상되어 78.62 AUC를 기록했으며, BatchNorm을 사용한 경우는 77.82 AUC였다.
• RevGNN는 ogbn-products에서 SGC와 SIGN을 능가하여 82.16%의 테스트 정확도를 기록했으며, 각각 74.87%와 77.60%를 기록한 반면, 메모리 사용량은 크게 낮았다.