GNNによるグラフ分類とREADOUT関数の理論と実装

グラフ分類の流れ: ノード埋め込みからREADOUTを経てラベル予測するまでの全体像

上の図は、GNNによるグラフ分類の全体の流れを示しています。入力グラフからメッセージパッシングでノード埋め込みを計算し、READOUT関数で1本のベクトルに集約してから分類器に渡すという3段構造が一目で分かります。この「集約」の部分こそが本記事のテーマです。

分子が「毒性を持つかどうか」を予測したいとしましょう。分子はたくさんの原子(ノード)と結合(エッジ)からできたグラフです。原子1個1個に「この原子は重要か」というラベルを付けたいわけではありません。知りたいのは、分子という1つのまとまり全体が、毒性ありか/なしかという、たった1つの答えです。ところがグラフ・ニューラルネットワーク(GNN)のメッセージパッシングが計算してくれるのは、原子ごとの特徴ベクトル(ノード埋め込み)の集まりです。原子100個の分子なら100本のベクトルが出てきます。この「ばらばらの100本のベクトル」を、どうやって「分子1個を表す1本のベクトル」に押し縮めればよいのでしょうか。

この押し縮めの操作を READOUT関数(あるいはグローバルプーリング)と呼びます。一見すると「全部足すか、平均を取るか、最大値を取るか、どれでもよさそう」に見えます。しかし実は、この選択がモデルの表現力を決定的に左右します。たとえば「水素原子が3個ついた構造」と「6個ついた構造」を区別したいとき、平均(mean)を取ってしまうと両者が同じ値になって見分けがつかなくなることがあります。READOUTの設計は、単なる実装上の都合ではなく、「どんなグラフを区別できるか」という GNN の能力の上限に直結しているのです。

この理論的な核心が、グラフ同型判定の古典的アルゴリズムである Weisfeiler-Lehman(WL)検査 との対応であり、その対応を最大限に活かすよう設計されたのが GIN(Graph Isomorphism Network) です。READOUTを理解することは、応用面でも「分子の物性予測」「化合物のスクリーニング」「ソーシャルネットワークのコミュニティ分類」「タンパク質の機能予測」といったグラフ単位の分類・回帰タスクすべての土台になります。本記事では、READOUTが満たすべき条件を数式で導出し、なぜsum集約が最も強いのかをWL検査と結び付けて示した上で、PyTorch GeometricでTUDatasetのグラフ分類を実装し、READOUTの種類による精度差を実際に観察します。

本記事の内容

  • グラフ分類の問題設定と、なぜノード埋め込みの集約が必要なのか
  • READOUT関数が満たすべき置換不変性の条件と、その数学的な意味
  • sum/mean/maxの単射性の違いをWL同型判定の観点から導出(GINがsumを選ぶ根拠)
  • PyTorch GeometricによるTUDatasetグラフ分類の実装
  • READOUTの種類ごとの精度差と、層を深くしたときの過平滑化の影響をグラフで比較

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が格段に深まります。READOUTはメッセージパッシングの「次の段」にあたる操作なので、まずメッセージパッシングを理解しておくことが重要です。

グラフ分類という問題設定

まず、私たちが解こうとしている問題をはっきりさせましょう。GNNには大きく2種類のタスクがあります。1つはノードレベルのタスクで、たとえば「ソーシャルネットワークの各ユーザーがボットかどうか」を1人ずつ判定します。もう1つが本記事で扱うグラフレベルのタスクで、グラフ全体に対して1つのラベルを付けます。

具体的には、グラフ $G = (V, E)$ が与えられたとき、その全体に対するラベル $y_G$(たとえば「毒性あり=1/なし=0」)を予測します。学習データは「グラフとそのラベルの組」の集合 $\{(G_1, y_1), (G_2, y_2), \dots, (G_N, y_N)\}$ で、各グラフのノード数 $|V|$ はバラバラです。分子なら原子3個のものもあれば50個のものもあります。

ここで本質的な困難が現れます。ニューラルネットワークの分類器(最後の全結合層)は、固定長のベクトルを入力として受け取ります。ところがメッセージパッシングが出力するのは、ノードごとの埋め込みベクトルの集まりで、その本数(=ノード数)はグラフごとに違うのです。

メッセージパッシングを $K$ 層適用した後、各ノード $v$ は埋め込みベクトル $\bm{h}_v^{(K)} \in \mathbb{R}^d$ を持ちます。グラフ $G$ 全体では、ノード集合 $V = \{v_1, \dots, v_n\}$ に対して

$$ \{\bm{h}_{v_1}^{(K)}, \bm{h}_{v_2}^{(K)}, \dots, \bm{h}_{v_n}^{(K)}\} $$

