第3章：MPNN実装

汎用メッセージパッシングフレームワーク：分子から結晶まで適用可能な統一的実装

3.1 MPNNフレームワークの詳細

Message Passing Neural Networks（MPNN）は、Gilmer et al.（2017）によって提案された汎用的なグラフニューラルネットワークフレームワークです。CGCNNが結晶材料に特化しているのに対し、MPNNは分子、タンパク質、結晶など、あらゆるグラフ構造化データに適用できます。

3.1.1 論文の主要な貢献（Gilmer et al., 2017）

Gilmerらの論文（Proceedings of the 34th International Conference on Machine Learning, PMLR 70, pp. 1263-1272）は、以下の重要な貢献をしました：

1. 統一フレームワーク：既存のGNN手法（GCN、GraphSAGE、GAT等）を包含する一般化（pp. 1264-1265）
2. 量子化学予測：QM9データセットで13種類の量子化学特性を高精度に予測（表1、p. 1269）
3. カスタマイズ性：Message、Update、Readout関数を自由に設計可能（pp. 1265-1266）

数学的定式化（論文式(1)-(3)、pp. 1265-1266）：

Message関数（式(1)）：

\[ m_v^{t+1} = \sum_{w \in \mathcal{N}(v)} M_t(\mathbf{h}_v^t, \mathbf{h}_w^t, \mathbf{e}_{vw}) \]

Update関数（式(2)）：

\[ \mathbf{h}_v^{t+1} = U_t(\mathbf{h}_v^t, m_v^{t+1}) \]

Readout関数（式(3)）：

\[ \hat{y} = R(\{\mathbf{h}_v^T \mid v \in G\}) \]

ここで：

• \( \mathbf{h}_v^t \): ノード \( v \) の第 \( t \) ステップでの隠れ状態
• \( \mathcal{N}(v) \): ノード \( v \) の近傍ノード集合
• \( \mathbf{e}_{vw} \): エッジ特徴量
• \( M_t \): Message関数（学習可能なニューラルネットワーク）
• \( U_t \): Update関数（GRU、LSTM、またはMLPを使用）
• \( R \): Readout関数（グラフレベルの表現を生成）

3.1.2 CGCNN vs MPNN：設計思想の違い

3.2 Message関数の実装パターン

3.2.1 シンプルなMessage関数

# Example 1: 基本的なMessage関数の実装
# Google Colab環境セットアップ
!pip install torch-geometric torch-scatter torch-sparse rdkit

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import MessagePassing

class SimpleMessageFunction(MessagePassing):
    """シンプルなMessage関数

    論文: Gilmer et al. (2017), ICML, pp. 1265-1266
    """
    def __init__(self, node_dim, edge_dim, message_dim):
        """
        Args:
            node_dim (int): ノード特徴量の次元
            edge_dim (int): エッジ特徴量の次元
            message_dim (int): メッセージの次元
        """
        super().__init__(aggr='add')  # 集約方法: 合計

        # Message生成のための全結合層
        self.message_net = nn.Sequential(
            nn.Linear(node_dim + node_dim + edge_dim, message_dim),
            nn.ReLU(),
            nn.Linear(message_dim, message_dim)
        )

    def forward(self, x, edge_index, edge_attr):
        """
        Args:
            x (Tensor): ノード特徴量 [num_nodes, node_dim]
            edge_index (Tensor): エッジリスト [2, num_edges]
            edge_attr (Tensor): エッジ特徴量 [num_edges, edge_dim]

        Returns:
            Tensor: 集約されたメッセージ [num_nodes, message_dim]
        """
        return self.propagate(edge_index, x=x, edge_attr=edge_attr)

    def message(self, x_i, x_j, edge_attr):
        """メッセージ生成（エッジごとに実行）

        Args:
            x_i (Tensor): 受信ノードの特徴量 [num_edges, node_dim]
            x_j (Tensor): 送信ノードの特徴量 [num_edges, node_dim]
            edge_attr (Tensor): エッジ特徴量 [num_edges, edge_dim]

        Returns:
            Tensor: メッセージ [num_edges, message_dim]
        """
        # 受信ノード、送信ノード、エッジを連結
        msg_input = torch.cat([x_i, x_j, edge_attr], dim=1)

        # MLPでメッセージ生成
        return self.message_net(msg_input)

