Google が TensorFlow-GNN 1.0 のリリースを正式に発表しました。大規模なグラフニューラルネットワークを構築するための動的かつインタラクティブなサンプリング

2005 年には画期的な著作「The Graph Neural Network Model」が出版され、グラフ ニューラル ネットワークが誰でも利用できるようになりました。

これまで、科学者がグラフ データを処理する方法は、データの前処理段階でグラフを一連の「ベクトル表現」に変換することでした。

CNN の登場により、情報損失というこの欠点は完全に変わりました。過去 20 年間、世代を追うごとにモデルが進化し続け、ML 分野の進歩を推進してきました。

本日、Google は大規模な GNN を構築するための本番環境でテスト済みのライブラリである TensorFlow GNN 1.0 (TF-GNN) のリリースを正式に発表しました。

TensorFlow でのモデリングとトレーニングの両方をサポートし、大規模なデータ ストアからの入力グラフの抽出もサポートします。

TF-GNN は、オブジェクトの種類と関係が異なるノードとエッジのセットによって表される異種グラフ用にゼロから構築されています。

現実世界のオブジェクトとその関係はさまざまな形で現れますが、TF-GNN の異種焦点により、それらを非常に自然に表現できます。

Google の科学者 Anton Tsitsulin 氏は、複雑なヘテロ構造モデリングが復活したと述べています。

TF-GNN 1.0がリリースされました

物体とそれらの相互関係は、私たちの世界のいたるところに存在します。

オブジェクトを理解する上での関係の重要性は、輸送ネットワーク、生産ネットワーク、ナレッジ グラフ、ソーシャル ネットワークなど、オブジェクト自体の属性を個別に見るのと同じくらい重要です。

離散数学とコンピュータ サイエンスでは、さまざまな不規則な方法でエッジによって任意に接続された「ノード」で構成されるグラフなどのネットワークを長い間形式化してきました。

ただし、ほとんどの機械学習アルゴリズムでは、ピクセルのグリッドや単語のシーケンスなど、入力オブジェクト間の規則的で均一な関係のみが許可され、関係がまったくないことも許可されます。

グラフ ニューラル ネットワーク (略して GNN) は、グラフの接続性 (初期のアルゴリズム DeepWalk や Node2Vec など) とさまざまなノードおよびエッジ入力の特徴の両方を活用する強力な手法です。

GNN は、グラフ全体 (この分子は特定の方法で反応するか?)、個々のノード (引用に基づいて、このドキュメントのトピックは何ですか?)、および潜在的なエッジ (この製品は別の製品と一緒に購入される可能性が高いですか?) について予測を行うことができます。

GNN は、グラフ上で予測を行うことに加えて、より一般的なニューラル ネットワークの使用例とのギャップを埋める強力なツールです。

グラフの離散的な関係情報を連続的にエンコードし、別のディープラーニング システムに自然に組み込むことができます。

Google は本日、大規模な GNN を構築するための実稼働テスト済みのライブラリである TensorFlow GNN 1.0 (TF-GNN) を正式に発表しました。

TensorFlow では、このようなグラフは tfgnn.GraphTensor 型のオブジェクトによって表されます。

これは、tf.data.Dataset、tf.function などで「ファーストクラス オブジェクト」として受け入れられる複合テンソル型 (Python クラス内のテンソルのコレクション) です。

グラフ構造と、ノード、エッジ、グラフ全体の特性の両方を保存できます。

GraphTensor のトレーニング可能な変換は、高レベル Keras API の Layers オブジェクトとして定義することも、tfgnn.GraphTensor プリミティブを使用して直接定義することもできます。

GNN: コンテキスト内のオブジェクトの予測

次に、TF-GNN をさらに説明するために、その典型的なアプリケーションの 1 つを見てみましょう。

相互参照テーブルによって定義された大規模データベース内のノードのクラスのプロパティを予測する

