""" ============================================================================= CS486 · 第 3 讲 · 代码 01：打开思维链模型——"一跳查找"的电路 ============================================================================= 【这个脚本要回答的问题】（讲义 §4） 第 1 讲的 cot_L2（2 层、4 头）能 100% 地写出 1–7 跳的推理链。我们从外面知道"它会了"， 这一讲打开它，看它**怎样**做每一跳。第 1 讲 §3.5 猜过一个机制： 第 1 层：每条陈述 "x>y" 右边的 y "看一眼"左边 2 个位置的 x，把"我的来源是 x"记下来（绑定）； 第 2 层：在当前字母 c 的位置，去找"来源是 c"的那个 y，把 y 搬过来（查找）。 这个猜测对吗？ 演示 1：注意力模式——在"要预测下一个推理字母"的位置上，每个头把注意力放在哪里？ 演示 2：第 1 层的头在题面里每条陈述的右边字母处，看向哪里？ 演示 3：一个例子的注意力"热图"（文字版）。 演示 4：头消融——关掉一个或几个头，逐步查找的准确率还剩多少？ 演示 5：因果干预——改动题面里的某一个字母，模型的预测会不会跟着改？ 【运行】 python3 01_lookup_circuit.py # 约 15 秒（需要第 1 讲缓存的 cot_L2；没有缓存会先训练约 35 秒） python3 01_lookup_circuit.py --quick # 冒烟测试：用一个只训练 60 步的模型，数字没有意义 【依赖】PyTorch + 本目录的 cot_lab.py + 第 1 讲的模块。 """ import argparse import os import random import sys import time import torch HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) import cot_lab as cl # noqa: E402 import hop_world as hw # noqa: E402 import tiny_reasoner as tr # noqa: E402 from console import pct, print_header, print_observations, print_table # noqa: E402 CATS = ("匹配陈述的左边", "匹配陈述的右边（答案）", "当前字母作为别处的右边", "跳数数字", "前一个位置", "自己", "其他") @torch.no_grad() def attention_at_steps(model, probs): """ 在每个"要预测下一个推理字母"的位置上，统计每个头的注意力落在哪几类位置上。 【思路】设当前字母是 c。题面里以 c 开头的陈述 "c>y" 叫"匹配陈述"： 它的左边（c 所在的位置）、右边（y 所在的位置——就是正确答案）； c 还会作为另一条陈述 "x>c" 的右边出现；再加上"前一个位置""自己"和"其他"。 陈述 j 的三个字符在题面里的位置是 4j、4j+1、4j+2（每条陈述加分号占 4 个字符）。 【返回】dict：(层, 头) → {类别: 平均注意力} """ L, H = model.cfg.n_layers, model.cfg.n_heads sums = {(l, h): dict.fromkeys(CATS, 0.0) for l in range(L) for h in range(H)} n = 0 for p in probs: inp, tgt = hw.format_example(p, "cot") _, attns = cl.forward_with_head_mask(model, torch.tensor([hw.encode(inp + tgt)])) for t, cur in cl.trace_positions(p): j = next(i for i, st in enumerate(p.statements) if st[0] == cur) jt = next(i for i, st in enumerate(p.statements) if st[2] == cur) digit = len(p.prompt) - 3 # 题面结尾是 ";k?s"，数字 k 在倒数第 3 个位置 pos = {CATS[0]: 4 * j, CATS[1]: 4 * j + 2, CATS[2]: 4 * jt + 2, CATS[3]: digit, CATS[4]: t - 1, CATS[5]: t} for (l, h), d in sums.items(): row = attns[l][0, h, t] used = 0.0 for c, q in pos.items(): d[c] += row[q].item() used += row[q].item() d[CATS[6]] += 1.0 - used n += 1 return {key: {c: v / n for c, v in d.items()} for key, d in sums.items()} @torch.no_grad() def layer0_at_targets(model, probs): """ 第 1 层（layer 0）的每个头，在题面里每条陈述右边字母的位置上，看向哪里。 【返回】dict：头 → {"左边 2 格（来源）", "左边 1 格（'>'）", "自己", "其他"} 的平均注意力 """ H = model.cfg.n_heads keys = ("左边 2 格（来源）", "左边 1 格（'>'）", "自己", "其他") sums = {h: dict.fromkeys(keys, 0.0) for h in range(H)} n = 0 for p in probs: _, attns = cl.forward_with_head_mask(model, torch.tensor([hw.encode(p.prompt)])) for j in range(len(p.statements)): t = 4 * j + 2 for h in range(H): row = attns[0][0, h, t] a2, a1, a0 = row[t - 2].item(), row[t - 1].item(), row[t].item() for kname, v in zip(keys, (a2, a1, a0, 1 - a2 - a1 - a0)): sums[h][kname] += v n += 1 return {h: {kname: v / n for kname, v in d.items()} for h, d in sums.items()} def demo1(model, probs): print_header("演示 1：在'要预测下一个推理字母'的位置上，每个头的注意力落在哪里（n = 8、k = 7，200 道题 × 7 步）") stats = attention_at_steps(model, probs) rows = [[f"第 {l + 1} 层 头 {h}"] + [pct(stats[(l, h)][c], 0) for c in CATS] for (l, h) in sorted(stats)] print_table(["头"] + list(CATS), rows, [14, 14, 22, 22, 8, 10, 6, 6]) print("（第 1 步的'前一个位置'就是 '?'，'自己'是起点字母；某一位置同时属于两类时（例如前一个位置恰好是数字）会被重复计入。）") return stats def demo2(model, probs): print_header("演示 2：第 1 层的头，在题面里每条陈述'右边字母'的位置上看向哪里") stats = layer0_at_targets(model, probs) keys = list(stats[0].keys()) rows = [[f"第 1 层 头 {h}"] + [pct(stats[h][k], 0) for k in keys] for h in sorted(stats)] print_table(["头"] + keys, rows, [14, 18, 18, 8, 8]) return stats def demo3(model, stats1, stats2): print_header("演示 3：一个例子——最'像查找'的头在起点位置的注意力（文字版热图）") lookup = max(stats1, key=lambda key: stats1[key][CATS[1]]) binder = max(stats2, key=lambda h: stats2[h]["左边 2 格（来源）"]) p = hw.make_eval_set(8, 3, 1, seed=13)[0] with torch.no_grad(): _, attns = cl.forward_with_head_mask(model, torch.tensor([hw.encode(p.prompt)])) t = len(p.prompt) - 1 row = attns[lookup[0]][0, lookup[1], t] glyph = "".join(" .:*#"[min(4, int(row[i].item() * 5))] for i in range(len(p.prompt))) print(f"题面 {p.prompt}") print(f"注意力 {glyph} ← 第 {lookup[0] + 1} 层 头 {lookup[1]}，在最后一个字符（起点 {p.start}）处") print(f"（' ' < 0.2 ≤ '.' < 0.4 ≤ ':' < 0.6 ≤ '*' < 0.8 ≤ '#'） 正确的下一个字母：{p.walk[1]}" f"（在陈述 {p.start}>{p.walk[1]} 的右边）") return lookup, binder @torch.no_grad() def demo4(model, probs, lookup_heads, binder): print_header("演示 4：头消融——关掉某些头之后，教师强制下逐步查找的准确率（n = 8、k = 7，200 道题）") L, H = model.cfg.n_layers, model.cfg.n_heads means = cl.head_output_means(model, probs[:100]) def run(heads): mask = torch.ones(L, H) for l, h in heads: mask[l, h] = 0 zero = cl.step_accuracy(model, probs, mask) mean = cl.step_accuracy(model, probs, replace={key: means[key] for key in heads}) return zero, mean full = cl.step_accuracy(model, probs) rows = [["不消融（原模型）", pct(full), pct(full)]] out = {"full": full, "single": {}} for l in range(L): for h in range(H): out["single"][(l, h)] = run([(l, h)]) rows.append([f"第 {l + 1} 层 头 {h}", pct(out["single"][(l, h)][0]), pct(out["single"][(l, h)][1])]) out["pair"] = run(lookup_heads) names = " 和 ".join(f"头 {h}" for _, h in lookup_heads) rows.append([f"第 2 层 {names}（两个查找头）", pct(out["pair"][0]), pct(out["pair"][1])]) print_table(["消融的头", "零消融（输出置 0）", "均值消融（输出换成平均值）"], rows, [30, 20, 26]) return out @torch.no_grad() def predict_first(model, prompts): """在每个题面末尾预测下一个 token（= 推理链的第一个字母）。""" return [o[:1] for o in tr.generate(model, prompts, 1)] def demo5(model, probs): print_header("演示 5：因果干预——只改题面里的一个字母，看推理链第一步的预测（300 道题，n = 8）") rng = random.Random(5) base, edit_tgt, edit_src, edit_other = [], [], [], [] for p in probs: used = set("".join(p.statements)) free = [c for c in hw.LETTERS if c not in used] z = rng.choice(free) j = next(i for i, st in enumerate(p.statements) if st[0] == p.start) other = next(i for i, st in enumerate(p.statements) if i != j and st[0] != p.start and st[2] != p.start) sts = list(p.statements) def prompt_with(new_sts): return ";".join(new_sts) + f";{p.k}?{p.start}" base.append(p.prompt) a = sts[:] a[j] = f"{p.start}>{z}" # (a) 改匹配陈述的右边：p.start > z edit_tgt.append((prompt_with(a), z)) b = sts[:] b[j] = f"{z}>{sts[j][2]}" # (b) 改匹配陈述的左边：不再有以起点开头的陈述 edit_src.append((prompt_with(b), sts[j][2])) c = sts[:] c[other] = f"{c[other][0]}>{z}" # (c) 对照：改一条无关陈述的右边 edit_other.append((prompt_with(c), p.walk[1])) pred0 = predict_first(model, base) pa = predict_first(model, [x for x, _ in edit_tgt]) pb = predict_first(model, [x for x, _ in edit_src]) pc = predict_first(model, [x for x, _ in edit_other]) n = len(probs) out = { "base": sum(x == p.walk[1] for x, p in zip(pred0, probs)) / n, "tgt_follow": sum(x == z for x, (_, z) in zip(pa, edit_tgt)) / n, "src_old": sum(x == y for x, (_, y) in zip(pb, edit_src)) / n, "other_same": sum(x == y for x, (_, y) in zip(pc, edit_other)) / n, } rows = [ ["不改", "预测 = 正确的下一个字母", pct(out["base"])], ["(a) 把匹配陈述的右边改成题里没有的字母 z", "预测 = z", pct(out["tgt_follow"])], ["(b) 把匹配陈述的左边改成 z（起点不再有出边）", "预测仍 = 原来的右边字母", pct(out["src_old"])], ["(c) 对照：改一条无关陈述的右边", "预测 = 原来的正确答案", pct(out["other_same"])], ] print_table(["干预", "检查", "比例"], rows, [46, 26, 8], ["<", "<", ">"]) return out def main(): ap = argparse.ArgumentParser(description="一跳查找的电路") ap.add_argument("--quick", action="store_true", help="冒烟测试") args = ap.parse_args() t0 = time.perf_counter() model = tr.get_model("cot_L2", quick=args.quick)[0] count = 20 if args.quick else 200 probs = hw.make_eval_set(8, 7, count, seed=9) s1 = demo1(model, probs) s2 = demo2(model, probs) lookup, binder = demo3(model, s1, s2) lookup_heads = sorted(sorted(((1, h) for h in range(model.cfg.n_heads)), key=lambda key: -s1[key][CATS[1]])[:2]) abl = demo4(model, probs, lookup_heads, binder) itv = demo5(model, hw.make_eval_set(8, 3, 30 if args.quick else 300, seed=10)) if args.quick: print("\n（--quick 冒烟测试：模型只训练了 60 步，上面的数字没有意义，不打印观察。）") print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒") return (_, ha), (_, hb) = lookup_heads sink = max(((1, h) for h in range(model.cfg.n_heads)), key=lambda key: s1[key][CATS[3]]) zs, ms = abl["single"][sink] print_observations([ f"第 2 层的头 {ha} 和头 {hb} 在推理位置上把 {pct(s1[(1, ha)][CATS[1]], 0)} 和 {pct(s1[(1, hb)][CATS[1]], 0)} 的注意力" "直接放在'匹配陈述的右边'——正确答案所在的位置。它们是'查找头'：用当前字母当查询，找到来源是它的那条陈述，" "把右边的字母搬过来。它们几乎不看当前字母本身在题面里出现的地方。", f"第 1 层的头 {binder} 在每条陈述右边字母的位置上，把 {pct(s2[binder]['左边 2 格（来源）'], 0)} 的注意力放在" "左边 2 格（这条陈述的来源字母）上：它是'绑定头'，让右边字母的表示里带上'我的来源是谁'；" "查找头比较的正是这个被绑定进来的来源。两层合起来就是第 1 讲猜测的'绑定 + 查找'——2 层恰好做 1 跳。", f"消融确认了这个电路：关掉绑定头，逐步准确率从 {pct(abl['full'])} 掉到 {pct(abl['single'][(0, binder)][0])}" f"（均值消融 {pct(abl['single'][(0, binder)][1])}）；两个查找头**单独**关掉时只降到 " f"{pct(abl['single'][(1, ha)][0])} / {pct(abl['single'][(1, hb)][0])}——它们互为备份；同时关掉降到 " f"{pct(abl['pair'][0])}（均值消融 {pct(abl['pair'][1])}）。只看单头消融，会以为头 {hb} 毫无作用。", f"注意力和消融并不总是一致：第 2 层头 {sink[1]} 的注意力有 {pct(s1[sink][CATS[3]], 0)} 落在跳数数字上，" f"落在匹配陈述上的是 {pct(s1[sink][CATS[0]] + s1[sink][CATS[1]], 0)}，零消融却让准确率降到 {pct(zs)}；换成均值消融则是 {pct(ms)}。它的输出几乎是一个常数，" "后面的层'习惯了'这个常数——置零破坏的是这个习惯，不是查找。零消融会高估这类头的作用，所以要和均值消融对照着看。", f"因果干预：把匹配陈述的右边改成一个题里没有的字母，{pct(itv['tgt_follow'])} 的预测跟着变成那个字母；" f"改一条无关陈述，{pct(itv['other_same'])} 的预测不变；把匹配陈述的左边改掉（起点不再有出边），" f"只有 {pct(itv['src_old'])} 的预测还是原来的右边字母——模型确实是'按来源找、照搬右边'。", ]) print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒") if __name__ == "__main__": main()