# 使用例
node_dim = 64
edge_dim = 10
message_dim = 64

msg_fn = SimpleMessageFunction(node_dim, edge_dim, message_dim)

# ダミーデータ
x = torch.randn(5, node_dim)  # 5ノード
edge_index = torch.tensor([[0, 1, 1, 2, 2, 3, 3, 4],
                            [1, 0, 2, 1, 3, 2, 4, 3]], dtype=torch.long)
edge_attr = torch.randn(8, edge_dim)

# Message関数実行
messages = msg_fn(x, edge_index, edge_attr)

print(f"Message関数:")
print(f"  入力ノード特徴量: {x.shape}")
print(f"  エッジ数: {edge_index.shape[1]}")
print(f"  出力メッセージ: {messages.shape}")
# 出力例:
# Message関数:
#   入力ノード特徴量: torch.Size([5, 64])
#   エッジ数: 8
#   出力メッセージ: torch.Size([5, 64])

3.2.2 エッジネットワークを用いたMessage関数

# Example 2: エッジネットワーク（Edge Network）を用いたMessage関数
class EdgeNetworkMessage(MessagePassing):
    """エッジネットワークを用いたMessage関数

    エッジ特徴量をニューラルネットワークで処理し、
    メッセージの重み付けに使用する高度な手法。
    """
    def __init__(self, node_dim, edge_dim, message_dim):
        super().__init__(aggr='add')

        # ノード特徴量の変換
        self.node_lin = nn.Linear(node_dim, message_dim)

        # エッジネットワーク（エッジ特徴量 → 重み）
        self.edge_net = nn.Sequential(
            nn.Linear(edge_dim, message_dim),
            nn.ReLU(),
            nn.Linear(message_dim, message_dim)
        )

    def forward(self, x, edge_index, edge_attr):
        # ノード特徴量の変換
        x = self.node_lin(x)
        return self.propagate(edge_index, x=x, edge_attr=edge_attr)

    def message(self, x_j, edge_attr):
        """エッジネットワークで重み付けされたメッセージ

        Args:
            x_j (Tensor): 送信ノードの特徴量 [num_edges, message_dim]
            edge_attr (Tensor): エッジ特徴量 [num_edges, edge_dim]

        Returns:
            Tensor: 重み付きメッセージ [num_edges, message_dim]
        """
        # エッジ特徴量から重みを生成
        edge_weight = self.edge_net(edge_attr)

        # 送信ノードの特徴量に重みを適用
        return x_j * edge_weight

# 使用例
edge_msg_fn = EdgeNetworkMessage(node_dim=64, edge_dim=10, message_dim=64)
messages_edge = edge_msg_fn(x, edge_index, edge_attr)

print(f"エッジネットワークMessage関数:")
print(f"  出力メッセージ: {messages_edge.shape}")
print(f"  パラメータ数: {sum(p.numel() for p in edge_msg_fn.parameters()):,}")

# 出力例:
# エッジネットワークMessage関数:
#   出力メッセージ: torch.Size([5, 64])
#   パラメータ数: 13,120

3.3 Update関数の実装パターン

3.3.1 GRU（Gated Recurrent Unit）を用いたUpdate

# Example 3: GRUを用いたUpdate関数
class GRUUpdate(nn.Module):
    """GRU（Gated Recurrent Unit）を用いたUpdate関数

    論文: Gilmer et al. (2017), ICML, p. 1266
    GRUは隠れ状態を時系列的に更新するRNNの一種。
    メッセージパッシングの各ステップで状態を更新する。
    """
    def __init__(self, hidden_dim):
        """
        Args:
            hidden_dim (int): 隠れ状態の次元
        """
        super().__init__()

        # PyTorchのGRU Cell
        self.gru = nn.GRUCell(hidden_dim, hidden_dim)

    def forward(self, h, m):
        """状態を更新

        Args:
            h (Tensor): 現在の隠れ状態 [num_nodes, hidden_dim]
            m (Tensor): 集約されたメッセージ [num_nodes, hidden_dim]

        Returns:
            Tensor: 更新された隠れ状態 [num_nodes, hidden_dim]
        """
        # GRUで状態更新
        # h^{t+1} = GRU(h^t, m^{t+1})
        return self.gru(m, h)

