""" ============================================================================= 第 2 讲 · 02：compute-optimal 的 test-time scaling —— 并行采样还是顺序修改？ ============================================================================= 对应讲义：lectures/lecture02_test_time_compute.md 第 4 节（Snell et al. 2024） ----------------------------------------------------------------------------- 【问题】 ----------------------------------------------------------------------------- 给你 N 次生成的预算，怎么花最好？两种极端： · 并行（parallel）：独立采样 N 次（best-of-N）—— 广泛地"探索"不同的解法； · 顺序（sequential）：采样 1 次，再连续修改 N − 1 次 —— 在一个解答上不断"打磨"。 以及介于两者之间的：c 条修改链、每条长 L（c × L = N）。最后从全部 N 个解答里用 verifier 挑一个 （weighted best-of-N：同一个答案的 verifier 概率相加，见 tts_tools.select_weighted）。 Snell et al. 2024 的核心发现：**最优的分配取决于题目的难度**。简单题适合顺序修改，难题需要并行探索。 所以先估计难度、再按难度分配计算（compute-optimal），比一律用 best-of-N 高效得多。 ----------------------------------------------------------------------------- 【三个演示】 ----------------------------------------------------------------------------- 演示 1 预算 N = 16：不同链长 L 在五个难度档上的准确率。 演示 2 compute-optimal：在开发集上为每个难度档选最好的 L，在测试集上和 best-of-N 比较； best-of-N 要多花多少预算才能追上？ 演示 3 FLOPs-matched：小模型 + test-time compute vs 参数多 14 倍的大模型（采样一次）。 推理负载越小（推理 token 远少于预训练 token），小模型能分到的推理预算越多。 【约定】所有生成都用 temperature = 0.7（本讲的"标准用法"）。verifier 是 tts_tools.verifier_score （正确答案平均高 2 分，噪声按答案固定），挑选用 weighted best-of-N。难度用 32 次采样估计的 pass@1 划分（"oracle 难度"， 需要知道标准答案；Snell et al. 也用了不需要答案的"预测难度"，见讲义）。 运行时间：约 15 秒（只用标准库）。 """ import random from tts_tools import (ToyLLM, difficulty_bins, estimate_pass1, make_problem_set, mean, pad, run_chain, select_weighted) TEMPERATURE = 0.7 N_BINS = 5 REPEATS = 3 # 每道题每种配置重复 3 次，减小随机波动 REPEATS_FLOPS = 8 # 演示 3 每个难度档只有约 30 道测试题，多重复几次 def solve(model, problem, rng, budget, chain_len): """ 用 budget 次生成解一道题：budget // chain_len 条修改链，每条长 chain_len， 最后从全部解答里用 weighted best-of-N 挑一个（同一答案的 verifier 概率相加）。 【参数】model；problem；rng；budget：生成次数（N）；chain_len：每条链的长度 L（1 ≤ L ≤ N） 【返回】bool：挑出的答案是否正确 【边界】budget 不是 chain_len 的整数倍时，多出来的预算不用（向下取整），保证不超预算。 """ pool = [] for _ in range(max(1, budget // chain_len)): pool.extend(run_chain(model, problem, rng, chain_len, TEMPERATURE)) return select_weighted(problem, pool).correct def evaluate_grid(model, problems, configs, rng): """ 对每个 (预算 N, 链长 L) 配置，算出每道题的平均准确率。 【参数】configs：(N, L) 的列表 【返回】dict：(N, L) → 每道题准确率的列表（和 problems 一一对应） """ table = {} for n, l in configs: table[(n, l)] = [mean(solve(model, p, rng, n, l) for _ in range(REPEATS)) for p in problems] return table def bin_mean(values, bins, b, subset): """第 b 档、且下标在 subset 里的题目的平均值。""" picked = [values[i] for i in subset if bins[i] == b] return mean(picked) def demo_ratio_by_difficulty(table, bins, everyone): print("=" * 78) print("演示 1：预算 N = 16 次生成，c 条链 × 每条长 L —— 五个难度档上的准确率") print("=" * 78) print(pad("配置", 22) + "".join(pad(f"档{b + 1}", 9, "right") for b in range(N_BINS)) + pad("全部", 9, "right")) for l in (1, 2, 4, 8, 16): label = f"{16 // l} 条 × L = {l}" + ("（并行）" if l == 1 else "（顺序）" if l == 16 else "") vals = table[(16, l)] print(pad(label, 22) + "".join(pad(f"{bin_mean(vals, bins, b, everyone):.3f}", 9, "right") for b in range(N_BINS)) + pad(f"{mean(vals):.3f}", 9, "right")) print("（档 1 最简单，档 5 最难）") print("\n观察：") print(" · 简单的档（1、2）：预算 16 已经足够，各种配置都接近 1.0。") print(" · 中等难度的档（3）：偏向顺序修改的配置（L = 8、16）明显好于纯并行 —— 这些题的错误多半是") print(" \"手滑\"，模型仔细复查一下就能修好，把预算花在打磨解答上更划算。") print(" · 最难的档（4、5）：纯顺序修改不比纯并行好，最好的是两者的混合（8 条 × L = 2）。") print(" 这些题的错误常常出在模型自己的\"误解\"上，复查时模型会坚持错误的做法，修改再多也改不过来，") print(" 需要多条独立的链去\"碰\"出不同的做法；但每条链修改一两次，仍能顺手修掉手滑。") print(" · Snell et al. 的结论正是如此：简单题适合纯顺序，难题存在一个最佳的顺序 / 并行比例。") print(" 没有一个对所有题都最好的比例，这就是\"按难度分配\"的动机。") print(" · 每档 60 道题 × 重复 3 次，相邻配置之间 0.03 左右的差别可能只是随机波动。\n") def demo_compute_optimal(table, bins, dev, test, budgets, ls_for): print("=" * 78) print("演示 2：compute-optimal（按难度档选 L）vs best-of-N") print("=" * 78) chosen = {} print("在开发集上为每个难度档选出的链长 L：") for n in budgets: chosen[n] = [max(ls_for[n], key=lambda l: bin_mean(table[(n, l)], bins, b, dev)) for b in range(N_BINS)] print(f" N = {n:>3d}：" + "，".join(f"档{b + 1} L = {chosen[n][b]}" for b in range(N_BINS))) bon_budgets = sorted(n for n, l in table if l == 1) bon = {n: mean(table[(n, 1)][i] for i in test) for n in bon_budgets} print("\n测试集上的准确率：") print(pad("预算 N", 9) + pad("best-of-N", 12, "right") + pad("compute-optimal", 18, "right") + " best-of-N 追上它需要的预算") for n in budgets: opt = mean(table[(n, chosen[n][bins[i]])][i] for i in test) catch = next((m for m in bon_budgets if bon[m] >= opt), None) catch_text = f"N = {catch}（{catch / n:.0f} 倍）" if catch else f"> {bon_budgets[-1]}" print(pad(str(n), 9) + pad(f"{bon[n]:.3f}", 12, "right") + pad(f"{opt:.3f}", 18, "right") + " " + catch_text) print("best-of-N 在更多预算下：" + "，".join(f"N = {m}：{bon[m]:.3f}" for m in bon_budgets)) print("\n观察：") print(" · 三个预算下，按难度选 L 的 compute-optimal 策略都比 best-of-N 准；best-of-N 要花 2 倍的预算") print(" 才能追上（Snell et al. 在 MATH 上报告的是约 4 倍，玩具世界里的差距小一些）。") print(" · 预算越大，优势越小（N = 64 时只差 0.009）：并行采样的 coverage 随预算增长，") print(" 最终几乎所有档都选了 L = 1。按难度分配在\"预算紧张\"时最有价值。") print(" · 选 L 用的是开发集，评估用的是测试集 —— 如果在同一批题上既选又评，结果会偏乐观。") print(" 开发集上选出的 L 本身也有噪声（例如 N = 16 时档 4、档 5 选了 L = 2，和演示 1 一致）。") print(" · 代价：要先知道每道题的难度。这里用的是\"oracle 难度\"（需要标准答案），") print(" 实际系统需要用 verifier 的分数等信号来\"预测\"难度，而预测本身也要花计算。\n") return chosen def demo_flops_matched(small, large, problems, bins, test, chosen, rng): print("=" * 78) print(f"演示 3：FLOPs-matched —— {small.name} 模型 + test-time compute vs {large.name} 大模型（采样一次）") print("=" * 78) ratio = large.size / small.size print("总计算量 = 预训练（≈ 6 × 参数量 × 预训练 token 数）+ 推理（≈ 2 × 参数量 × 推理 token 数 × 每题采样数）。") print("设 R = 推理 token 总数 / 预训练 token 数。让小模型的总计算量等于大模型（每题采样一次），") print(f"小模型每题能用的采样数 M = {ratio:.0f} + 3 × ({ratio:.0f} − 1) / R：") rows = [] for r in (0.5, 2.0, 10.0): m = int(ratio + 3 * (ratio - 1) / r) rows.append((r, m)) print(f" R = {r:>4}：M = {m}") large_acc = [mean(large.generate(p, rng, TEMPERATURE).correct for _ in range(16)) for p in problems] print() print(pad("", 22) + "".join(pad(f"档{b + 1}", 9, "right") for b in range(N_BINS)) + pad("全部", 9, "right")) print(pad(f"{large.name} 采样一次", 22) + "".join(pad(f"{bin_mean(large_acc, bins, b, test):.3f}", 9, "right") for b in range(N_BINS)) + pad(f"{mean(large_acc[i] for i in test):.3f}", 9, "right")) for r, m in rows: ref = 64 if m >= 40 else 16 # 借用最接近的预算上为每档选好的 L accs = [0.0] * len(problems) for i in test: l = chosen[ref][bins[i]] accs[i] = mean(solve(small, problems[i], rng, m, l) for _ in range(REPEATS_FLOPS)) print(pad(f"{small.name}，M = {m}（R = {r}）", 22) + "".join(pad(f"{bin_mean(accs, bins, b, test):.3f}", 9, "right") for b in range(N_BINS)) + pad(f"{mean(accs[i] for i in test):.3f}", 9, "right")) print("\n观察：") print(" · 档 1~4：三种推理负载下，10B 模型加 test-time compute 都追平或超过了 14 倍大的模型。") print(" · 最难的档 5：只有推理负载很低（R = 0.5，每题能用 92 次生成）时小模型才赢；") print(" R = 2 和 R = 10 时大模型领先 —— 小模型\"不知道\"的东西，有限的 test-time compute 补不上。") print(" · R 越大（模型上线后要服务的推理量越多），小模型每题能分到的推理预算越少，") print(" \"训练一个更大的模型\"就越划算。Snell et al. 的结论正是如此：") print(" test-time compute 和 pretraining compute 不是谁替代谁，要看题目难度和推理负载。\n") def main(): rng = random.Random(41) small, large = ToyLLM(1e10), ToyLLM(1.4e11) problems = make_problem_set(random.Random(40), n_problems=300, min_steps=3, max_steps=8) everyone = list(range(len(problems))) dev, test = everyone[:150], everyone[150:] pass1 = [estimate_pass1(small, p, rng, 32, TEMPERATURE) for p in problems] bins = difficulty_bins(pass1, N_BINS) print(f"{small.name} 模型在 temperature = {TEMPERATURE} 下，五个难度档的平均 pass@1：" + "，".join(f"档{b + 1} {bin_mean(pass1, bins, b, everyone):.2f}" for b in range(N_BINS)) + "\n") budgets = (4, 16, 64) ls_for = {4: (1, 2, 4), 16: (1, 2, 4, 8, 16), 64: (1, 2, 4, 8, 16, 32, 64)} configs = [(n, l) for n in budgets for l in ls_for[n]] + [(n, 1) for n in (8, 32, 128, 256)] table = evaluate_grid(small, problems, configs, rng) demo_ratio_by_difficulty(table, bins, everyone) chosen = demo_compute_optimal(table, bins, dev, test, budgets, ls_for) demo_flops_matched(small, large, problems, bins, test, chosen, rng) if __name__ == "__main__": main()