""" ============================================================================= 第 12 讲 · 代码 1：残差流、logit lens 与直接 logit 归因——第 1 讲的 base model 是怎样"想起"一个人住在哪里的 ============================================================================= 对应讲义：lectures/lecture12_mechanistic_interpretability.md 第 2 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 第 1 讲的 base model 读到 "kari lives in " 时，会给 kari 所住城市的首字母很高的概率。这个预测是在网络的哪一步形成的？ 是哪些部件（嵌入、哪个注意力头、哪个 MLP）把它"写"进去的？ 演示 1：残差流的视角：一个 2 层 Transformer 的输出 = 嵌入 + 注意力₀ + MLP₀ + 注意力₁ + MLP₁（逐项相加）。 演示 2：logit lens：把每一层之后的残差流直接接到最后的 LayerNorm 和输出层上，看"如果在这里停下，模型会预测什么"—— 60 个人，正确城市首字母的平均概率和 top-1 准确率怎样逐层变化。 演示 3：直接 logit 归因（direct logit attribution）：把"正确首字母 vs 其他 7 个城市首字母"的 logit 差 分解到每个部件、每个注意力头上——谁贡献最大。 演示 4：贡献最大的那几个头在看哪里：最后一个位置的注意力落在名字上吗？ ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 01_logit_lens.py # 约 3 秒；需要 PyTorch 依赖 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py（第 1 讲的 base model 只读加载，不做任何修改）。 【PyTorch 小知识】 - 我们不改模型，而是把它的前向"手写"一遍，把每一步的中间结果留下来——这是可解释性研究里最常用的办法 （更方便的工具有 TransformerLens 一类的库，用"钩子"自动保存中间结果）。 - F.layer_norm(x, (d,), weight, bias)：和 nn.LayerNorm 的计算一样，只是参数由调用者传入。 """ import math import os import sys import torch import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import pad, print_header # noqa: E402 from fm_world import ALPHABET, CITIES, build_world # noqa: E402 from tiny_gpt import CharTokenizer, load_or_train_base_model # noqa: E402 @torch.no_grad() def decompose(model, idx, ablate=(), mean_y=None): """ 手写 base model 的前向，返回最后一个位置上残差流的各个组成部分。 【思路】pre-LN Transformer 的每一层都是 x ← x + 某个部件(LN(x))，所以最后的残差流 = 所有部件输出之和： x_final = 嵌入 + Σ_层 (注意力输出 + MLP 输出) 注意力输出还可以再拆成"每个头的贡献 + 输出投影的偏置"：输出投影 proj 是线性的， proj(拼接的各头输出) = Σ_头 W_proj[:, 这个头的那几列] · 头的输出 + 偏置。 【参数】model：TinyGPT；idx：(1, T) 的字符编号；ablate：要消融的 (层, 头) 集合——这些头的输出在所有位置都被替换； mean_y：None 时"零消融"（置零）；否则是 dict (层, 头) → 这个头在整个数据集上的平均输出（dh 维），"均值消融" 【返回】dict： "parts" 有序 dict 名字 → 最后一个位置的向量（d,），各项之和 = 最后的残差流 "stages" list[(名字, 向量)]：嵌入之后、每层之后的残差流（logit lens 用） "attn" list：每一层的注意力权重 (H, T, T) """ T = idx.shape[1] x = model.tok_emb(idx) + model.pos_emb(torch.arange(T)) parts = {"嵌入（token + 位置）": x[0, -1].clone()} stages = [("嵌入之后", x[0, -1].clone())] attns, ys = [], [] for li, blk in enumerate(model.blocks): att_mod = blk.attn h = blk.ln1(x) B, T_, D = h.shape H = att_mod.n_head q, k, v = att_mod.qkv(h).split(D, dim=2) q, k, v = (t.view(B, T_, H, D // H).transpose(1, 2) for t in (q, k, v)) s = q @ k.transpose(-2, -1) / math.sqrt(D // H) a = s.masked_fill(~att_mod.mask[:T_, :T_], float("-inf")).softmax(-1) attns.append(a[0]) y = a @ v # (1, H, T, dh) for (al, ah) in ablate: if al == li: y = y.clone() y[:, ah] = 0.0 if mean_y is None else mean_y[(al, ah)] # 零消融 / 均值消融：这个头不再携带任何关于输入的信息 ys.append(y[0].detach()) Wp = att_mod.proj.weight # (D, D) dh = D // H for hh in range(H): parts[f"第 {li} 层 注意力头 {hh}"] = (Wp[:, hh * dh:(hh + 1) * dh] @ y[0, hh, -1]) parts[f"第 {li} 层 注意力输出偏置"] = att_mod.proj.bias.clone() x = x + att_mod.proj(y.transpose(1, 2).reshape(B, T_, D)) mlp_out = blk.mlp(blk.ln2(x)) parts[f"第 {li} 层 MLP"] = mlp_out[0, -1].clone() x = x + mlp_out stages.append((f"第 {li} 层之后", x[0, -1].clone())) total = sum(parts.values()) assert torch.allclose(total, x[0, -1], atol=1e-5) # 各部件之和 = 最后的残差流 out_logits = model.head(model.ln_f(x[0, -1])) return dict(parts=parts, stages=stages, attn=attns, final=x[0, -1], logits=out_logits, y=ys) def lens_logits(model, vec): """logit lens：把残差流向量直接接到最后的 LayerNorm 和输出层上。""" return model.head(model.ln_f(vec)) def prompts(world): """每个人一个提示 "\\n<名字> lives in "，答案是城市的首字母（8 个城市的首字母互不相同）。""" return [(p.name, f"\n{p.name} lives in ", p.city[0]) for p in world.people] def demo1_2_lens(model, tok, world): print_header("演示 1、2：残差流与 logit lens（60 个人的提示 \"<名字> lives in \"）") letters = sorted({c[0] for c in CITIES}) city_ids = [tok.stoi[c] for c in letters] stage_names = None probs, accs, city_mass = {}, {}, {} for name, prompt, ans in prompts(world): out = decompose(model, torch.tensor([tok.encode(prompt)])) stage_names = [n for n, _ in out["stages"]] for sname, vec in out["stages"]: p = lens_logits(model, vec).softmax(-1) probs.setdefault(sname, []).append(p[tok.stoi[ans]].item()) accs.setdefault(sname, []).append(float(p.argmax().item() == tok.stoi[ans])) city_mass.setdefault(sname, []).append(p[city_ids].sum().item()) print("最后的残差流 = 嵌入 + 第 0 层注意力 + 第 0 层 MLP + 第 1 层注意力 + 第 1 层 MLP（逐项相加，代码里断言过）\n") print(f"{pad('在这里停下', 14)}{pad('正确首字母的平均概率', 22, '>')}{pad('top-1 准确率', 14, '>')}{pad('8 个城市首字母的总概率', 24, '>')}") for s in stage_names: avg_p, acc, mass = (sum(v) / len(v) for v in (probs[s], accs[s], city_mass[s])) print(f"{pad(s, 14)}{pad(f'{avg_p:.3f}', 22, '>')}{pad(f'{acc:.3f}', 14, '>')}{pad(f'{mass:.3f}', 24, '>')}") print("\n观察 Observations：") print(" 1. 只看嵌入（相当于只知道'当前字符是空格、在第几个位置'），8 个城市首字母的总概率是 0：模型对城市一无所知。") print(" 2. 第 0 层之后，城市首字母的总概率升到 0.115，但正确的那个只有 0.014（约等于 0.115 ÷ 8）：") print(" '接下来可能是一个城市'开始出现，'是哪个城市'还完全不知道。") print(" 第 1 层之后，城市首字母的总概率 0.993、正确的那个 0.855（top-1 准确率 0.933）：具体的事实是在最后一层才被挑出来的。") print(" '类别'先出现，'具体事实'后出现——大模型的 logit lens 研究里也常看到类似的由粗到细。") print(" 3. logit lens 是一个粗糙的工具：中间层的残差流并不是为'直接接输出层'而训练的。") print(" 后来的 tuned lens（Belrose et al. 2023，Stella Biderman 是作者之一）给每一层学一个小的线性变换再接输出层，更可靠。") def demo3_attribution(model, tok, world): print_header("演示 3：直接 logit 归因——'正确首字母 − 其他 7 个城市首字母的平均'这个 logit 差，每个部件贡献多少") letters = sorted({c[0] for c in CITIES}) totals, n = {}, 0 for name, prompt, ans in prompts(world): out = decompose(model, torch.tensor([tok.encode(prompt)])) final = out["final"] # LayerNorm 的"冻结"线性化：减均值、除以整个残差流的标准差（对这一个样本固定），再乘增益 mu, sigma = final.mean(), final.var(unbiased=False).add(model.ln_f.eps).sqrt() W = model.head.weight # (V, D) direction = W[tok.stoi[ans]] - torch.stack([W[tok.stoi[c]] for c in letters if c != ans]).mean(0) for pname, vec in out["parts"].items(): contrib = (direction * model.ln_f.weight * (vec - vec.mean()) / sigma).sum().item() totals[pname] = totals.get(pname, 0.0) + contrib n += 1 del mu total_all = sum(totals.values()) / n print(f"{pad('部件', 26)}{pad('平均贡献（logit 差）', 22, '>')}{pad('占总量', 10, '>')}") for pname, v in sorted(totals.items(), key=lambda kv: -kv[1]): print(f"{pad(pname, 26)}{pad(f'{v / n:+.3f}', 22, '>')}{pad(f'{v / n / total_all:.0%}', 10, '>')}") print(f"{pad('合计', 26)}{pad(f'{total_all:+.3f}', 22, '>')}") print("\n观察 Observations：") print(" 1. 归因把一个预测拆成了各部件的'投票'。第 1 层的 MLP 一项就占了 98%；") print(" 所有注意力头的**直接**贡献都接近 0，嵌入也几乎不投票（它只知道当前字符是空格）。") print(" 2. 这是一个**线性化**的分析：LayerNorm 的除法被冻结成一个常数，部件之间的相互作用没有算进去") print(" （比如第 1 层的 MLP 能工作，是因为第 0 层先把名字的信息搬到了这个位置）。") print(" 要证明因果关系，还需要消融或修补（第 6 讲 §6 的因果追踪；代码 03 在 induction 模型上做）。") return totals def demo4_attention(model, tok, world): print_header("演示 4：每个注意力头在最后一个位置上看哪里，以及消融它之后模型还记得多少") ps = prompts(world) base_p = sum(decompose(model, torch.tensor([tok.encode(pr)]))["logits"].softmax(-1)[tok.stoi[a]].item() for _, pr, a in ps) / len(ps) # 每个头在整个数据集（60 个提示的所有位置）上的平均输出，用于"均值消融" acc = {} for _, pr, _ in ps: for li, y in enumerate(decompose(model, torch.tensor([tok.encode(pr)]))["y"]): for hh in range(y.shape[0]): acc.setdefault((li, hh), []).append(y[hh]) mean_y = {k: torch.cat(v).mean(0) for k, v in acc.items()} print(f"不消融：正确首字母的平均概率 {base_p:.3f}\n") print(f"{pad('注意力头', 14)}{pad('看名字的注意力', 16, '>')}{pad('零消融后', 12, '>')}{pad('均值消融后', 12, '>')}{pad('均值消融的变化', 16, '>')}") rows = [] for li in range(len(model.blocks)): for hh in range(model.blocks[li].attn.n_head): on_name, p_zero, p_mean = 0.0, 0.0, 0.0 for name, prompt, ans in ps: idx = torch.tensor([tok.encode(prompt)]) a = decompose(model, idx)["attn"][li][hh, -1] on_name += a[1:1 + len(name)].sum().item() # 位置 0 是开头的换行，名字从位置 1 开始 p_zero += decompose(model, idx, ablate={(li, hh)})["logits"].softmax(-1)[tok.stoi[ans]].item() p_mean += decompose(model, idx, ablate={(li, hh)}, mean_y=mean_y)["logits"].softmax(-1)[tok.stoi[ans]].item() k = len(ps) rows.append((li, hh, on_name / k, p_mean / k)) delta = f"{p_mean / k - base_p:+.3f}" print(f"{pad(f'第 {li} 层 头 {hh}', 14)}{pad(f'{on_name / k:.2f}', 16, '>')}{pad(f'{p_zero / k:.3f}', 12, '>')}" f"{pad(f'{p_mean / k:.3f}', 12, '>')}{pad(delta, 16, '>')}") gather = [] for hh in range(model.blocks[0].attn.n_head): tot = 0.0 for name, prompt, _ in ps: a = decompose(model, torch.tensor([tok.encode(prompt)]))["attn"][0][hh] tot += a[len(name), 1:len(name)].sum().item() # 名字最后一个字符的位置 → 名字前面的字符 gather.append(tot / len(ps)) print("\n第 0 层各头在'名字的最后一个字符'处，放在名字前面几个字符上的注意力：" + "，".join(f"头 {h} {g:.2f}" for h, g in enumerate(gather))) worst = min(rows, key=lambda r: r[3]) name, prompt, ans = ps[0] row = decompose(model, torch.tensor([tok.encode(prompt)]))["attn"][worst[0]][worst[1], -1] shown = " ".join(f"{ch}:{w:.2f}" for ch, w in zip(prompt, row.tolist()) if w >= 0.05 and ch != "\n") print(f"\n例子：{prompt!r}（答案 {ans}…）第 {worst[0]} 层头 {worst[1]} 在最后一个位置的注意力（只列 ≥ 0.05 的）：{shown}") print("\n观察 Observations：") print(" 1. 演示 3 里所有注意力头的直接贡献都接近 0，但消融揭示了它们的作用：每一个头被消融，正确首字母的概率都下降。") print(" 影响最大的是第 1 层头 0（−0.548）：它在最后一个位置把 85% 的注意力放在名字的字符上，") print(" 把'这个人是谁'的信息搬过来，第 1 层的 MLP 再用它查出城市。第 0 层的头在名字的最后一个字符处") print(" 看名字前面的字符（0.58–0.72）：先在名字内部把几个字符汇总起来。直接归因只能看到'最后一步'，消融才能看到间接作用。") print(" 2. 零消融和均值消融的结论基本一致。但要注意，这个 2 层模型里信息流是**分布式**的——") print(" 没有哪一个头是唯一的'负责人'，也有看起来不看名字、消融后却影响很大的头（第 1 层头 1，−0.362）；") print(" 要把它完全讲清楚，需要更细的路径修补（path patching），这里不展开。") print(" 3. 这和第 6 讲 §6 因果追踪的结论一致：事实的'回忆'发生在最后一个位置的 MLP 里，名字的信息要先被注意力搬过来。") print(" 4. 这个模型只有 2 层、8 个头，已经需要好几个部件配合；真实的大模型有几千个头、几十层，") print(" 这就是为什么机制可解释性要先从极小的模型、极简单的任务开始（代码 02、03）。") if __name__ == "__main__": world, tok = build_world(), CharTokenizer(ALPHABET) base, _ = load_or_train_base_model(world, tok) demo1_2_lens(base, tok, world) demo3_attribution(base, tok, world) demo4_attention(base, tok, world)