# 使用例
hidden_dim = 64
update_fn = GRUUpdate(hidden_dim)

# 現在の隠れ状態
h_current = torch.randn(5, hidden_dim)

# 集約されたメッセージ（Message関数の出力）
messages_agg = torch.randn(5, hidden_dim)

# Update実行
h_next = update_fn(h_current, messages_agg)

print(f"GRU Update関数:")
print(f"  現在の状態: {h_current.shape}")
print(f"  メッセージ: {messages_agg.shape}")
print(f"  更新後の状態: {h_next.shape}")
print(f"  状態の変化量: {torch.norm(h_next - h_current).item():.4f}")

# 出力例:
# GRU Update関数:
#   現在の状態: torch.Size([5, 64])
#   メッセージ: torch.Size([5, 64])
#   更新後の状態: torch.Size([5, 64])
#   状態の変化量: 5.2341

3.3.2 MLPを用いたシンプルなUpdate

# Example 4: MLPを用いたUpdate関数
class MLPUpdate(nn.Module):
    """MLPを用いたシンプルなUpdate関数

    GRUよりパラメータが少なく、計算も高速。
    """
    def __init__(self, hidden_dim):
        super().__init__()

        # 2層MLP
        self.mlp = nn.Sequential(
            nn.Linear(2 * hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )

    def forward(self, h, m):
        """状態を更新

        Args:
            h (Tensor): 現在の隠れ状態 [num_nodes, hidden_dim]
            m (Tensor): 集約されたメッセージ [num_nodes, hidden_dim]

        Returns:
            Tensor: 更新された隠れ状態 [num_nodes, hidden_dim]
        """
        # 現在の状態とメッセージを連結
        combined = torch.cat([h, m], dim=1)

        # MLPで新しい状態を計算
        h_new = self.mlp(combined)

        # 残差接続（オプション）
        return h_new + h

# 使用例
mlp_update_fn = MLPUpdate(hidden_dim=64)
h_next_mlp = mlp_update_fn(h_current, messages_agg)

print(f"MLP Update関数:")
print(f"  更新後の状態: {h_next_mlp.shape}")
print(f"  パラメータ数（MLP）: {sum(p.numel() for p in mlp_update_fn.parameters()):,}")
print(f"  パラメータ数（GRU）: {sum(p.numel() for p in update_fn.parameters()):,}")

# 出力例:
# MLP Update関数:
#   更新後の状態: torch.Size([5, 64])
#   パラメータ数（MLP）: 12,352
#   パラメータ数（GRU）: 24,768

3.4 Readout関数の実装パターン

3.4.1 Set2Set Readout

# Example 5: Set2Set Readout関数
from torch_geometric.nn import Set2Set

class Set2SetReadout(nn.Module):
    """Set2Set Readout関数

    論文: Vinyals et al. (2015) "Order Matters: Sequence to sequence for sets"
    Gilmer et al. (2017) ICML, p. 1266で推奨

    順序不変なグラフレベル表現を生成する高度な手法。
    Attention機構を用いて重要なノードを強調。
    """
    def __init__(self, hidden_dim, processing_steps=3):
        """
        Args:
            hidden_dim (int): ノード特徴量の次元
            processing_steps (int): Set2Setの処理ステップ数
        """
        super().__init__()

        # Set2Set層（PyTorch Geometric提供）
        self.set2set = Set2Set(hidden_dim, processing_steps=processing_steps)

        # 出力層
        self.fc = nn.Linear(2 * hidden_dim, 1)  # Set2Setは2倍の次元を出力

    def forward(self, x, batch):
        """グラフレベルの表現を生成

        Args:
            x (Tensor): ノード特徴量 [num_nodes, hidden_dim]
            batch (Tensor): バッチインデックス [num_nodes]

        Returns:
            Tensor: 予測値 [batch_size, 1]
        """
        # Set2Setでグラフ表現を生成
        graph_repr = self.set2set(x, batch)

        # 全結合層で予測
        return self.fc(graph_repr)

# 使用例
from torch_geometric.data import Batch, Data

# 複数のグラフをバッチ化
data_list = [
    Data(x=torch.randn(3, 64)),
    Data(x=torch.randn(4, 64)),
    Data(x=torch.randn(5, 64))
]
batch = Batch.from_data_list(data_list)

# Set2Set Readout
readout_fn = Set2SetReadout(hidden_dim=64, processing_steps=3)
predictions = readout_fn(batch.x, batch.batch)

print(f"Set2Set Readout:")
print(f"  バッチサイズ: {batch.num_graphs}")
print(f"  総ノード数: {batch.num_nodes}")
print(f"  予測値: {predictions.shape}")
print(f"  予測値の例: {predictions.squeeze().detach().numpy()}")

# 出力例:
# Set2Set Readout:
#   バッチサイズ: 3
#   総ノード数: 12
#   予測値: torch.Size([3, 1])
#   予測値の例: [-0.234, 0.567, -0.891]

3.5 完全なMPNNモデル

# Example 6: 完全なMPNNモデルの実装
class MPNN(nn.Module):
    """完全なMPNNモデル

    論文: Gilmer et al. (2017), ICML, pp. 1263-1272
    """
    def __init__(self,
                 node_features,
                 edge_features,
                 hidden_dim=64,
                 num_layers=3,
                 readout_steps=3):
        """
        Args:
            node_features (int): 入力ノード特徴量の次元
            edge_features (int): エッジ特徴量の次元
            hidden_dim (int): 隠れ層の次元
            num_layers (int): メッセージパッシングの層数
            readout_steps (int): Set2Setの処理ステップ数
        """
        super().__init__()

        # 入力の埋め込み
        self.node_embedding = nn.Linear(node_features, hidden_dim)

        # Message関数（複数層）
        self.message_layers = nn.ModuleList([
            EdgeNetworkMessage(hidden_dim, edge_features, hidden_dim)
            for _ in range(num_layers)
        ])

        # Update関数（GRU）
        self.update_layers = nn.ModuleList([
            GRUUpdate(hidden_dim)
            for _ in range(num_layers)
        ])

        # Readout関数（Set2Set）
        self.readout = Set2SetReadout(hidden_dim, processing_steps=readout_steps)

    def forward(self, data):
        """
        Args:
            data (Data): PyTorch Geometric Dataオブジェクト
                - x: ノード特徴量 [num_nodes, node_features]
                - edge_index: エッジリスト [2, num_edges]
                - edge_attr: エッジ特徴量 [num_edges, edge_features]
                - batch: バッチインデックス [num_nodes]

        Returns:
            Tensor: 予測値 [batch_size, 1]
        """
        # ノードの埋め込み
        h = self.node_embedding(data.x)

        # メッセージパッシング（複数層）
        for message_layer, update_layer in zip(self.message_layers, self.update_layers):
            # Message: 近傍から情報を集約
            m = message_layer(h, data.edge_index, data.edge_attr)

            # Update: 隠れ状態を更新
            h = update_layer(h, m)

        # Readout: グラフレベルの予測
        return self.readout(h, data.batch)

# モデル初期化
model = MPNN(
    node_features=11,  # QM9の原子特徴量（原子番号等）
    edge_features=4,   # 結合タイプ、距離等
    hidden_dim=64,
    num_layers=3,
    readout_steps=3
)

print(f"完全なMPNNモデル:")
print(f"  総パラメータ数: {sum(p.numel() for p in model.parameters()):,}")
print(f"  メッセージパッシング層数: 3")
print(f"  隠れ層次元: 64")
print(f"  Readout: Set2Set（3ステップ）")

# ダミーデータで動作確認
dummy_data = Data(
    x=torch.randn(10, 11),
    edge_index=torch.randint(0, 10, (2, 20)),
    edge_attr=torch.randn(20, 4),
    batch=torch.zeros(10, dtype=torch.long)
)

output = model(dummy_data)
print(f"\nモデル出力:")
print(f"  入力: {dummy_data.num_nodes}ノード、{dummy_data.num_edges}エッジ")
print(f"  出力: {output.shape}")

# 出力例:
# 完全なMPNNモデル:
#   総パラメータ数: 124,993
#   メッセージパッシング層数: 3
#   隠れ層次元: 64
#   Readout: Set2Set（3ステップ）
#
# モデル出力:
#   入力: 10ノード、20エッジ
#   出力: torch.Size([1, 1])

3.6 QM9データセットでの分子特性予測

3.6.1 QM9データセットの概要

QM9データセット（Ramakrishnan et al., 2014, Scientific Data, 1, 140022, pp. 1-7）は、量子化学計算による分子特性の大規模データベースです。134,000の有機分子（最大9重原子、C、H、O、N、F）について、DFT計算により13種類の量子化学特性が計算されています（pp. 3-4）。

主要な量子化学特性:

• HOMO: 最高被占軌道エネルギー（電子の供与能力）
• LUMO: 最低空軌道エネルギー（電子の受容能力）
• Gap: HOMO-LUMOギャップ（励起エネルギー、重要な電子的特性）
• μ: 双極子モーメント（分子の極性）
• α: 分極率（外部電場への応答）
• ZPVE: ゼロ点振動エネルギー

# Example 7: QM9データセットの読み込みとMPNN訓練
!pip install torch-geometric-temporal  # QM9データセット用

from torch_geometric.datasets import QM9
import torch
import torch.nn as nn
from torch.optim import Adam
from torch_geometric.loader import DataLoader
from sklearn.metrics import mean_absolute_error
import numpy as np

# QM9データセット読み込み
dataset = QM9(root='./data/qm9')

print(f"QM9データセット:")
print(f"  総分子数: {len(dataset):,}")
print(f"  ノード特徴量次元: {dataset[0].x.shape[1]}")
print(f"  エッジ特徴量次元: {dataset[0].edge_attr.shape[1]}")
print(f"  ターゲット特性数: {dataset[0].y.shape[1]}")

# サンプル分子の確認
sample_mol = dataset[0]
print(f"\nサンプル分子:")
print(f"  原子数: {sample_mol.num_nodes}")
print(f"  結合数: {sample_mol.num_edges}")
print(f"  HOMO-LUMOギャップ: {sample_mol.y[0, 4].item():.4f} eV")
print(f"  双極子モーメント: {sample_mol.y[0, 0].item():.4f} Debye")

# データを訓練・検証・テストに分割
# QM9の標準分割: 110,000 / 10,000 / 13,885
train_dataset = dataset[:110000]
val_dataset = dataset[110000:120000]
test_dataset = dataset[120000:]

# DataLoader作成
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

print(f"\nデータ分割:")
print(f"  訓練: {len(train_dataset):,}分子")
print(f"  検証: {len(val_dataset):,}分子")
print(f"  テスト: {len(test_dataset):,}分子")

# 出力例:
# QM9データセット:
#   総分子数: 130,831
#   ノード特徴量次元: 11
#   エッジ特徴量次元: 4
#   ターゲット特性数: 19
#
# サンプル分子:
#   原子数: 5
#   結合数: 8
#   HOMO-LUMOギャップ: 0.2586 eV
#   双極子モーメント: 0.0000 Debye
#
# データ分割:
#   訓練: 110,000分子
#   検証: 10,000分子
#   テスト: 10,831分子

3.6.2 HOMO-LUMOギャップ予測の訓練

# Example 8: HOMO-LUMOギャップ予測の訓練
def train_qm9_model(model, train_loader, val_loader,
                    target_idx=4,  # HOMO-LUMOギャップ
                    epochs=50, lr=0.001, device='cuda'):
    """QM9データセットでMPNNを訓練

    Args:
        model (nn.Module): MPNNモデル
        train_loader (DataLoader): 訓練データ
        val_loader (DataLoader): 検証データ
        target_idx (int): 予測する特性のインデックス（4: HOMO-LUMOギャップ）
        epochs (int): エポック数
        lr (float): 学習率
        device (str): デバイス

    Returns:
        dict: 訓練履歴
    """
    model = model.to(device)
    optimizer = Adam(model.parameters(), lr=lr)
    criterion = nn.L1Loss()  # Mean Absolute Error

    history = {'train_loss': [], 'val_loss': [], 'val_mae': []}

    for epoch in range(epochs):
        # ===== 訓練フェーズ =====
        model.train()
        train_loss = 0.0

        for batch in train_loader:
            batch = batch.to(device)
            optimizer.zero_grad()

            # 予測（ターゲット特性のみ）
            pred = model(batch)
            target = batch.y[:, target_idx].unsqueeze(1)

            # 損失計算
            loss = criterion(pred, target)

            # バックプロパゲーション
            loss.backward()
            optimizer.step()

            train_loss += loss.item() * batch.num_graphs

        train_loss /= len(train_loader.dataset)

        # ===== 検証フェーズ =====
        model.eval()
        val_loss = 0.0
        y_true, y_pred = [], []

        with torch.no_grad():
            for batch in val_loader:
                batch = batch.to(device)
                pred = model(batch)
                target = batch.y[:, target_idx].unsqueeze(1)

                loss = criterion(pred, target)
                val_loss += loss.item() * batch.num_graphs

                y_true.extend(target.cpu().numpy())
                y_pred.extend(pred.cpu().numpy())

        val_loss /= len(val_loader.dataset)
        val_mae = mean_absolute_error(y_true, y_pred)

        # 履歴に記録
        history['train_loss'].append(train_loss)
        history['val_loss'].append(val_loss)
        history['val_mae'].append(val_mae)

        # 進捗表示
        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs}:")
            print(f"  Train Loss: {train_loss:.4f} eV")
            print(f"  Val Loss: {val_loss:.4f} eV")
            print(f"  Val MAE: {val_mae:.4f} eV")

    return history

