拒绝采样

0. 任务目标

拒绝采样（Rejection Sampling, RS）是第 9 章的主题，也是最简单的一种「用奖励改进策略」的方法：让模型对每个 prompt 生成 $N$ 条回复，用奖励模型给每条打分，挑出好的那些，再拿它们做一轮 SFT。没有 RL 循环，没有优势估计，没有 KL 惩罚——它就是「采样 + 筛选 + 模仿」。

正因为简单，它在工业界用得极多（Llama 2/3 的后训练流水线里都有它）。但也正因为简单，它有一个特别容易被忽略的方法论陷阱，而这个陷阱就是本组作业的全部重点：

所以这一组作业的交付物不是「一个准确率数字」，而是一对准确率数字之差。学会这件事之后，你看任何后训练论文里的消融表格，视角都会变。

1. 流水线与运行命令

参考实现把第 9 章的四个阶段完整实现了一遍，任务是 GSM8K（小学数学应用题）：

1. Stage 1 · 生成：用 Qwen/Qwen3-1.7B 对每个训练 prompt 生成 $N$ 条解题过程（num_completions_per_prompt: 8）。
2. Stage 2 · 打分：用 nvidia/AceMath-7B-RM 给每一条打分。
3. Stage 3a · 选择：按某种策略挑出一个子集。四种策略的差别全在这一步，前两个阶段完全一样。
4. Stage 3b · SFT：在选中的 (prompt, completion) 对上微调同一个 Qwen3-1.7B。
5. Stage 3c · 评测：在 GSM8K 测试切片上量 exact-match 准确率，记为 test_accuracy。

默认切片是 1000 条训练 / 200 条测试，生成用 temperature: 0.8、top_p: 0.95、max_new_tokens: 512，SFT 用 lr 5e-6、2 个 epoch，评测用 eval_temperature: 0.0（贪心，保证可复现）。

先建缓存，再跑四个配置

cd _src/code/

# (1) 只做一次：Stage 1 + Stage 2，产出一份 JSONL 缓存
python -m rejection_sampling.preprocess \
    --config rejection_sampling/configs/top_per_prompt.yaml

# (2) 四个选择策略各训一次（train.py 首次调用会自动跑 preprocess，之后命中缓存）
python -m rejection_sampling.train --config rejection_sampling/configs/top_per_prompt.yaml
python -m rejection_sampling.train --config rejection_sampling/configs/random_per_prompt.yaml
python -m rejection_sampling.train --config rejection_sampling/configs/top_k_overall.yaml
python -m rejection_sampling.train --config rejection_sampling/configs/random_k_overall.yaml

缓存文件每行是一个 prompt 的全部打分结果：

{"question": "...", "answer": "72",
 "completions": ["...", "..."], "rewards": [0.71, 0.33, ...]}

四种选择策略

两个 *_k_overall 配置的 selection.top_k 必须写成同一个值（默认 1000），否则样本预算不对齐，对照就无效了。

2. 为什么必须有 random 对照组

这一节是本组作业的灵魂，重要到值得占一整节。

问题的形状

拒绝采样的完整操作是「生成 → 打分 → 选择 → SFT」。你想检验的假设是「用奖励做选择比不用好」，但你实际做的动作里包含了至少四个可能独立产生效果的成分：

只跑一个 top_per_prompt，然后拿它和「原始模型」比，你得到的是这四项的总和。而论文里、简历里、周报里想写的那句话——「我们的奖励模型带来了 X 个点的提升」——对应的只是最后一项。

对照组把式子改成了减法

设 $A_{\text{base}}$ 是原始模型的准确率，$A_{\text{rand}}$ 是随机选择后训练的准确率，$A_{\text{top}}$ 是奖励选择后训练的准确率。那么：

换句话说：随机对照组把一个绝对值问题变成了差值问题。而差值对很多系统性误差是免疫的——两组共享同一份候选池、同一套训练超参、同一个评测脚本，任何影响两边的因素都会在相减时抵消掉。这正是对照实验的全部威力所在。

3. 对照组必须锁死哪些变量

一个合格的对照组必须只改变你想检验的那一个变量，其余全部保持一致。这里想检验的变量是「选择时用不用奖励」，所以必须锁死下面五项：

锁死之后，两个配置之间唯一的差别就落在选择那一行代码上：max(rewards) 换成了 rng.randrange(len(completions))。这就是「受控对照」在代码层面的样子。

还有两个容易漏掉的坑

其一，评测脚本本身要固定。GSM8K 的 exact-match 依赖答案抽取逻辑（从生成文本里把最终数字抠出来）。如果你在两次运行之间改了这个正则，得到的差值毫无意义。配置里 eval_temperature: 0.0 就是出于同样的考虑——贪心解码去掉了评测阶段的随机性。

