MLエンジニアへの道 #36 - GNNパイプライン

Last Edited: 1/23/2025

このブログ記事では、GNN APIを用いて異なる下流タスクに対応するパイプラインを構築します。

ML

これまで、GNNに関する議論では、下流タスクの詳細については避け、下流タスクのための潜在ノード表現の作成に焦点を当ててきました。 しかし、下流タスクは、適切なデータ準備、モデル選択、評価手法の決定、およびGNNパイプラインの構築において重要です。 基本的なモデルやGNN APIの機能を使ったパイプライン構築を学んだ今、下流タスクとそれに基づくパイプラインについて議論します。

グラフレベルのタスク

グラフレベルのタスクは、最終的な潜在ノード埋め込みを集約したり、新しいグラフを生成したりして、グラフのクラスや値を予測することを指します。 具体例として、分子の毒性の二値分類、材料科学の化学特性予測、薬物候補生成(ドラッグディスカバリー)が挙げられます。 これらのタスクでは、複数のグラフとそれに関連付けられたラベルを準備し、データセットを訓練用、検証用、テスト用に分割して教師あり学習を行います。

上記は、TF-GNNおよびPyGを使用してグラフレベルのタスク向けのデータセットを生成する方法を示しています。 PyGの実装は、前回の記事で取り上げた内容とほぼ同じであり、私たちが最も馴染みのある設定に沿っています。 モデルがトレーニング中にテストデータセットを見ないため、これらのモデルは 帰納的(inductive) 設定 でトレーニングされているといいます。 帰納的設定でトレーニングされたモデルは、未見のデータに対しても一般化できます。

ノードレベルおよびエッジレベルのタスク

ノードレベルおよびエッジレベルのタスクでは、ノードやエッジのクラスや値を予測します。ノードレベルタスクの例として、 金融取引ネットワークにおける不正アカウント検出や、引用ネットワークにおける学術論文の発表会場予測があります。 一方、エッジレベルタスクの例として、推薦システムでの評価予測や、SNSアカウント間の将来のリンク予測があります。 これらのタスクは主に単一のグラフを対象とし、複数のグラフを扱うケースは少ないです。単一のグラフを扱う場合、 各ノードやエッジをデータポイントとして扱い、データセットを分割するためにマスクを使用します。

トレーニング中にモデルがグラフ全体の構造を観察して推論を行うため、単一グラフをマスク処理して扱う場合、 モデルは 推移的(transductive) 設定でトレーニングされるといいます。これまでに紹介したモデルは隣接行列を使用し、 グラフ全体の構造が既知であることを前提としています。そのため、これらのモデルは単一グラフの埋め込みを生成する際に、 新しいグラフへの一般化能力が低い傾向にあります。

大規模な動的グラフ

グラフが時間とともに変化しない場合、新しいグラフに一般化できないモデルでも問題はありません。 しかし、グラフが動的に変化する場合、これは問題となります。また、従来のモデルは隣接行列を使用するため、 大規模なグラフでは処理速度が遅くなります。そのため、ユーザーアイテム二部グラフを用いた評価予測のような大規模かつ動的なグラフでは、 ルートノード周辺のサブグラフをサンプリングし、各サブグラフの埋め込みを使用してルートノードの予測を行うようにモデルをトレーニングします。

この方法を使用すると、モデルを帰納的(inductive)な設定でトレーニングでき、 動的に変化するグラフの未見データやサブグラフに一般化することが可能になります。 また、モデルのサイズが小さくなり、処理速度も向上します。このようにタスクに応じて、 適切なデータ準備、モデル、トレーニング設定、評価方法を選択する必要があります。

結論

この記事では、グラフレベル、ノードレベル、およびエッジレベルのタスク、帰納的および推移的設定、 ならびにGNN APIを使用した適切なGNNパイプラインについて説明しました。改めて強調すると、 取り組んでいる問題の種類を認識し、それに応じた適切なGNNパイプラインを選択することが重要です。 GNN APIの使用方法に関する詳細情報については、以下の公式ドキュメントを参照することをお勧めします。

リソース