GNNの実装はもはや難しくありません。この記事では、効率的なGNNとスケーラブルなグラフ表現学習の最新の進歩についてまとめています。

グラフ ニューラル ネットワークは、現実世界に適用する場合、メモリ制限、ハードウェア制限、信頼性制限など、多くの課題に直面します。この記事では、ケンブリッジ大学の博士課程の学生である Chaitanya K. Joshi が、データの準備、効率的なアーキテクチャ、学習パラダイムという 3 つの側面から、研究者がこれらの問題を克服するために行った進歩についてレビューします。

効率的でスケーラブルなグラフ表現学習のためのツールボックス。

この記事の目的は、効率的なグラフ ニューラル ネットワークとスケーラブルなグラフ表現学習に関する主要なアイデアの概要を説明することです。また、グラフ ニューラル ネットワークを現実世界に合わせて拡張し、リアルタイムのシナリオに適用できるようにするデータ準備、GNN アーキテクチャ、学習パラダイムの主要な進歩についても紹介します。

具体的な内容は以下のとおりです。

グラフニューラルネットワークが直面する現実世界の課題

グラフ ニューラル ネットワークは、グラフ、セット、3D ポイント クラウドなどの不規則なデータ構造で動作できる新しいディープラーニング アーキテクチャです。近年、GNN の開発は、ソーシャル ネットワーク、推奨システム、生物医学的発見など、さまざまな分野に広がっています。ただし、実際のアプリケーションでは、GNN の構築には次のような課題があります。

メモリ制限

現実世界のネットワークは非常に大規模かつ複雑になることがあります。たとえば、Facebook には約 30 億のアクティブ アカウントがあり、いいね、コメント、共有などのさまざまな方法でやり取りするため、アカウントをノードとして構成されるグラフに無数のエッジが作成されます。

Facebook の全ユーザーとそのやり取りの表など、現実世界のグラフ ネットワークは非常に大きくて扱いにくいため、そのような巨大なグラフを GPU メモリに収めてニューラル ネットワークをトレーニングするのは不可能な場合があります。

ハードウェアの制限

グラフは本質的にスパースなオブジェクトであり、GNN は効率的でスケーラブルな計算のためにそのスパース性を活用する必要があります。しかし、現代の GPU は行列に対する集中的な演算を処理するように設計されているため、これは言うほど簡単ではありません。スパース行列用のカスタマイズされたハードウェア アクセラレータは、GNN の適時性とスケーラビリティを大幅に向上させることができますが、その設計方法については未解決の問題が残っています。

最新の GPU は密な行列演算に適していますが、グラフは本質的にスパースな構造です。隣接行列が非常に疎でない限り、GNN を実装するときに、グラフを単純に密行列として扱い、マスクを使用して接続されていないノードを識別する方が高速になることがよくあります。

信頼性の制限

巨大なグラフを処理する簡単な方法は、グラフを小さなサブグラフに分割し、ミニバッチ勾配降下法を使用して GNN をトレーニングすることです (各サブグラフはデータのミニバッチとして機能します)。しかし、このアプローチの最大の問題は、サンプルが独立している標準的な機械学習データセットとは異なり、ネットワーク データのリレーショナル構造によってサンプル間に統計的な依存関係が生まれることです。したがって、サブグラフが完全なグラフのセマンティクスを保持し、GNN のトレーニングに信頼できる勾配を提供することを保証することは簡単な作業ではありません。

画像全体のセマンティクスと勾配情報の両方を保持するサンプリング手順を設計するにはどうすればよいでしょうか?

巨大な画像の処理

サブサンプリング技術

巨大なグラフを GNN に適合させようとする既存の論文では、大きなグラフを管理しやすいサブグラフに分割するグラフのサブサンプリングに焦点が当てられています。 ClusterGCN はクラスタリング アルゴリズムを使用してグラフをクラスター化されたサブグラフに分割し、各クラスターを単一のデータ ミニバッチとして扱って GNN をトレーニングします。これは実際にはうまく機能しますが、特にクラスターが類似のラベルを持つ意味のあるクラスターを形成することが多い同種グラフの場合はうまく機能します。

出典: GraphSAINT: グラフサンプリングに基づく帰納的学習法

GraphSAINT は、サブグラフのミニバッチを構築するための、より一般的な確率的グラフ サンプラーを提案します。可能なサンプリング方式には、均一ノード/エッジ サンプリングとランダム ウォーク サンプリングが含まれます。ただし、前のセクションで説明した信頼性の問題 (セマンティック情報と勾配情報) により、サブサンプリング方法では、フルイメージでのトレーニングと比較してモデルのパフォーマンスが制限される可能性があります。

履歴ノードの埋め込み

