[論文レビュー] GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings

Matthias Fey, Jan Eric Lenssen+2 more|arXiv (Cornell University)|Jun 10, 2021

Advanced Graph Neural Networks被引用数 16

ひとこと要約

GNNAutoScale (GAS) は、履歴埋め込みを用いて計算グラフをプルーニングすることにより、任意のメッセージパッシンググラフニューラルネットワーク (GNN) の大規模グラフにおけるスケーラブルな学習を可能にし、入力サイズにかかわらず GPU メモリ使用量を一定に保つ。モデルの表現力は完全に維持され、 ogbn-products や Reddit などの大規模ベンチマークで最先端の性能を達成し、VR-GCN や MVS-GNN といった先行手法を上回る。

ABSTRACT

We present GNNAutoScale (GAS), a framework for scaling arbitrary message-passing GNNs to large graphs. GAS prunes entire sub-trees of the computation graph by utilizing historical embeddings from prior training iterations, leading to constant GPU memory consumption in respect to input node size without dropping any data. While existing solutions weaken the expressive power of message passing due to sub-sampling of edges or non-trainable propagations, our approach is provably able to maintain the expressive power of the original GNN. We achieve this by providing approximation error bounds of historical embeddings and show how to tighten them in practice. Empirically, we show that the practical realization of our framework, PyGAS, an easy-to-use extension for PyTorch Geometric, is both fast and memory-efficient, learns expressive node representations, closely resembles the performance of their non-scaling counterparts, and reaches state-of-the-art performance on large-scale graphs.

研究の動機と目的

• GPU メモリ制限と近隣ノードの爆発的増加のため、大規模グラフにおける深く表現力のある GNN の学習のスケーラビリティの課題に対処すること。
• エッジのサブサンプリングや非学習可能な伝搬による性能低下を避けるために、メッセージパッシング GNN の完全な表現力を維持しつつ、スケーラブルな学習を実現すること。
• スケーラビリティとメッセージパッシングメカニズムを分離することで、特定の GNN アーキテクチャに依存せず、GCNII や PNA などの多様な GNN オペレータに適用可能となるように一般化すること。
• 履歴埋め込みによる近似誤差の理論的バウンディングを提示し、実用的な手法でこれらのバウンディングをタイトにすること。
• CPU に保存された履歴を用いて、1 つの GPU のみで大規模データセット (例: ogbn-products) におけるフルバッチ GNN の学習を可能にすること。

提案手法

• GAS は、現在のミニバッチに含まれるノードおよびその直接の 1 ホップ近隣ノードに限定して計算グラフをプルーニングし、GPU メモリ使用量を入力サイズに依存しない定数に保つ。
• 履歴埋め込み（過去の学習イテレーションにおけるノード表現）を用いて、ミニバッチ外のノードをメモリ効率の良い代替手段として使用する。
• 履歴埋め込みは各イテレーションごとに更新され、非アクティブなノードからの情報伝搬を保証し、完全な計算グラフを保存せずにトポロジカルな依存関係を維持する。
• 埋め込みの陳腐化度と関数のリプシッツ連続性に基づく理論的近似誤差バウンディングを提供し、実用的な戦略でこれらのバウンディングをタイトにする。
• PyGAS は、PyTorch Geometric の拡張であり、最小限のコード変更で既存の GNN コードベースにシームレスに統合可能である。
• この手法はモデル設計とは直交しており、GCNII や PNA などの深く表現力のあるアーキテクチャを含む、任意のメッセージパッシング GNN に適用可能である。

実験結果

リサーチクエスチョン

• RQ1 エッジのサブサンプリングや表現力の損失なしに、任意のメッセージパッシング GNN を大規模グラフにスケーリングすることは可能か？
• RQ2 履歴埋め込みを用いることで生じる理論的近似誤差は何か？ そして、実用的にどのようにしてこれを最小化できるか？
• RQ3 履歴埋め込みは、フルバッチ GNN の表現力を維持しつつ、GPU メモリ使用量を一定に保つことができるか？
• RQ4 GAS フレームワークは、深く表現力のある GNN に適用した場合、大規模グラフベンチマークで最先端の性能を達成するか？
• RQ5 アーキテクチャの変更なしに、多様な GNN アーキテクチャに一般化可能か？

主な発見

• GAS を用いることで、ogbn-products（2.4M ノード、61.9M エッジ）のような大規模グラフにおいて、1 つの GPU のみでフルバッチ GNN の学習が可能となり、GPU メモリ使用量は入力サイズに依存しない。
• ogbn-products データセットにおいて、PNA-GAS は 79.91% の精度を達成し、以前の最先端手法である GraphSAINT（79.08%）および OOM で失敗したフルバッチ PNA を上回った。
• Reddit において、PNA-GAS は 97.17% の精度を達成し、GraphSAINT（97.00%）および VR-GCN（94.50%）を上回り、大規模なノード分類タスクにおける優れた性能を示した。
• GCNII-GAS モデルは Reddit で 96.77% の精度を達成し、OOM で失敗したフルバッチ GCNII を上回った。深層モデルのスケーラブルな学習が可能であることを示した。
• ogbn-products では、履歴の保存に要するメモリ消費量が 1 レイヤーあたり約 2GB に留まり、CPU RAM に保存可能であり、学習は効率的かつスケーラブルであった。
• オープンソース実装である PyGAS を用いることで、最小限のコード変更で深く表現力のある GNN を大規模グラフで学習可能となり、非スケーリング手法に近い性能を達成した。