Day 134 — カスタムアーキテクチャ設計④ — Multi-Head Self-Attentionブロックの自作設計

📚 背景知識（読んでから問題へ）

🪄 今日の比較の設計（5種類のブロック × num_tokens 2パターン）

AttnMLPBlockはResidualMLPBlockの1回のResidual加算を2回繰り返す構造（Attention sub-layer + FFN sub-layer）と捉えると理解しやすい。今日はさらにnum_tokens（トークンへの分割数）を4と8の2パターンで比較し、粒度とパラメータ数のトレードオフを見る。

🎯 問題

Day131〜133と同系統の合成分類データ・共通骨格（DeepNet）を使い、これまでのResidualMLPBlock・GatedMLPBlock・SEMLPBlockに加えて新しくAttnMLPBlockを実装し、比較します。

import numpy as np
from sklearn.datasets import make_classification

# --- 合成データ生成（Day131〜133と同一） ---
X, y = make_classification(
    n_samples=2400, n_features=20, n_informative=12,
    n_classes=5, n_clusters_per_class=1, random_state=42,
    flip_y=0.03,
)

TRAIN_LEN = 1800   # 先頭1800件をtrain、残り600件をtest
HIDDEN = 64        # ブロック内部の次元（全ブロック共通）
# 今日比較する5種類:
#   (a) PlainMLP  : Day131と同じ（正規化・Skipなし、劣化問題の対照群）
#   (b) ResMLP    : Day131のResidualMLPBlock（H(x) = F(x) + x）
#   (c) GatedMLP  : Day132のGatedMLPBlock（H(x) = T(x)F(x) + (1-T(x))x、gate_bias_init=-3）
#   (d) SEMLP     : Day133のSEMLPBlock（H(x) = x + S(F(x)) * F(x)、reduction=4）
#   (e) AttnMLP   : 今日新規実装。dimをnum_tokens個のトークンに分割しMulti-Head Self-Attentionを適用

💡 ヒント

import torch
import torch.nn as nn

class AttnMLPBlock(nn.Module):
    def __init__(self, dim, num_tokens=8, num_heads=4):
        super().__init__()
        assert dim % num_tokens == 0
        self.num_tokens = num_tokens
        self.token_dim = dim // num_tokens
        self.norm1 = nn.LayerNorm(dim)
        self.mha = nn.MultiheadAttention(
            embed_dim=self.token_dim, num_heads=num_heads, batch_first=True
        )
        self.norm2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, dim))
        self.last_attn = None   # 可視化用

    def forward(self, x):
        b = x.shape[0]
        h = self.norm1(x)
        tokens = h.view(b, self.num_tokens, self.token_dim)   # (batch, num_tokens, token_dim)
        attn_out, attn_w = self.mha(
            tokens, tokens, tokens, need_weights=True, average_attn_weights=True
        )
        self.last_attn = attn_w.mean(dim=0).detach()    # (num_tokens, num_tokens)
        attn_out = attn_out.reshape(b, -1)               # (batch, dim)
        x = x + attn_out                                # Attention sub-layerのResidual

        h2 = self.norm2(x)
        f = self.ffn(h2)
        return x + f                                # FFN sub-layerのResidual

✅ 模範解答

設計要素↔コード変数の対応表

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import make_classification
from torch.utils.data import TensorDataset, DataLoader

torch.manual_seed(0)

# --- サンプルデータ生成（Day131〜133と同一） ---
X, y = make_classification(
    n_samples=2400, n_features=20, n_informative=12,
    n_classes=5, n_clusters_per_class=1, random_state=42, flip_y=0.03,
)
TRAIN_LEN = 1800
HIDDEN = 64

X_mean, X_std = X[:TRAIN_LEN].mean(axis=0), X[:TRAIN_LEN].std(axis=0)
X_scaled = (X - X_mean) / X_std

X_train = torch.tensor(X_scaled[:TRAIN_LEN], dtype=torch.float32)
y_train = torch.tensor(y[:TRAIN_LEN], dtype=torch.long)
X_test = torch.tensor(X_scaled[TRAIN_LEN:], dtype=torch.float32)
y_test = torch.tensor(y[TRAIN_LEN:], dtype=torch.long)