GNNAutoScale (GAS) は、大規模なグラフに GNN を適用するための基本的なサブサンプリング手法の有望な代替手段です。 GAS は、Chen らによる過去のノード埋め込みに関する研究に基づいており、以前のトレーニングで取得したノード埋め込みを現在のトレーニングに再利用します。

出典: 「GNNAutoScale: 履歴埋め込みによるスケーラブルで表現力豊かなグラフ ニューラル ネットワーク」

GAS フレームワークには 2 つの主要コンポーネントがあります。まず、最初の部分ではノードのミニバッチを構築し (高速ランダム サブサンプリングを実行)、GNN 計算グラフを整理して、ミニバッチ内のノードとその 1 ホップ隣接ノードのみを保持します。つまり、GAS のスケールは GNN の深さに依存しません。次に、GNN 集約にノード埋め込みのミニバッチが必要になるたびに、GAS は CPU に保存されている履歴埋め込みからそれらを取得します。同時に、現在のミニバッチ ノードの履歴埋め込みも継続的に更新されます。

2 番目の部分は、サブサンプリングとの主な違いです。これにより、GNN は情報の表現を最大化し、現在のミニバッチ データを履歴埋め込みと組み合わせて、完全な近傍情報を取得して利用しながら、大規模グラフへのスケーラビリティを確保できます。

GAS の作者は、そのアイデアを人気のある PyTorch ジオメトリ ライブラリにも統合しました。これにより、ほとんどのメッセージ パッシング GNN を非常に大きなグラフでトレーニングできるようになり、GPU メモリ要件が削減され、ほぼ完全なバッチ パフォーマンス (つまり、完全なグラフでトレーニングする場合のパフォーマンス) が維持されます。

大規模グラフにスケーリングするためのツールキットには、他にも次のようなアイデアがあります。

[CVPR2020]   L ² -GCN: グラフ畳み込みネットワークのレイヤーごとの効率的な学習トレーニング。Yuning You、Tianlong Chen、Zhangyang Wang、Yang Shen。
[KDD2020]  近似 PageRank を使用したグラフ ニューラル ネットワークのスケーリング。Aleksandar Bojchevski、Johannes Klicpera、Bryan Perozzi、Amol Kapoor、Martin Blais、Benedek Rózemberczki、Michal Lukasik、Stephan Günnemann。
[ICLR 2021]  テンソル関数によるグラフトラバーサル: スケーラブルな学習のためのメタアルゴリズム。Elan Markowitz、Keshav Balasubramanian、Mehrnoosh Mirtaheri、Sami Abu-El-Haija、Bryan Perozzi、Greg Ver Steeg、Aram Galstyan。
[ニュールIPS2021]  グラフ ニューラル ネットワークの深さと範囲の分離。Hanqing Zeng、Muhan Zhang、Yinglong Xia、Ajitesh Srivastava、Andrey Malevich、Rajgopal Kannan、Viktor Prasanna、Long Jin、Ren Chen。

スケーラブルでリソース効率の高いGNNアーキテクチャ

グラフ拡張MLP

スケーラブルな GNN を開発する場合、グラフの関係構造を考慮せずに、単純な MLP を少数のノードで実行するという直感に反する考え方があります。

Wu らによる Simplified Graph Convolutional Network (SGC) は、このアイデアを提案した最初の研究でした。 SGC は本質的に、Kipf と Welling によるバニラ GCN の「脱構築」であり、(高価だが学習不可能な) スパース近傍特徴集約ステップと、それに続く ReLU 非線形ステップ (安価で学習可能) の線形投影ステップを切り離します。大規模なグラフを処理する場合、特徴の集約は CPU 上で効率的に事前計算され (スパース操作では優れたパフォーマンスを発揮)、その後、「構造強化」ノード特徴がバッチ処理され、GPU でトレーニングされた MLP に渡されます。

出典: グラフ畳み込みネットワークの簡素化

Rossi らによる SIGN: Scalable Inception Graph Neural Networks は、コンピューター ビジョンの Inception ネットワークからインスピレーションを得て、複数の事前計算ステップを実行して各グラフ ノードのより優れた初期構造的特徴を取得することにより、SGC のアイデアをさらに一歩進めようとしています。

このアーキテクチャ ファミリの他の興味深い開発には、構造拡張 MLP の理論的限界を調査した Chen らの研究や、ノード機能に対して MLP を実行した後の同様のラベル伝播アプローチが、均質グラフ上のより扱いにくい GNN よりも優れていることを示す Huang らの論文などがあります。

効率的なグラフ畳み込み層

