[論文レビュー] TF-GNN: Graph Neural Networks in TensorFlow

Oleksandr Ferludin, Arno Eigenwillig+25 more|arXiv (Cornell University)|Jul 7, 2022

Advanced Graph Neural Networks被引用数 15

ひとこと要約

TF-GNN は、4段階の API ハイアラルキーを備えたスケーラブルなグラフニューラルネットワーク（GNN）を実装するオープンソースのTensorFlowライブラリであり、異種・多関係的グラフデータをサポートする。研究者や開発者が、TensorFlowの低レベル演算から高レベルの Keras モデルまで、さまざまな抽象化レベルで GNN を構築・学習・デプロイできる。OGBN-MAG において、より小さな最適化された MPNN モデルが、はるかに大きな複雑なアーキテクチャを上回る最先端の性能を達成した。

ABSTRACT

TensorFlow-GNN (TF-GNN) is a scalable library for Graph Neural Networks in TensorFlow. It is designed from the bottom up to support the kinds of rich heterogeneous graph data that occurs in today's information ecosystems. In addition to enabling machine learning researchers and advanced developers, TF-GNN offers low-code solutions to empower the broader developer community in graph learning. Many production models at Google use TF-GNN, and it has been recently released as an open source project. In this paper we describe the TF-GNN data model, its Keras message passing API, and relevant capabilities such as graph sampling and distributed training.

研究の動機と目的

• 複雑な異種関係的データ上で GNN を学習するためのスケーラブルで生産向けのフレームワークの不足に対処すること。
• 初心者から熟練の研究者まで、あらゆるスキルレベルの開発者を対象とした階層的な API を提供すること。
• 分散学習、TPUアクセラレーション、複数モodalの事前学習モデルへのアクセスを含む、TensorFlowエコシステムへのシームレスな統合を実現すること。
• 実世界の応用における GNN の効率的なモデル開発、ハイパーパramータチューニング、デプロイを可能にすること。
• 単純でよく最適化された GNN モデルが、ベンチマークのグラフ学習タスクにおいて、より複雑なアーキテクチャを上回ることを示すこと。

提案手法

• TF-GNN は4段階の API を実装している：異種グラフ表現のためのデータレベル、GNN レイヤー構築のための Keras メッセージパッシング API、メッセージパッシング操作の定義のための GraphUpdate、学習と推論のためのオーケストレーター。
• エッジ中心、ノード中心、グラフ中心の GNN ビルディングブロックをサポートしており、複雑な関係構造の柔軟なモデリングを可能にする。
• ノードおよびエッジ特徴量がグラフの近傍にわたって適用される学習可能な関数により更新される、メッセージパッシングフレームワークを採用している。
• TensorFlowエコシステムと統合されており、分散学習、TPUアクセラレーション、推論用の SavedModel エクスポートを可能にしている。
• 迅速なプロトタイピングを可能にする、VanillaMPNN、GraphSAGE、GATv2、MultiHeadAttention などのプリビルドコンポーネントを含むモデルライブラリを提供している。
• Vizier との統合により、モデル容量、ドロップアウト、正則化、メッセージ次元などのハイパーパramータチューニングが自動最適化可能である。

実験結果

リサーチクエスチョン

• RQ1 TensorFlow 上に、異種的・多関係的グラフデータをサポートする統合的かつスケーラブルな GNN フレームワークを構築できるか？
• RQ2 階層的な API デザインは、機械学習の経験レベルが異なるユーザーを効果的に支援できるか？
• RQ3 単純でよく最適化された MPNN モデルは、大規模なグラフベンチマークにおいて、より複雑なトランスフォーマー基盤のアーキテクチャを上回れるか？
• RQ4 TensorFlow エコシステムとの統合は、GNN の学習効率とデプロイ準備度をどの程度向上できるか？
• RQ5 自動ハイパーパramータチューニングは、アーキテクチャの複雑さを伴わずに、高性能な GNN 設定を効果的に発見できるか？

主な発見

• Vizier を用いたハイパーパramータチューニングで学習された、589万パラメータの単純な MPNN モデルが、OGBN-MAG で検証精度 0.5149 を達成し、はるかに大きな 2680万パラメータの異種グラフトランスフォーマー（HGT）モデルを上回った。
• Vizier 最適化済みの MPNN モデルはテスト精度 0.5027 を達成し、HGT のテスト精度 0.4982 を上回り、ハイパーパラメータチューニングの有効性がアーキテクチャの複雑さを上回ることを示した。
• TF-GNN フレームワークは、Google 内の複数のチームで GNN モデルの生産デプロイを可能にし、そのスケーラビリティと使いやすさを検証した。
• エッジ中心およびグラフ中心のメッセージパッシングをサポートすることで、ノード中心のフレームワークで取り組みにくいモデル（例：グラフトランスフォーマー）をネイティブに表現可能となった。
• TensorFlow エコシステムとの統合により、TPU 上での効率的な学習と、推論用の SavedModel へのスムーズなエクスポートが実現され、デプロイの障壁が低減した。
• モデルライブラリと高レベル API により開発時間が大幅に短縮され、GNN の迅速な実験とデプロイが可能になった。