Day 127 — SOTA手法の追跡・論文実装② — 汎化性能を高めるOptimizerの工夫（SAM: Sharpness-Aware Minimization）

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

⛰️ SAMの2段階更新（今日のワークフロー）

今日はこのワークフローを、少ない訓練データで過学習が起きやすい合成データを題材に体験する。

🎯 問題

Day126と同様の合成分類データ（ただし今回は訓練データを少なめにして過学習が起きやすい状況を作る）に対し、通常のAdamとSAM+Adamをそれぞれ実装して訓練し、テスト精度と「損失地形の鋭さ（sharpness）」の両方を比較します。

import numpy as np
from sklearn.datasets import make_classification

# --- 合成データ生成（5クラス分類。訓練データを意図的に少なくして過学習を誘発） ---
X, y = make_classification(
    n_samples=1200, n_features=20, n_informative=12,
    n_classes=5, n_clusters_per_class=1, random_state=42,
    flip_y=0.05,
)

TRAIN_LEN = 300   # 訓練データはわずか300件

💡 ヒント

import torch

def sam_step(model, loss_fn, x, y, base_optimizer, rho=0.05, eps=1e-12):
    # --- ① 1回目の forward/backward（現在地の勾配を取得） ---
    base_optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()

    # --- ② e_hat を計算してパラメータに加算（摂動） ---
    grad_norm = torch.norm(
        torch.stack([p.grad.norm(2) for p in model.parameters() if p.grad is not None])
    )
    e_hats = {}
    with torch.no_grad():
        for p in model.parameters():
            if p.grad is None: continue
            e_hat = rho * p.grad / (grad_norm + eps)
            p.add_(e_hat)
            e_hats[p] = e_hat

    # --- ③ 摂動後の地点で 2回目の forward/backward ---
    base_optimizer.zero_grad()
    loss2 = loss_fn(model(x), y)
    loss2.backward()

    # --- ④ 元のパラメータへ戻してから、摂動後の勾配で更新 ---
    with torch.no_grad():
        for p in model.parameters():
            if p in e_hats: p.sub_(e_hats[p])
    base_optimizer.step()
    return loss.item()

✅ 模範解答

記号↔変数の対応表

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

torch.manual_seed(0)

# --- サンプルデータ生成（問題文と同じ。TRAIN_LENを小さくして過学習を誘発） ---
X, y = make_classification(
    n_samples=1200, n_features=20, n_informative=12,
    n_classes=5, n_clusters_per_class=1, random_state=42, flip_y=0.05,
)
TRAIN_LEN = 300

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)

class MLP(nn.Module):
    def __init__(self, in_dim=20, hidden=32, num_classes=5):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, num_classes),
        )
    def forward(self, x):
        return self.net(x)

loss_fn = nn.CrossEntropyLoss()

# --- ステップ1: SAMの数式2本をそのままコードに翻訳 ---
def sam_step(model, loss_fn, x, y, base_optimizer, rho=0.05, eps=1e-12):
    """e_hat(w) = rho * grad(L(w)) / ||grad(L(w))||_2
       w_{t+1} = w_t - eta * grad(L(w + e_hat(w)))|_{w=w_t}"""
    base_optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()

    grad_norm = torch.norm(
        torch.stack([p.grad.norm(2) for p in model.parameters() if p.grad is not None])
    )
    e_hats = {}
    with torch.no_grad():
        for p in model.parameters():
            if p.grad is None: continue
            e_hat = rho * p.grad / (grad_norm + eps)
            p.add_(e_hat)
            e_hats[p] = e_hat

    base_optimizer.zero_grad()
    loss2 = loss_fn(model(x), y)
    loss2.backward()

    with torch.no_grad():
        for p in model.parameters():
            if p in e_hats: p.sub_(e_hats[p])
    base_optimizer.step()
    return loss.item()

# --- ステップ2/3: baseline(Adam) と SAM+Adam の学習ループ ---
def train_baseline(epochs=100, lr=0.01):
    torch.manual_seed(0)
    model = MLP()
    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 = loss_fn(model(xb), yb)
            loss.backward()
            optimizer.step()
    return model

