banner
约 4,500 字
15 分钟

On-Policy Distillation 技术阅读

摘要

完整解释学生采样、固定教师概率评分、逐token优势与更新流程,核对数学蒸馏、知识与行为恢复、单提示实验及成本口径,说明方法怎样迁移到自己的模型。

On-policy distillation 让学生模型生成自己的轨迹,再让固定教师计算这些轨迹中每个 token 的概率,依据教师与学生的概率差训练学生。它在学生实际会访问的状态上提供密集信号,补充只模仿教师答案和只看末端奖励的限制。

Thinking Machines 的这篇文章把方法用于数学推理和领域助手:前者比较不同后训练方式,后者先学习内部文档,再用旧模型恢复被破坏的指令行为。核心问题是采样来自谁、教师如何评分、哪些参数更新,以及已经学到的知识与行为能否同时保留。实验给出有条件的收益,也留下教师偏差、起点能力和成本口径等边界。

信息

内容

原文

On-Policy Distillation

作者与日期

Kevin Lu 与 Thinking Machines Lab,2025-10-27

类型与 DOI

技术研究文章;10.64434/tml.20251026,DOI 日期标记与页面发布日期不同

阅读范围

英文全文,归档 PDF 21 页

关键对象

学生当前策略、固定教师、学生轨迹上的逐 token log probability

主要证据

Qwen3 数学蒸馏、领域知识与 IFEval 对照、同架构 RL 教师再蒸馏、单提示实验

来源与实现

原文,Tinker cookbook

先分清三种训练信号

监督微调可以使用教师生成的完整答案,学生在教师轨迹上学习下一个 token。通常称为 off-policy,是因为训练路径由另一个策略产生。学生在实际生成中犯错后,会进入训练数据没有覆盖的前缀,后续错误可能累积。

On-policy RL 由学生自己生成,然后按实际结果给奖励。如果一条长推理只有最终正确或错误,它能告诉模型哪条完整轨迹更好,但没有直接提供每一步的局部质量信息。

On-policy distillation 同样让学生生成,教师在这些学生前缀上提供 token 分布信息。因此学生既遇到自己的状态,又获得比一次末端分数更细的监督。

方法

主要轨迹来源

训练信号

需要的额外条件

教师答案 SFT

教师或标注数据

目标序列上的 token 监督

合适的答案、格式与覆盖

末端奖励 RL

学生

结果奖励与 advantage

可用奖励、有效探索与采样

本文 on-policy distillation

学生

教师与学生逐 token 概率差

教师能对学生序列计算 log probability

这里的密集指信号随 token 提供,不能理解成每个 token 都有正确性的人工标签。教师条件概率包含它对当前前缀的行为判断,仍可能受前缀错误或教师能力限制影响。

一次更新怎样走通

下面的总览按原文 Implementation 与 Pseudocode 整理。教师只参与前向评分,学生负责采样和参数更新。

按原文整理的 on-policy 蒸馏流程
按原文整理的 on-policy 蒸馏流程

依据:原文 Implementation,归档 PDF 第 5—7 页。(图由AI辅助绘制)

准备提示与两个模型。 任务提示来自训练集合,学生当前模型承担生成,固定教师承担概率计算。教师可以是更大模型,也可以是希望恢复行为的旧版本模型。教师大小与适用能力是两件事。

学生生成并保存采样概率。 对每个提示生成一条或多条轨迹,记录学生在各 token 的 log probability。轨迹包含实际由学生选择的前缀,后续教师评分必须以同一前缀为条件。

教师重放输入,计算学生 token 的概率。 教师接收提示与学生序列,计算每个已经生成 token 的 log probability。它无需逐 token 自己重新生成一遍答案,也不需要梯度更新。长序列评分可以以批处理前向计算实现,性能取决于模型与服务。

构造逐 token advantage。 使用教师 log probability 减去学生采样 log probability,形成当前 token 的反馈。保存采样时的概率很重要:若采样与训练模型版本已经不同,更新需要对应的重要性采样处理。

更新学生,再重新采样。 原文沿用 RL 的重要性采样损失执行反向,只有学生参数更新。下一轮轨迹来自更新后的学生,教师仍固定,训练继续针对学生现在会访问的前缀。

这种流程不需要单独生成 teacher completion 数据集,但每次学生轨迹仍需要教师评分。不能把它称为完全不使用教师算力,也不能用一段普通聊天 API 文本评分自动替代 token log probability。

Reverse KL 与逐 token 信号的含义

记学生策略为 π_θ、教师为 π_T,共同前缀为 h。原文选择反向 KL:

