""" ============================================================================= 第 12 讲 · 代码 4：grokking——先背下来，很久之后才"顿悟"；以及顿悟之后网络用的是傅里叶变换 ============================================================================= 对应讲义：lectures/lecture12_mechanistic_interpretability.md 第 5 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- Power et al. (2022) 发现：在模加法这样的小算法任务上，网络会先把训练集背下来（训练准确率 100%，测试准确率接近 0）， 继续训练很久之后，测试准确率才突然升到接近 100%——"grokking"。Nanda et al. (2023) 打开了这样一个网络， 发现它用离散傅里叶变换和三角恒等式做加法，并且找到了在"顿悟"之前就在平稳变化的进度指标。 演示 1：一个小 MLP 学 (a + b) mod 59，只用 40% 的数对训练（AdamW，权重衰减 1.0，全 batch，30,000 步）： 训练 / 测试准确率、参数范数，每 2000 步一次。 演示 2：嵌入矩阵的傅里叶分析：能量集中在几个频率上吗？这个"集中度"是突然变化还是逐渐变化？ 演示 3：因果检验：只保留嵌入里的几个关键频率（去掉其余），测试准确率还剩多少？反过来，只去掉关键频率呢？ 演示 4：两个对照：没有权重衰减；训练数据只有 30%。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 04_grokking.py # 约 60 秒（训练 3 个小 MLP，各 30,000 步）；只需要 PyTorch 【numpy / PyTorch 小知识】 - 离散傅里叶变换：把一个长度为 p 的序列 f(0), …, f(p − 1) 写成若干个频率 k 的余弦、正弦之和： f(a) = c₀ + Σ_k [ C_k cos(2πka/p) + S_k sin(2πka/p) ] 频率 k 的"能量"是 C_k² + S_k²。这里把嵌入矩阵的每一列（一个维度在 p 个数字上的取值）都这样分解， 再把所有维度的能量加起来，就得到"嵌入在各个频率上的能量"。 """ import math import os import sys import torch import torch.nn as nn import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) import tiny_gpt # noqa: E402,F401 （导入时设置 torch.set_num_threads(1)） from console import pad, print_header # noqa: E402 P, D, WIDTH, STEPS = 59, 64, 128, 30000 class ModAddMLP(nn.Module): """ 学 (a + b) mod P 的小网络：两个数各查一个嵌入（同一张表），拼起来，过一层 ReLU 隐藏层，输出 P 个类别。 【思路】这是 grokking 研究里常用的最小模型之一（Nanda et al. 用的是 1 层 Transformer，本质上相似）。 嵌入表 W_E（P × D）是我们要做傅里叶分析的对象。 """ def __init__(self): super().__init__() self.E = nn.Embedding(P, D) self.W1 = nn.Linear(2 * D, WIDTH) self.W2 = nn.Linear(WIDTH, P) def forward(self, x, E=None): E = self.E.weight if E is None else E h = torch.cat([E[x[:, 0]], E[x[:, 1]]], dim=-1) return self.W2(F.relu(self.W1(h))) def split(frac, seed=0): """把全部 P × P 个数对随机分成训练集（比例 frac）和测试集。返回 (训练输入, 训练答案, 测试输入, 测试答案)。""" pairs = torch.tensor([(a, b) for a in range(P) for b in range(P)]) perm = torch.randperm(len(pairs), generator=torch.Generator().manual_seed(seed)) n = int(frac * len(pairs)) tr, te = pairs[perm[:n]], pairs[perm[n:]] return tr, (tr[:, 0] + tr[:, 1]) % P, te, (te[:, 0] + te[:, 1]) % P def fourier_power(E): """ 嵌入矩阵 E（P × D）在频率 k = 1 … (P − 1)/2 上的能量（所有维度加起来）。 【返回】长度 (P − 1)/2 的张量；第 k − 1 个元素是频率 k 的能量 """ a = torch.arange(P, dtype=torch.float32) out = [] for k in range(1, (P - 1) // 2 + 1): c = torch.cos(2 * math.pi * k * a / P) @ E s = torch.sin(2 * math.pi * k * a / P) @ E out.append((c.pow(2) + s.pow(2)).sum()) return torch.stack(out) def keep_frequencies(E, freqs, keep=True): """ 只保留（keep=True）或只去掉（keep=False）嵌入里频率在 freqs 中的那些傅里叶分量（常数项总是保留）。 【思路】f(a) = c₀ + Σ_k (2/P)[C_k cos + S_k sin]，其中 C_k = Σ_a f(a) cos(2πka/P)。按频率把这些项加回去或减掉。 """ a = torch.arange(P, dtype=torch.float32) out = E.mean(0, keepdim=True).expand_as(E).clone() for k in range(1, (P - 1) // 2 + 1): if (k in freqs) != keep: continue cos, sin = torch.cos(2 * math.pi * k * a / P), torch.sin(2 * math.pi * k * a / P) out += (2 / P) * (torch.outer(cos, cos @ E) + torch.outer(sin, sin @ E)) return out def train(frac=0.4, wd=1.0, seed=0, log_every=2000, snapshots=False): """ 全 batch 训练（AdamW，lr 1e-3，betas (0.9, 0.98)，和 Nanda et al. 相同的设置）。 【返回】(model, history)：history 每条是 dict(step, train_acc, test_acc, norm, top5, freqs) top5 是嵌入能量里最大的 5 个频率所占的比例，freqs 是那 5 个频率（snapshots=True 时才算） """ tr, ytr, te, yte = split(frac, seed) torch.manual_seed(seed) model = ModAddMLP() opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=wd, betas=(0.9, 0.98)) history = [] for step in range(STEPS + 1): if step % log_every == 0: with torch.no_grad(): rec = dict(step=step, train_acc=(model(tr).argmax(-1) == ytr).float().mean().item(), test_acc=(model(te).argmax(-1) == yte).float().mean().item(), norm=math.sqrt(sum(p.pow(2).sum().item() for p in model.parameters()))) if snapshots: pw = fourier_power(model.E.weight) top = torch.topk(pw, 5) rec["top5"] = (top.values.sum() / pw.sum()).item() rec["freqs"] = sorted((top.indices + 1).tolist()) history.append(rec) if step == STEPS: break loss = F.cross_entropy(model(tr), ytr) opt.zero_grad() loss.backward() opt.step() return model, history, (te, yte) def demo1_2(model, history): print_header(f"演示 1、2：学 (a + b) mod {P}，40% 的数对做训练集——准确率、参数范数、嵌入的傅里叶集中度") print(f"{pad('步', 7, '>')}{pad('训练准确率', 12, '>')}{pad('测试准确率', 12, '>')}{pad('参数范数', 10, '>')}" f"{pad('最强 5 个频率占的能量', 22, '>')}{pad('是哪 5 个频率', 26, '>')}") for r in history: print(f"{pad(r['step'], 7, '>')}{pad(format(r['train_acc'], '.3f'), 12, '>')}" f"{pad(format(r['test_acc'], '.3f'), 12, '>')}{pad(format(r['norm'], '.1f'), 10, '>')}" f"{pad(format(r['top5'], '.3f'), 22, '>')}{pad(str(r['freqs']), 26, '>')}") print(f"（嵌入有 {(P - 1) // 2} 个频率；如果能量均匀分布，5 个频率只占 {5 / ((P - 1) // 2):.3f}）") print("\n观察 Observations：") print(" 1. 训练准确率在 2000 步之内就到了 1.000，测试准确率却在 0 附近停留了好几千步：网络先把训练集**背**了下来。") print(" 2. 测试准确率从 4000–6000 步开始爬升（0.003 → 0.013 → 0.067），14000 步到 0.920，20000 步到 1.000——") print(" 这就是 grokking：背下来很久之后才'顿悟'。") print(" 3. 但'顿悟'并不是突然发生的：嵌入能量向最强 5 个频率集中的比例，从初始的 0.20 开始一路平稳上升——") print(" 测试准确率还接近 0 的时候（4000 步）已经到 0.24，测试准确率到 1 之后还在继续升（30000 步 0.61）；") print(" 最强的那几个频率从 10000 步起基本稳定。Nanda et al. 正是用这类'进度指标'说明：泛化电路在背后平稳地形成，") print(" 只是在它强到足以压过'背诵'之前，测试准确率看不出来。") print(" 4. 参数范数先升后降（2000 步 71.4 → 30000 步 51.8）：背诵阶段权重变大，之后权重衰减慢慢把用于背诵的部分'清理'掉——") print(" Nanda et al. 把训练分成三个阶段：背诵、电路形成、清理。") def demo3_causal(model, test): print_header("演示 3：因果检验——只保留 / 只去掉嵌入里的关键频率，测试准确率还剩多少") te, yte = test pw = fourier_power(model.E.weight.detach()) order = torch.argsort(pw, descending=True) with torch.no_grad(): base = (model(te).argmax(-1) == yte).float().mean().item() print(f"原模型：测试准确率 {base:.3f}") for n in (3, 5, 8): key = {int(i) + 1 for i in order[:n]} keep = (model(te, keep_frequencies(model.E.weight, key, True)).argmax(-1) == yte).float().mean().item() drop = (model(te, keep_frequencies(model.E.weight, key, False)).argmax(-1) == yte).float().mean().item() print(f"最强的 {n} 个频率 {sorted(key)}：只保留它们 → {keep:.3f}；只去掉它们 → {drop:.3f}") print(f"（随机猜的准确率是 1/{P} = {1 / P:.3f}）") print("\n观察 Observations：") print(" 1. 只保留最强的 5 个频率（把嵌入里其余 24 个频率的成分全部去掉），测试准确率还有 0.996；保留 8 个，1.000。") print(" 反过来，去掉这 8 个频率，准确率掉到 0.052（随机水平是 0.017）。网络做加法用的就是这几个频率——这是因果证据。") print(" 只去掉最强的 3 个频率，准确率还有 0.934：信息分散在 8 个左右的频率上，有一定的冗余。") print(" 2. Nanda et al. 的解释：把 a 表示成角度 ωa（ω = 2πk/P）的余弦和正弦，用三角恒等式算出 cos(ω(a + b)) 和 sin(ω(a + b))，") print(" 再对每个候选答案 c 算 cos(ω(a + b − c))——它在 c = a + b mod P 时最大。'加法'变成了'在圆上转动'。") print(" 他们的 1 层 Transformer（P = 113）用了 5 个关键频率；我们的小 MLP（P = 59）用的频率不同、也更分散，但做的是同一类事。") def demo4_controls(): print_header("演示 4：两个对照（同样 30,000 步）") for title, kw in [("没有权重衰减（wd = 0），40% 的数对", dict(wd=0.0)), ("有权重衰减，但只有 30% 的数对", dict(frac=0.3))]: _, h, _ = train(log_every=10000, **kw) print(f"{title}：" + "；".join(f"第 {r['step']} 步 训练 {r['train_acc']:.3f} / 测试 {r['test_acc']:.3f} / 范数 {r['norm']:.1f}" for r in h)) print("\n观察 Observations：") print(" 1. 没有权重衰减时，网络背下训练集之后就没有动力去找更简单的解：参数范数一路涨到 294，测试准确率始终是 0。") print(" 权重衰减让'用很多参数背诵'变得昂贵，'用几个频率的通用算法'变得划算——又是第 11 讲的最小描述长度。") print(" 2. 训练数据只有 30% 时，30,000 步结束时测试准确率只有 0.021，似乎刚刚开始爬升：数据越少，背诵越便宜，顿悟来得越晚。") print(" grokking 对超参数非常敏感——这也是为什么它在大模型的真实训练里很少被直接看到。") if __name__ == "__main__": model, history, test = train(snapshots=True) demo1_2(model, history) demo3_causal(model, test) demo4_controls()