# 使用例（実データがあれば）
# model_qm9 = MPNN(
#     node_features=11,
#     edge_features=4,
#     hidden_dim=64,
#     num_layers=3,
#     readout_steps=3
# )
#
# history = train_qm9_model(
#     model=model_qm9,
#     train_loader=train_loader,
#     val_loader=val_loader,
#     target_idx=4,  # HOMO-LUMOギャップ
#     epochs=50,
#     lr=0.001,
#     device='cuda' if torch.cuda.is_available() else 'cpu'
# )

print(f"訓練関数の定義完了")
print(f"期待される性能（論文値、Gilmer et al. 2017, 表1, p. 1269）:")
print(f"  HOMO-LUMOギャップ MAE: 0.043 eV")
print(f"  双極子モーメント MAE: 0.033 Debye")
print(f"  分極率 MAE: 0.092 Bohr³")

3.7 CGCNN vs MPNN：定量的比較

3.7.1 結晶 vs 分子での性能差

出典:

• CGCNN: Xie & Grossman (2018), Physical Review Letters, 120, 145301, 表I, p. 4
• MPNN: Gilmer et al. (2017), ICML, 表1, p. 1269

3.7.2 アーキテクチャの違いが性能に与える影響

CGCNNが結晶で優れる理由:

1. 周期境界条件: 無限に繰り返される結晶構造を適切に扱う
2. エッジゲート機構: 原子間距離に応じた適応的な重み付け
3. ドメイン特化設計: 結晶材料の特性（配位環境、長距離相互作用）に最適化