DKL(πθ∥πT)=Ea∼πθ(⋅∣h)[log⁡πθ(a∣h)−log⁡πT(a∣h)].D_{\mathrm{KL}}(\pi_\theta\|\pi_T) =\mathbb E_{a\sim\pi_\theta(\cdot\mid h)} \left[\log\pi_\theta(a\mid h)-\log\pi_T(a\mid h)\right].

采样来自学生,计算的是学生选择的 token 在两种策略中的概率差。原文将逐 token advantage 设置成负的采样 log ratio:

At=log⁡πT(yt∣ht)−log⁡πsample(yt∣ht).A_t=\log\pi_T(y_t\mid h_t)-\log\pi_{\mathrm{sample}}(y_t\mid h_t).

若教师更支持某个学生已生成 token,这一项可能为正;若学生过度偏爱教师不支持的 token,可能为负。单个采样项允许正负变化,完整分布上的 KL 才具有非负性质。把每个 token 的 log ratio 都当成精确且非负的 KL,会误解颜色图与 advantage。

一个教学算例:同一前缀下学生选某 token 的概率为 0.2,教师为 0.6,信号为 ln(0.6/0.2)≈1.10;另一个 token 学生为 0.1、教师为 0.01,信号约为 -2.30。数值仅解释概率比,未来自作者实验。

原文为简化使用 discount factor=0,即 token 只用当前概率信号,不将未来 token 的反馈累计为当前 advantage。它区别于长时程回报归因,也影响短片段训练为什么可以进行。学生仍通过各位置更新学习整体行为,但实现没有显式求每一步对最终任务结果的全部影响。

为什么错误终点可能没有明显惩罚

原文用 SimpleBench 的冰块问题举例,学生忽视题目物理前提,把问题按纯数量算术处理。教师对早期使路线偏离的 token 给出较低概率;在已经进入错误前缀以后,最终错误数字可能反而很符合该前缀,因此没有最大的局部惩罚。

原文图 6:学生轨迹的教师概率反馈
原文图 6:学生轨迹的教师概率反馈

来源:原文 Figure 6,归档 PDF 第 6 页。颜色展示该实例的 token 信号,图中最终答案需结合此前错误路线理解。

这说明教师评分是条件行为信号。它能定位部分路线分叉,也可能顺着错误前缀给出高概率的后续。若任务必须满足实际环境约束,仍应保留结果测试、事实核验或外部奖励。

作者用相对教师的分布一致性解释这种目标不容易通过传统奖励捷径优化,但教师分布不等于现实正确性。如果教师错误、偏好不适合任务或不会使用当前工具,低 KL 仍可能得到不理想结果。对自己的项目应分别测与教师的相似度和真实任务效果。

教学伪代码与实现依赖

下面依据原文伪代码整理,变量表达流程,未在项目中运行,也不承诺是当前库可直接复制的接口。

Python
# 教师参数固定,采样与更新由学生完成
rollouts = sample_student(prompts)
student_logprobs = rollouts.sampling_logprobs
teacher_logprobs = score_teacher(rollouts.prompts, rollouts.tokens)

# 对齐相同前缀、相同目标 token;仅目标生成位置参与训练
advantages = teacher_logprobs - student_logprobs
update_student_with_importance_sampling(
    rollouts,
    sampling_logprobs=student_logprobs,
    advantages=advantages,
)

实现时需核对 tokenizer、chat template、角色与特殊 token,确保两套概率对应同一文本位置。不同词表不能简单把学生 token ID 当成教师的相同 ID;不同分词方式下如何对齐,需要额外方法,不能假装差两个数组就完成了跨模型蒸馏。

还要区分提示部分和生成部分,正确处理 padding、终止 token、截断与训练 mask。教师概率和采样概率都应视为该次轨迹的固定数据,反向传播不能通过教师;训练中的当前学生 log probability 用于损失与采样修正。

原文使用 compute_logprobs,这是对给定序列评分的能力要求。文本生成接口即使能返回少量自己输出的概率,也不一定能对任意学生序列重放评分。应根据实际服务确认这一点,再估计实现成本。

数学推理:起点、实测与外推分别看

作者用 Qwen3-8B-Base 作学生,先在 OpenThoughts-3 上做教师轨迹 SFT,数据 completion 来自 QwQ-32B。后续 on-policy 阶段的教师是 Qwen3-32B,两阶段教师并非同一个模型。

400K prompts 的全参数 SFT 检查点在 AIME’24 上达到约 60%,作为后训练起点。原文再比较增加 SFT、RL 与 on-policy distillation 的代价。

原文图 7:增加 SFT 数据的收益
原文图 7:增加 SFT 数据的收益

来源:Figure 7,归档 PDF 第 9 页。实线是已有数据范围,虚线包含对更大训练规模的外推。