たとえば、コンピュータサイエンス (CS) の引用データベース arxiv では、1 対多の引用と多対 1 の引用関係があり、各論文の主題領域を予測できます。

ほとんどのニューラル ネットワークと同様に、GNN は多数のラベル付きサンプル (数百万単位) のデータセットでトレーニングされますが、各トレーニング ステップは、はるかに小さいバッチのトレーニング サンプル (数百など) で構成されます。

数百万のサンプルに拡張するために、GNN は基礎となるグラフの比較的小さなサブグラフのストリームでトレーニングされます。各サブグラフには、中央のラベル付きノードの GNN 結果を計算し、モデルをトレーニングするのに十分な生データが含まれています。

このプロセスはサブグラフ サンプリングとも呼ばれ、GNN トレーニングにとって非常に重要です。

既存のツールのほとんどは、バッチでサンプリングを実行して、トレーニング用の静的サブグラフを生成します。

TF-GNN は、動的かつインタラクティブなサンプリングを通じてこれを改善するツールを提供します。

サブグラフ サンプリングは、大きなグラフから小さくて実用的なサブグラフを抽出し、GNN トレーニングの入力例を作成するプロセスです。

TF-GNN 1.0 では、関連するすべてのスケールの動的またはバッチサブグラフ サンプリング (Colab ノートブックでのインタラクティブ サンプリング) を構成するための柔軟な Python API が導入されています。

具体的には、単一のトレーニング ホストのメイン メモリに保存されている小さなデータセットの「効率的なサンプリング」、または Apache Beam を介してネットワーク ファイル システムに保存されている大規模なデータセット (最大数億のノードと数十億のエッジ) の分散サンプリングです。

これらの同じサンプリングされたサブグラフ上で、GNN のタスクはルート ノードの非表示 (または潜在) 状態を計算することです。非表示状態は、ルート ノードの近傍に関する関連情報を集約してエンコードします。

一般的なアプローチは「メッセージ パッシング ニューラル ネットワーク」です。

メッセージ パッシングの各ラウンドで、ノードは着信エッジに沿って隣接ノードからメッセージを受信し、これらのエッジから自身の隠し状態を更新します。

n ラウンド後、ルート ノードの非表示状態は、n エッジ内のすべてのノードの集約情報を反映します (下の図に示すように、n = 2)。メッセージと新しい隠し状態は、ニューラル ネットワークの隠し層によって計算されます。

異種グラフでは、異なるタイプのノードとエッジに対して個別にトレーニングされた隠し層を使用することが理にかなっていることがよくあります。

この図は、各ステップでノード状態が外部ノードから内部ノードに伝播され、新しいノード状態が内部ノードで計算される単純な「メッセージ パッシング ニューラル ネットワーク」を示しています。ルートノードに到達すると、最終的な予測を行うことができます。

トレーニングのセットアップは、ラベル付きノードを持つ GNN の隠し状態の上に出力層を配置し、損失を計算し (予測誤差を測定するため)、バックプロパゲーションを介してモデルの重みを更新することによって行われます。これは、ニューラル ネットワークのトレーニングで一般的です。

GNN は、教師ありトレーニングに加えて、教師なしトレーニングも可能であり、ノードとその特徴の離散グラフ構造の連続表現 (または埋め込み) を計算できます。

これらの表現は、他の ML システムでもよく使用されます。

このようにして、グラフによってエンコードされた個別の関係情報は、より一般的なニューラル ネットワークの使用例に組み込むことができます。 TF-GNN は、異種グラフ上の教師なし目標のきめ細かい指定をサポートします。

GNNアーキテクチャの構築

TF-GNN ライブラリは、さまざまな抽象化レベルでの GNN の構築とトレーニングをサポートします。

最高レベルでは、ユーザーは、Kera レイヤーで表現される、ライ​​ブラリにバンドルされている定義済みモデルのいずれかを使用できます。

