[論文レビュー] LMC: Fast Training of GNNs via Subgraph Sampling with Provable Convergence

Zhihao Shi, Xize Liang+1 more|arXiv (Cornell University)|Feb 2, 2023

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 の一次静止点への収束を達成すること。

提案手法

• 勾配誤差をバイアス（破棄されたメッセージ由来）と分散（確率的サンプリング由来）に分解することで、不偏なミニバッチ勾配を提案。
• 部分グラフサンプリング中に破棄されたメッセージを再取得・補正できるメッセージ伝達フレームワークを用いてバックワードパスを再定式化。
• 不完全な最新メッセージと履歴メッセージを組み合わせた補正機構を導入し、バイアス補正を効率的に行う。
• 局所的からグローバルな次数比に基づくスコア関数を用いて、スタリネスと正確性のバランスを取る学習可能スケーリング係数 $\beta_i = \text{score}(i)\alpha$ を導入。
• メッセージ補正を通じた正確な勾配推定を維持することで、一次静止点への収束を保証。
• サンプリングされた部分グラフ上で直接実行可能なため、さまざまな GNN アーキテクチャに広く適用可能。

実験結果

リサーチクエスチョン

• RQ1 ミニバッチ外のメッセージを破棄するにもかかわらず、部分グラフ単位のサンプリング手法が GNN 学習において保証された収束を達成できるか？
• RQ2 バックワードパスでメッセージを破棄することによって生じるバイアスを、高い計算コストを伴わずに効果的に補正できるか？
• RQ3 最新メッセージと履歴メッセージの両方を勾配補正に用いることで、収束速度とモデル精度にどのような影響を与えるか？
• RQ4 小さなバッチサイズ下での、提案手法とフルバッチ学習の性能比較は？
• RQ5 大規模グラフ学習タスクにおいて、収束保証を維持しながらも高い効率性を維持できるか？

主な発見

• LMC は GNN の一次静止点への収束を達成し、保証された収束を有する最初の部分グラフ単位のサンプリング手法です。
• Ogbn-arxiv や FLICKR、REDDIT、PPI などの大規模ベンチマークにおいて、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% の正確度を達成し、すべてのベースラインを上回る最高性能を記録しました。