という $n$ 本のベクトルの集まりが得られます。この $n$ 本を、グラフごとに本数が違っても必ず1本の固定長ベクトル $\bm{h}_G \in \mathbb{R}^d$ に変換しなければなりません。この変換こそがREADOUT関数です。

$$ \bm{h}_G = \mathrm{READOUT}\left(\{\bm{h}_v^{(K)} : v \in V\}\right) $$

そしてこの $\bm{h}_G$ を全結合層+softmaxに通して、グラフラベルを予測します。ここまでで「なぜ集約が必要か」がはっきりしました。次に、この集約関数がどんな条件を満たさなければならないのかを考えていきます。

READOUTとは — 順番に依存してはいけない

READOUTの設計に入る前に、一番大事な制約を直感的につかみましょう。グラフのノードには、本来「1番目のノード」「2番目のノード」という決まった順番はありません。私たちがデータをコンピュータに載せるとき、便宜上 $v_1, v_2, \dots$ と番号を振りますが、その番号の付け方は完全に恣意的です。同じ分子でも、原子に番号を振る順番を変えれば、ノードの並び順は入れ替わります。

ところが、番号の振り方を変えただけで「毒性あり」だった予測が「毒性なし」に変わってしまったら、それは明らかにおかしいですよね。同じグラフは、ノードの番号付けによらず、必ず同じグラフ表現 $\bm{h}_G$ にならなければならないのです。この性質を置換不変性(permutation invariance) と呼びます。

イメージとしては、「クラスの平均身長」を考えると分かりやすいです。出席番号順に並べても、名前のあいうえお順に並べても、平均身長は同じ値になります。平均という操作は「順番に依存しない」からです。一方で「1番目の生徒の身長」は、並べ方を変えれば変わってしまいます。READOUTには、平均のような「順番に依存しない」操作だけが許されるのです。

これを数式で表します。ノードの並び替え(置換)を表す関数を $\pi$ とします。$\pi$ はノードの番号を入れ替える操作で、たとえば $\pi(1)=3, \pi(2)=1, \dots$ のように対応づけます。READOUT関数 $f$ が置換不変であるとは、任意の置換 $\pi$ に対して

$$ \begin{equation} f\left(\{\bm{h}_{v_1}, \bm{h}_{v_2}, \dots, \bm{h}_{v_n}\}\right) = f\left(\{\bm{h}_{v_{\pi(1)}}, \bm{h}_{v_{\pi(2)}}, \dots, \bm{h}_{v_{\pi(n)}}\}\right) \end{equation} $$

が成り立つことです。左辺と右辺で、入力されるベクトルの「中身」はまったく同じで、「並び順」だけが違います。それでも出力が同じでなければならない、というのがこの式の主張です。

置換不変性の直感: ノードの順番を変えても同じグラフ表現が得られる

左側がノードをA→B→C→Dの順に並べた場合、右側がC→A→D→Bと並べ替えた場合です。どちらも同じ4ノードのグラフであり、同じ埋め込みベクトルの集合を持ちます。sumをREADOUTとして使えば、足す順番に関係なく合計は必ず同じになり、置換不変性が自動的に満たされることが視覚的に確認できます。

この条件を最もシンプルに満たすのが、要素ごとの可換な集約です。集合の要素を1つずつ取り込む二項演算 $\oplus$ が、足し算のように順番を入れ替えても結果が変わらない(可換かつ結合的)なら、

$$ \bm{h}_G = \bm{h}_{v_1} \oplus \bm{h}_{v_2} \oplus \dots \oplus \bm{h}_{v_n} $$

は自動的に置換不変になります。この $\oplus$ として実際によく使われるのが、sum(総和)/mean(平均)/max(要素ごとの最大値) の3つです。次のセクションでは、これら3つを具体的に定義し、それぞれの性質を比べていきます。

3つの基本的なREADOUTの定義

3つの代表的なREADOUTを定義します。いずれもノード埋め込みの集合 $\{\bm{h}_v\}_{v \in V}$(簡単のため層の上付き添字 $(K)$ は省略)を入力とし、$d$ 次元ベクトル $\bm{h}_G$ を返します。これらが「順番に依存しない」演算であることに注目してください。

sum(総和プーリング)

$$ \begin{equation} \bm{h}_G^{\mathrm{sum}} = \sum_{v \in V} \bm{h}_v \end{equation} $$

全ノードのベクトルを単純に足し合わせます。足し算は順番を入れ替えても結果が同じなので、置換不変性は明らかです。

mean(平均プーリング)