TF-GNN には、研究文献からの少数のモデル コレクションに加えて、慎重にキュレーションされたモデリングの選択肢を提供する、高度に構成可能なモデル テンプレートが付属しています。

Google では、これらの選択肢が社内の多くの質問に対する強力な基準となることを発見しました。テンプレートは GNN レイヤーを実装します。ユーザーは Kera レイヤーから初期化するだけで済みます。

 import tensorflow_gnn as tfgnn from tensorflow_gnn.models import mt_albis def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec): """Builds a GNN as a Keras model.""" graph = inputs = tf.keras.Input(type_spec=graph_tensor_spec) # Encode input features (callback omitted for brevity). graph = tfgnn.keras.layers.MapFeatures( node_sets_fn=set_initial_node_states)(graph) # For each round of message passing...for _ in range(2): # ... create and apply a Keras layer. graph = mt_albis.MtAlbisGraphUpdate( units=128, message_dim=64, attention_type="none", simple_conv_reduce_type="mean", normalization_type="layer", next_state_type="residual", state_dropout_rate=0.2, l2_regularizatinotallow=1e-5, )(graph) return tf.keras.Model(inputs, graph)

最も低いレベルでは、ユーザーは、ノードからそのすべての出力エッジにデータをブロードキャストしたり、すべての入力エッジからノードにデータを集中させたりするなど、グラフ内でデータを渡すためのプリミティブに基づいて、GNN モデルをゼロから作成できます。

TF-GNN のグラフ データ モデルは、特徴や隠れた状態に関しては、ノード、エッジ、および入力グラフ全体を平等に扱います。

したがって、MPNN のようなノード中心のモデルを直接表現できるだけでなく、より一般的な形式のグラフ ネットワークも表現できます。

これは、コア TensorFlow 上のモデリング フレームワークとして Kera を使用して実行できます (ただし、必ずしも実行する必要はありません)。

研修手配

上級ユーザーはカスタム モデル トレーニングを自由に実行できますが、TF-GNN Runner は一般的なケースで Kera モデルのトレーニングを調整する簡潔な方法も提供します。

簡単な呼び出しは次のようになります。

 from tensorflow_gnn import runner runner.run( task=runner.RootNodeBinaryClassification("papers", ...), model_fn=model_fn, trainer=runner.KerasTrainer(tf.distribute.MirroredStrategy(), model_dir="/tmp/model"), optimizer_fn=tf.keras.optimizers.Adam, epochs=10, global_batch_size=128, train_ds_provider=runner.TFRecordDatasetProvider("/tmp/train*"), valid_ds_provider=runner.TFRecordDatasetProvider("/tmp/validation*"), gtspec=..., )

Runner は、分散トレーニングや Cloud TPU 上の固定形状の tfgnn.GraphTensor パディングなど、ML Pain に対する既成のソリューションを提供します。

単一タスクのトレーニング（上記参照）に加えて、複数（2 つ以上）のタスクの共同トレーニングもサポートします。

たとえば、教師なしタスクを教師ありタスクと混合して、アプリケーション固有の帰納的バイアスを持つ最終的な連続表現 (または埋め込み) を形成できます。呼び出し元は、タスク パラメータをタスク マップに置き換えるだけです。

 from tensorflow_gnn import runner from tensorflow_gnn.models import contrastive_losses runner.run( task={ "classification": runner.RootNodeBinaryClassification("papers", ...), "dgi": contrastive_losses.DeepGraphInfomaxTask("papers"), }, ... )

さらに、TF-GNN Runner には、モデル属性のための統合された勾配実装が含まれています。

統合勾配出力は、観測された GraphTensor と同じ接続性を持つ GraphTensor ですが、その特徴は勾配値に置き換えられ、勾配値が大きいほど、小さい勾配値よりも GNN 予測に大きく貢献します。

要約すると、Google は TF-GNN が TensorFlow における GNN の大規模な応用を促進し、この分野におけるさらなる革新を促進することを期待しています。