GNNの一般化メッセージパッシング|エッジとグラフレベルを統合した更新設計
標準的なGNNのメッセージパッシングはノード埋め込みのみを更新します。しかしエッジや流体場の大域的な状態を同時に管理したい場面では、この枠組みでは不十分です。本記事では、エッジ・ノード・グラフの3階層を統一的に更新するBattaglia et al. [2018] の一般化メッセージパッシングを解説します。
ノードレベルGNNの限界と一般化メッセージパッシングの動機
標準GNNが扱えない情報
これまでの記事で解説してきたGNNは、メッセージパッシングを通じてノード埋め込み を更新します。最終的な出力はノード埋め込み であり、ノード分類・リンク予測・グラフ分類(前の記事で解説したプーリングを経由)に使うことができます。
しかし実際の問題では、この枠組みが扱いにくい状況があります。
エッジに対して直接予測を行いたい場合がその一例です。配管ネットワークでは各管路(エッジ)の流量や圧力損失を予測したいことがあります。標準GNNでは最終層のノード埋め込みからエッジの端点の情報を取り出して予測するしかなく、エッジの「状態」をメッセージパッシング中に逐次更新する手段がありません。
また、グラフ全体のグローバルな情報をメッセージパッシング中に各ノードへ伝えたい場合も扱いにくいです。マルチフィジクス解析では全体圧力・温度・レイノルズ数のようなグローバルな条件が各節点の更新に影響しますが、標準GNNではこの情報を各ノードへ届ける自然な仕組みがありません。
3種類の埋め込みを同時に管理する
こうした問題に対処するため、Battaglia et al. [2018] は各メッセージパッシングのイテレーションでエッジ埋め込み・ノード埋め込み・グラフ埋め込みの3種類を同時に管理・更新する一般化メッセージパッシング(generalized message passing)を提案しました。
エッジ・ノード・グラフの統合更新|一般化メッセージパッシングの定式化
3つの埋め込みとその役割
一般化メッセージパッシングでは、各メッセージパッシングイテレーション において3種類の埋め込みが管理されます。
- :エッジ の埋め込み
- :ノード の埋め込み
- :グラフ全体の埋め込み
各イテレーションでは、エッジ更新 → ノードへの集約 → ノード更新 → グラフ更新の順で処理が進みます。
エッジの更新
最初にエッジ埋め込みを更新します。エッジ の新しい埋め込みは、前のイテレーションのエッジ埋め込み、両端のノード埋め込み、およびグラフ全体の埋め込みを入力として計算します。
は任意の微分可能な関数(MLPなど)です。エッジが端点ノードの情報を取り込んで自身を更新するため、端点に生じた変化がエッジの状態に反映されます。またグラフ全体の埋め込み を入力とすることで、グローバルな条件変化もエッジに伝播します。
流体工学への対応で考えると、管路ネットワークでは式(1)が「ある管路の流状態を、接続ノードの圧力・速度情報とシステム全体の運転条件から更新する」操作に相当します。
ノードへの集約
更新されたエッジ埋め込みを使って、各ノード へのメッセージを集約します。
には、これまでの記事で解説した和・平均・最大値・アテンションなどの集約関数を使えます。式(2)では近傍のノード埋め込みではなくエッジ埋め込みを集約する点が標準GNNとの違いです。エッジの状態が集約されることで、ノード埋め込みにエッジの情報が自然に統合されます。
ノードの更新
集約されたメッセージとグラフ全体の埋め込みを使って、ノード埋め込みを更新します。
グラフ全体の埋め込み が各ノードの更新に直接入力されることで、グローバルな情報がすべての節点に伝播します。 には基本GNN・スキップ接続・ゲート付き更新(GRU/LSTM)など既存の手法をそのまま適用できます。
グラフ全体の更新
最後に、全ノードと全エッジの埋め込みを使ってグラフ全体の埋め込みを更新します。
は全ノードと全エッジの埋め込みの集合を入力とするため、前の記事で解説したグラフプーリングの手法がそのまま使えます。グラフ全体の埋め込みは次のイテレーションで式(1)・式(3)に利用されるため、グローバルな情報が繰り返しエッジとノードに伝播します。
4式の更新順序と情報の流れ
式(1)〜式(4)の更新順序と各式の依存関係をまとめます。
| ステップ | 更新対象 | 主な入力 |
|---|---|---|
| 式(1) エッジ更新 | 、端点ノード、 | |
| 式(2) ノード集約 | 更新済みエッジ埋め込み | |
| 式(3) ノード更新 | 、集約メッセージ、 | |
| 式(4) グラフ更新 | 、全ノード、全エッジの埋め込み |
エッジが先に更新され、その結果がノード集約に使われる点がこの枠組みの核心です。エッジとノードとグラフの情報が双方向に影響し合いながら更新されます。
流体シミュレーションへの応用と設計上の指針
非構造格子上でのエッジ分類
非構造格子を用いたCFDでは、格子エッジは「境界エッジ(壁面・流入・流出)」と「内部エッジ」に分類されます。この分類はシミュレーションの境界条件を決定する重要な情報です。
標準GNNでは最終層のノード埋め込みから間接的にエッジの状態を推定するしかありませんが、一般化メッセージパッシングではエッジ埋め込み に対して直接分類損失を定義できます。各メッセージパッシングイテレーションを経て更新されたエッジ埋め込みに分類器を適用するだけで、エッジレベルの予測タスクが成立します。
グラフ埋め込みによるグローバル条件の管理
熱流体解析では、系全体のレイノルズ数・マッハ数・プラントル数といた無次元パラメータが流れ場全体の挙動を支配します。これらのグローバルなパラメータをグラフ埋め込み にエンコードしておくと、式(1)と式(3)を通じてすべてのエッジとノードの更新にグローバル条件が反映されます。
例えば圧縮性流れと非圧縮性流れでは同じ格子トポロジーでも更新の様子が異なります。グラフ埋め込みにマッハ数の情報を持たせることで、1つのGNNモデルが複数の流れ条件に対応できるようになります。
標準GNNとの設計の使い分け
一般化メッセージパッシングの4式すべてを実装すると設計の自由度が増しますが、実装コストも上がります。タスクに応じて必要な部分だけを取り入れた設計が実践的です。
| タスク | 推奨設計 |
|---|---|
| ノード分類・回帰のみ | 標準GNN(エッジ・グラフ埋め込みは不要) |
| エッジの分類・回帰が必要 | 式(1)(2)(3)を実装(グラフ埋め込みは省略可) |
| グローバル条件を活用したい | 式(1)(3)(4)を実装(グラフ埋め込みを管理) |
| エッジ・ノード・グラフすべてが必要 | 式(1)〜式(4)すべてを実装 |
・ には本シリーズで解説してきた基本GNN・GAT・スキップ接続・ゲート付き更新などをそのまま使えます。 には和・平均・最大値・アテンションなどの集約関数が適用でき、 にはグラフプーリングの手法が対応します。この枠組みはChapter 5全体で解説してきた手法を統合する設計指針として機能します。
さらに深く学ぶための書籍
本記事のテーマをより深く学びたい方には、以下の書籍をおすすめします。
- Graph Representation Learning(William L. Hamilton, Morgan & Claypool, 2020):本記事の内容の原典です。一般化メッセージパッシングの定式化とその設計上の意義を丁寧に解説しており、GNNを体系的に学ぶ出発点として最適な一冊です。