准确率长期卡在 70%？把 weight_decay 用成 L2、给 Norm/偏置也衰减、与 AdamW 混用的三连坑

# adamw_decay_pitfalls.py
import argparse, math, torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(0)

class TinyTransformer(nn.Module):
    def __init__(self, d=128, ncls=5, nhead=4, nl=2):
        super().__init__()
        self.embed = nn.Linear(20, d)
        encoder_layer = nn.TransformerEncoderLayer(d_model=d, nhead=nhead, dim_feedforward=256, batch_first=True, norm_first=True)
        self.enc = nn.TransformerEncoder(encoder_layer, num_layers=nl)
        self.pool = nn.AdaptiveAvgPool1d(1)
        self.head = nn.Linear(d, ncls)
    def forward(self, x):              # x: [B,20] → 伪装成 5 个 token
        x = x.view(x.size(0), 5, 4)    # [B,5,4]
        x = self.embed(x)              # [B,5,d]
        x = self.enc(x)                # [B,5,d]
        x = x.transpose(1,2)           # [B,d,5]
        x = self.pool(x).squeeze(-1)   # [B,d]
        return self.head(x)

def make_loader(n=6000, bs=128):
    X = torch.randn(n, 20)
    W = torch.randn(20, 5); b = torch.randn(5)
    y = (X @ W + b).argmax(dim=1)      # 线性可分 + 轻微扰动
    ds = torch.utils.data.TensorDataset(X, y)
    return torch.utils.data.DataLoader(ds, batch_size=bs, shuffle=True, drop_last=True)

def split_decay(model):
    decay, no_decay = [], []
    for n, p in model.named_parameters():
        if p.requires_grad is False: 
            continue
        is_norm = any(k in n.lower() for k in ["norm", "bn"])
        if n.endswith(".bias") or is_norm:
            no_decay.append(p)
        else:
            decay.append(p)
    return [
        {"params": decay, "weight_decay": 0.01},
        {"params": no_decay, "weight_decay": 0.0},
    ]

def run(bug=True, steps=400):
    device = "cpu"
    model = TinyTransformer().to(device)
    train_loader = make_loader(n=6000, bs=128)
    val_loader   = make_loader(n=2000, bs=256)

    if bug:
        # 错误：Adam + weight_decay（耦合 L2）且不排除 Norm/偏置
        opt = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=1e-2)
    else:
        # 正确：AdamW（解耦）+ param group 排除 Norm/偏置
        opt = torch.optim.AdamW(split_decay(model), lr=3e-4, weight_decay=1e-2)

    def one_epoch(loader, train=True):
        model.train(train)
        total, correct, loss_sum, n = 0, 0, 0.0, 0
        for x, y in loader:
            logits = model(x.to(device))
            loss = F.cross_entropy(logits, y.to(device))
            if train:
                opt.zero_grad(set_to_none=True)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                opt.step()
            with torch.no_grad():
                pred = logits.argmax(1)
                correct += (pred.cpu() == y).sum().item()
                total   += y.numel()
                loss_sum += float(loss.item()) * y.size(0)
                n += y.size(0)
        return loss_sum / n, correct / total

    # 训练若干 epoch 并周期评估
    for ep in range(1, 11):
        tr_loss, tr_acc = one_epoch(train_loader, train=True)
        va_loss, va_acc = one_epoch(val_loader, train=False)
        # 监控：Norm 权重的均值范数
        with torch.no_grad():
            norm_ws = [p for n,p in model.named_parameters() if ("norm" in n.lower()) and p.ndim>=1]
            norm_mean = float(torch.stack([w.norm() for w in norm_ws]).mean()) if norm_ws else 0.0
        tag = "BUG" if bug else "FIX"
        print(f"[{tag}] epoch={ep:02d} train_acc={tr_acc*100:5.1f}% val_acc={va_acc*100:5.1f}% norm|w|≈{norm_mean:.4f}")

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--bug", choices=["on","off"], default="on")
    args = ap.parse_args()
    print("== 错误设置 =="); run(bug=True)
    print("\n== 正确设置 =="); run(bug=False)

optimizer = torch.optim.AdamW(
    [
        {"params": decay_params, "weight_decay": wd},
        {"params": no_decay_params, "weight_decay": 0.0},
    ],
    lr=lr, betas=(0.9, 0.999), eps=1e-8
)

def build_param_groups(model, wd=0.01):
    decay, no_decay = [], []
    for n, p in model.named_parameters():
        if p.requires_grad is False: continue
        if n.endswith(".bias") or "norm" in n.lower() or "bn" in n.lower():
            no_decay.append(p)
        else:
            decay.append(p)
    return [
        {"params": decay, "weight_decay": wd},
        {"params": no_decay, "weight_decay": 0.0},
    ]

def assert_param_groups(optimizer):
    has_decay = any(g.get("weight_decay", 0) > 0 for g in optimizer.param_groups)
    has_nodecay = any(g.get("weight_decay", 0) == 0 for g in optimizer.param_groups)
    assert has_decay and has_nodecay, "缺少衰减/不衰减分组，请检查"

def summarize_decay(model, optimizer):
    names = {id(p): n for n,p in model.named_parameters()}
    for i, g in enumerate(optimizer.param_groups):
        wd = g.get("weight_decay", 0)
        cnt = len(g["params"])
        some = [names.get(id(p), "<?>") for p in g["params"][:5]]
        print(f"[group {i}] wd={wd} count={cnt} e.g. {some}")

def norm_weight_health(model):
    with torch.no_grad():
        norms = [p.norm().item() for n,p in model.named_parameters() if "norm" in n.lower()]
        if norms:
            print(f"[health] norm|w| mean={sum(norms)/len(norms):.4f}")