[論文レビュー] PyG 2.0: Scalable Learning on Real World Graphs

Matthias Fey, J. Sunil+11 more|ArXiv.org|Jul 22, 2025

Advanced Graph Neural Networks被引用数 5

ひとこと要約

PyG 2.0 は PyTorch Geometric の大規模な更新で、スケーラビリティ、ヘテロジニアス、時系列グラフのサポート、説明性を強化し、実世界の大規模グラフでの効率的な学習と LLM との統合・関係型ディープラーニングを可能にします。

ABSTRACT

PyG (PyTorch Geometric) has evolved significantly since its initial release, establishing itself as a leading framework for Graph Neural Networks. In this paper, we present Pyg 2.0 (and its subsequent minor versions), a comprehensive update that introduces substantial improvements in scalability and real-world application capabilities. We detail the framework's enhanced architecture, including support for heterogeneous and temporal graphs, scalable feature/graph stores, and various optimizations, enabling researchers and practitioners to tackle large-scale graph learning problems efficiently. Over the recent years, PyG has been supporting graph learning in a large variety of application areas, which we will summarize, while providing a deep dive into the important areas of relational deep learning and large language modeling.

研究の動機と目的

• 実世界の大規模かつヘテロジニアスなグラフに対する、スケーラブルなグラフ学習を動機づけ、可能にする。
• ヘテロジニアスおよび時系列グラフに対応するアーキテクチャ的・アルゴリズム的改善、効率的なデータ読み込み、最適化されたメッセージ伝搬を導入する。
• 堅牢な説明性ツールとストレージバックエンドおよびGPUアクセラレーションとのシームレスな統合を提供する。
• Relational Deep Learning および LLMs の Retrieval-augmented generation との統合などの応用を実証する。

提案手法

• メモリ効率とデータ局所性のために CSR/CSC 形式へのキャッシュを用いた最適化されたメッセージ伝搬のため EdgeIndex を導入する。
• 集約を第一級の概念とし、さまざまな集約演算子（mean、max、sum、学習可能/カスタム変種）へのプラグアンドプレー対応を提供する。
• edgeタイプごとに層を複製し、型特異的集約を用いた双方向メッセージ伝搬を実行する変換による native ヘテロジニアスグラフサポートを提供する。
• カーネル融合と torch.compile を用いたエンドツーエンドのコンパイルをサポートし、2–3倍の実行時間短縮を達成する。
• FeatureStore と GraphStore のリモートバックエンドを備えたスケーラブルなグラフ基盤と、サブグラフのためのサンプリングスタック（層ごとの剪定および時系列サブグラフサンプリングを含む）を提供する。
• cuGraph/cuGraph-PyG および WholeGraph を統合し、GPU 加速データ読み込みと分散ストレージを提供して、GPU間で線形スケーリングを実現する。

実験結果

リサーチクエスチョン

• RQ1 PyG を、ビリオン級のスケール、ヘテロジニアス、時系列グラフ上で効率的に GNN を訓練するように拡張するには？
• RQ2 End-to-end グラフ学習のためのプラグアンドプレー可能なスケーラビリティ、ストレージの柔軟性、および加速計算を可能にするアーキテクチャの primitive は何か？
• RQ3 ユーザーフレンドリーで即利用可能な形で、ヘテロジニアス GNN に説明性を統合するにはどうすればよいか？
• RQ4 PyG 2.0 は Relational Deep Learning および Graph-RAG のような高度な応用を大規模言語モデルでどう実現するか？

主な発見

• EdgeIndex ベースの加速はメモリのボトルネックを軽減し、大規模グラフでの効率的なメッセージ伝搬を実現する。
• 集約を第一級原理とすることで、多様な集約戦略の柔軟なプラグアンドプレー利用を可能にし、学習可能なものや型破りなものを含む。
• ヘテロジニアスグラフのサポートは native で自動化されており、エッジタイプごとの GNN複製と双方向メッセージ伝搬のための FX ベース変換を通じて実現する。
• Torch.compile ベースのカーネル融合は、一般的な GNN アーキテクチャのフォワード/バックワードパスで2–3倍の速度向上をもたらす。
• 層ごとの剪定と時系列サンプリングを含むサブグラフサンプリングは、計算とメモリのオーバーヘッドを抑えつつ非常に大規模なグラフ上での GNN 訓練を拡張する。
• cuGraph-PyG および WholeGraph の統合は、マルチ GPU 環境でのGPU加速、ほぼ線形スケーリングを提供する。