$$ \begin{equation} \bm{h}_G^{\mathrm{mean}} = \frac{1}{|V|} \sum_{v \in V} \bm{h}_v \end{equation} $$

総和をノード数 $|V|$ で割ります。これは「ノード1個あたりの平均的な特徴」を表します。

max(最大プーリング)

$$ \begin{equation} \bm{h}_G^{\mathrm{max}} = \left[\max_{v \in V} h_{v,1}, \ \max_{v \in V} h_{v,2}, \ \dots, \ \max_{v \in V} h_{v,d}\right]^\top \end{equation} $$

各次元 $j$ ごとに、全ノードの中での最大値を取ります。「どこか1つのノードでも強く反応した特徴」を拾い上げる操作です。

これら3つは一見すると「似たようなもの」に見えますが、決定的な違いがあります。それは、入力の集合が違えば必ず出力も違うという性質、すなわち単射性(injectivity) をどの程度満たすか、という点です。単射性とは「異なる入力を必ず異なる出力に写す」性質のことです。READOUTが単射でなければ、本来区別すべき2つのグラフが同じ表現に潰れてしまい、分類器がどう頑張っても見分けられなくなります。

なぜ単射性がそれほど重要なのか。次のセクションで、簡単な数値例を使って3つのREADOUTの「見分ける力」の差を体感してから、それをグラフ同型判定の理論へとつなげていきます。

なぜsumが強いのか — 単射性を数値例で見る

抽象的な議論の前に、小さな数値例で「どのREADOUTが何を区別できないか」を体感しましょう。ここでは話を最も単純にするため、ノードの特徴が1次元のスカラーだとします。

例1:mean は「個数」を区別できない

2つのノード集合を考えます。

  • 集合 $A = \{1, 1, 1\}$(特徴値1のノードが3個)
  • 集合 $B = \{1, 1\}$(特徴値1のノードが2個)

これは、たとえば「ある官能基が3個ついた分子」と「2個ついた分子」のような状況に対応します。それぞれのREADOUTを計算してみましょう。

  • sum: $A \to 3$、$B \to 2$ → 区別できる
  • mean: $A \to 1$、$B \to 1$ → 区別できない
  • max: $A \to 1$、$B \to 1$ → 区別できない

meanとmaxは、ノードの「個数」の情報を捨ててしまっています。同じ種類のノードが何個あろうと、平均や最大は変わらないからです。一方sumは、足し合わせるたびに値が増えるので、個数の違いがそのまま出力に残ります。

例2:max は「分布」を区別できない

  • 集合 $C = \{1, 2, 3\}$
  • 集合 $D = \{3, 3, 3\}$

  • sum: $C \to 6$、$D \to 9$ → 区別できる

  • mean: $C \to 2$、$D \to 3$ → 区別できる
  • max: $C \to 3$、$D \to 3$ → 区別できない

maxは最大値以外の情報をすべて捨てるため、最大値が同じなら中身がどれだけ違っても同じ出力になります。

sum/mean/maxの区別能力: 同じ色のセルは区別不能を示す比較表

この表の赤いセルが「区別できない」(出力が同じになってしまう)ケース、緑のセルが「区別できる」ケースです。左のパネルでは個数の違いにmeanとmaxが失敗し、中央パネルでは分布の違いにmaxが失敗しています。右パネルではsumだけがどちらの集合も正しく区別できており、理論的な序列が数値の上でも確認できます。

この2つの例から、「区別できる入力の集合の広さ」には明確な序列があることが見えてきます。sumは個数も分布も保持し、meanは分布の形は見るが個数を見ず、maxは最大値だけを見ます。多重集合(重複を許す集合)を区別する能力という意味では、

$$ \text{sum} \ \succ \ \text{mean} \ \approx \ \text{max} $$

という順序が成り立つのです($\succ$ は「より表現力が高い」を表します)。

なぜこの「多重集合を区別する力」がそれほど大事なのでしょうか。実はこれが、グラフが同型かどうかを判定する古典アルゴリズム、WL検査の心臓部とぴったり対応しているのです。次のセクションで、その対応を明らかにしていきます。

Weisfeiler-Lehman検査との対応

2つのグラフが「同じ形」かどうかを判定する問題をグラフ同型判定と呼びます。ノードの番号を付け替えるだけで一方が他方に重なるなら、その2つは同型です。この判定を効率よく近似する古典アルゴリズムが 1次元Weisfeiler-Lehman検査(1-WL、色洗練アルゴリズム) です。

WL検査の手続きは驚くほどメッセージパッシングに似ています。まず各ノードに初期の「色」(ラベル)を与えます。次に、各ノードの色を「自分の現在の色」と「隣接ノードの色の多重集合」を組み合わせて新しい色に更新します。これを繰り返し、色の分布が変化しなくなるまで続けます。最終的に2つのグラフの色のヒストグラムが違えば、それらは確実に非同型だと判定できます。

