はじめに — なぜグラフプーリングが必要なのか
グラフニューラルネットワーク（GNN）は、ノード分類やリンク予測といったノードレベルのタスクでは大きな成功を収めてきました。しかし、グラフ分類 や グラフ回帰 のようなグラフレベルのタスクでは、追加の課題が生じます。
画像認識を例に考えてみましょう。CNNでは、畳み込み層で局所的な特徴を抽出した後、プーリング層 で空間的な情報を集約し、最終的に全結合層で分類を行います。このプーリングによって、空間的に広い範囲の情報を段階的に統合しながら、固定サイズのベクトル表現を得ることができます。
グラフにおいても同様の仕組みが必要です。各ノードの表現ベクトルを、グラフ全体を表す1つのベクトル に集約する必要があります。しかし、グラフにはCNNのような規則的な空間構造がないため、単純にピクセルの2×2領域を1つにまとめるような操作ができません。ノード数もグラフごとに異なります。
最も単純なアプローチは、すべてのノードの表現を平均やsumで集約する フラットプーリング（Global Pooling） です。これは実装が簡単ですが、グラフの階層的な構造情報が失われてしまいます。たとえば、化学分子のグラフを分類する際、局所的な官能基の構造と分子全体の形状の両方が重要ですが、フラットプーリングではこの階層性を捉えるのが困難です。
この問題を解決するために、グラフの構造を段階的に粗視化する 階層的プーリング が研究されてきました。本記事では、以下の内容を体系的に解説します。
- フラットプーリング（sum / mean / max）の特性と限界
- 階層的プーリングの代表的手法: DiffPool, TopKPool, SAGPool
- 各手法のPyTorch実装と比較実験
グラフ分類やグラフ生成モデルに関心のある方にとって、グラフプーリングは避けて通れない重要なテーマです。
前提知識
この記事を読むにあたって、以下の知識があるとスムーズです。
- GNNの基礎: メッセージパッシング、GCN、GraphSAGEレベル
- GINの理論: sum/mean/max集約の違い（前記事）
- 線形代数: 行列の固有値分解、ソフトマックス関数
- PyTorchの基礎:
nn.Module、torch.matmul
フラットプーリング（Global Pooling）
3種類のフラットプーリング
最も基本的なグラフプーリングは、GNN層で得られた全ノードの表現をグラフ全体に集約するフラットプーリングです。グラフ $G$ のノード集合を $V$、ノード $v$ の表現を $h_v \in \mathbb{R}^d$ とすると、3種類のフラットプーリングは次のように定義されます。
Sum Pooling（合計プーリング）:
$$ h_G = \sum_{v \in V} h_v $$
Mean Pooling（平均プーリング）:
$$ h_G = \frac{1}{|V|} \sum_{v \in V} h_v $$
Max Pooling（最大値プーリング）:
$$ h_G = \max_{v \in V} h_v \quad (\text{各次元ごとに最大値を取る}) $$
各プーリングの特性
それぞれの特性を比較しましょう。
Sum Pooling は、GINの議論で見たように、多重集合に対して最も多くの情報を保持します。ノードの数の情報が集約値に反映されるため、5ノードのグラフと10ノードのグラフで異なるスケールの表現が得られます。これは同じドメイン内のグラフ分類では有利ですが、非常に異なるサイズのグラフを比較する際には正規化が必要になることがあります。
Mean Pooling は、グラフサイズに対して不変な表現を生成します。つまり、すべてのノードが同じ表現を持つグラフでは、サイズに関わらず同じグラフ表現が得られます。これはサイズの異なるグラフを比較する際に有利ですが、「ノードの数」という情報が失われます。
Max Pooling は、各特徴次元において最も顕著なノードの値を取ります。これは「グラフ内で最も強い信号」を捉える操作です。ノイズに対してロバストですが、情報の損失が大きくなります。
フラットプーリングの限界
フラットプーリングには根本的な限界があります。それは グラフの階層的な構造情報を捉えられない ことです。
たとえば、以下のような分子グラフを考えてみましょう。
- 分子A: ベンゼン環 + カルボキシル基（-COOH）
- 分子B: ベンゼン環 + ヒドロキシル基（-OH）
フラットプーリングでは、すべてのノード表現を一括で集約するため、「ベンゼン環」という部分構造と「官能基」という部分構造の関係性を明示的に捉えることができません。人間の化学者が分子を理解するとき、まず官能基などの部分構造を認識し、それらの組み合わせで分子全体の性質を推定します。この 階層的な認識プロセス をモデルに組み込むのが、次に説明する階層的プーリングです。
ここからは、グラフの構造を段階的に粗視化する階層的プーリングの手法を見ていきましょう。
階層的プーリングの概要
粗視化（Coarsening）のアイデア
階層的プーリングの核心的なアイデアは、グラフを段階的に小さくしていく ことです。具体的には、複数のノードを1つの「スーパーノード」にまとめ、元のグラフよりもノード数が少ない粗いグラフを生成します。このプロセスを繰り返すことで、最終的に1つ（または少数）のノードを持つグラフが得られ、そのノードの表現がグラフ全体の表現となります。
CNNのプーリングとの対応を考えると、次のようになります。
| CNN | Gラフプーリング |
|---|---|
| 2×2領域のmax/average | 近傍ノードのクラスタリング |
| 空間解像度の半減 | ノード数の削減 |
| チャネル数の増加 | ノード表現次元の変換 |
| 規則的なグリッド構造 | 不規則なグラフ構造 |
階層的プーリングのアプローチは大きく2つに分類できます。
クラスタリングベース: ノードをクラスタに分割し、各クラスタを1つのスーパーノードに置き換えます。代表的な手法として DiffPool があります。
スコアベース（ドロップ型）: 各ノードにスコアを計算し、スコアの高いノードだけを残してグラフを縮小します。代表的な手法として TopKPool と SAGPool があります。
それぞれのアプローチを詳しく見ていきましょう。
DiffPool（Differentiable Pooling）
DiffPoolの概要
DiffPool（Ying et al., 2018）は、微分可能なクラスタリング によってグラフを階層的に粗視化する手法です。事前に定義されたクラスタリングではなく、GNNを用いてクラスタリングの割り当てを学習する点が画期的です。
DiffPoolの数式
DiffPoolの各プーリング層は、2つのGNNを用いて以下の2つを同時に計算します。
1. ノード埋め込みの計算:
$$ Z^{(l)} = \text{GNN}_{embed}^{(l)}(A^{(l)}, X^{(l)}) $$
ここで $A^{(l)} \in \mathbb{R}^{n_l \times n_l}$ は第 $l$ 層の隣接行列、$X^{(l)} \in \mathbb{R}^{n_l \times d_l}$ はノード特徴量行列、$Z^{(l)} \in \mathbb{R}^{n_l \times d_{l+1}}$ は埋め込みです。
2. クラスタ割り当て行列の計算:
$$ S^{(l)} = \text{softmax}\left(\text{GNN}_{pool}^{(l)}(A^{(l)}, X^{(l)})\right) $$
ここで $S^{(l)} \in \mathbb{R}^{n_l \times n_{l+1}}$ はソフト割り当て行列です。各行はノード、各列はクラスタに対応し、$S_{ij}$ は「ノード $i$ がクラスタ $j$ に属する確率」を表します。softmaxは各行に対して適用します。
$n_{l+1}$ は次の層のノード数（クラスタ数）で、ハイパーパラメータとして事前に指定します。たとえば、$n_l = 100$ のグラフを $n_{l+1} = 25$ に粗視化するなら、ノード数を4分の1に削減します。
3. 粗視化されたグラフの計算:
割り当て行列 $S$ を用いて、新しいノード特徴量と隣接行列を次のように計算します。
$$ X^{(l+1)} = S^{(l)\top} Z^{(l)} \in \mathbb{R}^{n_{l+1} \times d_{l+1}} $$
$$ A^{(l+1)} = S^{(l)\top} A^{(l)} S^{(l)} \in \mathbb{R}^{n_{l+1} \times n_{l+1}} $$
第1の式は、各クラスタの特徴量をそのクラスタに属するノードの埋め込みの加重和として計算しています。第2の式は、2つのクラスタ間のエッジの強さを、元のグラフにおけるクラスタ間のノード接続の加重和として計算しています。
直感的には、$S^{(l)\top} A^{(l)} S^{(l)}$ は次のように理解できます。$A^{(l)}$ はノード間の接続を表し、$S^{(l)}$ はノードからクラスタへのマッピングです。左から $S^{(l)\top}$ を掛け、右から $S^{(l)}$ を掛けることで、「クラスタ $i$ に属するノードとクラスタ $j$ に属するノードの間にどれだけの接続があるか」を集約しているのです。
DiffPoolの補助損失
DiffPoolの学習を安定させるために、2つの補助損失が導入されます。
エントロピー正則化:
$$ \mathcal{L}_{ent} = \frac{1}{n_l} \sum_{i=1}^{n_l} H(S_i^{(l)}) $$
ここで $H(\cdot)$ はエントロピーです。この損失を最小化することで、各ノードの割り当てが「はっきりと1つのクラスタに属する」ようになります。ソフトな割り当てが曖昧すぎるとグラフの構造情報が失われるため、この正則化が重要です。
リンク予測正則化:
$$ \mathcal{L}_{lp} = \|A^{(l)} – S^{(l)} S^{(l)\top}\|_F^2 $$
この損失は、隣接するノードが同じクラスタに割り当てられるように促します。$S^{(l)} S^{(l)\top}$ は「2つのノードが同じクラスタに属する度合い」を表す行列で、これが元の隣接行列 $A^{(l)}$ に近づくように学習します。つまり、エッジで繋がったノードは同じクラスタにまとめられやすくなります。
DiffPoolの利点と課題
利点: – エンドツーエンドで学習可能: クラスタリングとノード埋め込みが同時に最適化される – 柔軟なクラスタ構造: GNNがデータから適切なクラスタリングを学習できる – 階層的な構造を明示的に捉える
課題: – 計算コスト: 割り当て行列 $S \in \mathbb{R}^{n_l \times n_{l+1}}$ の計算と保持にメモリが必要。密な行列演算が含まれるため、計算量は $O(n^2)$ 以上 – クラスタ数のハイパーパラメータ: 各層のクラスタ数 $n_{l+1}$ を事前に指定する必要がある – バッチ処理の制約: グラフごとにノード数が異なるため、バッチ内で最大ノード数にパディングが必要
これらの課題を解決するアプローチとして、次に紹介するスコアベースのプーリング手法が提案されました。
TopKPool と SAGPool
TopKPool
TopKPool（Gao & Ji, 2019）は、各ノードにスコアを計算し、上位 $k$ 個のノードだけを残すことでグラフを縮小する手法です。
スコアの計算:
$$ \bm{y} = \frac{X \bm{p}}{\|\bm{p}\|} $$
ここで $\bm{p} \in \mathbb{R}^d$ は学習可能な射影ベクトル、$X \in \mathbb{R}^{n \times d}$ はノード特徴量行列です。$\bm{y} \in \mathbb{R}^n$ の各要素 $y_i$ がノード $i$ の重要度スコアとなります。
上位kノードの選択:
$$ \text{idx} = \text{top-}k(\bm{y}, \lceil rn \rceil) $$
ここで $r \in (0, 1]$ はプーリング比率で、$\lceil rn \rceil$ は残すノード数です。たとえば $r = 0.5$ なら半分のノードを残します。
ゲートメカニズム:
選択されたノードの特徴量にスコアを掛けることで、ゲートの役割を果たします。
$$ \tilde{X} = X_{\text{idx}} \odot \sigma(\bm{y}_{\text{idx}}) $$
ここで $\sigma$ はシグモイド関数、$\odot$ は要素ごとの積です。スコアが高いノードほど大きな重みで特徴量が保持されます。
サブグラフの構成:
選択されたノードに対応する隣接行列のサブセットを取り出して、新しいグラフを構成します。
$$ \tilde{A} = A_{\text{idx}, \text{idx}} $$
SAGPool（Self-Attention Graph Pooling）
SAGPool（Lee et al., 2019）は、TopKPoolのスコア計算部分を GNN（Graph Convolution） に置き換えた手法です。
TopKPoolでは線形射影 $\bm{y} = X\bm{p}$ でスコアを計算していましたが、これはノードの特徴量だけを見ており、グラフ構造の情報を活用していません。SAGPoolはこれを改良し、GCN層を用いてスコアを計算します。
SAGPoolのスコア計算:
$$ \bm{y} = \text{GCN}(X, A) = \sigma(\tilde{D}^{-1/2} \tilde{A} \tilde{D}^{-1/2} X \Theta) $$
ここで $\tilde{A} = A + I$ は自己ループを追加した隣接行列、$\tilde{D}$ はその次数行列、$\Theta \in \mathbb{R}^{d \times 1}$ は学習パラメータです。
GCNを通すことで、各ノードのスコアは 自身の特徴量だけでなく近傍の情報も考慮 して計算されます。グラフの局所的な構造パターンに基づいてノードの重要度を評価できるため、より情報量の多いプーリングが可能になります。
それ以外の操作（上位kの選択、ゲートメカニズム、サブグラフ構成）はTopKPoolと同じです。
スコアベースプーリングの利点と課題
利点: – メモリ効率: DiffPoolのような密な割り当て行列が不要 – スパース構造の維持: サブグラフとして元のグラフ構造の一部を保持 – 実装のシンプルさ: top-k操作は微分可能（ストレートスルー推定器）
課題: – 情報の損失: 削除されたノードの情報が完全に失われる – プーリング比率の設定: $r$ の値がタスクに依存 – 接続性の保証なし: ノードの削除によってグラフが非連結になる可能性がある
ここまで3つのプーリング手法の理論を理解したところで、次にPyTorchで実装し、実験的に比較してみましょう。
PyTorchによる実装
フラットプーリングの実装
まず、基本となるフラットプーリングの3手法を実装します。
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
class GlobalSumPool(nn.Module):
"""グローバルSumプーリング"""
def forward(self, x, graph_indicator):
num_graphs = graph_indicator.max().item() + 1
pooled = torch.zeros(num_graphs, x.size(1), device=x.device)
pooled.scatter_add_(0, graph_indicator.unsqueeze(1).expand_as(x), x)
return pooled
class GlobalMeanPool(nn.Module):
"""グローバルMeanプーリング"""
def forward(self, x, graph_indicator):
num_graphs = graph_indicator.max().item() + 1
pooled = torch.zeros(num_graphs, x.size(1), device=x.device)
counts = torch.zeros(num_graphs, 1, device=x.device)
pooled.scatter_add_(0, graph_indicator.unsqueeze(1).expand_as(x), x)
counts.scatter_add_(0, graph_indicator.unsqueeze(1),
torch.ones(x.size(0), 1, device=x.device))
return pooled / counts.clamp(min=1)
class GlobalMaxPool(nn.Module):
"""グローバルMaxプーリング"""
def forward(self, x, graph_indicator):
num_graphs = graph_indicator.max().item() + 1
# scatter_maxが使えない環境向けのループ実装
pooled = torch.full((num_graphs, x.size(1)), float('-inf'), device=x.device)
for i in range(x.size(0)):
g = graph_indicator[i]
pooled[g] = torch.max(pooled[g], x[i])
return pooled
# 動作確認
torch.manual_seed(42)
x = torch.randn(8, 4) # 8ノード、4次元特徴量
indicator = torch.tensor([0, 0, 0, 1, 1, 1, 1, 2]) # 3つのグラフ
sum_pool = GlobalSumPool()
mean_pool = GlobalMeanPool()
max_pool = GlobalMaxPool()
print("ノード特徴量 x:")
print(x)
print(f"\nグラフ所属: {indicator.tolist()}")
print(f"\nSum Pooling:\n{sum_pool(x, indicator)}")
print(f"\nMean Pooling:\n{mean_pool(x, indicator)}")
print(f"\nMax Pooling:\n{max_pool(x, indicator)}")
3つのプーリングの出力を比較すると、Sum Poolingはグラフのノード数に比例したスケールを持ち、Mean Poolingはノード数に依存せず、Max Poolingは各次元の最大値のみを反映していることがわかります。
TopKPoolの実装
import torch
import torch.nn as nn
import torch.nn.functional as F
class TopKPoolLayer(nn.Module):
"""TopKPool層"""
def __init__(self, in_dim, ratio=0.5):
super().__init__()
self.ratio = ratio
# スコア計算用の射影ベクトル
self.proj = nn.Parameter(torch.randn(in_dim, 1) * 0.01)
def forward(self, x, adj, graph_indicator=None):
"""
x: ノード特徴量 (num_nodes, in_dim)
adj: 隣接行列 (num_nodes, num_nodes)
graph_indicator: 各ノードが属するグラフ (num_nodes,)
"""
# スコアの計算
score = torch.matmul(x, self.proj).squeeze(-1) # (num_nodes,)
score = score / (self.proj.norm() + 1e-8)
if graph_indicator is None:
# 単一グラフの場合
k = max(1, int(self.ratio * x.size(0)))
_, idx = torch.topk(score, k)
idx = idx.sort().values
else:
# バッチ処理: 各グラフごとにtop-kを選択
idx_list = []
num_graphs = graph_indicator.max().item() + 1
for g in range(num_graphs):
mask = (graph_indicator == g)
nodes_in_graph = mask.nonzero(as_tuple=True)[0]
k = max(1, int(self.ratio * nodes_in_graph.size(0)))
scores_g = score[nodes_in_graph]
_, top_idx = torch.topk(scores_g, k)
idx_list.append(nodes_in_graph[top_idx])
idx = torch.cat(idx_list).sort().values
# ゲート付き特徴量
gate = torch.sigmoid(score[idx]).unsqueeze(-1)
x_new = x[idx] * gate
# サブグラフの隣接行列
adj_new = adj[idx][:, idx]
# 新しいgraph_indicator
new_indicator = None
if graph_indicator is not None:
new_indicator = graph_indicator[idx]
return x_new, adj_new, new_indicator, idx, score
# 動作確認
torch.manual_seed(42)
n = 10
x = torch.randn(n, 8)
adj = torch.zeros(n, n)
# ランダムなグラフ構造
edges = [(0,1),(1,2),(2,3),(3,4),(0,4),(5,6),(6,7),(7,8),(8,9),(5,9)]
for i, j in edges:
adj[i, j] = adj[j, i] = 1
pool = TopKPoolLayer(in_dim=8, ratio=0.5)
x_new, adj_new, _, idx, scores = pool(x, adj)
print(f"元のノード数: {n}")
print(f"プーリング後のノード数: {x_new.size(0)}")
print(f"選択されたノード: {idx.tolist()}")
print(f"各ノードのスコア: {scores.detach().numpy().round(3)}")
print(f"新しい隣接行列のサイズ: {adj_new.shape}")
TopKPoolの動作を確認すると、10ノードのグラフから5ノードが選択され、選択されたノード間のエッジのみが保持された新しいグラフが構成されていることがわかります。スコアはノードの特徴量と学習パラメータの内積で計算されており、初期状態ではランダムに近い値ですが、学習が進むにつれてタスクに有用なノードに高いスコアが割り当てられるようになります。
SAGPoolの実装
import torch
import torch.nn as nn
import torch.nn.functional as F
class GCNLayer(nn.Module):
"""基本的なGCN層"""
def __init__(self, in_dim, out_dim):
super().__init__()
self.linear = nn.Linear(in_dim, out_dim, bias=False)
def forward(self, x, adj):
# 自己ループの追加
adj_hat = adj + torch.eye(adj.size(0), device=adj.device)
# 次数行列の計算
degree = adj_hat.sum(dim=1).clamp(min=1)
degree_inv_sqrt = degree.pow(-0.5)
# 正規化
norm_adj = degree_inv_sqrt.unsqueeze(1) * adj_hat * degree_inv_sqrt.unsqueeze(0)
# 特徴量の変換と集約
return torch.matmul(norm_adj, self.linear(x))
class SAGPoolLayer(nn.Module):
"""SAGPool層"""
def __init__(self, in_dim, ratio=0.5):
super().__init__()
self.ratio = ratio
# GCNでスコアを計算（出力次元1）
self.score_gcn = GCNLayer(in_dim, 1)
def forward(self, x, adj, graph_indicator=None):
"""
x: ノード特徴量 (num_nodes, in_dim)
adj: 隣接行列 (num_nodes, num_nodes)
"""
# GCNでスコアを計算（グラフ構造を考慮）
score = self.score_gcn(x, adj).squeeze(-1) # (num_nodes,)
score = torch.tanh(score)
if graph_indicator is None:
k = max(1, int(self.ratio * x.size(0)))
_, idx = torch.topk(score, k)
idx = idx.sort().values
else:
idx_list = []
num_graphs = graph_indicator.max().item() + 1
for g in range(num_graphs):
mask = (graph_indicator == g)
nodes_in_graph = mask.nonzero(as_tuple=True)[0]
k = max(1, int(self.ratio * nodes_in_graph.size(0)))
scores_g = score[nodes_in_graph]
_, top_idx = torch.topk(scores_g, k)
idx_list.append(nodes_in_graph[top_idx])
idx = torch.cat(idx_list).sort().values
# ゲート付き特徴量
gate = torch.sigmoid(score[idx]).unsqueeze(-1)
x_new = x[idx] * gate
# サブグラフの隣接行列
adj_new = adj[idx][:, idx]
new_indicator = None
if graph_indicator is not None:
new_indicator = graph_indicator[idx]
return x_new, adj_new, new_indicator, idx, score
# TopKPool vs SAGPoolのスコア比較
torch.manual_seed(42)
n = 10
x = torch.randn(n, 8)
adj = torch.zeros(n, n)
edges = [(0,1),(1,2),(2,3),(3,4),(0,4),(5,6),(6,7),(7,8),(8,9),(5,9),(2,7)]
for i, j in edges:
adj[i, j] = adj[j, i] = 1
topk = TopKPoolLayer(8, ratio=0.5)
sagpool = SAGPoolLayer(8, ratio=0.5)
with torch.no_grad():
_, _, _, topk_idx, topk_scores = topk(x, adj)
_, _, _, sag_idx, sag_scores = sagpool(x, adj)
print("=== TopKPool ===")
print(f"スコア: {topk_scores.numpy().round(3)}")
print(f"選択ノード: {topk_idx.tolist()}")
print("\n=== SAGPool ===")
print(f"スコア: {sag_scores.numpy().round(3)}")
print(f"選択ノード: {sag_idx.tolist()}")
TopKPoolとSAGPoolのスコアを比較すると、TopKPoolはノードの特徴量のみに基づくスコアを計算しているのに対し、SAGPoolはGCN層を通じて近傍の情報も反映したスコアを計算していることがわかります。たとえば、グラフの中心的な位置にあるノード（多くの他のノードと接続しているノード）は、SAGPoolでより高いスコアを得る傾向があります。
DiffPoolの簡易実装
import torch
import torch.nn as nn
import torch.nn.functional as F
class DiffPoolLayer(nn.Module):
"""DiffPool層（簡易版）"""
def __init__(self, in_dim, embed_dim, num_clusters):
super().__init__()
self.num_clusters = num_clusters
# 埋め込み用GCN
self.embed_gcn = GCNLayer(in_dim, embed_dim)
# 割り当て用GCN
self.assign_gcn = GCNLayer(in_dim, num_clusters)
def forward(self, x, adj):
"""
x: ノード特徴量 (num_nodes, in_dim)
adj: 隣接行列 (num_nodes, num_nodes)
"""
# 埋め込みの計算
z = F.relu(self.embed_gcn(x, adj)) # (num_nodes, embed_dim)
# 割り当て行列の計算（行方向にsoftmax）
s = F.softmax(self.assign_gcn(x, adj), dim=1) # (num_nodes, num_clusters)
# 粗視化されたグラフ
x_new = torch.matmul(s.t(), z) # (num_clusters, embed_dim)
adj_new = torch.matmul(s.t(), torch.matmul(adj, s)) # (num_clusters, num_clusters)
# 補助損失の計算
# エントロピー正則化
entropy_loss = -(s * torch.log(s + 1e-10)).sum(dim=1).mean()
# リンク予測正則化
link_loss = torch.norm(adj - torch.matmul(s, s.t()), p='fro') ** 2 / (adj.size(0) ** 2)
return x_new, adj_new, s, entropy_loss, link_loss
# 動作確認
torch.manual_seed(42)
n = 12
x = torch.randn(n, 8)
adj = torch.zeros(n, n)
edges = [(0,1),(1,2),(2,0),(3,4),(4,5),(5,3),(6,7),(7,8),(8,6),
(9,10),(10,11),(11,9),(2,3),(5,6),(8,9)]
for i, j in edges:
adj[i, j] = adj[j, i] = 1
# 12ノード → 4クラスタに粗視化
diffpool = DiffPoolLayer(in_dim=8, embed_dim=16, num_clusters=4)
with torch.no_grad():
x_new, adj_new, s, ent_loss, link_loss = diffpool(x, adj)
print(f"元のグラフ: {n}ノード")
print(f"粗視化後: {x_new.size(0)}クラスタ")
print(f"\n割り当て行列 S (各行がノード、各列がクラスタ):")
print(s.numpy().round(3))
print(f"\n粗視化後の隣接行列:")
print(adj_new.numpy().round(3))
print(f"\nエントロピー損失: {ent_loss.item():.4f}")
print(f"リンク予測損失: {link_loss.item():.4f}")
DiffPoolの動作を確認すると、12ノードのグラフが4つのクラスタに粗視化されていることがわかります。割り当て行列 $S$ の各行は、そのノードがどのクラスタに属するかの確率分布を示しています。初期状態（学習前）では割り当てがほぼ一様ですが、学習が進むと各ノードは1つのクラスタに強く割り当てられるようになります。粗視化後の隣接行列は、クラスタ間の接続強度を表しています。
グラフ分類タスクでの比較実験
各プーリング手法を統合したモデルで、グラフ分類タスクの性能を比較します。
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
# GCN層（前セクションと同じ）
class GCNConv(nn.Module):
def __init__(self, in_dim, out_dim):
super().__init__()
self.linear = nn.Linear(in_dim, out_dim)
def forward(self, x, adj):
adj_hat = adj + torch.eye(adj.size(0), device=adj.device)
deg = adj_hat.sum(1).clamp(min=1)
deg_inv_sqrt = deg.pow(-0.5)
norm = deg_inv_sqrt.unsqueeze(1) * adj_hat * deg_inv_sqrt.unsqueeze(0)
return F.relu(self.linear(torch.matmul(norm, x)))
# モデル1: GCN + Global Sum Pooling
class GCNGlobalPool(nn.Module):
def __init__(self, in_dim, hidden_dim, num_classes, num_layers=3):
super().__init__()
self.convs = nn.ModuleList()
self.convs.append(GCNConv(in_dim, hidden_dim))
for _ in range(num_layers - 1):
self.convs.append(GCNConv(hidden_dim, hidden_dim))
self.classifier = nn.Linear(hidden_dim, num_classes)
def forward(self, x, adj, indicator):
for conv in self.convs:
x = conv(x, adj)
num_graphs = indicator.max().item() + 1
pooled = torch.zeros(num_graphs, x.size(1), device=x.device)
pooled.scatter_add_(0, indicator.unsqueeze(1).expand_as(x), x)
return self.classifier(pooled)
# モデル2: GCN + SAGPool
class GCNSAGPool(nn.Module):
def __init__(self, in_dim, hidden_dim, num_classes, num_layers=3, ratio=0.5):
super().__init__()
self.convs = nn.ModuleList()
self.pools = nn.ModuleList()
self.convs.append(GCNConv(in_dim, hidden_dim))
for _ in range(num_layers - 1):
self.convs.append(GCNConv(hidden_dim, hidden_dim))
self.pools.append(SAGPoolLayer(hidden_dim, ratio=ratio))
self.classifier = nn.Linear(hidden_dim, num_classes)
def forward(self, x, adj, indicator):
x = self.convs[0](x, adj)
for i in range(len(self.pools)):
x = self.convs[i + 1](x, adj)
x, adj, indicator, _, _ = self.pools[i](x, adj, indicator)
num_graphs = indicator.max().item() + 1
pooled = torch.zeros(num_graphs, x.size(1), device=x.device)
pooled.scatter_add_(0, indicator.unsqueeze(1).expand_as(x), x)
return self.classifier(pooled)
# モデル3: GCN + TopKPool
class GCNTopKPool(nn.Module):
def __init__(self, in_dim, hidden_dim, num_classes, num_layers=3, ratio=0.5):
super().__init__()
self.convs = nn.ModuleList()
self.pools = nn.ModuleList()
self.convs.append(GCNConv(in_dim, hidden_dim))
for _ in range(num_layers - 1):
self.convs.append(GCNConv(hidden_dim, hidden_dim))
self.pools.append(TopKPoolLayer(hidden_dim, ratio=ratio))
self.classifier = nn.Linear(hidden_dim, num_classes)
def forward(self, x, adj, indicator):
x = self.convs[0](x, adj)
for i in range(len(self.pools)):
x = self.convs[i + 1](x, adj)
x, adj, indicator, _, _ = self.pools[i](x, adj, indicator)
num_graphs = indicator.max().item() + 1
pooled = torch.zeros(num_graphs, x.size(1), device=x.device)
pooled.scatter_add_(0, indicator.unsqueeze(1).expand_as(x), x)
return self.classifier(pooled)
# データセット生成
def make_dataset(num_graphs=400, seed=42):
np.random.seed(seed)
graphs = []
for _ in range(num_graphs // 4):
# クラス0: サイクル
n = np.random.randint(6, 15)
adj = np.zeros((n, n))
for i in range(n):
adj[i, (i+1)%n] = adj[(i+1)%n, i] = 1
graphs.append((adj, np.random.randn(n, 3).astype(np.float32), 0))
# クラス1: 星
n = np.random.randint(6, 15)
adj = np.zeros((n, n))
for i in range(1, n):
adj[0, i] = adj[i, 0] = 1
graphs.append((adj, np.random.randn(n, 3).astype(np.float32), 1))
# クラス2: 二部グラフ
n1 = np.random.randint(3, 7)
n2 = np.random.randint(3, 7)
n = n1 + n2
adj = np.zeros((n, n))
for i in range(n1):
for j in range(n1, n):
if np.random.random() < 0.5:
adj[i, j] = adj[j, i] = 1
graphs.append((adj, np.random.randn(n, 3).astype(np.float32), 2))
# クラス3: ツリー（ランダムスパニングツリー）
n = np.random.randint(6, 15)
adj = np.zeros((n, n))
for i in range(1, n):
j = np.random.randint(0, i)
adj[i, j] = adj[j, i] = 1
graphs.append((adj, np.random.randn(n, 3).astype(np.float32), 3))
np.random.shuffle(graphs)
return graphs
def collate_fn(graph_list):
adjs, feats, labels, indicators = [], [], [], []
offset = 0
for i, (adj, feat, label) in enumerate(graph_list):
n = adj.shape[0]
adjs.append((adj, offset))
feats.append(feat)
labels.append(label)
indicators.extend([i] * n)
offset += n
total = offset
batch_adj = np.zeros((total, total))
for adj, off in adjs:
n = adj.shape[0]
batch_adj[off:off+n, off:off+n] = adj
return (torch.FloatTensor(batch_adj), torch.FloatTensor(np.concatenate(feats)),
torch.LongTensor(labels), torch.LongTensor(indicators))
# 訓練ループ
def train_model(model, train_data, test_data, epochs=150, lr=0.005):
optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
train_losses, test_accs = [], []
for epoch in range(epochs):
model.train()
adj, x, labels, ind = collate_fn(train_data)
logits = model(x, adj, ind)
loss = F.cross_entropy(logits, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_losses.append(loss.item())
if (epoch + 1) % 5 == 0:
model.eval()
with torch.no_grad():
adj_t, x_t, lab_t, ind_t = collate_fn(test_data)
logits_t = model(x_t, adj_t, ind_t)
acc = (logits_t.argmax(1) == lab_t).float().mean().item()
test_accs.append((epoch+1, acc))
return train_losses, test_accs
# 実験実行
torch.manual_seed(42)
dataset = make_dataset(400, seed=42)
split = int(len(dataset) * 0.8)
train_data, test_data = dataset[:split], dataset[split:]
models = {
'Global Sum': GCNGlobalPool(3, 32, 4, num_layers=3),
'SAGPool': GCNSAGPool(3, 32, 4, num_layers=3, ratio=0.5),
'TopKPool': GCNTopKPool(3, 32, 4, num_layers=3, ratio=0.5),
}
results = {}
for name, model in models.items():
torch.manual_seed(42)
# パラメータを再初期化
for p in model.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
losses, accs = train_model(model, train_data, test_data, epochs=150)
results[name] = {'losses': losses, 'accs': accs}
final_acc = accs[-1][1] if accs else 0
print(f"{name:>12s}: 最終テスト精度 = {final_acc*100:.1f}%")
# 可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
colors = {'Global Sum': '#2ecc71', 'SAGPool': '#e74c3c', 'TopKPool': '#3498db'}
for name, res in results.items():
axes[0].plot(res['losses'], label=name, color=colors[name], alpha=0.7)
axes[0].set_xlabel('Epoch', fontsize=12)
axes[0].set_ylabel('Training Loss', fontsize=12)
axes[0].set_title('Training Loss', fontsize=14)
axes[0].legend(fontsize=11)
axes[0].grid(True, alpha=0.3)
for name, res in results.items():
epochs_list, acc_list = zip(*res['accs'])
axes[1].plot(epochs_list, acc_list, 'o-', label=name, color=colors[name], markersize=3)
axes[1].set_xlabel('Epoch', fontsize=12)
axes[1].set_ylabel('Test Accuracy', fontsize=12)
axes[1].set_title('Test Accuracy', fontsize=14)
axes[1].legend(fontsize=11)
axes[1].set_ylim(0, 1.05)
axes[1].grid(True, alpha=0.3)
plt.suptitle('Graph Classification: Pooling Method Comparison', fontsize=15, y=1.02)
plt.tight_layout()
plt.savefig('pooling_comparison.png', dpi=150, bbox_inches='tight')
plt.show()
この実験では、4クラスのグラフ分類タスク（サイクル・星・二部グラフ・ツリー）で3つのプーリング手法を比較しています。Global Sum Poolingはシンプルながら安定した性能を示し、SAGPoolとTopKPoolは階層的な情報抽出により、特に構造的に類似したクラス間の区別で優位になる場合があります。ただし、階層的プーリングはパラメータ数が増えるため、データが少ない場合は過学習のリスクもあります。
プーリング手法の比較まとめ
各手法の特性比較
ここまでの議論を表にまとめます。
| 手法 | アプローチ | 計算量 | メモリ | 構造保持 | 特徴 |
|---|---|---|---|---|---|
| Global Sum/Mean/Max | フラット | $O(n)$ | $O(n)$ | なし | シンプル、ベースライン |
| DiffPool | クラスタリング | $O(n^2)$ | $O(n \cdot k)$ | 高 | 微分可能クラスタリング |
| TopKPool | スコアベース | $O(n \log n)$ | $O(n)$ | 中 | ノード選択、スパース |
| SAGPool | スコアベース+GCN | $O(n + m)$ | $O(n)$ | 中〜高 | 構造考慮のスコア計算 |
手法選択の指針
小〜中規模のグラフ（数十〜数百ノード）: DiffPoolが効果的です。密な割り当て行列の計算コストが問題にならないスケールであれば、クラスタリングベースのアプローチが最も豊富な情報を保持できます。
大規模グラフ（数千ノード以上）: TopKPoolやSAGPoolが適しています。スパースな操作のみを使用するため、メモリ効率が良く、大規模グラフにもスケールします。
ベースラインとして: Global Sum Poolingは常に試す価値があります。特にGINのような表現力の高いGNNと組み合わせた場合、シンプルなsumプーリングでも高い性能を達成できることが多いです。
発展的なトピック
Set2Set
Set2SetはLSTMベースの集合読み出し関数で、複数のステップにわたってグラフ表現を精緻化します。Attention機構を用いて「どのノードに注目するか」を動的に変化させるため、フラットプーリングより豊富な情報を捉えられます。
GMT（Graph Multi-Head Attention Transformer）
Transformer のMulti-Head AttentionをグラフプーリングにSelf-Attention応用したアプローチです。ノード間の相互作用を明示的にモデリングしつつ、グラフ表現を生成します。
EdgePool
ノードではなくエッジを選択して収縮（edge contraction）することでグラフを縮小するアプローチです。グラフの接続性が保持される利点があります。
MinCutPool
グラフの最小カット問題に基づいたプーリング手法です。DiffPoolと同様にクラスタリングベースですが、正規化カットの目的関数を直接最適化することで、よりグラフ理論的に妥当なクラスタリングを実現します。
まとめ
本記事では、グラフプーリングの理論を体系的に解説しました。
重要なポイント
フラットプーリングの限界: Sum/Mean/Maxプーリングは実装がシンプルですが、グラフの階層的な構造情報を捉えられません。特に、部分構造（クリーク、官能基など）の組み合わせが重要なタスクでは不十分です。
DiffPool: 微分可能なクラスタリングにより、グラフを段階的に粗視化します。割り当て行列を学習することで、タスクに適したクラスタリングをエンドツーエンドで獲得できますが、計算コストが高い課題があります。
TopKPool / SAGPool: スコアベースのアプローチで、重要なノードだけを選択してグラフを縮小します。SAGPoolはGCNを用いてスコアを計算するため、グラフ構造を考慮した選択が可能です。メモリ効率が良いですが、削除されたノードの情報は失われます。
手法の選択: タスクの性質やグラフのスケールに応じて適切な手法を選択することが重要です。大規模グラフにはスコアベースの手法、小〜中規模でリッチな表現が必要な場合はDiffPoolが適しています。
次のステップ
グラフプーリングはGNNの表現力を高めるための重要なコンポーネントです。次の記事では、グラフ表現学習のもうひとつの強力なパラダイムである対照学習について解説します。