GNNのメッセージパッシングはノードごとの埋め込みを出力しますが、分子特性予測や流体場の大域的な性質予測のようにグラフ全体を1つの表現で捉えたい場合があります。本記事では、ノード埋め込みの集合をグラフ全体の埋め込みに変換するグラフプーリングの手法を解説します。セットプーリング（平均・和・アテンション集約）とグラフコースニング（クラスタリングによる階層的集約）の2つのアプローチを扱います。

グラフプーリングとは何か｜ノード埋め込みからグラフ全体の表現へ

ノードレベル埋め込みとグラフレベル埋め込みの違い

GNNのメッセージパッシングを $K$ 回繰り返すと、各ノード $u$ の最終的な埋め込み $\mathbf{z}_u \in \mathbb{R}^d$ が得られます。ノード分類やリンク予測のように個々のノードやエッジに対して予測を行うタスクでは、これで十分です。

しかし予測のターゲットがグラフ全体の性質である場合、ノード埋め込みをそのまま使うことはできません。グラフ全体の表現 $\mathbf{z}_\mathcal{G} \in \mathbb{R}^d$ を得るために、ノード埋め込みの集合 $\{\mathbf{z}_u \mid u \in \mathcal{V}\}$ をまとめて1つのベクトルに変換する操作が必要になります。これをグラフプーリング（graph pooling）と呼びます。

流体力学の文脈では、非構造格子上のGNNが各節点の速度・圧力を埋め込みとして出力した後、グラフ全体の表現として圧力損失・揚力係数・流量といったマクロ的な量を予測する場面がグラフプーリングに対応します。

グラフプーリングの2つのアプローチ

グラフプーリングのアプローチは大きく2種類に分けられます。

アプローチ	概要
セットプーリング	グラフ構造を無視し、ノード埋め込みの集合に対してプーリング関数を適用する
グラフコースニング	グラフ構造を利用してノードをクラスタに割り当て、段階的にグラフを粗大化する

セットプーリングはシンプルで汎用的ですが、グラフのトポロジー情報を利用しません。グラフコースニングはCNNの階層的プーリングに対応する考え方であり、グラフ構造を活かした表現学習が可能です。

セットプーリング｜平均・和とアテンションベースの集約

単純な和・平均によるプーリング

最も基本的なセットプーリングは、全ノードの埋め込みの和または平均を取ることです。

$$\mathbf{z}_\mathcal{G} = \frac{\sum_{v \in \mathcal{V}} \mathbf{z}_v}{f_n(|\mathcal{V}|)} \tag{1}$$

ここで $f_n(|\mathcal{V}|)$ は正規化関数であり、$f_n = 1$ とすれば和、$f_n = |\mathcal{V}|$ とすれば平均になります。式(1)はノード数の異なるグラフにも適用でき、ノード順序に依存しない（置換不変な）操作です。

グラフが比較的小規模で、ノードの数や位置に強い意味的な違いがない場面（例：小分子の分類）では、平均プーリングだけでも実用的な精度が得られることが知られています。

流体格子への応用では、格子点の埋め込みを平均することは「グラフ全体の平均的な状態」を表現することに相当します。圧力損失のように空間平均と相関する量の予測には効果的ですが、局所的な特徴（衝撃波の位置など）は失われます。

アテンションベースのプーリング

単純な平均・和では、各ノードが等しい重みで集約されます。あるノードがグラフレベルの予測に強く関与し、別のノードはほとんど関与しないような場合、重み付きの集約が有利です。

Vinyals et al. [2015] が提案した手法は、LSTMとアテンションを組み合わせた反復的なプーリングを行います。クエリベクトル $\mathbf{q}_t$ （各ノードの埋め込みをどのような基準で注目するかを表すベクトル）を介して各ノードへの注目度を計算し、コンテキストベクトル $\mathbf{o}_t$ に集約します。この操作を $t = 1, \ldots, T$ の $T$ ステップ繰り返します。

第 $t$ ステップでは、まず前ステップのコンテキストベクトル $\mathbf{o}_{t-1}$ とクエリベクトル $\mathbf{q}_{t-1}$ を入力としてLSTMがクエリベクトルを更新します。

$$\mathbf{q}_t = \text{LSTM}(\mathbf{o}_{t-1},\, \mathbf{q}_{t-1}) \tag{2}$$

このように、ステップごとにクエリベクトルを更新することで、ノード埋め込みを様々な視点で注目できます。次に、クエリベクトル $\mathbf{q}_t$ と各ノード埋め込み $\mathbf{z}_v$ からアテンションスコア（どの程度クエリベクトル $\mathbf{q}_t$ を注視するか）を計算します。

$$e_{v,t} = f_a(\mathbf{z}_v,\, \mathbf{q}_t),\quad \forall v \in \mathcal{V} \tag{3}$$

ここで $f_a : \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}$ はアテンション関数（内積など）です。アテンションスコアをソフトマックスで正規化して重みを得ます。

$$a_{v,t} = \frac{\exp(e_{v,t})}{\sum_{u \in \mathcal{V}} \exp(e_{u,t})},\quad \forall v \in \mathcal{V} \tag{4}$$

重み $a_{v,t}$ でノード埋め込みを集約してコンテキストベクトルを得ます。

$$\mathbf{o}_t = \sum_{v \in \mathcal{V}} a_{v,t}\, \mathbf{z}_v \tag{5}$$

$T$ ステップ分のコンテキストベクトルを連結してグラフ全体の埋め込みとします。