ここで決定的に重要なのが、「隣接ノードの色の多重集合を使う」という部分です。隣に同じ色のノードが何個あるか、という個数の情報がWL検査の判別力の源泉です。たとえば「赤い隣人が3人」と「赤い隣人が2人」を区別できなければ、WL検査は本来非同型なグラフを同型と誤判定してしまいます。

WL検査とGNNメッセージパッシングの対応関係図

左列のWL検査と右列のGNNメッセージパッシングが、ステップごとに1対1で対応しています。「初期色/ラベル」「隣接の多重集合を収集」「ハッシュ/MLPで新表現を生成」「色の分布/全層連結で判定」という4ステップが鏡のように対応しており、GINがWL検査を数値計算として実現したものであることが分かります。この対応を最大化するためにsumとMLPが選ばれた必然性が、図から直感的につかめます。

GNNのメッセージパッシングをWL検査と同じだけ強くするには、各層の集約(aggregation)とREADOUTが、この多重集合を単射に写す必要があります。これがGINの理論的主張です。Xuらの定理(簡略版)は次のように述べられます。

メッセージパッシングGNNがグラフを区別する能力は、たかだか1-WL検査と同等である。そして、層内集約とREADOUTがいずれも多重集合上で単射な関数であるとき、GNNは1-WL検査と等しい判別力を達成する。

前のセクションで見た通り、多重集合(個数を含む集合)を区別できるのはsumだけでした。meanは個数を、maxは分布を潰してしまいます。したがって、WL検査と同じ判別力を得るには、集約とREADOUTにsumを使うのが理にかなっているのです。これがGINがsum集約を採用する直接の根拠です。

GINでは、層内のノード更新を次の式で行います。

$$ \begin{equation} \bm{h}_v^{(k)} = \mathrm{MLP}^{(k)}\!\left((1 + \epsilon^{(k)}) \cdot \bm{h}_v^{(k-1)} + \sum_{u \in \mathcal{N}(v)} \bm{h}_u^{(k-1)}\right) \end{equation} $$

ここで $\mathcal{N}(v)$ はノード $v$ の隣接ノード集合、$\epsilon^{(k)}$ は自分自身の重みを調整する学習可能なスカラー(または固定値0)、$\mathrm{MLP}$ は多層パーセプトロンです。隣接ノードの集約にsumを使っている点に注目してください。さらにMLPを後段に置くことで、「sumで集約した多重集合」を単射に近い形で別の表現に写すことができます。なぜMLPが必要かというと、sumだけでは異なる多重集合がたまたま同じ和になる可能性が残るためで、MLPの非線形変換がそれを分離する役割を担います。

そしてグラフ全体のREADOUTにも、同じ理由からsumを使います。GINでは特に、各層の出力をすべて連結してから読み出す工夫がよく使われます。

$$ \begin{equation} \bm{h}_G = \big\Vert_{k=0}^{K} \ \sum_{v \in V} \bm{h}_v^{(k)} \end{equation} $$

ここで $\big\Vert$ はベクトルの連結を表します。なぜ全層を連結するかというと、浅い層は局所的な構造(小さな部分グラフ)を、深い層は広域的な構造を捉えており、両方の情報をグラフ表現に含めたいからです。この設計により、GINは理論上1-WL検査と同等の判別力を持つGNNとなります。

GINのアーキテクチャ: 各層でSUM集約とMLPを経て全層の出力を連結する

GINのアーキテクチャ図では、左の入力グラフから始まり、各GINConv層でSUM集約とMLPを適用してノード埋め込みを更新していきます。各層の出力からSUM READOUTでグラフ表現を計算し(図中の下向き点線矢印)、最終的にそれらを全部連結することで、局所構造から大域構造まで複数スケールの情報を1本のベクトルに凝縮していることが分かります。

ここまでで、READOUTの理論的な背景が一通り揃いました。sumがなぜ強いのか、それがWL検査の何と対応するのかが見えました。では、この差は実際のデータでどれくらい効いてくるのでしょうか。次のセクションから、PyTorch Geometricを使って実装し、自分の目で確かめていきます。

Python実装:TUDatasetでグラフ分類

ここからは実装に移ります。使うのは、グラフ分類のベンチマークとして定番の TUDataset に含まれる MUTAG というデータセットです。MUTAGは188個の化合物グラフからなり、各グラフが「変異原性あり/なし」の2クラスに分類されます。ノードは原子、エッジは化学結合を表し、ノードには原子種を表すone-hot特徴が付いています。まず必要なライブラリを確認しておきましょう。PyTorch GeometricはPyTorchの拡張ライブラリで、グラフデータの扱いとGNN層を提供します。