train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)

# --- ステップ1&2: AttnMLPBlockの自作設計（num_tokens・num_heads可変） ---
class AttnMLPBlock(nn.Module):
    def __init__(self, dim, num_tokens=8, num_heads=4):
        super().__init__()
        assert dim % num_tokens == 0
        self.num_tokens = num_tokens
        self.token_dim = dim // num_tokens
        self.norm1 = nn.LayerNorm(dim)
        self.mha = nn.MultiheadAttention(
            embed_dim=self.token_dim, num_heads=num_heads, batch_first=True
        )
        self.norm2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, dim))
        self.last_attn = None

    def forward(self, x):
        b = x.shape[0]
        h = self.norm1(x)
        tokens = h.view(b, self.num_tokens, self.token_dim)
        attn_out, attn_w = self.mha(
            tokens, tokens, tokens, need_weights=True, average_attn_weights=True
        )
        self.last_attn = attn_w.mean(dim=0).detach()
        attn_out = attn_out.reshape(b, -1)
        x = x + attn_out

        h2 = self.norm2(x)
        f = self.ffn(h2)
        return x + f

class SEMLPBlock(nn.Module):
    """Day133と同じ実装（比較対象）"""
    def __init__(self, dim, reduction=4):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.fc1 = nn.Linear(dim, dim)
        self.act = nn.ReLU()
        self.norm2 = nn.LayerNorm(dim)
        self.fc2 = nn.Linear(dim, dim)
        reduced = max(1, dim // reduction)
        self.excitation = nn.Sequential(
            nn.Linear(dim, reduced), nn.ReLU(), nn.Linear(reduced, dim), nn.Sigmoid(),
        )
        self.last_scale = None

    def forward(self, x):
        h = self.norm1(x)
        h = self.act(self.fc1(h))
        h = self.norm2(h)
        f = self.fc2(h)
        s = self.excitation(f)
        self.last_scale = s.mean(dim=0).detach()
        return x + s * f

class GatedMLPBlock(nn.Module):
    """Day132と同じ実装（比較対象）"""
    def __init__(self, dim, gate_bias_init=-3.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.fc1 = nn.Linear(dim, dim)
        self.act = nn.ReLU()
        self.norm2 = nn.LayerNorm(dim)
        self.fc2 = nn.Linear(dim, dim)
        self.gate_fc = nn.Linear(dim, dim)
        nn.init.constant_(self.gate_fc.bias, gate_bias_init)

    def forward(self, x):
        h = self.norm1(x)
        h = self.act(self.fc1(h))
        h = self.norm2(h)
        f = self.fc2(h)
        t = torch.sigmoid(self.gate_fc(x))
        return t * f + (1 - t) * x

class ResidualMLPBlock(nn.Module):
    """Day131と同じ実装（比較対象）"""
    def __init__(self, dim):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.fc1 = nn.Linear(dim, dim)
        self.act = nn.ReLU()
        self.norm2 = nn.LayerNorm(dim)
        self.fc2 = nn.Linear(dim, dim)

    def forward(self, x):
        h = self.norm1(x)
        h = self.act(self.fc1(h))
        h = self.norm2(h)
        h = self.fc2(h)
        return x + h

class PlainBlock(nn.Module):
    """Day131と同じ実装（劣化問題の対照群）"""
    def __init__(self, dim):
        super().__init__()
        self.fc1 = nn.Linear(dim, dim)
        self.act1 = nn.ReLU()
        self.fc2 = nn.Linear(dim, dim)
        self.act2 = nn.ReLU()

    def forward(self, x):
        h = self.act1(self.fc1(x))
        h = self.act2(self.fc2(h))
        return h

# --- ステップ3: 1ブロックあたりの総パラメータ数の比較 ---
def count_params(module):
    return sum(p.numel() for p in module.parameters())

res_extra = count_params(ResidualMLPBlock(HIDDEN))
gated_extra = count_params(GatedMLPBlock(HIDDEN))
se_extra = count_params(SEMLPBlock(HIDDEN, reduction=4))
attn8_block = AttnMLPBlock(HIDDEN, num_tokens=8, num_heads=4)
attn8_extra = count_params(attn8_block)
attn8_mha_only = count_params(attn8_block.mha)   # Attention sub-layer自体のみ

print("--- 1ブロックあたりの総パラメータ数 ---")
print(f"ResidualMLPBlock:            {res_extra} 個")
print(f"GatedMLPBlock:                {gated_extra} 個")
print(f"SEMLPBlock(reduction=4):      {se_extra} 個")
print(f"AttnMLPBlock(num_tokens=8):   {attn8_extra} 個（うちMHA自体は {attn8_mha_only} 個）")
# 出力例:
# ResidualMLPBlock:             8,576 個
# GatedMLPBlock:                12,736 個
# SEMLPBlock(reduction=4):      10,704 個
# AttnMLPBlock(num_tokens=8):   8,864 個（うちMHA自体は 288 個）
# → AttnMLPBlockの総パラメータ数はGatedMLPBlockより少ない。
#   Attention sub-layer自体(288個)は驚くほど軽量で、大半(8320個)はFFN sub-layerが占めている

# --- ステップ4: 共通骨格を持つDeepNet ---
class DeepNet(nn.Module):
    def __init__(self, in_dim=20, hidden=HIDDEN, num_classes=5, n_blocks=20,
                 mode="attn", gate_bias_init=-3.0, reduction=4, num_tokens=8, num_heads=4):
        super().__init__()
        self.input_layer = nn.Linear(in_dim, hidden)
        if mode == "plain":
            self.blocks = nn.ModuleList([PlainBlock(hidden) for _ in range(n_blocks)])
        elif mode == "res":
            self.blocks = nn.ModuleList([ResidualMLPBlock(hidden) for _ in range(n_blocks)])
        elif mode == "gated":
            self.blocks = nn.ModuleList(
                [GatedMLPBlock(hidden, gate_bias_init=gate_bias_init) for _ in range(n_blocks)]
            )
        elif mode == "se":
            self.blocks = nn.ModuleList(
                [SEMLPBlock(hidden, reduction=reduction) for _ in range(n_blocks)]
            )
        elif mode == "attn":
            self.blocks = nn.ModuleList(
                [AttnMLPBlock(hidden, num_tokens=num_tokens, num_heads=num_heads) for _ in range(n_blocks)]
            )
        else:
            raise ValueError(mode)
        self.output_layer = nn.Linear(hidden, num_classes)

    def forward(self, x):
        h = self.input_layer(x)
        for block in self.blocks:
            h = block(h)
        return self.output_layer(h)

def evaluate(model):
    model.eval()
    with torch.no_grad():
        preds = model(X_test).argmax(dim=1)
    return (preds == y_test).float().mean().item()

def train_and_measure(mode, n_blocks=20, epochs=60, lr=1e-3,
                       gate_bias_init=-3.0, reduction=4, num_tokens=8, num_heads=4):
    torch.manual_seed(0)
    model = DeepNet(n_blocks=n_blocks, mode=mode, gate_bias_init=gate_bias_init,
                     reduction=reduction, num_tokens=num_tokens, num_heads=num_heads)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_loader:
            optimizer.zero_grad()
            loss = F.cross_entropy(model(xb), yb)
            loss.backward()
            optimizer.step()
    acc = evaluate(model)
    return model, acc

# --- ステップ4: PlainMLP・ResMLP・GatedMLP・SEMLP・AttnMLPの5者比較（N_BLOCKS=20） ---
configs = [
    ("plain", {}),
    ("res", {}),
    ("gated", {"gate_bias_init": -3.0}),
    ("se", {"reduction": 4}),
    ("attn", {"num_tokens": 8, "num_heads": 4}),
]
for mode, kwargs in configs:
    model, acc = train_and_measure(mode, n_blocks=20, **kwargs)
    print(f"{mode:6s}: acc={acc:.4f}")
# 出力例（乱数・実行環境により変動）:
# plain : acc=0.2117
# res   : acc=0.7217
# gated : acc=0.7350
# se    : acc=0.7417
# attn  : acc=0.7508
# → AttnMLPが5者中最も高精度。トークン間の関連度を明示的に計算する分、
#   SEブロック（チャネル単独の重要度のみ）より一段柔軟な情報の混ぜ合わせができている

# --- ステップ5: num_tokens(4 vs 8)の比較 ---
print("\n--- num_tokensの比較（N_BLOCKS=20） ---")
for nt in [4, 8]:
    _, acc = train_and_measure("attn", n_blocks=20, num_tokens=nt, num_heads=4)
    block = AttnMLPBlock(HIDDEN, num_tokens=nt, num_heads=4)
    print(f"num_tokens={nt}: acc={acc:.4f}  token_dim={HIDDEN // nt}  総パラメータ数={count_params(block)}")
# 出力例の傾向:
# num_tokens=4: acc=0.7422  token_dim=16  総パラメータ数=9,376
# num_tokens=8: acc=0.7508  token_dim= 8  総パラメータ数=8,864
# → num_tokens=8(より細かい粒度)の方がパラメータ数が少ないにもかかわらず精度が高い。
#   token_dimが小さくなってもトークン数が増えることで、より細かい関連性を捉えられている

# --- ステップ6: Attention重み行列の可視化 ---
attn_model, _ = train_and_measure("attn", n_blocks=20, num_tokens=8, num_heads=4)
attn_model.eval()
with torch.no_grad():
    _ = attn_model(X_test)
last_block_attn = attn_model.blocks[-1].last_attn   # shape: (8, 8)
for i in range(8):
    row = " ".join(f"{last_block_attn[i, j]:.3f}" for j in range(8))
    print(f"token{i}: {row}")
# 出力例の傾向:
# 対角成分(token_i - token_i)が最も高い値を持つ行が多い(自己参照が支配的)一方、
# 一部のトークンペア(例: token2-token5)で対角以外の値が突出して高くなる傾向が見られる
# → 特定のトークン同士が強く関連づけられていることをネットワークが学習している

📊 結果の可視化（パラメータ内訳とAttention重み行列）

🪜 Step-by-Step 解説

1なぜ`dim`をそのままAttentionに入力せず「トークン化」するのか

tokens = h.view(b, self.num_tokens, self.token_dim)   # (batch, num_tokens, token_dim)

2Attention sub-layerとFFN sub-layerが「2段のResidual」を持つ理由

3Attention sub-layer自体のパラメータ数が意外なほど少ない理由

# GatedMLPBlockのgate_fc: dim -> dim（フル結合）                → dim*dim + dim = 4,160 個
# AttnMLPBlockのmha: embed_dim=token_dim(=8)のQKV+出力射影      → 288 個

4Attention重み行列が示す「学習された関連づけ」

🧮 数学・統計の補足（文系向け）

🏆 Kaggleでの実践的な使い方

よく使われるコンペカテゴリ: ☑ 表形式データ（Tabular） / ☑ 自然言語処理（NLP） / ☑ 画像認識（CV） / ☐ 時系列（Time Series）

⚠️ よくある誤解・ミス

🚀 次のステップ

• 発展: 今日はトークンへの分割を「連続した次元を機械的にグループ化する」という単純な方法で行いましたが、実際のFT-Transformerでは各特徴量（列）を個別に埋め込みベクトルに変換してから「1特徴量=1トークン」として扱います。AttnMLPBlockを拡張し、入力層の直後に各特徴量を独立した小さな埋め込みに変換する層を挟み、「特徴量そのもの」をトークンとしてAttentionにかける版を実装し、精度がどう変わるか試してみましょう
• 次回予告: Phase6テーマ2「カスタムアーキテクチャ設計」5日目（テーマ2の最終日）。Day131〜134で作った4種類のブロック（Residual・Gated・SE・Attention）をすべて組み合わせたハイブリッドブロックを設計し、どの組み合わせが最も効果的かをアブレーション実験（機構を1つずつ外して効果を測る実験）で検証します

Phase 6 の学習マップ（全20テーマ）

Day131でResidual Connection、Day132でGated MLP、Day133でSqueeze-and-Exciteを実装。今日はトークン間の関連度を動的に計算するMulti-Head Self-Attentionブロックを実装した。次回はこれら4機構を統合したハイブリッドブロックへ進む。

📝 自己評価（解いた後に記入）