[論文レビュー] TGL: A General Framework for Temporal GNN Training on Billion-Scale Graphs

Hongkuan Zhou, Da Zheng+4 more|arXiv (Cornell University)|Mar 28, 2022

Advanced Graph Neural Networks被引用数 11

ひとこと要約

TGL は、100億スケールのグラフ上で大規模な時系列 GNN のトレーニングを統合的かつスケーラブルに行うためのフレームワークであり、時系列サンプラー、ノードメモリ、マイルーブ、メッセージパッシングエンジンを備え、多様な TGNN アーキテクチャの効率的トレーニングを可能にする。13倍の高速化を達成し、4つの GPU で100億エッジのグラフを1〜10時間でトレーニング可能であり、ランダムチャンクスケジューリングや時系列 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 アーキテクチャを統一したトレーニングフレームワークに統合すること。
• 100億エッジのグラフ上で、陳腐化するノードメモリや非効率なサンプリングといった、大バッチトレーニングの課題を克服すること。
• 小さなベンチマークをはるかに超える大規模な実世界データセット上で TGNN の評価を可能にすること。
• 産業および学術的利用に耐える生産環境対応で、マルチ GPU 対応のシステムを提供すること。

提案手法

• 動的グラフにおける時系列順エッジの効率的格納とアクセスを可能にする時系列 CSR データ構造の設計。
• 時系列順序とマルチコア処理を活用して、近隣エッジサンプリングを高速化する並列時系列サンプラーの開発。
• トレーニングバッチをより小さな時系列的に整合性のあるチャンクにシャッフルすることで、大バッチトレーニングにおける依存関係の損失を低減するランダムチャンクスケジューリング技術の導入。
• ノードメモリモジュールとマイルーブの統合により、時間ステップにわたる動的ノード表現の維持と更新を実現。
• 上記のコンポONENTを設定ファイルによる構成で組み合わせることで、多様な TGNN バリエーションをサポートするメッセージパッシングエンジンの設計。
• NCCL と CUDA ストリームを用いたデータ転送と計算のオーバーラップにより、マルチ GPU トレーニングの実装。

実験結果

リサーチクエスチョン

• RQ1 統合フレームワークは、100億規模の動的グラフ上で、多様な時系列 GNN アーキテクチャを効率的にトレーニングできるか？
• RQ2 ノードメモリが陳腐化する状況において、大バッチトレーニングをどのようにスケーラブルかつ正確に実現できるか？
• RQ3 最適化されたデータ構造と並列サンプリングにより、大規模な時系列グラフでどの程度のパフォーマンス向上が達成できるか？
• RQ4 実世界の100億エッジデータセット上で、時系列 GNN は小さなスケールのベンチマークと比較してどの程度の性能を示すか？
• RQ5 マルチ GPU トレーニングは、大規模な時系列 GNN のトレーニングをどの程度スケーリングできるか？

主な発見

• TGL は、リンク予測およびノード分類タスクにおいて、オープンソースのベースラインと比較して平均13倍の高速化を達成しながら、精度を維持または向上させた。
• 時系列並列サンプラーは、マルチコア CPU でベースラインサンプラーと比較して173倍の高速化を達成した。
• 4 GPU のマシンでは、TGL は100億以上の時系列エッジを含む1エポックのトレーニングを1〜10時間で完了した。
• ランダムチャンクスケジューリングにより、時系列依存性の損失が軽減され、ベースライン手法が失敗するような大規模バッチサイズ（例：4800）でも収束が可能になった。
• マルチ GPU トレーニングでは、JODIE で4 GPU で2.74倍、TGN で4 GPU で2.25倍の高速化を達成し、優れたスケーラビリティを示した。
• メモリベースのモデル（例：TGN）は、GDELT データセットを1エポックトレーニングするのに30分未満で完了したが、非メモリベースのモデルは3時間以上を要した。