16.1 off-policy 蒸馏与 on-policy 蒸馏
设 为教师, 为学生,二者都是同一 token 词表(vocabulary)上的分布。要在智能体轨迹上让 接近 ,有两种自然的做法。 off-policy 蒸馏(off-policy distillation)从教师采样回合,并用 15.1 节的带掩码 SFT 损失在这些回合上微调学生。由于环境项相互抵消,它的期望损失为 即轨迹分布之间的前向 KL(forward KL)加上一个常数。前向 KL 倾向于 mode-covering:教师赋予较高概率的 token 若被学生忽略,学生就会受到惩罚。问题出在下标上。期望是对教师的回合取的,所以学生只在教师访问过的前缀上接受训练。 on-policy 蒸馏(on-policy distillation)从学生采样回合,在每个学生前缀 上向教师询问其分布,并在那里最小化逐 token 的散度。本书的默认选择是逐 token 的反向 KL, 它具有 mode-seeking 的倾向:学生在教师概率很小的地方放置概率质量会受到惩罚。现在前缀 来自学生,包括它自身错误所产生的观测,所以教师的纠正恰好出现在学生部署时会到达的位置。on-policy 蒸馏是 DAgger 的软标签版本;DAgger 是一种模仿学习算法,它在学习者自己访问的状态上向专家查询(Ross, Gordon, and Bagnell, 2010)。该论文的动机同样适用:在序列预测中,未来的输入取决于过去的输出,这破坏了监督学习背后的独立同分布假设。 只有当教师知道在学生的前缀上(包括糟糕的前缀)该怎么做时,on-policy 蒸馏才有用。因此我们构建一个见过错误的教师。with_mistake 包装 oracle,使它在某一轮犯一个给定的错误,然后就像那一轮没有发生过一样继续。在得到的示范中,出错的那一轮被从损失中掩码掉,所以教师学到的是从错误中恢复,而不是犯这个错误。
clamp 示范中,错误出现在第 1 轮:读取之后,专家应用了 fix_range_end,它没有改变源代码(no change),随后专家继续执行正确的 swap_minmax、一次测试和一次提交。只有出错调用的四个 token 被从损失中掩码掉。全部 180 条示范的奖励都是 1。
教师是一个标准配置的 TinyAgent(157k 参数),先在弱行为上训练以掌握格式,再在 oracle 示范和恢复示范上训练。另一个干净教师(clean teacher)只看过 20 条 oracle 示范;16.4 节用它说明恢复数据为什么重要。学生是一个更小的 TinyAgent,宽度 32、只有一层,只在弱行为上训练过,对应一个懂格式但没有技能的小型预训练模型。
16.2 作为稠密信号的逐 token 反向 KL
结果 RL 每个回合给出一个数。on-policy 蒸馏为每个动作 token 给出一个数 ,并且无论回合是否成功都能得到。要使用它,我们需要它的梯度。 固定一个前缀 ,记 ,。于是 ,利用 , 因为第二个求和是 。这是单个 token 上的策略梯度(policy gradient),其奖励为 ,即负的对数比,只是要最小化而不是最大化。由此有两种实现。采样 token 形式只使用学生实际生成的 token 以及教师对它的对数概率,把 当作普通策略梯度损失中的逐 token 优势(advantage)。它每个 token 只需要教师的一个数。全词表形式精确计算对 的求和,并让自动微分对其求导。它在单个 token 上没有采样噪声,但在每个位置都需要教师的完整分布。 这个推导在固定前缀下成立。完整目标 还通过学生访问哪些前缀而依赖于 ,这会额外引入一个得分函数(score function)项,其中每个 token 的学习信号还包括后续 token 的 KL(逐 token KL 的剩余回报,return-to-go)。实际做法会丢掉它。Thinking Machines 的实现使用“折扣因子为零”(“a discount factor of zero”),所以每个 token 只使用当前位置的 KL 信号(16.6 节)。GKD 式训练在学生生成的序列上计算逐 token 散度,并且不对采样过程求导。我们遵循同样的约定,并明确说明它是一个近似:它优化的是学生已访问前缀上的局部分布差异,没有显式优化早期动作对未来 KL 的影响。 在 logit 层面,两个 KL 方向也不同。对学生的 logits 求导, 反向 KL 在一个 token 上的梯度与学生自己的概率 成正比。如果学生几乎不给教师偏好的 token 分配概率质量,反向 KL 几乎拉不动它;它主要是压低学生当前偏好的东西。前向 KL 则无论如何都以 的力度把教师的 token 拉上来。下一个代码块通过枚举验证采样 token 恒等式,并针对一个错过了教师众数的学生打印两种 logit 梯度。distill_loss 保存全词表版本。它接收同一位置上的学生与教师 logits(与全书一样,对齐到被预测的位置),把教师从计算图中分离,在任一方向上计算逐位置散度,并在动作掩码上取平均。
16.3 实验:蒸馏与结果 RL 的比较
现在我们训练学生基座的三个副本,各训练 40 步,批次结构相同:每步 4 个训练任务,每个任务 8 个回合(每步 32 个回合)。- off-policy 蒸馏:教师采样 32 个回合,学生在其上走一步
sft_loss。 - on-policy 蒸馏:学生采样 32 个回合,教师在一次前向传播中为每个前缀打分,学生走一步
distill_loss(反向 KL)。 - 结果 RL:学生采样,使用带逐任务组基线的 REINFORCE 和隐藏测试奖励,与 15.5 节相同。第 9 章给出完整的讨论。
证据。 本书测得:单个随机种子,评估分别使用 80 和 40 个回合,所以低于约 0.1 的差异属于噪声。“40 步时蒸馏 ≫ 结果 RL”这一排序远超噪声。教师使用了恢复示范;16.4 节测量这一点有多重要。
16.4 暴露偏差与误差累积
平均通过率掩盖了策略失败的方式。假设当学生的历史看起来与教师的相同时,学生在第 轮以概率 犯错,并且在第一次犯错之后,它就处在前缀空间中一个从未训练过的区域,在那里失败的概率接近 1。那么在 轮中,它的成功率约为 ,而处于分布之外的期望轮数按 而不是 增长:一次早期错误会毁掉之后的每一轮。这是 Ross, Gordon, and Bagnell (2010) 对行为克隆(behavior cloning)的经典分析,而区分 off-policy 训练与 on-policy 训练的,正是第二个假设:学生在自己犯错之后会做什么?forced_rollout 用给定的动作列表开启一个回合,这些动作经由与真实 rollout 相同的 token 路径注入,随后让模型继续。我们注入三种错误类型,并测量后续部分的成功率。
- 只在干净示范上训练的教师自己也无法恢复。 干净教师在每个干净开局上都成功(1.00),而在注入错误之后只在 0.00 和 0.10 的任务上成功。教师本身也是在某个分布上训练出来的策略,它在从未见过的前缀上的行为并不是专业能力。on-policy 蒸馏的效果上限取决于教师在学生的前缀上表现如何。
- off-policy 学生继承的是教师的干净行为,而不是它的恢复能力。 完整教师能从两种注入的错误中恢复(1.00),但它自己的样本几乎从不包含错误,所以 off-policy 蒸馏从未向学生展示过恢复。off-policy 学生在干净开局上的成功率为 0.80,犯错之后为 0.03 和 0.00,与干净教师的模式相同。这就是被测量出来的暴露偏差。
- on-policy 学生能从它自己会犯的错误中恢复,而且只限于这些错误。 在一次无效编辑之后,它的成功率为 0.47,而 off-policy 学生为 0.00,因为学生基座会猜测操作,在训练中做过这类编辑,于是教师在那里纠正了它们。它无法从没有终止的调用中恢复(0.00),因为在格式正确的行为上训练的学生基座几乎从不产生这种调用,所以该前缀从未出现在它的训练批次中。on-policy 蒸馏覆盖的是学生自己的错误分布,而不是所有错误。
16.5 全词表与 top-k 教师分布
distill_loss 需要教师在每个位置上的完整分布:每个批次 个数,对于 Qwen3 的 151,936 词表,fp32 下每个 token 约 0.6 MB(3.6 节)。当教师运行在单独的推理服务上时,这些数据必须经过网络传输。因此,系统只为每个位置发送教师的前 个条目,或者只发送被采样 token 的对数概率(16.2 节的采样 token 形式)。截断会改变目标,而改变的方式取决于如何处理缺失的概率质量:
- 重归一化 top- 把两个分布都限制在教师的前 个 token 上并重新归一化。它忽略学生放在该集合之外的任何概率质量,所以一个把概率质量分散到教师排除的 token 上的学生不会因此受到惩罚。当 时,损失恒为 0。
- top- 加尾部桶(tail bucket)对学生和教师都保留教师的 个 token,并把其余所有 token 合并成一个“其他”结果。粗化后分布之间的散度是完整 KL 的下界(数据处理不等式),而学生仍会因为放在教师前 个之外的总概率质量受到惩罚。
tail_bucket 返回 个结果上归一化的对数概率,所以 distill_loss 可以直接把它们当作 logits 使用。下面的循环在若干个 上比较两种截断方式。
陷阱。 先在教师上计算 top-,再把教师其余的 logits 设为 去调用全词表损失,那么只要学生在前 个之外有概率质量,反向 KL 就是无穷大,在浮点运算中会变成 nan 梯度。要么显式地把学生限制在相同的支撑集上,要么保留一个尾部桶。
16.6 近期工作
以下条目是来自公开来源的文献报告结果,查阅日期为 2026-10-13。本书没有复现其中任何一项。- GKD(Agarwal et al., “On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes”,v1 2023-06-23;查阅了摘要)。指出了标准蒸馏中训练与推理之间的不一致,并用教师反馈在学生自己生成的输出上训练学生。它允许使用前向 KL 以外的散度,以适应无法表示教师的学生,并且可以与 RL 微调结合。
- MiniLLM(Gu et al.,v1 2023-06-14;查阅了摘要)。用反向 KL 取代前向 KL,使学生不会高估教师的低概率区域,并以 on-policy 方式优化它。它报告在 120M 到 13B 参数的指令遵循任务上,暴露偏差更低、校准更好。
- Qwen3 Technical Report(Qwen Team,v1 2025-05-14;查阅了第 4.5 节和表 21)。描述了为较小模型进行的由强到弱蒸馏,分两个阶段:先在教师输出上做 off-policy 蒸馏,再做 on-policy 蒸馏,其中“学生模型生成 on-policy 序列用于微调”(“the student model generates on-policy sequences for fine-tuning”),并以最小化与教师 logits 的 KL 散度为目标进行训练。表 21 从同一个经过 off-policy 蒸馏的 Qwen3-8B 检查点(AIME’24 55.0)出发,比较了耗费 17,920 GPU 小时的 RL(67.6)与耗费 1,800 GPU 小时的 on-policy 蒸馏(74.4)。
- On-Policy Distillation(Kevin Lu and Thinking Machines Lab,2025-10-27;查阅了全文)。使用 16.2 节的采样 token 形式:逐 token 优势是负的反向 KL,由教师对学生所采样 token 的对数概率计算得到,并采用“折扣因子为零”(“a discount factor of zero”)。文章说明它不使用 top- logit 蒸馏。从在教师推理轨迹上做过 off-policy SFT 的 Qwen3-8B-Base 出发(AIME’24 上为 60%),on-policy 蒸馏报告在约 150 步内达到 70%,并估计相对于外推的 SFT 可节省 9 到 30 倍的计算量,具体取决于是否计入教师采样。
16.7 实践中的设计问题
问题:教师与学生的分词器不一致。
现象: 蒸馏损失很大且不下降,或者损失下降而学生变差。
原因: distill_loss 按索引逐一比较分布。如果教师与学生的词表顺序不同,索引 对应的是不同的字符串;如果它们切分文本的方式不同,位置 对应的是不同的前缀。
可选方案: (a) 使用来自同一模型家族、共享分词器(tokenizer)的教师和学生,这是常见选择。(b) 对于字符串相同但 id 不同的词表,把教师条目映射到学生 id,对一一映射而言这是精确的。(c) 对于切分方式不同的情况,在序列层面蒸馏(在教师文本上做 off-policy SFT,由学生重新分词),或者对 token 片段做显式对齐,这两种做法都会失去逐 token 信号。
建议做法: (a)。下一个代码块展示前两种失败模式。
问题:教师成本。 现象: on-policy 蒸馏的每一步比 RL 的每一步更慢,或者教师服务成为集群的瓶颈。 原因: on-policy 蒸馏在每一步都要对每个学生 token 运行一次教师前向传播,而教师通常比学生大得多。off-policy 蒸馏每个回合只为教师付一次费用,其数据可以跨轮次、跨学生复用。RL 完全不需要教师,但需要环境和验证器。 可选方案: (a) 只对动作 token 打分(工具 token 不需要目标),在观测占据上下文主体时节省最多。(b) 返回带尾部桶的 top-,而不是完整分布(16.5 节)。(c) 使用采样 token 形式,每个 token 一个数。(d) 先做 off-policy 蒸馏,只在最后阶段做 on-policy 蒸馏,Qwen3 报告和 Thinking Machines 的文章都是这样做的。 建议做法: (d) 加 (a) 和 (c)。下一个代码块根据 16.3 节的 token 计数估计每一步的前向传播成本。
问题:教师在某些任务类型上是错的。 现象: 学生的平均分接近教师,而评估中的某个切片对两者都是零。 原因: 蒸馏损失衡量的是与教师的一致性,而不是正确性。只要教师自信地出错,学生就会被训练成以同样的方式自信地出错,而且用的正是让蒸馏高效的那种稠密信号。 可选方案: (a) 蒸馏前用验证器逐切片验证教师,并在教师失败的切片上丢弃教师信号。(b) 混入结果奖励,。(c) 在该切片上修正或替换教师。 建议做法: (a),并把被丢弃的切片当作单独的 RL 问题或数据问题来处理。我们注入这一失败:一个教师副本在zero任务类型上使用hardcode_tests的示范上做了微调,这是一个在某类任务上用可被投机利用的奖励训练出来的教师的简化示例。
zero 上则完全失败。我们用三种方式把它 on-policy 地蒸馏到学生基座中:普通的反向 KL;反向 KL 加上在同一批次上权重为 的 REINFORCE 损失;以及用上面的逐任务类型验证来门控(gated)同样的混合,使失败任务类型的回合只得到 RL 信号而没有教师信号。
zero 上复制了这个漏洞利用:它在那里的 32 个样本中有 27 个使用 hardcode_tests,通过率为 0。(它在 compare 上 0.22 的分数是一个独立的、单随机种子下由短预算造成的效应;KL + RL 运行在那里达到 0.59。)
以相同权重混入结果 RL 项,并不能在 40 步内修复该任务类型,投机行为反而变得更普遍(32 个中有 29 个)。原因在于组基线:在 zero 任务上,学生的八个样本全是教师的投机行为,全部失败,优势全为 0,所以 RL 项没有可以推动的方向,而 KL 项一直把它拉向投机行为。
在该任务类型上关闭教师信号消除了投机行为,学生在 zero 上最常见的行为变成了根本不做编辑(32 个中有 11 个),但通过率仍是 0,而且这一次运行在 minmax 上也丢失了准确率(0.44):没有教师信号,组内也没有任何成功样本,就没有什么能教会它 add_zero_guard。图 16.4 汇总了这四个策略。
如果教师在某类任务上系统性地出错,单靠模仿它不能纠正这些错误。它必须先由独立检查发现,然后用一个能够探索的信号来解决:在这个切片上,这意味着付出 RL 的成本并保证足够的探索以找到一次成功(练习 6),或者提供新的示范。
小结
- off-policy 蒸馏是在教师回合上做 SFT,在教师前缀上最小化轨迹分布之间的前向 KL。on-policy 蒸馏在学生生成的前缀上最小化逐 token 散度,默认使用反向 KL。它就是带软教师的 DAgger。
- 在固定前缀下,:这是一个奖励为 的单 token 策略梯度。实际做法忽略未来前缀对当前 token 的依赖(折扣为 0)。反向 KL 的 logit 梯度随学生自己的概率缩放,所以当学生错过教师的众数时,起步会变慢。
- 从同一个小型学生出发,在每步 32 个回合、共 40 步的条件下,两种蒸馏方法都远远胜过结果 RL(本书测得)。两者之间的差别体现在恢复能力上,而不是平均通过率:off-policy 学生继承了教师的干净行为,无法从错误中恢复;on-policy 学生能从它自己犯的错误中恢复,但无法从它从未犯过的错误中恢复。
- 教师自身必须在学生的前缀上胜任。只在干净示范上训练的教师无法从注入的错误中恢复,因此也不可能教会恢复。
- top- 教师分布会使目标产生偏差。重归一化 top- 丢弃了学生在教师支撑集之外的概率质量,并在 时给出零损失。尾部桶只多花一个数,在每个 下都与完整梯度保持接近。
- 蒸馏信号衡量的是一致性,而不是正确性。教师在一个任务类型上的错误被原样传递给了学生,而额外的 RL 项和门控都没能在预算内修复它,因为该任务类型上的每个组奖励方差都为零。
练习
- 实现 on-policy 蒸馏的采样 token 形式(逐 token 优势 ,从计算图中分离,用于 REINFORCE 损失),并在相同步数下比较它与
distill_loss的学习曲线和梯度方差。采样形式在每个前缀上需要多少个样本,才能使其与真实梯度的余弦达到 16.5 节中全词表梯度的水平? - 改变一个假设:让学生与教师一样大。off-policy 蒸馏和 on-policy 蒸馏是否仍然殊途同归?然后把学生变得更小(
d_model=16),比较前向 KL 与反向 KL,并用mass-covering 和 mode-seeking来解释其中的差异。 - 加回实际做法丢掉的剩余回报项:用某个 token 之后各 token 的逐 token KL 之和(折扣 )为该 token 的得分函数加权,并在 16.4 节的恢复测量上比较 、 和 。
- 在 on-policy 蒸馏过程中,以每轮 0.1 的概率向学生的 rollout 注入没有终止的调用这一错误(并把注入的 token 从损失中掩码掉),然后重新测量恢复能力。把这一做法与 DART 式的噪声注入以及 16.4 节发现的局限联系起来。
- 在教师温度为 0.5、1 和 2 时(在 softmax 之前把教师 logits 除以温度)运行 top- 测量。要使梯度余弦达到 0.99 所需的 如何随教师的熵变化?
- 修复 16.7 节的
zero任务类型:门控掉教师,并只在被门控的任务上提高学生的 rollout 温度。什么样的温度能让组奖励方差足够频繁地非零,从而让 RL 找到add_zero_guard?这对其他任务类型有什么影响? - 设计评审:你必须用一个同家族的 32B 教师、一个沙箱化的基于测试的验证器以及固定的 GPU 预算,对一个 1.5B 参数的编程智能体做后训练。列出各个阶段(数据、off-policy 蒸馏、on-policy 蒸馏、RL),说明你会测量什么来决定何时从一个阶段进入下一个阶段,以及如何在学生复制之前发现一个在某个代码仓库或某种语言上出错的教师。