残念ながら、同次グラフ上のノード分類以外の多くのタスクでは、SGC/SIGN のようなモデルよりも表現力豊かな GNN アーキテクチャが必要になる場合があります。これは通常、グラフ上で分類または推論タスクを実行するときに発生します。この場合、最もパフォーマンスの高い GNN は通常、ノードとエッジの両方の機能を活用します。

このようなモデルは、GNN アーキテクチャ設計のメッセージ パッシング スタイル (Petar Veličković によって普及) に従っており、各エッジを異なる方法で処理するという点で「異方性」と見なすことができます (一方、バニラ GCN は、同じ学習可能な重みが各エッジに適用されるという点で「等方性」です)。ただし、各レイヤーのノードとエッジで潜在空間埋め込みを維持すると、GNN の推論の待ち時間とメモリ要件が大幅に増加します。

出典: 異方性グラフニューラルネットワークは必要か?

Tailor らによる Efficient Graph Convolution (EGC) は、このジレンマを解決しようと試みます。彼らは基本的な GCN 設計 (中規模サイズのグラフのバッチに適切にスケーリング) から始めて、GCN のスケーラビリティを維持しながら、表現力が最大限に高い畳み込み GNN バージョンを設計しました。驚くべきことに、Open Graph ベンチマークのグラフ分類タスクでは、より複雑で扱いにくいベースライン モデルよりも優れたパフォーマンスを発揮します。

出典: 異方性グラフニューラルネットワークは必要か?

EGC レイヤーは PyTorch Geometric にも統合されており、プラグアンドプレイの代替として使用して GNN のパフォーマンスとスケーラビリティを向上させることができます。

Li らによるもう 1 つの興味深いアイデアは、コンピューター ビジョンにおける効率的な ConvNet (可逆接続、グループ畳み込み、重み付け、バランスのとれたモデル) を活用して、GNN のメモリとパラメーターの効率を向上させるというものです。彼らのフレームワークは、（前例のない）1000 以上のレイヤーを持つ GNN をトレーニングすることができ、Open Graph ベンチマークの大規模ノード分類タスクで優れたパフォーマンスを実現します。

GNN 圧縮学習パラダイム

データ準備技術と効果的なモデルアーキテクチャに加えて、学習モデル、つまりモデルのトレーニング方法によっても、GNN のパフォーマンスが大幅に向上し、レイテンシが短縮されます。

知識の蒸留はパフォーマンスを向上させる

知識蒸留 (KD) は、高性能だがリソースを大量に消費する教師モデルからリソース効率の高い生徒に知識を転送する一般的なニューラル ネットワーク学習パラダイムです。 KD の概念は、もともと Hinton らによって提案されました。KD は、教師モデルの出力ロジットを標準的な教師あり学習損失と一致させるように生徒をトレーニングします。

Yang らによる最近の研究とそれに続く Zhang らによる研究では、リソース効率の高い GNN を開発するためのこの単純なロジットベースの KD アイデアの威力が実証されました。彼らは表現力豊かな GNN を教師モデルとしてトレーニングし、ノード機能とグラフ構造の相関性が高い場合に展開を容易にするために、KD を使用して知識を MLP 学生に転送しました。この考え方は、前のセクションの SGC や SIGN などのグラフ拡張 MLP に拡張することができ、MLP のようなモデルのパフォーマンスを大幅に向上させると同時に、実稼働システムに簡単に導入できるようになります。

出典: グラフレスニューラルネットワーク: 蒸留により古い MLP に新しいトリックを教える

コンピューター ビジョンでは、実践者は、潜在的な埋め込み空間を揃える損失関数を介して、表現知識を教師から生徒に転送することで、ロジット ベースの KD を超えようとします。

出典: グラフニューラルネットワークの表現知識蒸留について

Yang らによる先駆的な研究では、教師のノード埋め込み空間からローカルなトポロジ構造を保存するように生徒をトレーニングすることにより、GNN の表現蒸留を初めて調査しました。彼らはこのアプローチをローカル構造保存 (LSP) と呼んでいます。これは、生徒が教師のノードの埋め込み空間に存在する直接の近隣ノードのペアワイズ類似性を模倣するように促すためです。

最近、Joshi らは LSP 研究を拡張し、すぐ隣の物体を超えた潜在的な相互作用を保持することで、グローバルなトポロジカル構造を考慮しました。彼らは、GNN の新しい表現蒸留の目的を 2 つ提案しました。(1) 明示的なアプローチであるグローバル構造保存は、LSP を拡張してすべてのペアの類似性を考慮します。(2) 暗黙的なアプローチであるグラフ対照表現蒸留は、対照学習を使用して、共有表現空間で生徒ノードの埋め込みを教師ノードの埋め込みと揃えます。

出典: グラフニューラルネットワークの表現知識蒸留について