MPNNが分子で優れる理由:

1. Set2Set Readout: 分子サイズに不変な柔軟な表現学習
2. GRU Update: 複雑な電子構造を捉える時系列的な状態更新
3. カスタマイズ性: 分子の特性（芳香族性、結合次数等）に応じた柔軟な設計

3.7.3 計算コストの比較

MPNNの計算コストが高い理由:

• GRU Updateは再帰的計算が必要（並列化困難）
• Set2Set Readoutは複数ステップの処理が必要
• エッジネットワークはCGCNNのゲート機構より複雑

3.8 まとめ

この章では、MPNNの汎用フレームワークとQM9データセットでの分子特性予測を学びました：

1. MPNNフレームワーク: Message、Update、Readoutの3段階による汎用的な設計
2. Message関数: シンプルなMLPからエッジネットワークまで多様な実装
3. Update関数: GRU（時系列的更新）vs MLP（シンプル）のトレードオフ
4. Readout関数: Set2Setによる柔軟なグラフレベル表現学習
5. QM9予測: HOMO-LUMOギャップ（MAE 0.043 eV）、双極子モーメント（MAE 0.033 Debye）
6. CGCNN vs MPNN: 結晶特化 vs 汎用フレームワークのトレードオフ

次章では、組成ベース特徴量（Magpie）とGNN（CGCNN/MPNN）の定量的比較を、Matbenchベンチマークで実施します。