继续按观察到的近似 log-linear 趋势推算,作者估计约 2M prompts 可能达到 70%。这是外推,依赖曲线继续成立;不能写成该团队已经训练完成并测得 2M 的结果。

从同一 400K SFT 起点做 on-policy distillation,作者在约 150 步达到约 70%。不同 LoRA 容量也得到改进,部分设置仍低于全参数训练。密集监督提高学习效率,无法消除所有容量和优化差异。

原文图 9:on-policy 蒸馏的增量收益
原文图 9:on-policy 蒸馏的增量收益

来源:Figure 9,归档 PDF 第 10 页。横轴表示作者计量的增量计算,三个面板使用不同参数更新设置。

文章另引用 Qwen3 技术报告:SFT 起点为 55.0,RL 为 67.6、17,920 GPU hours,on-policy distillation 为 74.4、约 1,800 GPU hours。该表来自另一团队及其配方,不能与 Thinking Machines 的 60→70 直接拼成同一实验。

9—30 倍来自哪种比较

作者分别计算教师与学生 FLOPs。若已有 SFT 数据集可直接使用,或生成成本在多次训练中摊销,增加 SFT 的比较不再收取完整教师生成成本;on-policy 蒸馏仍计入教师对新轨迹的评分,得到约 9 倍的计算效率改善。

若新任务尚无教师数据,需要重新生成大规模 completion,把教师生成也计入 off-policy 成本,估计改善约 30 倍。教师前向概率评分可有效并行,实际 GPU hours 比例又可能与 FLOPs 不同。作者在特定口径下讨论约 18 倍的 GPU hours 改善,不能将几种数字同时当成独立实测承诺。

结果或比较

证据性质

使用时的条件

SFT-400K 约 60%

作者实测起点

Qwen3-8B-Base 与指定数据、训练设置

SFT-2M 约 70%

曲线外推

假设扩展趋势不提前停滞

OPD 约 150 步到 70%

作者后训练实验

同一 SFT 起点、Qwen3 教师与概率评分

9—30 倍

成本模型与不同计费边界

是否计入离线数据教师生成、评分与并行效率

Qwen3 表中的 74.4

原文引用的另一报告

原 Qwen3 配方,不能与上述起点混同

对自己的训练,应记录离线数据生成、学生采样、教师评分、学生反向、评价与基础设施开销。已有数据免费可用与每次新生成数据,属于不同预算问题。

领域助手:学知识与保行为为什么会互相影响

第二个任务使用已经后训练的 Qwen3-8B,将公司内部文档加入训练。评价分别检查内部知识问答与 IFEval。原模型没有见过内部文档,再大也不能凭参数记忆准确回答所有内部事实,所以首先需要提供知识学习信号。

作者同时准备背景聊天数据:使用 Tulu3 prompts,让训练前的 Qwen3-8B 重新生成回答,混入中训练以尽量保留原行为。与更大教师生成的背景数据相比,原模型自己的分布在其设置中更适合保留原能力。数据来自更强模型,不保证对当前目标更合适。

增加文档比例提高知识成绩,却降低指令遵循;至少约 30% 背景聊天可减轻退化,但原文所测比例都没有完全保持原 IFEval。LoRA 限制参数更新仍不足以阻止遗忘,知识学习也较少。

原文图 10:领域数据与背景数据比例
原文图 10:领域数据与背景数据比例

来源:Figure 10,归档 PDF 第 13 页。两个维度分别是内部知识与指令行为,不能用其中一个代替整体模型质量。

用旧模型恢复行为

在完成知识中训练后,作者对 Tulu3 prompts 做 on-policy distillation,教师是训练前的 Qwen3-8B。这个教师与学生规模相同,目标是恢复原有助手行为。此阶段不使用内部文档,训练信号来自旧模型的条件分布。

模型阶段

内部 QA

IFEval

原 Qwen3-8B

18%

85%

100% 内部文档中训练

43%

45%

70% 文档混合中训练

36%

79%

70% 混合后再蒸馏

41%

83%

来源:原文 Figure 13,归档 PDF 第 16 页。后一阶段接近恢复指令行为,同时在该内部 QA 上未损失新知识,甚至有正向变化;IFEval 83% 仍与原 85% 有差距。

这个实验展示知识与行为可以通过阶段分工处理,不证明任意多轮持续学习都不会遗忘。内部数据与评价细节公开程度有限,读者无法仅凭文章完整复现该任务。自己的模型应在领域、通用指令和旧任务回归上持续测量,特别检查教师不掌握的新知识是否受到蒸馏压制。

同架构教师实验:密集信号是否加快学习

作者从 Qwen3-8B-Base 开始,用 rank=128 的 LoRA 在 DeepMath 上做 RL,将得到的模型作为教师,再将其行为 on-policy 蒸馏回相同 Base 起点。

