[논문 리뷰] GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings

Matthias Fey, Jan Eric Lenssen외 2명|arXiv (Cornell University)|2021. 06. 10.

Advanced Graph Neural Networks인용 수 16

한 줄 요약

GNNAutoScale (GAS)은 역사적 임bedding을 사용하여 계산 그래프를 정제함으로써, 입력 크기와 관계없이 일정한 GPU 메모리 사용을 유지하면서도 임의의 메시지 전달 그래프 신경망(GNN)의 확장 가능한 훈련을 가능하게 한다. 이는 전체 모델 표현력을 유지하면서도, ogbn-products 및 Reddit과 같은 대규모 벤치마크에서 최신 기술 수준의 성능을 달성한다. 이는 VR-GCN 및 MVS-GNN과 같은 이전 방법들을 능가한다.

ABSTRACT

We present GNNAutoScale (GAS), a framework for scaling arbitrary message-passing GNNs to large graphs. GAS prunes entire sub-trees of the computation graph by utilizing historical embeddings from prior training iterations, leading to constant GPU memory consumption in respect to input node size without dropping any data. While existing solutions weaken the expressive power of message passing due to sub-sampling of edges or non-trainable propagations, our approach is provably able to maintain the expressive power of the original GNN. We achieve this by providing approximation error bounds of historical embeddings and show how to tighten them in practice. Empirically, we show that the practical realization of our framework, PyGAS, an easy-to-use extension for PyTorch Geometric, is both fast and memory-efficient, learns expressive node representations, closely resembles the performance of their non-scaling counterparts, and reaches state-of-the-art performance on large-scale graphs.

연구 동기 및 목표

• GPU 메모리 제한과 이웃 노드 폭발 문제로 인해 깊고 표현력 있는 GNN을 대규모 그래프에서 훈련시키는 데 발생하는 확장성 문제를 해결하기 위해.
• 간선 샘플링이나 비학습 가능한 전파로 인한 성능 저하를 피하면서도, 메시지 전달 GNN의 전체 표현력을 유지하면서도 확장 가능한 훈련을 가능하게 하기 위해.
• 메시지 전달 메커니즘과 확장성을 분리함으로써 특정 GNN 아키텍처를 초월해 일반화할 수 있도록 하여, GCNII 및 PNA와 같은 다양한 GNN 연산자에 적용할 수 있도록 하기 위해.
• 역사적 임bedding으로 인한 근사 오차에 대한 이론적 경계를 제공하고, 실용적인 방법으로 이러한 경계를 강화할 수 있도록 하기 위해.
• 단일 GPU만을 사용하면서도 CPU에 저장된 역사 자료를 활용하여 대규모 데이터셋(예: ogbn-products)에서 전체 배치 GNN 훈련을 가능하게 하기 위해.

제안 방법

• GAS는 현재 미니배치에 속하지 않는 노드를 대체하기 위해 이전 훈련 반복에서 생성된 역사적 임bedding을 메모리 효율적인 대체 수 Mitt로 사용한다.
• GNN 계산 그래프를 현재 미니배치에 속한 노드와 그들의 직접적인 1단계 이웃 노드로만 유지함으로써 GPU 메모리 소비를 GNN 깊이와 관계없이 일정하게 유지한다.
• 역사적 임bedding은 각 반복마다 업데이트되며, 비활성화된 노드로부터 정보를 전파함으로써 전체 계산 그래프를 저장하지 않더라도 위상적 의존성을 유지한다.
• 임bedding의 노후화 정도와 함수의 리프시츠 연속성에 기반한 이론적 근사 오차 경계를 제공하며, 실용적인 전략을 통해 이러한 경계를 강화한다.
• PyGAS는 PyTorch Geometric의 확장 기능으로, 최소한의 코드 변경으로 기존 GNN 코드베이스에 원활하게 통합할 수 있도록 한다.
• 이 방법은 모델 설계와 독립적이므로, GCNII와 같은 깊은 모델이나 PNA와 같은 표현력 있는 아키텍처를 포함한 모든 메시지 전달 GNN에 적용 가능하다.

실험 결과

연구 질문

• RQ1 간선 샘플링을 하지 않거나 모델 표현력을 훼손하지 않고도, 임의의 메시지 전달 GNN을 대규모 그래프에 대해 확장 가능하게 훈련시킬 수 있는가?
• RQ2 역사적 임bedding을 사용할 경우 발생하는 이론적 근사 오차는 무엇이며, 실무에서 이를 최소화할 수 있는가?
• RQ3 역사적 임bedding은 전체 배치 GNN의 표현력을 유지하면서도 일정한 GPU 메모리 사용을 가능하게 할 수 있는가?
• RQ4 GAS 프레임워크는 깊고 표현력 있는 GNN에 적용했을 때 대규모 그래프 벤치마크에서 최신 기술 수준의 성능을 달성하는가?
• RQ5 아키텍처 수정 없이도 다양한 GNN 아키텍처로 일반화될 수 있는가?

주요 결과

• GAS는 ogbn-products(240만 개 노드, 6190만 개 간선)와 같은 대규모 그래프에서 단일 GPU만을 사용하여 전체 배치 GNN 훈련을 가능하게 하며, 입력 크기와 관계없이 GPU 메모리 사용이 일정하다.
• ogbn-products 데이터셋에서 PNA-GAS는 79.91%의 정확도를 달성하여 이전 최신 기술 수준의 방법인 GraphSAINT(79.08%)와 전체 배치 PNA(메모리 초과로 실패)를 능가했다.
• Reddit에서 PNA-GAS는 97.17%의 정확도를 기록하여 GraphSAINT(97.00%)와 VR-GCN(94.50%)를 모두 능가했으며, 대규모 노드 분류에서 뛰어난 성능을 보였다.
• GCNII-GAS 모델은 Reddit에서 96.77%의 정확도를 달성하여 전체 배치 GCNII(메모리 초과로 실패)를 능가했으며, 깊은 모델도 대규모로 훈련시킬 수 있음을 보여주었다.
• ogbn-products에서 역사 자료를 저장하는 데 약 2GB의 메모리 소비만 발생하여 CPU RAM에 저장하는 데 가능했으며, 훈련은 여전히 효율적이고 확장 가능했다.
• 오픈소스 구현인 PyGAS는 깊고 표현력 있는 GNN을 대규모 그래프에서 최소한의 코드 변경으로 훈련시킬 수 있게 하였으며, 비확장 가능한 대안과 거의 유사한 성능을 달성했다.