[논문 리뷰] TGL: A General Framework for Temporal GNN Training on Billion-Scale Graphs

Hongkuan Zhou, Da Zheng외 4명|arXiv (Cornell University)|2022. 03. 28.

Advanced Graph Neural Networks인용 수 11

한 줄 요약

TGL은 빌리언 스케일 그래프에서 대규모 시간적 GNN 훈련을 위한 통합 프레임워크로, 시간 샘플러, 노드 메모리, 마일박스, 메시지 전달 엔진을 통해 다양한 TGNN 아키텍처의 효율적 훈련을 가능하게 한다. 4개의 GPU에서 빌리언 엣지 그래프를 1~10시간 내에 훈련시키며 최대 13배의 성능 향상을 달성했으며, 랜덤 청크 스케줄링과 Temporal-CSR와 같은 신규 기법을 통해 확장성과 성능을 향상시켰다.

ABSTRACT

Many real world graphs contain time domain information. Temporal Graph Neural Networks capture temporal information as well as structural and contextual information in the generated dynamic node embeddings. Researchers have shown that these embeddings achieve state-of-the-art performance in many different tasks. In this work, we propose TGL, a unified framework for large-scale offline Temporal Graph Neural Network training where users can compose various Temporal Graph Neural Networks with simple configuration files. TGL comprises five main components, a temporal sampler, a mailbox, a node memory module, a memory updater, and a message passing engine. We design a Temporal-CSR data structure and a parallel sampler to efficiently sample temporal neighbors to formtraining mini-batches. We propose a novel random chunk scheduling technique that mitigates the problem of obsolete node memory when training with a large batch size. To address the limitations of current TGNNs only being evaluated on small-scale datasets, we introduce two large-scale real-world datasets with 0.2 and 1.3 billion temporal edges. We evaluate the performance of TGL on four small-scale datasets with a single GPU and the two large datasets with multiple GPUs for both link prediction and node classification tasks. We compare TGL with the open-sourced code of five methods and show that TGL achieves similar or better accuracy with an average of 13x speedup. Our temporal parallel sampler achieves an average of 173x speedup on a multi-core CPU compared with the baselines. On a 4-GPU machine, TGL can train one epoch of more than one billion temporal edges within 1-10 hours. To the best of our knowledge, this is the first work that proposes a general framework for large-scale Temporal Graph Neural Networks training on multiple GPUs.

연구 동기 및 목표

• 대규모 동적 그래프에서 시간적 GNN을 훈련하기 위한 일반적이고 확장 가능한 프레임워크의 부족을 해결하기 위해.
• 간단한 설정 파일을 통해 다양한 TGNN 아키텍처를 단일 훈련 프레임워크로 통합하기 위해.
• 노드 메모리 노후화 및 빌리언 엣지 그래프에서의 비효율적 샘플링과 같은 대규모 배치 훈련의 과제를 극복하기 위해.
• 작은 벤치마크를 넘어서 실제 대규모 실세계 데이터셋에서 TGNN의 평가를 가능하게 하기 위해.
• 산업 및 학술 용도로 사용 가능한 프로덕션 수준의 다중 GPU 호환 시스템을 제공하기 위해.

제안 방법

• 동적 그래프 내 시간 순서로 정렬된 엣지를 효율적으로 저장하고 액세스하기 위한 Temporal-CSR 데이터 구조 설계.
• 시간 순서와 다중 코어 처리를 활용해 이웃 샘플링을 가속화하는 병렬 시간 샘플러 개발.
• 학습 배치를 더 작은 시간적으로 일관된 청크로 재배치함으로써 대규모 배치 훈련 중 의존성 손실을 줄이기 위한 랜덤 청크 스케줄링 기법 도입.
• 노드 메모리 모듈과 마일박스를 통합해 시간 단계 간 동적으로 변화하는 노드 표현을 유지하고 업데이트하기 위해.
• 설정 파일을 통한 구성으로 위의 구성 요소를 조합하여 다양한 TGNN 변종을 지원하는 메시지 전달 엔진 설계.
• NCCL 및 CUDA 스트림을 활용해 데이터 전송과 계산을 겹치는 다중 GPU 훈련 구현.

실험 결과

연구 질문

• RQ1 통합 프레임워크는 빌리언 스케일 동적 그래프에서 다양한 시간적 GNN 아키텍처를 효율적으로 훈련시킬 수 있는가?
• RQ2 노드 메모리가 노후화될 경우 대규모 배치 훈련을 어떻게 확장성 있고 정확하게 유지할 수 있는가?
• RQ3 최적화된 데이터 구조와 병렬 샘플링을 통해 대규모 시간적 그래프에서 어떤 성능 향상을 달성할 수 있는가?
• RQ4 작은 스케일 벤치마크를 넘어서 실제 빌리언 엣지 실세계 데이터셋에서 TGNN의 성능은 어떻게 되는가?
• RQ5 다중 GPU 훈련은 대규모 시간적 GNN의 훈련을 어느 정도로 확장시킬 수 있는가?

주요 결과

• TGL은 링크 예측 및 노드 분류 작업에서 정확도를 유지하거나 향상시키면서도 오픈소스 기반 베이스라인 대비 평균 13배의 성능 향상을 달성했다.
• 시간적 병렬 샘플러는 기준 샘플러 대비 다중 코어 CPU에서 173배의 성능 향상을 기록했다.
• 4개의 GPU를 탑재한 머신에서 TGL은 10억 개 이상의 시간 엣지를 가진 한 에포크를 1~10시간 내에 훈련시켰다.
• 랜덤 청크 스케줄링은 시간 의존성 손실을 완화하여 기존 방법이 실패하는 대규모 배치 크기(예: 4800)에서도 수렴 가능하게 했다.
• 다중 GPU 훈련은 JODIE에서 4개 GPU로 2.74배의 성능 향상을, TGN에서 4개 GPU로 2.25배의 성능 향상을 기록하여 강력한 확장성을 입증했다.
• 메모리 기반 모델인 TGN은 GDELT 데이터셋에서 한 에포크를 30분 이내에 훈련시키지만, 메모리 기반이 아닌 모델은 3시간 이상 소요된다.