演習問題

Easy（基礎確認）

Medium（応用）

import torch
import torch.nn as nn
from torch.optim import Adam
from torch_geometric.datasets import QM9
from torch_geometric.loader import DataLoader

# QM9データセット読み込み
dataset = QM9(root='./data/qm9')
train_dataset = dataset[:110000]
val_dataset = dataset[110000:120000]

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)

# MPNNモデル初期化
model = MPNN(
    node_features=11,
    edge_features=4,
    hidden_dim=64,
    num_layers=3,
    readout_steps=3
)

# 訓練
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)
optimizer = Adam(model.parameters(), lr=0.001)
criterion = nn.L1Loss()

for epoch in range(50):
    model.train()
    train_loss = 0.0

    for batch in train_loader:
        batch = batch.to(device)
        optimizer.zero_grad()

        # HOMO-LUMOギャップ予測（インデックス4）
        pred = model(batch)
        target = batch.y[:, 4].unsqueeze(1)

        loss = criterion(pred, target)
        loss.backward()
        optimizer.step()

        train_loss += loss.item() * batch.num_graphs

    train_loss /= len(train_loader.dataset)

    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1}: Train Loss = {train_loss:.4f} eV")

# 検証
model.eval()
val_preds, val_targets = [], []
with torch.no_grad():
    for batch in val_loader:
        batch = batch.to(device)
        pred = model(batch)
        target = batch.y[:, 4].unsqueeze(1)

        val_preds.extend(pred.cpu().numpy())
        val_targets.extend(target.cpu().numpy())