$$\mathbf{z}_\mathcal{G} = \mathbf{o}_1 \oplus \mathbf{o}_2 \oplus \cdots \oplus \mathbf{o}_T \tag{6}$$

ここで $\oplus$ はベクトルの連結を表します。初期値 $\mathbf{q}_0$ と $\mathbf{o}_0$ はともにゼロベクトルとします。

アテンションスコア

このアテンションベースのプーリングはグラフ分類タスクで広く採用されています。異なるクエリベクトルを使うことで、多角的な視点でノードを注目するため、単純な平均よりも情報量の多いグラフ表現が得られます。

グラフコースニング｜クラスタリングによる階層的プーリング

グラフコースニングのアイデア

CNNでは、畳み込み後にmax-poolingやaverage-poolingを繰り返すことで特徴マップを段階的に縮小します。グラフコースニング（graph coarsening）はこれに対応するグラフ版の操作で、グラフ上のノードをクラスタにまとめながら段階的にグラフを粗大化します。

セットプーリングとの大きな違いは、グラフコースニングがグラフのトポロジー情報（隣接関係）を活用する点です。クラスタに属するノード同士の接続関係を保ちながら粗大化するため、グラフの構造的な特徴が保存されます。

割り当て行列とグラフの粗大化

グラフコースニングでは、まずクラスタリング関数 $f_c$ によって各ノードの $c$ クラスタへの帰属の強さを表す割り当て行列 $\mathbf{S} \in \mathbb{R}^{|\mathcal{V}| \times c}$ を計算します。

$$\mathbf{S} = f_c(\mathcal{G},\, \mathbf{Z}) \tag{7}$$

各要素$\mathbf{S}[u, i] \geq 0$ はノード $u$ とクラスタ $i$ の結びつきの強さを表します。割り当て行列 $\mathbf{S}$ を使って、粗大化後の隣接行列 $\mathbf{A}^\text{new} \in \mathbb{R}^{c \times c}$ を計算します。

$$\mathbf{A}^\text{new} = \mathbf{S}^\top \mathbf{A}\, \mathbf{S} \tag{8}$$

$\mathbf{A}^\text{new}[i, j]$ は粗大化後のクラスタ $i$ とクラスタ $j$ の間のエッジ強度を表します。ポイントは、式(8)によって隣接行列がクラスタ間の隣接のようなものを表す行列が得られている点です。同様に、粗大化後のノード特徴行列 $\mathbf{X}^\text{new} \in \mathbb{R}^{c \times d}$ は次のように得られます。

$$\mathbf{X}^\text{new} = \mathbf{S}^\top \mathbf{X} \tag{9}$$

$\mathbf{X}^\text{new}$ の各行は、そのクラスタに帰属するノードの特徴の加重平均を表します。粗大化されたグラフ $(\mathbf{A}^\text{new}, \mathbf{X}^\text{new})$ に対して再びGNNのメッセージパッシングを適用し、さらに粗大化するというプロセスを繰り返すことで、最終的に十分小さくなったグラフにセットプーリングを適用してグラフ全体の埋め込みを得ます。

クラスタリング関数の設計と実装上の制約

グラフコースニングを一気通貫で学習するためには、クラスタリング関数 $f_c$ が微分可能でなければなりません。微分できないと誤差逆伝播法が使えないからです。これはスペクトルクラスタリングなどの古典的なアルゴリズムが直接使えないことを意味します。

実用的なアプローチとして、 $f_c$ としてGNNを使う、つまり別のGNNで割り当て行列 $\mathbf{S}$ を予測する方法があります [Ying et al., 2018]。入力グラフ $(\mathcal{G}, \mathbf{Z})$ を別のGNNに与え、ソフトマックスで正規化した出力を割り当て行列 $\mathbf{S}$ として使います。この方法なら、グラフ構造と特徴量の両方を利用してクラスタリングが行われ、主タスクの損失と一緒に学習できます。

粗大化するアプローチの実装上の注意点をまとめると次のとおりです。

課題	内容
微分可能性	一気通貫で学習するためには、$f_c$ は微分可能でなければならない
学習の不安定性	セットプーリングより学習が不安定になりやすい
計算コスト	粗大化後の密な隣接行列 $\mathbf{A}^\text{new}$ の計算に $O(

クラスタリングではなくノードを選択して除去するアプローチも提案されており [Cangea et al., 2018; Gao and Ji, 2019]、計算量の観点から有利な場面があります。

流体シミュレーションへの応用

非構造格子上で動作するGNNにグラフプーリングを組み合わせることで、格子全体の大域的特性を予測するサロゲートモデルを構築できます。

平均プーリング（式(1)）は実装が最もシンプルで、流量・平均圧力・マクロ的な熱流束のような空間平均量の予測に適しています。アテンションベースのプーリング（式(2)〜式(6)）は、格子上の特定の領域（壁面近傍や衝撃波付近など）が大域的な性質に強く影響する場合に有効です。

グラフコースニングは多スケール構造を持つ問題との相性が良いです。例えば境界層を細かく解像した格子と、主流域を粗く解像した格子が共存する非構造格子では、コースニングによってマルチスケールの情報を段階的に集約できます。

さらに深く学ぶための書籍

本記事のテーマをより深く学びたい方には、以下の書籍をおすすめします。

• Graph Representation Learning（William L. Hamilton, Morgan & Claypool, 2020）：本記事の内容の原典です。グラフプーリングの各手法の数学的定式化と直感的な解釈を丁寧に解説しており、GNNを体系的に学ぶ出発点として最適な一冊です。