""" ============================================================================= 第 5 讲 · 代码 3：先写草稿再回答——scratchpad / chain-of-thought 为什么有用 ============================================================================= 对应讲义：lectures/lecture05_emergence_icl.md 第 5 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 第 1 讲的 base model 没学会两位数加法（新题 0.09）。Nye et al. 2021（Show Your Work）和 Wei et al. 2022（Chain-of-Thought）给出的办法是：**不要让模型一步写出答案，先让它写出中间步骤**。 这个文件在 6 位数加法上做一个"同样的算力，只改写法"的对照实验： 演示 1：同样的模型（2 层、64 维，约 10 万参数）、同样的训练算力（每步 1024 个 token、1500 步）， 三种写法：直接写答案 / 先写草稿 / 先写同样长的一串填充符 "."。直接写答案和草稿各训练 2 个种子： 结果稳定吗？草稿有用，是因为"多了几个位置"，还是因为草稿的内容？ 演示 2：一个只有 1 层、32 维（约 1.5 万参数）的模型：直接写答案几乎全错；先写草稿呢？ ——"每多写一个 token，就多一次完整的前向计算"：草稿把串行计算摊到了很多个 token 上。 演示 3：自洽性（self-consistency）：在一个"学到一半"的草稿模型上，采样多条草稿再对答案投票， 和贪心解码比一比。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 03_scratchpad_cot.py # 第一次约 1.5 分钟（训练 7 个小模型并缓存到 checkpoints/）；之后约 10 秒 需要 PyTorch。依赖本目录的 seq_tasks.py，以及 code/lecture01 的 tiny_gpt.py、console.py（只读导入）。 【Python 小知识】 - collections.Counter(列表).most_common(1)：数每个元素出现几次，返回出现最多的那个（和它的次数）。 - zip(*samples)：把"10 个样本 × 300 道题"的列表"转置"成"300 道题 × 10 个样本"。 - copy.deepcopy(model.state_dict())：把当前的参数完整复制一份。不复制的话，存下来的只是"引用"， 训练继续进行时它会跟着变。 """ import copy import os import sys from collections import Counter import torch 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 seq_tasks import (ADD_ALPHABET, CharTokenizer, GPTConfig, addition_answers, addition_batch, # noqa: E402 addition_block_size, addition_example, addition_lengths, addition_prompts, config_dict, greedy_fill, load_or_train, new_model, random_pairs, sample_fill, train_masked_lm) TOK = CharTokenizer(ADD_ALPHABET) MAIN = dict(n_layer=2, n_head=4, d_model=64) # 和第 1 讲 base model 同样的形状 TINY = dict(n_layer=1, n_head=2, d_model=32) N_DIGITS = 6 TOKENS_PER_STEP = 1024 # 每步处理的 token 数：所有写法都一样 → 训练算力相同 STEPS = 1500 CURVE_EVERY = 100 SEEDS = (0, 1) TEST_SEED = 99 # ============================================================================= # 训练与评测 # ============================================================================= def batch_size_for(n_digits, mode): """ 同样的算力下，每步放几道题。 【思路】训练一步的计算量大约正比于"这一步处理的 token 数"。直接写答案的题短（6 位数 21 个 token）， 草稿的题长（57 个 token），所以每步放的题数不同：B = 1024 // 序列长度。 结果是：直接写答案的模型**看到的题目多将近 3 倍**——这对它是有利的，比较才算公平。 """ return TOKENS_PER_STEP // addition_block_size(n_digits, mode) def get_model(n_digits, mode, shape, seed=0, snapshot=False): """ 训练（或读取缓存的）一个加法模型，并记录学习曲线。 【参数】 n_digits, mode : 几位数、哪种写法（见 seq_tasks.addition_example） shape : MAIN 或 TINY seed : 种子 s：模型初始化用 s，出题顺序用 1 + s（换一个种子 = 换一次"运气"） snapshot : True 时，保存训练中途**第一个**贪心准确率落在 [0.2, 0.9] 的检查点 （演示 3 需要一个"学到一半"的模型；这个规则在看结果之前就定好了） 【返回】(model, info)；info 里有 curve（每 100 步在 100 道题上的贪心整题准确率）、snapshot（dict 或 None）、 n_params、train_seconds、problems_seen """ cfg = GPTConfig(vocab_size=len(ADD_ALPHABET), block_size=addition_block_size(n_digits, mode), **shape) bsz = batch_size_for(n_digits, mode) meta = dict(config=config_dict(cfg), task="addition", n_digits=n_digits, mode=mode, steps=STEPS, batch_size=bsz, lr=3e-3, model_seed=seed, data_seed=1 + seed, curve_every=CURVE_EVERY, snapshot=snapshot) name = f"add{n_digits}_{mode}_L{shape['n_layer']}d{shape['d_model']}_s{seed}" def build_and_train(): model = new_model(cfg, seed=seed) curve, snap = [], {} eval_a, eval_b = random_pairs(100, n_digits, torch.Generator().manual_seed(7)) def callback(step, m): acc = exact_match(m, eval_a, eval_b, n_digits, mode) curve.append((step, acc)) if snapshot and not snap and 0.2 <= acc <= 0.9: snap.update(step=step, acc=acc, state_dict=copy.deepcopy(m.state_dict())) train_masked_lm(model, lambda g: addition_batch(TOK, n_digits, mode, bsz, g), steps=STEPS, seed=1 + seed, callback=callback, callback_every=CURVE_EVERY) return model, dict(curve=curve, snapshot=snap or None, problems_seen=bsz * STEPS) model, info, _ = load_or_train(name, meta, build_and_train) return model, info def decode_answers(out_ids, n_digits): """模型写出的整段目标 → 最后 n+1 个字符（答案的位置；草稿写错格式时也照样取这几位来比较）。""" return [TOK.decode(row)[-(n_digits + 1):] for row in out_ids.tolist()] def exact_match(model, a, b, n_digits, mode): """贪心解码，答案一字不差才算对。返回准确率。""" _, t_len = addition_lengths(n_digits, mode) outs = decode_answers(greedy_fill(model, addition_prompts(TOK, a, b, n_digits, mode), t_len), n_digits) return sum(o == t for o, t in zip(outs, addition_answers(a, b, n_digits))) / len(a) def evaluate(model, n_digits, mode, n=1000): """ 在 n 道新题上评测（默认 1000 道；题目用固定种子 TEST_SEED 生成。训练时的题目是现场随机出的， 6 位数加法有 10^12 种题，和测试题重合的概率可以忽略）。 【返回】dict：em（整题准确率）、per_digit（答案每一位的准确率，从最高位到最低位） """ a, b = random_pairs(n, n_digits, torch.Generator().manual_seed(TEST_SEED)) _, t_len = addition_lengths(n_digits, mode) outs = decode_answers(greedy_fill(model, addition_prompts(TOK, a, b, n_digits, mode), t_len), n_digits) truth = addition_answers(a, b, n_digits) per_digit = [sum(o[i] == t[i] for o, t in zip(outs, truth)) / n for i in range(n_digits + 1)] return dict(em=sum(o == t for o, t in zip(outs, truth)) / n, per_digit=per_digit) def first_step_above(curve, threshold=0.5): """学习曲线上第一次达到 threshold 的步数；从没达到时返回 "—"。""" for step, acc in curve: if acc >= threshold: return step return "—" # ============================================================================= # 演示 1：同样的算力，三种写法，多个种子 # ============================================================================= def demo_equal_compute(): print_header("演示 1：同样的模型、同样的训练算力，只改\"写法\"——直接写答案 / 先写填充符 / 先写草稿") for mode in ("direct", "filler", "scratch"): p, t = addition_example(347, 285, 4, mode) print(f" {pad(mode, 8)} {p}{t}") print(f"（上面是 4 位数的样子；下面的实验都是 {N_DIGITS} 位数加法。）") print(f"\n模型：2 层、64 维；每步 {TOKENS_PER_STEP} 个 token，{STEPS} 步（三种写法的训练算力相同）。测试：1000 道新题。\n") print(pad("写法", 9) + pad("种子", 6) + pad("每步题数", 10, ">") + pad("共见过题数", 12, ">") + pad("整题准确率", 12, ">") + pad("首次≥0.5 的步数", 18, ">") + " 答案每一位的准确率（最高位 → 最低位）") runs = [("direct", s) for s in SEEDS] + [("filler", 0)] + [("scratch", s) for s in SEEDS] results = {} for mode, seed in runs: model, info = get_model(N_DIGITS, mode, MAIN, seed=seed, snapshot=(mode == "scratch" and seed == 0)) r = evaluate(model, N_DIGITS, mode) results[(mode, seed)] = (r, info) print(pad(mode, 9) + pad(seed, 6) + pad(batch_size_for(N_DIGITS, mode), 10, ">") + pad(f"{info['problems_seen']:,}", 12, ">") + pad(f"{r['em']:.3f}", 12, ">") + pad(first_step_above(info["curve"]), 18, ">") + " " + " ".join(f"{v:.2f}" for v in r["per_digit"])) print(f"\n学习曲线（每 {CURVE_EVERY} 步在 100 道题上测一次整题准确率）：") steps = [s for s, _ in results[runs[0]][1]["curve"]] print(pad("写法/种子", 12) + "".join(pad(s, 6, ">") for s in steps)) for mode, seed in runs: print(pad(f"{mode}/{seed}", 12) + "".join(pad(f"{acc:.2f}", 6, ">") for _, acc in results[(mode, seed)][1]["curve"])) print(""" 观察 Observations： · 草稿：两个种子都在 1000 道新题上全对（1.000），学习曲线在第 500–600 步越过 0.5，第 900–1000 步到 1.00。 · 直接写答案：种子 0 在第 900–1000 步之间"顿悟"，最后 0.828；种子 1 到第 1500 步仍是 0.000。 逐位看种子 1：最高两位和最低两位都对（1.00 0.95 … 0.99 1.00），中间三位约 0.1，就是 10 个数字里瞎猜—— 中间位要知道低位传上来的进位，而模型还没学会在一次前向传播里把进位链算完。 同样的配置、只换随机种子，一个 0.83、一个 0：这种"平台期什么时候结束"的偶然性，正是代码 01 演示 3 说的"真实的跳变"。 · 公平起见：直接写答案的模型见过 72,000 道题，草稿模型只见过 25,500 道——草稿模型用更少的题学得更快、更稳。 · 填充符对照：先写 35 个 "." 再写答案，和草稿一样长、多出的"计算位置"一样多，结果是 0.000， 逐位看和直接写答案的种子 1 一样卡在平台上。所以草稿有用，靠的是**内容**：每个中间数字都只依赖附近 几个已经写出来的数字，模型每一步只需要做一件很小的事。Wei et al. 2022 的消融（只给"同样长度的点"， variable compute only）在真实模型上也没有提升。""") return results # ============================================================================= # 演示 2：1 层的模型 # ============================================================================= def demo_tiny_model(): print_header(f"演示 2：只有 1 层、32 维的模型（约 1.5 万参数）——直接写答案 vs 先写草稿（{N_DIGITS} 位数）") print(pad("写法", 10) + pad("参数量", 10, ">") + pad("整题准确率", 12, ">") + " 答案每一位的准确率（最高位 → 最低位）") for mode in ("direct", "scratch"): model, info = get_model(N_DIGITS, mode, TINY, seed=0) r = evaluate(model, N_DIGITS, mode) print(pad(mode, 10) + pad(f"{info['n_params']:,}", 10, ">") + pad(f"{r['em']:.3f}", 12, ">") + " " + " ".join(f"{v:.2f}" for v in r["per_digit"])) print(""" 观察 Observations： · 1 层的模型直接写答案：0.000，只有最高两位学会了；先写草稿：1000 道题全对。 · 为什么？一个 L 层的 Transformer，从读到输入到写出一个 token，中间只有 L 次"注意力 + MLP"。 直接写答案时，中间各位要在这 L 步之内把整条进位链算完；写草稿时，每写一个 token 就多一次完整的前向计算， 而且这次计算能读到前面写下的中间结果——**草稿把串行计算摊到了很多个 token 上** （讲义 §5.2："每个答案能用的串行计算量"）。 · 这里的草稿模型是**在草稿格式上训练过的**（Nye et al. 2021 的做法），所以 1.5 万参数也够用。 真实大模型的 chain-of-thought prompting 不训练，只在提示里给示例——那要求模型本来就会写像样的推理步骤， 所以 Wei et al. 发现它只在约 1000 亿参数以上才有用。两件事不要混淆（讲义 §5.3）。""") # ============================================================================= # 演示 3：自洽性 # ============================================================================= def vote(answers): """ 多数投票：出现次数最多的答案。 【参数】answers：一道题的若干个答案字符串 【返回】得票最多的那个 【边界】并列时取最先出现的那个（Counter.most_common 对次数相同的元素保持插入顺序），结果是确定的。 """ return Counter(answers).most_common(1)[0][0] def demo_self_consistency(results): print_header("演示 3：自洽性（self-consistency）——采样多条草稿，对最终答案投票") _, info = results[("scratch", 0)] snap = info["snapshot"] if snap is None: print("（训练中途没有出现准确率在 [0.2, 0.9] 之间的检查点，跳过本演示）") return cfg = GPTConfig(vocab_size=len(ADD_ALPHABET), block_size=addition_block_size(N_DIGITS, "scratch"), **MAIN) model = new_model(cfg, seed=0) model.load_state_dict(snap["state_dict"]) model.eval() print(f"使用演示 1 里种子 0 的草稿模型训练到第 {snap['step']} 步的检查点（当时在 100 道题上的贪心准确率 {snap['acc']:.2f}）。") n, n_samples = 300, 10 a, b = random_pairs(n, N_DIGITS, torch.Generator().manual_seed(TEST_SEED)) truth = addition_answers(a, b, N_DIGITS) prompts = addition_prompts(TOK, a, b, N_DIGITS, "scratch") _, t_len = addition_lengths(N_DIGITS, "scratch") greedy = decode_answers(greedy_fill(model, prompts, t_len), N_DIGITS) greedy_acc = sum(o == t for o, t in zip(greedy, truth)) / n print(f"\n{n} 道新题，贪心解码的整题准确率：{greedy_acc:.3f}\n") print(pad("温度", 8) + pad("单个样本", 10, ">") + pad("投票 5 个", 10, ">") + pad("投票 10 个", 11, ">") + pad("10 个里至少一个对", 18, ">") + pad("答案种类（平均）", 18, ">")) for temp in (0.7, 1.0): g = torch.Generator().manual_seed(123) samples = [decode_answers(sample_fill(model, prompts, t_len, temp, g), N_DIGITS) for _ in range(n_samples)] per_problem = list(zip(*samples)) # 每道题的 10 个答案 single = sum(s == t for s, t in zip(samples[0], truth)) / n v5 = sum(vote(ans[:5]) == t for ans, t in zip(per_problem, truth)) / n v10 = sum(vote(ans) == t for ans, t in zip(per_problem, truth)) / n any_ok = sum(t in ans for ans, t in zip(per_problem, truth)) / n kinds = sum(len(set(ans)) for ans in per_problem) / n print(pad(temp, 8) + pad(f"{single:.3f}", 10, ">") + pad(f"{v5:.3f}", 10, ">") + pad(f"{v10:.3f}", 11, ">") + pad(f"{any_ok:.3f}", 18, ">") + pad(f"{kinds:.1f}", 18, ">")) print(""" 观察 Observations： · 投票确实有用：单个样本（T = 0.7）只有 0.290，明显低于贪心的 0.553——采样会引入"偶然的笔误"； 投票 5 个 0.397、10 个 0.483，一路回升。不同样本的笔误出在不同的地方，多数票把它们抵消掉了。 这就是 self-consistency 的基本机制。 · 但在这个玩具里，投票**没有超过贪心解码**（10 个样本投票 0.483 < 贪心 0.553）。Wang et al. 2022 在大模型上 看到的是投票明显好于贪心。差别在哪？他们的题目（数学应用题）有**很多条不同但都正确的推理路径**， 正确答案从许多路径汇聚而来，错误答案则五花八门；贪心只走一条路，第一步走偏就回不来 （第 1 讲 §2.3：贪心不等于最可能）。我们的加法草稿只有**一条**正确路径，贪心几乎就是每一步的众数—— 投票能做的，只是把采样带来的额外错误"投回去"。 · "10 个里至少一个对"是 0.830，远高于任何单一策略：如果有一个可靠的验证器（比如把答案代回去检查）， 多次采样 + 验证能做得更好——这是 2023 年之后"推理时计算"（test-time compute）的出发点（讲义 §7）。""") if __name__ == "__main__": res = demo_equal_compute() demo_tiny_model() demo_self_consistency(res)