from sklearn.metrics import mean_absolute_error
mae = mean_absolute_error(val_targets, val_preds)
print(f"Validation MAE: {mae:.4f} eV")
# 期待値: 約0.043 eV（論文値）

Hard（発展）

import torch
import torch.nn as nn
from torch_geometric.nn import MessagePassing

class BondTypeMessage(MessagePassing):
    """結合タイプを明示的に扱うMessage関数

    結合タイプ（単結合=1、二重結合=2、三重結合=3、芳香族=4）
    ごとに異なるMLPを使用してメッセージを生成。
    """
    def __init__(self, node_dim, message_dim, num_bond_types=4):
        """
        Args:
            node_dim (int): ノード特徴量の次元
            message_dim (int): メッセージの次元
            num_bond_types (int): 結合タイプの種類数
        """
        super().__init__(aggr='add')

        # 結合タイプごとのMLP
        self.bond_mlps = nn.ModuleList([
            nn.Sequential(
                nn.Linear(2 * node_dim, message_dim),
                nn.ReLU(),
                nn.Linear(message_dim, message_dim)
            )
            for _ in range(num_bond_types)
        ])

        # 結合タイプのone-hot埋め込み
        self.num_bond_types = num_bond_types

    def forward(self, x, edge_index, bond_type):
        """
        Args:
            x (Tensor): ノード特徴量 [num_nodes, node_dim]
            edge_index (Tensor): エッジリスト [2, num_edges]
            bond_type (Tensor): 結合タイプ [num_edges]（0-indexed）

        Returns:
            Tensor: 集約されたメッセージ [num_nodes, message_dim]
        """
        return self.propagate(edge_index, x=x, bond_type=bond_type)

    def message(self, x_i, x_j, bond_type):
        """結合タイプに応じたメッセージ生成

        Args:
            x_i (Tensor): 受信ノード [num_edges, node_dim]
            x_j (Tensor): 送信ノード [num_edges, node_dim]
            bond_type (Tensor): 結合タイプ [num_edges]

        Returns:
            Tensor: メッセージ [num_edges, message_dim]
        """
        # ノードを連結
        combined = torch.cat([x_i, x_j], dim=1)

        # 結合タイプごとにメッセージを生成
        messages = []
        for i in range(self.num_bond_types):
            # 該当する結合タイプのエッジを抽出
            mask = (bond_type == i)
            if mask.any():
                # 該当MLPでメッセージ生成
                msg_i = self.bond_mlps[i](combined[mask])
                messages.append((mask, msg_i))

        # 全メッセージを統合
        output = torch.zeros(combined.shape[0], messages[0][1].shape[1],
                             device=combined.device)
        for mask, msg in messages:
            output[mask] = msg

        return output