Joshi らによる Open Graph Benchmark データセットの実験では、表現蒸留を使用して軽量 GNN をトレーニングすると、実験パフォーマンスと、ノイズの多いデータや分布外データに対する堅牢性が大幅に向上することが示されています。

GNN の他の KD 手法には、教師なし蒸留 (自己蒸留とも呼ばれる) やデータなし蒸留などがあります。

低精度GNNの量子化

量子化を考慮したトレーニング (QAT) は、もう 1 つの一般的なニューラル ネットワーク学習パラダイムです。従来のニューラル ネットワーク モデルの重みとアクティベーションは 32 ビット浮動小数点数 FP32 として保存されますが、QAT は INT8 や INT4 などの精度の低い整数の重みとアクティベーションを使用してモデルをトレーニングします。低精度モデルは、パフォーマンスの低下を犠牲にするものの、推論の待ち時間において大きな利点があります。

Tailor らによる DegreeQuant は、GNN 専用の QAT 手法を提案しました。コンピューター ビジョン CNN 用に設計された汎用 QAT を GNN に適用すると、通常、量子化後のパフォーマンスが非常に低下します。 DegreeQuant は、基礎となるデータのグラフ構造を量子化プロセスに巧みに統合することでこの問題を軽減することを目指しています。DegreeQuant は、多数の隣接ノード (高次数) を持つノードが QAT 中に不安定性を引き起こす可能性があることを示し、QAT を実行するときに高次数ノードをランダムにマスクすることを提案しています。これにより、FP32 モデルと比較して GNN の QAT がより安定し、INT8 のパフォーマンス低下が最小限に抑えられます。

GNN を量子化する他のアイデアとしては、Zhao らのニューラル アーキテクチャ検索や、Bahri らの 2 値化法 (量子化の極端な例) を活用することが挙げられます。一般に、QAT のパフォーマンス低下の一部は、前のセクションの Bahri らと Joshi らによる論文「Graph Contrastive Representation Distillation」に示されているように、知識蒸留によって回復できます。

結論と展望

この論文では、効率的なグラフ ニューラル ネットワークとスケーラブルなグラフ表現学習に焦点を当てています。まず、現実世界の GNN の理論的およびエンジニアリング上の課題を特定します。

• 巨大なグラフ - メモリ制限
• スパースコンピューティング - ハードウェアの制限
• 図のサブサンプリング - 信頼性の限界

次に、効率的でスケーラブルな GNN を開発するためのツールとして役立つ可能性のある 3 つの重要なアイデアを紹介します。

1. データ準備 - 履歴ノード埋め込み検索を通じて、大規模グラフのサンプリングから CPU-GPU までのハイブリッド トレーニングを実現します。 2. 効率的なアーキテクチャ - 大規模ネットワークへのスケーリングのためのグラフ拡張 MLP と、バッチ処理されたグラフ データに対するリアルタイム推論のための効率的なグラフ畳み込み設計。 3. 学習パラダイム - 量子化を考慮したトレーニング (低精度のモデルの重みとアクティベーション) と知識蒸留 (表現力豊かな教師モデルを使用して GNN をより効率的にする) を組み合わせて、推論の待ち時間とパフォーマンスを最大化します。

効率的でスケーラブルなグラフ表現学習のためのツールボックス。

近い将来、研究コミュニティは、GNN ネットワーク向けの効率的でスケーラブルなツールボックスの開発を継続し、直接統合を通じて PyTorch Geometric や DGL などの GNN ライブラリに登場することが期待されます。また、現実世界のグラフやリアルタイム アプリケーションで GNN が機能する成功事例がますます増えていくと期待しています。

長期的には、グラフ データ + GNN が難解な新興研究分野から、機械学習の研究とアプリケーションのための標準的なデータ + モデル パラダイム (2D 画像 + CNN、またはテキスト + Transformer のように) に移行することを期待しています。したがって、PyTorch や TensorFlow などの標準フレームワークへの GNN のより深い統合、GNN 専用のハードウェア アクセラレータの開発、グラフ データ向けのより洗練されたハードウェアとソフトウェアの共同設計が期待されます。

実際、グラフ データと GNN から大きな商業的価値を引き出している企業では、こうした取り組みがすでに行われている可能性があります。

この記事で取り上げたトピックについてさらに詳しく知りたい場合は、次の研究をご覧ください。

Abadal らによる広範な調査では、GNN の基礎からグラフ表現学習用のハードウェア アクセラレータの設計 (この論文では取り上げられていません) まで、あらゆる内容が網羅されています。

出典: グラフ ニューラル ネットワークの高速化に関する調査: アルゴリズムの観点。

この記事で言及されている原論文を参照してください。