[논문 리뷰] LMC: Fast Training of GNNs via Subgraph Sampling with Provable Convergence

Zhihao Shi, Xize Liang외 1명|arXiv (Cornell University)|2023. 02. 02.

Advanced Graph Neural Networks인용 수 5

한 줄 요약

LMC는 그래프 신경망(GNN) 학습을 위한 새로운 부분그래프 기반 샘플링 방법으로, 메시지 전달 구조를 활용해 백워드 프로파게이션 중 기각된 메시지를 보상함으로써 증명 가능한 수렴성을 보장한다. 이는 최신 및 이전 메시지를 효율적으로 보완함으로써 미니배치 기울기 추정을 정확하게 수행함으로써 기존 최고 수준의 방법들보다 더 빠른 수렴과 뛰어난 성능을 달성한다.

ABSTRACT

The message passing-based graph neural networks (GNNs) have achieved great success in many real-world applications. However, training GNNs on large-scale graphs suffers from the well-known neighbor explosion problem, i.e., the exponentially increasing dependencies of nodes with the number of message passing layers. Subgraph-wise sampling methods -- a promising class of mini-batch training techniques -- discard messages outside the mini-batches in backward passes to avoid the neighbor explosion problem at the expense of gradient estimation accuracy. This poses significant challenges to their convergence analysis and convergence speeds, which seriously limits their reliable real-world applications. To address this challenge, we propose a novel subgraph-wise sampling method with a convergence guarantee, namely Local Message Compensation (LMC). To the best of our knowledge, LMC is the {\it first} subgraph-wise sampling method with provable convergence. The key idea of LMC is to retrieve the discarded messages in backward passes based on a message passing formulation of backward passes. By efficient and effective compensations for the discarded messages in both forward and backward passes, LMC computes accurate mini-batch gradients and thus accelerates convergence. We further show that LMC converges to first-order stationary points of GNNs. Experiments on large-scale benchmark tasks demonstrate that LMC significantly outperforms state-of-the-art subgraph-wise sampling methods in terms of efficiency.

연구 동기 및 목표

• 기존 부분그래프 기반 샘플링 방법들이 기울기 추정의 부정확성으로 인해 GNN 학습에서 수렴 불안정성과 느린 수렴을 겪는 문제를 해결하기 위해.
• 대규모 그래프에서 GNN의 이웃 폭발 문제를 해결하기 위해 효율적인 미니배치 학습을 가능하게 하며, 증명 가능한 수렴 보장을 제공하기 위해.
• 백워드 프로파게이션 중 기각된 메시지를 보상함으로써 기하급수적인 계산 비용을 유발하지 않도록 하는 방법을 설계하기 위해.
• 재귀적 이웃 확장을 피하고 국소적 이웃 상호작용만을 사용하여 GNN의 1차 정류점에 수렴하는 방법을 설계하기 위해.

제안 방법

• 기울기 오차를 편향(기각된 메시지에서 기인)과 분산(스토케스틱 샘플링에서 기인)으로 분해하여 편향 없는 미니배치 기울기를 제안한다.
• 부분그래프 샘플링 중 기각된 메시지를 복구하고 보상하기 위해 메시지 전달 프레임워크를 재구성하여 백워드 프로세스를 재정의한다.
• 부족한 최신 메시지와 이전 메시지를 조합하여 기울기 편향을 효율적으로 보정하는 보상 메커니즘을 도입한다.
• 이전 메시지의 낙관성과 정확성의 균형을 맞추기 위해 $\beta_i = \text{score}(i)\alpha$로 정의되는 학습 가능한 스케일링 인자를 사용하며, 이는 국소 대 전반적 도수 비율 기반의 스코어 함수에 기반한다.
• 메시지 보상으로 정확한 기울기 추정을 유지함으로써 1차 정류점에 수렴함을 보장한다.
• 샘플된 부분그래프에 직접적으로 다양한 GNN 아키텍처를 적용함으로써 광범위한 적용 가능성을 확보한다.

실험 결과

연구 질문

• RQ1 미니배치 외부의 메시지를 기각함에도 불구하고 부분그래프 기반 샘플링 방법이 GNN 학습에서 증명 가능한 수렴을 달성할 수 있는가?
• RQ2 백워드 프로파게이션에서 메시지를 기각함으로써 발생하는 편향을 높은 계산 비용 없이 효과적으로 보상할 수 있는가?
• RQ3 기울기 보정에서 최신 메시지와 이전 메시지를 모두 사용할 경우 수렴 속도와 모델 정확도에 어떤 영향을 미치는가?
• RQ4 작은 배치 크기에서 전반배치 학습과 비교해 본다면 제안된 방법의 성능는 어떠한가?
• RQ5 대규모 그래프 학습 과제에서 수렴 보장을 확보하면서도 높은 효율성을 유지할 수 있는가?

주요 결과

• LMC는 GNN의 1차 정류점에 수렴하며, 증명 가능한 수렴성을 보장하는 첫 번째 부분그래프 기반 샘플링 방법이다.
• Ogbn-arxiv, FLICKR, REDDIT, PPI와 같은 대규모 벤치마크에서 LMC는 GAS 대비 최대 30%, FM 대비 최대 40%까지 에포크당 학습 시간을 단축시켰다.
• 작은 배치 크기(예: 1 또는 40)에서도 LMC는 전반배치 GD와 유사한 예측 성능를 달성하지만, 기존 베이스라인들은 상당히 뒤처진다.
• LMC는 전방 및 백워드 프로세스에서 샘플된 모든 메시지를 100% 활용하며, 전반배치 메모리 사용량과 동일하게 메시지 손실 없이 처리한다.
• 제거 실험 결과에 따르면, $\alpha = 0.4$ 이며 $\text{score}(i) = 2x - x^2$일 경우 작은 배치 크기에서 최고의 성능를 기록하며 낙관성 영향을 최소화한다.
• 배치 크기 1, $\alpha = 0.4$ 조건에서 Ogbn-arxiv에서 71.65%의 정확도를 기록하며 모든 베이스라인을 압도하는 최고의 성능를 달성한다.