# 使用例
node_dim = 64
message_dim = 64

# 結合タイプを考慮したMessage関数
bond_msg = BondTypeMessage(node_dim, message_dim, num_bond_types=4)

# ダミーデータ
x = torch.randn(5, node_dim)
edge_index = torch.tensor([[0, 1, 1, 2, 2, 3],
                            [1, 0, 2, 1, 3, 2]], dtype=torch.long)
bond_type = torch.tensor([0, 0, 1, 1, 3, 3], dtype=torch.long)  # 単結合、二重結合、芳香族

# Message関数実行
messages = bond_msg(x, edge_index, bond_type)

print(f"結合タイプ考慮Message関数:")
print(f"  入力ノード: {x.shape}")
print(f"  結合タイプ: {bond_type}")
print(f"  出力メッセージ: {messages.shape}")
print(f"  パラメータ数: {sum(p.numel() for p in bond_msg.parameters()):,}")

学習目標の確認

このchapterを完了すると、以下を説明できるようになります：

基本理解

• ✅ MPNNの3段階（Message/Update/Readout）を説明できる
• ✅ CGCNN vs MPNNの設計思想の違いを理解している
• ✅ QM9データセットの量子化学特性を説明できる
• ✅ Set2Set Readoutの動作原理を理解している

実践スキル

• ✅ MPNNのMessage、Update、Readout関数をスクラッチ実装できる
• ✅ QM9データセットでHOMO-LUMOギャップを予測できる（MAE < 0.05 eV目標）
• ✅ GRU UpdateとMLP Updateを実装し、性能を比較できる
• ✅ Set2Set Readoutを実装し、分子サイズ不変な表現を学習できる

応用力

• ✅ CGCNN vs MPNNの使い分けを定量的に評価できる
• ✅ カスタムMessage関数を設計し、ドメイン知識を組み込める
• ✅ 論文の性能（HOMO-LUMOギャップ MAE 0.043 eV）を再現するための条件を理解している

次のステップ

次章では、組成ベース特徴量（Magpie）とGNN（CGCNN/MPNN）の定量的比較を、Matbenchベンチマークで実施します。予測精度、計算コスト、データ要件、解釈性の4軸で徹底分析し、実践的な手法選択の意思決定能力を養います。

参考文献

1. Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., & Dahl, G. E. (2017). "Neural Message Passing for Quantum Chemistry." Proceedings of the 34th International Conference on Machine Learning, PMLR 70, pp. 1263-1272.
2. Ramakrishnan, R., Dral, P. O., Rupp, M., & von Lilienfeld, O. A. (2014). "Quantum chemistry structures and properties of 134 kilo molecules." Scientific Data, 1, 140022, pp. 1-7.
3. Schütt, K. T., Kindermans, P. J., Sauceda, H. E., Chmiela, S., Tkatchenko, A., & Müller, K. R. (2017). "SchNet: A continuous-filter convolutional neural network for modeling quantum interactions." Advances in Neural Information Processing Systems, 30, pp. 991-1001.
4. Fey, M., & Lenssen, J. E. (2019). "Fast Graph Representation Learning with PyTorch Geometric." ICLR Workshop on Representation Learning on Graphs and Manifolds, pp. 1-9.
5. Vinyals, O., Bengio, S., & Kudlur, M. (2015). "Order Matters: Sequence to sequence for sets." arXiv preprint arXiv:1511.06391, pp. 1-11.
6. Xie, T., & Grossman, J. C. (2018). "Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties." Physical Review Letters, 120(14), 145301, pp. 1-6.
7. Wu, Z., Ramsundar, B., Feinberg, E. N., Gomes, J., Geniesse, C., Pappu, A. S., Leswing, K., & Pande, V. (2018). "MoleculeNet: a benchmark for molecular machine learning." Chemical Science, 9(2), pp. 513-530.