原文图 14:RL 学到的策略再蒸馏
原文图 14:RL 学到的策略再蒸馏

来源:Figure 14,归档 PDF 第 17 页。两种路线的起点与目标教师相关,图中比较梯度步和训练 KL。

蒸馏在不到约 10 个梯度步恢复相应 AIME 表现,而 RL 约需 70 步,作者总结为 7—10 倍步数差。再结合较短轨迹与较小 batch,讨论约 50—100 倍计算差。它包含不同训练长度、batch 与已存在教师的条件,不能直接说明教师的完整发现成本也消失了。

作者将 RL 的重要成本解释为寻找和检验策略,而蒸馏在策略已经存在后学习它,绕过发现过程。这个解释有助于理解训练分工。数学上两者仍通过参数更新学习,不应把语义策略搜索解释成脱离权重优化的另一套物理过程。

信息量讨论也应限定范围。末端标量奖励的策略梯度信号较稀疏,逐 token 教师概率更密集;带过程反馈、模型式 RL 或环境观察监督的其他算法,可以有不同信息结构。

一条提示的实验到底说明什么

文章只选一个随机数学 prompt,连续训练 20 步,每步 256 个学生 rollout,共 5120 条评分轨迹,仍近似达到相应教师的 AIME 表现。

原文图 15:单提示、多轨迹蒸馏
原文图 15:单提示、多轨迹蒸馏

来源:Figure 15,归档 PDF 第 18 页。不同曲线的提示复用、每批样本与训练 KL 口径见原图说明。

一条 prompt 仍产生大量不同学生前缀,教师提供逐 token 监督;底座已经拥有广泛基础,教师也经过先前 RL。因此实验不能概括为只用一条普通训练答案,就从零学会整个数学领域。

这个结果支持在所测条件下,反复采样同一提示仍可包含有用行为信息。它没有证明任意任务都只需一个提示,也没有解决提示覆盖、教师偏差和未见领域的泛化。训练 KL 下降与真实任务进步仍需分开记录。

自身样本 SFT 与持续 on-policy 训练的差别

原文还观察,先从 Qwen3-32B 的分布生成一份固定数据,再对模型自身做 SFT,仍会降低 IFEval。理论分布一致不代表有限数据批次的梯度为零,学生更新后,固定旧轨迹也逐渐不再来自当前策略。

On-policy distillation 每轮重新采样学生,教师固定;当两者条件分布相同,概率差信号趋近零,因此行为恢复有不同约束。实际数值精度、概率引擎、有限批次和训练目标仍会影响稳定性,不能将该机制写成绝无退化的保证。

对持续学习,可以交替做新领域知识训练与旧行为恢复,但应保留旧任务评价和新知识评价。若教师在新领域不知道答案,蒸馏数据的选择尤其重要:旧行为 prompts 与领域知识 prompts 的任务目标不同,混在一起可能改变想保留的能力。

给自己的微调与 Agent 训练选路线

已有高质量教师答案时,先用 SFT 建立格式、领域与行为基础。学生能够采到有价值轨迹、教师能对这些轨迹评分时,可以比较 on-policy 蒸馏。任务依赖真实环境正确性或需要超过教师的结果时,还应研究环境奖励与 RL。

选择教师时检查目标能力、工具格式与语言分布。恢复旧行为可使用旧检查点;学习某项新能力则需要对应更强教师。模型规模只能作为候选条件,不能替代任务评价。

Agent 轨迹涉及工具消息、观察和行动参数,普通数学文本蒸馏不自动适用。哪些位置参加 loss、教师是否能看懂环境结果、错误恢复路径是否被采到,以及真实任务是否完成,都需要单独定义。对 API 受限模型,要先确认任意序列概率评分能力,再决定实现。

建议建立三个对照:固定教师数据 SFT、学生轨迹 on-policy 蒸馏、任务奖励学习。保持底座、评价、预算尽量可比,同时记录真实任务成功、领域知识、旧行为回归、生成长度、教师评分成本与学生训练成本。本文没有在你的模型上运行这些实验,原文的效率数字应作为研究动机,最终判断仍来自这些对照。

来源与复读位置

  1. Thinking Machines 原文,2025-10-27。训练信号比较与实现见归档 PDF 第 1—7 页,数学实验与成本计量见第 8—11 页。

  2. 领域助手、旧模型教师与知识/行为对照见第 12—16 页;匹配教师、单提示与持续学习实验见第 16—20 页。

  3. Tinker cookbook包含作者实现入口。本文伪代码按原文整理,未进行运行或复现;Qwen3 表属于原文引用结果,SFT-2M 为明确标注的外推。

END