其二，别忘了记录基线。$A_{\text{base}}$（完全不训练、直接评测原始模型）不在四个配置里，但它是第 2 节那张差值表的第一行。建议手动补一次：它能告诉你「多采样 + 再训一轮」这部分收益有多大，而这往往比 RM 的净贡献大得多——这个对比本身就很有教育意义。

4. 该盯什么指标

训练与评测阶段

建议整理成这样一张表（数字栏留给你自己填）：

先跑诊断，再跑训练

这条建议能省掉大量算力。rejection_sampling/diagnostics.py 直接读缓存文件，不需要训练任何模型：

python -m rejection_sampling.diagnostics \
    --cache rejection_sampling/output/rollouts/<hash>.jsonl \
    --out-dir rejection_sampling/output/diagnostics

它给出三个视图加一个汇总数：

1. 奖励直方图（正确 vs 错误回复）——两个分布分不分得开，就是奖励模型的原始分辨力。如果重叠得几乎完全，后面的训练大概率也不会有 gap，可以直接省下来。
2. 行内选择命中率——在可判定的 prompt 上（即那些既有正确又有错误回复的行），$\argmax(\text{reward})$ 选中正确回复的比例，对比随机基线的比例。这是在训练之前就能拿到的「RM 净贡献」估计，比跑完整训练便宜几百倍。
3. best-of-$N$ 扫描——$N$ 从 1 到 8，「按奖励取前 $N$ 条里至少有一条正确」的比例，对比「随机取 $N$ 条」。两条线的间距如果很快饱和，说明再加 $N$ 也没用。
4. decidable_fraction——可判定 prompt 的占比，也就是「行内选择有可能起作用」的那部分数据有多少。

5. 自选消融

消融一：扫 $N$（每题采几条）

改 num_completions_per_prompt（4 / 8 / 16），其余不动。这是第 9 章那条「收益随 $N$ 增长而递减」曲线的实测版本。

观察目标有两个：(1) 候选池变大，是不是真的让「最好的那条」变好了？(2) 还是只让奖励模型有更多机会被噪声骗到（$N$ 越大，越可能出现一条「奖励很高但其实是错的」的回复——这是 best-of-$N$ 版本的过度优化）？

消融二：top_per_prompt 与全局 top-$K$ 的结构差异

这两个策略的差别不是「哪个更好」，而是它们在优化不同的东西：

建议顺便统计一件事：$K$ 取当前值时，有多少 prompt 完全没有代表？这个数直接量化了上面那条偏置。$K$ 越小，覆盖越偏。

再进一步的变体（值得一试）：先做组内归一化再全局取 top-$K$。对每一行做 $\tilde r_{i,j} = (r_{i,j} - \mu_i)/\sigma_i$，然后在归一化后的分数上取全局 top-$K$。如果这个变体明显好过原始的 top_k_overall，你就实证了「跨 prompt 的奖励尺度不可比」确实是个真问题——顺便也把自己推到了 GRPO 的门口（HW3 的组内 z-score 正是同一个思路）。

消融三：把 decidable_fraction 当自变量

前面说过，行内选择的效果上限被可判定比例卡死。那就主动去改变它：

• 换一个更弱的策略模型（同时调小 max_train_samples 省算力）。弱模型在 GSM8K 上全对的行会变少，decidable_fraction 应该显著升高。如果奖励模型有真信号，那么 top 与 random 的 gap 也应该跟着变大。这是一个非常干净的因果检验，也回答了一个关键问题：RS 是在「拯救糟糕的生成」，还是只是在「一堆本来就还行的答案里挑挑拣拣」？
• 换一个更难的数据集切片，效果类似。
• 迭代两轮：用第一轮 RS 训出的模型重新生成、重新打分、再训一次。预期是第二轮增益明显小于第一轮，同时 decidable_fraction 下降（模型变强了，全对的行变多了）。这是第 9 章「迭代式 RS」的最小演示，也能让你亲眼看到为什么奖励模型需要跟着策略一起更新。

这三个消融的共同点是：它们都不是在调训练超参，而是在改变「奖励信号有没有发挥空间」。这正是本组作业希望你养成的直觉。

本组小结

交付物清单

• 四个配置的完整日志与 test_accuracy，以及成对整理的对照表。
• rollout 缓存的哈希文件名 —— 用来证明四次训练确实共享了同一个候选池。
• diagnostics 的三张图和 decidable_fraction。
• 一段结论，明确区分「多采样 + 再训一轮的贡献」和「奖励模型的净贡献」两部分。