pip install torch torch_geometric

最初に、データセットを読み込んで中身を観察します。グラフごとにノード数が違うことを実際に確認するのが目的です。

import torch
from torch_geometric.datasets import TUDataset

# データセットの読み込み(初回は自動ダウンロード)
dataset = TUDataset(root='data/TUDataset', name='MUTAG')

print(f'グラフ数: {len(dataset)}')
print(f'クラス数: {dataset.num_classes}')
print(f'ノード特徴の次元: {dataset.num_node_features}')

# 各グラフのノード数の分布を確認
node_counts = [g.num_nodes for g in dataset]
print(f'ノード数: 最小 {min(node_counts)} / 最大 {max(node_counts)} / 平均 {sum(node_counts)/len(node_counts):.1f}')

このコードを実行すると、グラフ数は188、クラス数は2、ノード特徴の次元は7(原子種7種類のone-hot)と表示されます。そしてノード数は最小10、最大28、平均約18と、グラフごとにばらつきがあることが分かります。この「ノード数がバラバラ」という事実こそが、固定長に集約するREADOUTを必要とする理由でした。

次に、グラフをミニバッチにまとめる仕組みを準備します。PyTorch Geometricでは、複数のグラフを1つの大きな「非連結グラフ」として連結し、どのノードがどのグラフに属するかを batch ベクトルで管理します。この batch ベクトルがあるおかげで、READOUTを「グラフごとに」適用できるのです。

import torch
from torch_geometric.loader import DataLoader

# 訓練・テストに分割(80% / 20%)
torch.manual_seed(0)
dataset = dataset.shuffle()
n_train = int(len(dataset) * 0.8)
train_dataset = dataset[:n_train]
test_dataset = dataset[n_train:]

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# バッチの中身を1つ確認
batch = next(iter(train_loader))
print(f'バッチ内のグラフ数: {batch.num_graphs}')
print(f'バッチ内の総ノード数: {batch.num_nodes}')
print(f'batchベクトルの中身(先頭20個): {batch.batch[:20].tolist()}')

出力を見ると、バッチには32個のグラフが含まれ、それらのノードがすべて連結されて総ノード数は数百程度になります。batch ベクトルは [0, 0, 0, ..., 1, 1, ..., 2, ...] のように、各ノードがどのグラフ(0番目、1番目、…)に属するかを示しています。READOUT関数はこの batch ベクトルを見て、「同じ番号のノードだけを集約する」ことでグラフごとの表現を作ります。

それでは、READOUTを切り替えられるGINベースのモデルを定義します。PyTorch Geometricは global_add_pool(sum)、global_mean_pool(mean)、global_max_pool(max)という3つのREADOUTを用意しているので、引数で選べるようにします。

import torch
import torch.nn.functional as F
from torch.nn import Linear, Sequential, BatchNorm1d, ReLU
from torch_geometric.nn import GINConv, global_add_pool, global_mean_pool, global_max_pool


def make_mlp(in_dim, hidden):
    # GINConvの内部MLP(多重集合を単射に写す役割)
    return Sequential(
        Linear(in_dim, hidden), BatchNorm1d(hidden), ReLU(),
        Linear(hidden, hidden), ReLU(),
    )


class GIN(torch.nn.Module):
    def __init__(self, in_dim, hidden, num_classes, num_layers=3, readout='sum'):
        super().__init__()
        self.convs = torch.nn.ModuleList()
        self.convs.append(GINConv(make_mlp(in_dim, hidden)))
        for _ in range(num_layers - 1):
            self.convs.append(GINConv(make_mlp(hidden, hidden)))
        self.lin = Linear(hidden, num_classes)
        # READOUT関数の選択
        self.pool = {'sum': global_add_pool,
                     'mean': global_mean_pool,
                     'max': global_max_pool}[readout]

    def forward(self, x, edge_index, batch):
        # メッセージパッシング(ノード埋め込みの計算)
        for conv in self.convs:
            x = conv(x, edge_index)
        # READOUT:ノード埋め込み → グラフ表現(batchごとに集約)
        hg = self.pool(x, batch)
        # 分類器
        return self.lin(hg)

このモデルでは、num_layers 個のGINConv層でノード埋め込みを計算した後、self.pool でグラフ全体の表現に集約し、最後の線形層でクラスを予測します。readout 引数に 'sum''mean''max' のいずれかを渡すだけで、READOUTを切り替えられる設計です。GINConvの内部にMLPを置いているのは、前のセクションで述べた「sum集約を単射に近づける」役割を担わせるためです。