def train_sam(epochs=100, lr=0.01, rho=0.05):
    torch.manual_seed(0)
    model = MLP()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_loader:
            sam_step(model, loss_fn, xb, yb, optimizer, rho=rho)
    return model

model_baseline = train_baseline()
model_sam = train_sam(rho=0.05)

# --- ステップ4: test精度の比較 ---
def evaluate(model):
    model.eval()
    with torch.no_grad():
        preds = model(X_test).argmax(dim=1)
    return (preds == y_test).float().mean().item()

acc_base = evaluate(model_baseline)
acc_sam = evaluate(model_sam)
print(f"baseline (Adam)      test accuracy = {acc_base:.3f}")
print(f"SAM+Adam (rho=0.05)  test accuracy = {acc_sam:.3f}")
# 出力例（乱数・実行環境により変動）:
# baseline (Adam)      test accuracy = 0.681
# SAM+Adam (rho=0.05)  test accuracy = 0.712

# --- ステップ5: 簡易Sharpness測定（重みへのランダム摂動でtrain lossの増分を見る） ---
def measure_sharpness(model, magnitudes=(0.0, 0.02, 0.05, 0.1), n_trials=5, seed=0):
    model.eval()
    torch.manual_seed(seed)
    base_loss = loss_fn(model(X_train), y_train).item()
    results = {}
    for mag in magnitudes:
        increases = []
        for _ in range(n_trials):
            original = {p: p.data.clone() for p in model.parameters()}
            with torch.no_grad():
                for p in model.parameters():
                    p.add_(torch.randn_like(p) * mag)
            perturbed_loss = loss_fn(model(X_train), y_train).item()
            with torch.no_grad():
                for p in model.parameters():
                    p.data.copy_(original[p])
            increases.append(perturbed_loss - base_loss)
        results[mag] = float(np.mean(increases))
    return results

sharpness_base = measure_sharpness(model_baseline)
sharpness_sam = measure_sharpness(model_sam)
# 出力例の傾向:
# baseline: magnitude=0.10 で loss増分 +0.85 前後 → 鋭い谷
# SAM:      magnitude=0.10 で loss増分 +0.31 前後 → 平らな谷

📊 Sharpnessの可視化（摂動の大きさ vs train lossの増分）

🪜 Step-by-Step 解説

1SAMの数式2本を、勾配計算2回のコードに対応づける

grad_norm = torch.norm(torch.stack([p.grad.norm(2) for p in model.parameters() if p.grad is not None]))
e_hat = rho * p.grad / (grad_norm + eps)
p.add_(e_hat)

2なぜ「摂動→勾配計算→元に戻す→本番更新」という順序が必要か

with torch.no_grad():
    for p in model.parameters():
        if p in e_hats: p.sub_(e_hats[p])
base_optimizer.step()

3rho=0のときSAMが通常のAdamと一致することを確認する

4Sharpness測定の設計思想

p.add_(torch.randn_like(p) * mag)
perturbed_loss = loss_fn(model(X_train), y_train).item()

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

🏆 Kaggleでの実践的な使い方

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

⚠️ よくある誤解・ミス

🚀 次のステップ

• 発展: 今日実装したsam_stepのrhoを[0.01, 0.05, 0.1, 0.2]のように複数試し、test精度とSharpnessがrhoに対してどう変化するかを折れ線でプロットしてみましょう。「摂動の大きさとしてどこが最適か」という感覚を、数値で掴めるようになります
• 次回予告: Day 128 — SOTA手法の追跡・論文実装③。引き続きPhase6テーマ1の枠内で、今日と同じ「論文の数式→実装」ワークフローを別の題材（Exponential Moving Average（EMA）による重みの安定化）に適用し、論文を読んで実装に落とし込む力をもう一段鍛えます

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

Day126でLabel Smoothing（損失関数の改造）、Day127でSAM（Optimizerの改造）と、「論文→実装」ワークフローを異なる題材で2回反復。次はさらに別のSOTA手法へ。

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