続いて、学習と評価のループを書きます。3種類のREADOUTそれぞれでモデルを訓練し、テスト精度を比較します。

import torch

def train_one_epoch(model, loader, optimizer):
    model.train()
    total_loss = 0
    for data in loader:
        optimizer.zero_grad()
        out = model(data.x, data.edge_index, data.batch)
        loss = F.cross_entropy(out, data.y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs
    return total_loss / len(loader.dataset)


@torch.no_grad()
def test_accuracy(model, loader):
    model.eval()
    correct = 0
    for data in loader:
        out = model(data.x, data.edge_index, data.batch)
        pred = out.argmax(dim=1)
        correct += int((pred == data.y).sum())
    return correct / len(loader.dataset)

訓練と評価の関数を用意したので、いよいよ3つのREADOUTを比較します。同じ条件(層数・隠れ次元・エポック数)で訓練し、テスト精度を並べます。

import torch

results = {}
for readout in ['sum', 'mean', 'max']:
    torch.manual_seed(0)
    model = GIN(dataset.num_node_features, hidden=32,
                num_classes=dataset.num_classes, num_layers=3, readout=readout)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

    best_acc = 0
    for epoch in range(1, 101):
        train_one_epoch(model, train_loader, optimizer)
        acc = test_accuracy(model, test_loader)
        best_acc = max(best_acc, acc)
    results[readout] = best_acc
    print(f'READOUT={readout:>4s}  ベストテスト精度: {best_acc:.3f}')

print('\n=== まとめ ===')
for k, v in results.items():
    print(f'{k:>4s}: {v:.3f}')

このコードを実行すると、典型的には sum が最も高い精度(おおむね0.84前後)を示し、meanmax がそれよりやや低い結果になることが多くなります。MUTAGの分類では、変異原性に関わる特定の部分構造(官能基)が「何個あるか」が効くため、個数情報を保持するsumが有利になるのです。これはまさに、前半で導いた「sumは多重集合の個数を区別できるが、meanとmaxは捨ててしまう」という理論が、実データの精度差として現れた瞬間です。ただしデータセットによっては、グラフごとのノード数の違いが大きいときに、サイズの影響を打ち消すmeanの方が安定することもあるため、READOUTはタスクに応じて選ぶべきものです。

sum/mean/max READOUTの精度と特性の比較バーチャートと横棒グラフ

左の棒グラフはMUTAGデータセットでの典型的なテスト精度を示しており、sumが最も高く(0.842)、mean(0.789)、max(0.780)が続きます。右の横棒グラフでは、各READOUTが「個数保持」「分布保持」「最大値検出」「グラフサイズへの頑健性」の4特性をどれだけ持つかを比較しています。sumは個数と分布を両方保持する一方でグラフサイズ変動に弱く、meanはサイズに頑健だが個数を失うという、タスクごとの使い分けの指針が見えてきます。

精度の差が見えたところで、結果を棒グラフにして比較しやすくしておきましょう。

import matplotlib.pyplot as plt

plt.figure(figsize=(7, 5))
names = list(results.keys())
accs = [results[n] for n in names]
colors = ['#1f77b4', '#ff7f0e', '#2ca02c']
plt.bar(names, accs, color=colors)
for i, a in enumerate(accs):
    plt.text(i, a + 0.005, f'{a:.3f}', ha='center', fontsize=12)
plt.ylabel('Test Accuracy')
plt.ylim(0.5, 1.0)
plt.title('READOUT comparison on MUTAG (GIN, 3 layers)')
plt.grid(axis='y', alpha=0.3)
plt.tight_layout()
plt.savefig('readout_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

この棒グラフから、3つのREADOUTの精度差が一目で読み取れます。sumの棒が最も高く、meanとmaxがそれに続くという序列が視覚的に確認できれば、理論で導いた「単射性の序列 sum ≻ mean ≈ max」が実データでも再現されたことになります。差が小さい場合もありますが、その理由は MUTAG が比較的小規模で、3つのREADOUTいずれでもある程度の判別ができるためです。より複雑なデータセットほど差が開きやすくなります。

なお、GINのREADOUTとして紹介したsum/mean/maxは「全ノードを均等に扱う」という点で共通しています。より洗練されたアプローチとして、Attention Readoutがあります。ノードごとに「このタスクにとってどれほど重要か」をスコアリングし、重要度の高いノードを大きく重み付けして集約する方法です。

Attention Readout: ノードの重要度をアテンションスコアで重み付けして集約する

図では、ノードh_1のアテンションスコアが最も高く(0.75)、h_2が中程度(0.20)、h_3がほぼ無視(0.05)されています。最終的なグラフ表現はこの重みを掛けた加重和になるため、分類に効くノードの情報を強調し、ノイズになるノードを抑制できます。分子グラフなら「毒性に関わる活性サイト付近の原子」を自動的に重視するよう学習が進むイメージです。

階層的プーリング:グラフを段階的に粗視化する

グローバルREADOUTは「全ノードを一気に1つに集約する」アプローチです。しかし、グラフには「部分グラフ」「コミュニティ」「モジュール」などの中間的な構造が存在することがあり、それを一段階で消してしまうと情報が失われます。この問題に対処するのが階層的プーリングです。

階層的プーリング: DiffPoolやTopKでノード数を段階的に削減してグラフを粗視化する

図では、8ノードの元グラフを段階的に4ノード→2ノードへと粗視化していく過程を示しています。各プーリング層では、よく似た特徴を持つノード群をひとまとめにして「スーパーノード」に縮約しながらエッジ構造を保ちます。最終的な2ノードのグラフに対してグローバルREADOUTを適用することで、元のグラフが持つ階層的な構造情報が保持されたまま分類器に渡されます。代表的な手法としては、クラスタ割当を学習可能な行列で行うDiffPoolと、アテンションスコアが高い上位$k$個だけを残すTopK(SAGPool)があります。

過平滑化:層を深くしすぎるとどうなるか

ここまではREADOUTの種類に注目してきましたが、グラフ分類の精度を左右するもう1つの重要な要因があります。それが過平滑化(over-smoothing) です。メッセージパッシングは「隣のノードの情報を混ぜる」操作なので、層を重ねるほど各ノードの埋め込みが周囲と似てきます。層を深くしすぎると、最終的にすべてのノードがほぼ同じベクトルに収束してしまい、ノード間の区別がつかなくなるのです。

直感的には、教室で全員が少しずつ意見を交換していくと、最初はバラバラだった意見がだんだん似通っていき、十分長く議論すると全員がほぼ同じ意見に収束してしまう、という現象に似ています。多様性が失われると、グラフの細かな構造の違いがREADOUTの段階ですでに消えてしまい、どんなに優れたREADOUTを使っても区別できなくなります。

過平滑化: 浅い層では多様なノード表現が深い層では均一に収束する様子

左の散布図(L=2、浅い層)では、3つのクラスに対応するノード埋め込みが空間上で明確に分離しています。右の散布図(L=8、深い層)では、同じノードたちの埋め込みが中央付近に密集し、クラスのクラスタがほぼ消滅しています。この均一化こそが過平滑化の本質で、READOUTがどんなに優れていても、集約する前のノード埋め込みがすでに区別不能になっていれば、分類性能は頭打ちになってしまいます。

これを定量的に見るため、ノード埋め込みの「多様性」を測る指標を導入します。最も簡単なのは、バッチ内の全ノード埋め込みの分散(各次元の分散の平均)です。過平滑化が進むと、この分散がゼロに近づいていきます。層数を変えながらこの分散と精度がどう変化するかを観察してみましょう。

import torch

@torch.no_grad()
def embedding_variance(model, loader):
    # 最終層のノード埋め込みの分散(多様性の指標)を計算
    model.eval()
    all_emb = []
    for data in loader:
        x = data.x
        for conv in model.convs:
            x = conv(x, data.edge_index)
        all_emb.append(x)
    emb = torch.cat(all_emb, dim=0)
    return emb.var(dim=0).mean().item()

この関数は、メッセージパッシング後のノード埋め込みを集めて、その分散の平均を返します。分散が大きいほどノードが多様で、小さいほど過平滑化が進んでいることを意味します。次に、層数を1から8まで変えながら、テスト精度と埋め込み分散を記録します。

import torch

depths = [1, 2, 3, 4, 6, 8]
acc_by_depth = []
var_by_depth = []

for L in depths:
    torch.manual_seed(0)
    model = GIN(dataset.num_node_features, hidden=32,
                num_classes=dataset.num_classes, num_layers=L, readout='sum')
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    best = 0
    for epoch in range(1, 81):
        train_one_epoch(model, train_loader, optimizer)
        best = max(best, test_accuracy(model, test_loader))
    acc_by_depth.append(best)
    var_by_depth.append(embedding_variance(model, test_loader))
    print(f'層数={L}  精度={best:.3f}  埋め込み分散={var_by_depth[-1]:.4f}')

実行すると、層数が増えるにつれて埋め込み分散が徐々に小さくなっていく傾向が見られます。精度のほうは、層数2〜3あたりでピークを迎え、層を深くしすぎると頭打ちまたは低下する、という典型的なパターンが多く観察されます。これは、浅すぎると遠くのノードの情報が届かず、深すぎると過平滑化で多様性が失われる、という2つの効果のバランスで最適な層数が決まることを示しています。

最後に、層数に対する精度と埋め込み分散の関係を2軸グラフで可視化します。

import matplotlib.pyplot as plt

fig, ax1 = plt.subplots(figsize=(8, 5))

color1 = '#1f77b4'
ax1.set_xlabel('Number of layers')
ax1.set_ylabel('Test Accuracy', color=color1)
ax1.plot(depths, acc_by_depth, 'o-', color=color1, label='accuracy')
ax1.tick_params(axis='y', labelcolor=color1)
ax1.grid(alpha=0.3)

ax2 = ax1.twinx()
color2 = '#d62728'
ax2.set_ylabel('Node embedding variance', color=color2)
ax2.plot(depths, var_by_depth, 's--', color=color2, label='variance')
ax2.tick_params(axis='y', labelcolor=color2)

plt.title('Depth vs. accuracy and over-smoothing (MUTAG, sum readout)')
fig.tight_layout()
plt.savefig('oversmoothing.png', dpi=150, bbox_inches='tight')
plt.show()

この2軸グラフからは、2つの重要なことが読み取れます。第一に、赤い破線(埋め込み分散)が層数の増加とともに低下していくこと、つまり過平滑化が確かに進行していることです。第二に、青い実線(精度)が中間の層数でピークを取り、その後は分散の低下とともに伸び悩む、もしくは下がる傾向が見えることです。両者を重ねて見ることで、「過平滑化による多様性の喪失が精度の頭打ちの一因である」という因果が視覚的に裏付けられます。実務では、この観察に基づいて、過平滑化を緩和する残差接続(residual connection)やジャンプ接続(jumping knowledge)、各層出力を連結するGINの読み出し方式などを使い、深いモデルでも多様性を保つ工夫をします。

過平滑化の緩和策: 残差接続とジャンプ知識ネットワークの構造比較

左の残差接続では、各GINConv層の入出力を足し合わせることで、前の層の情報を次の層に直接受け渡します。これにより「深く積んでも浅い層の局所情報が失われにくい」という効果があります。右のジャンプ知識ネットワーク(Jumping Knowledge)では、全層の出力を最終段でまとめて連結し、READOUTに渡します。浅い層(局所構造)から深い層(大域構造)まで、それぞれが捉えた情報を全て活用できるため、層数によらず安定した表現が得られます。GINが各層のSUM集約をすべて連結する設計も、この発想に基づいています。

まとめ

本記事では、ノード埋め込みをグラフ全体の表現に集約するREADOUT関数について、理論から実装まで解説しました。

  • グラフ分類の核心:メッセージパッシングが出力する「本数がバラバラのノード埋め込み」を、固定長のグラフ表現に変換するのがREADOUTの役割
  • 置換不変性:ノードの番号付けによらず同じグラフ表現を返すことが必須条件で、sum/mean/maxのような可換な集約がこれを満たす
  • 単射性の序列:sumは多重集合の個数も分布も保持し、meanは個数を、maxは分布を捨てる。区別できる入力の広さは sum ≻ mean ≈ max
  • WL検査との対応:GNNの判別力はたかだか1-WL検査と同等で、集約とREADOUTにsumを使うことでその上限に到達できる。これがGINがsumを採用する根拠
  • 実装と検証:PyTorch GeometricでMUTAGを分類し、sumが最も高い精度を示すことを確認。さらに層を深くすると過平滑化で埋め込み分散が低下し、精度が頭打ちになることをグラフで観察

READOUTは、単なる「最後の集約ステップ」ではなく、GNNが「どんなグラフを区別できるか」という能力の上限を決める設計の要でした。sum集約とMLPを組み合わせたGINが理論上最強の判別力を持つこと、そして深さと過平滑化のトレードオフがあることを理解しておくと、実際のグラフ分類タスクでモデルを設計する際の確かな指針になります。

次のステップとして、以下の記事も参考にしてください。

GNN: メッセージパッシングニューラルネットワーク(MPNN)を解説する
READOUTの前段であるノード埋め込みの計算過程を詳しく解説。集約・更新・読み出しの3ステップからGCN/GAT/GINのメッセージパッシングを統一的に理解する。
グラフプーリングの理論: グラフレベルの表現を学習する
DiffPoolやSAGPool(TopK)など階層的プーリングの仕組みを解説。ノード数を段階的に削減してグラフ全体の表現を獲得する方法をPyTorch Geometricの実装とともに理解する。