很多人把“知识蒸馏”理解成“把一个大模型的结论搬给一个小模型”,但真正到了强化学习场景,事情没有那么简单。
最近有一类论文在较真一个问题:On-Policy 蒸馏到底是不是真的在蒸馏?这个问题的杀伤力在于,它直接挑战了训练流程里的命名。很多工作声称自己在蒸馏教师模型,但细看训练目标后发现,教师只是被当成一个给分机器,学生学到的东西本质上来自策略梯度,而不是来自教师的知识迁移。
这篇文章想把这个争议拆开。我会先理清“知识蒸馏”和“On-Policy”这两个概念在数学上分别约束了什么,再分析什么样的训练流程配得上“蒸馏”这个名字,什么样的流程只是“用教师分数做强化学习”。最后会讨论一种被称为 OPSA 的无需额外监督的 On-Policy 蒸馏思路,并给出可以落地的理解框架与伪代码。
读完你会得到三个判断标准:教师的什么信息进入了学生、学生从哪里采样数据、监督信号的维度是 token 级还是 scalar 级。这三个标准能帮你快速识别项目里那些“名为蒸馏、实为强化”的方案。
1. 先看争议的本质:教师到底教了什么
“蒸馏”这个词,在自然语言和日常比喻里被用得很泛。社区里有人讨论怎么把一本书“蒸馏”成笔记,也有人问怎么把某个 Skill“蒸馏”给小模型。这些说法很形象,但也带来了一个隐患:把结果的浓缩当成了知识的传递。
在模型层面,知识蒸馏有相对严格的数学形式。教师模型不是只告诉学生“这个答案对不对”,而是告诉学生“在每一个候选词上,我认为概率是多少”。学生看到的是一张完整的概率分布,而不是一个离散的打分。
那么在 On-Policy 蒸馏里发生了什么?
常见的做法是:从当前学生策略中采样一批回答,用教师模型对这些回答做评价,再把评价分数作为强化信号,更新学生策略。这种流程确实涉及了教师模型,也确实是在线策略更新,但教师给出的信息是“你这个回答值 0.8 分”或者“这个回答比另一个回答好”。学生收到的是一个标量,不是一份逐 token 的“教案”。
问题的核心就在于这里:监督信号的维度,决定了我们能不能把它称作蒸馏。
如果教师给的是完整 logits 或完整的纠正后文本,学生可以精确地知道“哪个词应该被改成哪个词”。如果教师只给一个分数,学生只知道“这次采样好还是不好”,至于好在哪里、该往哪个方向改词,学生需要自己去探索。后者更接近强化学习,而不是知识蒸馏。
论文里之所以要把这个区别挑明,是因为很多工作在标题和摘要里都用了 distillation,实验设置却是标准的策略优化。命名本身不会让方法失效,但会误导读者对机制的理解。你不知道自己到底在做什么,就很难判断训练为什么失败。
2. 先弄清楚两个基础概念:知识蒸馏与 On-Policy 优化
2.1 知识蒸馏的经典定义
在分类任务中,知识蒸馏通常用一个温度参数软化教师的输出概率,然后让学生去拟合这个软标签。
设教师模型输出的 logits 为 z_t,学生模型输出的 logits 为 z_s,温度系数为 τ,则软化后的概率为:
p_i = exp(z_i / τ) / sum_j exp(z_j / τ)蒸馏损失可以写成:
L_KD = KL(p_t || p_s)这个公式的重点不是交叉熵本身,而是教师输出中包含的“暗知识”。对于一张图片,教师可能以 0.7 的概率认为是猫,0.2 的概率认为是狗,0.1 的概率认为是狐狸。这种不确定性分布包含了类别之间的相似性信息,而离散标签无法表达这些。
在生成模型中,蒸馏可以发生在两个层级。
第一层是 token 级 logits 蒸馏。教师在每个生成位置的词表分布都会透传给到学生,这是一种高信息密度的监督。
第二层是序列级蒸馏。教师先生成一段完整文本 y_t,学生用这段文本做交叉熵训练。这里丢失了教师在每个位置的概率不确定性,但仍然保留了完整的正确答案,学生可以学到“在什么输入下应该输出什么内容”。
如果把教师当成分数器,只回复一个 reward,那么信息密度会进一步降低。
2.2 On-Policy 的准确含义
On-Policy 是强化学习里的概念。它指训练数据必须由当前正在更新的策略产生。
策略记为 π_θ,训练时从 π_θ 中采样一批动作 a,然后用这些动作产生的奖励去更新 θ。更新完成后,策略变成了 π_θ',下一轮又要重新从 π_θ' 采样。
与之相对,Off-Policy 允许使用旧策略或外部策略产生的经验数据。比如从固定数据集里读取一批动作和奖励,多次更新模型。
On-Policy 的价值在于避免分布偏移。旧数据可能来自一个完全不同的策略,用它更新当前策略时会产生偏差。尤其在语言模型中,如果训练数据和当前策略的采样风格差异过大,学生模型的输出分布会被拉向一个并不匹配的区域。
2.3 两个概念叠在一起后的歧义
当“蒸馏”和“On-Policy”组合在一起时,可能产生两种完全不同的理解。
第一种理解:在策略采样环节,样本直接来自学生当前策略。学生的回答被交给教师做 token 级监督,教师返回的是修正后的完整回答或对应位置的 logits。
第二种理解:在策略优化环节,样本来自学生当前策略,但教师的监督只是一个分数,学生用这个分数作为做策略梯度更新。
这两种流程都叫 On-Policy Distillation,但机制完全不同。前者更像“教师带学生做练习”,后者更像“教师当考官但不出辅导题”。
论文的核心关切集中在第二种。因为在这种流程里,真正起作用的是策略梯度公式,而不是蒸馏损失。
3. 判别标准:什么样的流程算“真蒸馏”
判断一个训练流程是否是蒸馏,不要看它的项目名,也不要看它有没有调用教师模型,而要看训练目标里信息的流向。
3.1 教师的 logits 作为监督:真蒸馏
每个位置教师都会给出完整的词表概率。学生需要让自己的概率分布靠近教师的概率分布。
这种做法的条件是教师和学生共享词表,且需要保存教师的 logits 或在线前向传播。
它的优点是监督信息细。教师对最佳词、次佳词、风险词都有判断,学生可以从中学到丰富的语言偏好。缺点在于训练和推理开销高,大模型场景下 logits 非常大。
3.2 教师的生成文本作为监督:序列蒸馏,符合蒸馏的直观定义
教师对 prompt 生成一段文本 y_t,学生在给定 y_t 时做最大似然训练。
这种做法的信息量低于 logits,因为教师对非采样词的概率没有被显式传递。但从知识迁移的角度看,学生知道了一个完整且可执行的答案,这个答案可以在实际推理中复现出来。
如果教师生成文本时做了采样或 beam search,那么这个文本通常代表教师比较偏好的路径。学生学到的是一条行为轨迹,因此可以叫蒸馏。
3.3 教师只输出一个标量分数:不是蒸馏,是奖励信号
假设学生采样出回答 y_s,教师给出一个分数 r。
训练目标可以写成:
L = - r(y_s) * log π_θ(y_s | x)这个公式和 REINFORCE 没有本质区别。即使我们把 r 叫作 teacher score,把流程叫作 reward-conditioned distillation,模型的更新方向依然是由分数对采样概率的加权决定的。
这种流程能工作吗?能。在很多场景里,教师确实可以起到奖励模型的作用,驱动学生策略向高质量回答偏移。但它不能被称为蒸馏,因为没有把教师的“决策过程”迁移给学生。
教师知识被压缩成了一个标量,学生只能知道“这个效果好”,而不知道“教师为什么认为这个效果好”。长此以往,学生会在已有采样空间里做排序优化,而不是从教师的知识分布中获取新内容。
因此论文提出了一个干净的判断:如果教师的输出空间是 Reward 或 Ranking,那么它是在做策略优化;如果教师的输出空间是 Token 分布或完整 Token 序列,那么它才有资格叫蒸馏。
3.4 对比表格
| 教师输出形式 | 信息维度 | 训练损失本质 | 是否适合叫蒸馏 |
|---|---|---|---|
| 完整 logits | 词表级概率分布 | KL 散度 / 软交叉熵 | 非常适合 |
| 完整回答文本 | Token 序列 | 交叉熵 | 适合 |
| 修正后的重写文本 | Token 序列 | 交叉熵 | 属于蒸馏与纠错结合 |
| 标量奖励分数 | 一个实数 | REINFORCE / PPO | 应叫奖励学习 |
| 成对偏好排序 | 序关系 | Rank Loss / DPO | 应叫偏好优化 |
这个表格的价值在于,它把所有“教师参与”的方案都排了一遍。只有前三种配得上知识蒸馏的名字。
4. On-Policy 场景下被忽略的偏差:教师只能评价学生能想到的答案
如果把教师当成分数器,On-Policy 采样还有一个更隐蔽的问题:学生永远只从自己的策略里采样,教师也因此只能看到学生已经能够表达的内容。
假设学生的策略目前只会输出三种模板:
- 模板 A:直接给结论,缺少推导。
- 模板 B:先给解释再给结论,但解释太长。
- 模板 C:用列表分点说明,但不完整。
On-Policy 流程从这三种模板中采样,教师给分,学生根据分数调整三种模板的概率。这个流程可以帮助学生在 A、B、C 之间做出选择,但学生永远无法通过这个流程学到模板 D:先写摘要再给例子最后给公式。
因为模板 D 根本不在采样空间里。“蒸馏”这个词隐含了教师应该把没见过的新东西教给学生,而只靠 On-Policy 采样和教师打分,这件事无法发生。
换个角度说,从学生策略采样相当于让教师批改学生已经写完的卷子。批改能提高学生对已有题型的熟练度,却很难让学生凭空掌握教师掌握的新题型。
真正要让教师知识进入学生,需要一个“外挂”通道:教师直接生成高质量回答,然后把这段回答作为训练目标。这就是 OPSA 方法希望解决的核心问题。
5. OPSA 的设计思路:如何在无需额外监督的情况下保留两种通道
OPSA 并不是完全放弃 On-Policy。它保留了“从学生当前策略采样”的稳定性,同时引入了教师的示范通道。这里的“无需监督”主要有两层含义。
第一,不需要额外训练一个奖励模型。很多蒸馏方案会把教师模型或者专门的 Reward Model 当成评分器,但 Reward Model 本身就是额外的监督来源。OPSA 希望只使用一个冻结的教师模型完成评价和示范。
第二,不需要人工标注偏好数据。传统偏好优化需要人类标注员对多个回答打分。OPSA 避免增加这一步。
在具体设计上,OPSA 对每个 prompt 样本做了两路采样:
- 学生当前策略 π_θ 生成回答 y_s。
- 参考教师模型 π_ref 生成回答 y_t 或者一组候选回答。
教师模型除了生成回答,还会给出学生回答的条件对数概率 log p_ref(y_s | x)。这个对数概率不来自额外 Reward Model,而是来自模型自身的 next-token 预测能力。用更高能力的教师去计算当前学生回答的单句概率,可以作为一条不需要人类标注的参考信号。
这里的关键是作者并不把教师对数概率直接当作奖励去更新策略,而是当作一个监督方式选择的开关。
当教师对学生的回答 y_s 给出较高认可度时,说明学生已经落在安全区域,那么这一轮可以做较保守的更新,例如 KL 正则或轻微的策略平滑,避免采样噪声造成剧烈波动。
当教师对学生回答 y_s 的认可度明显较低时,说明学生的回答不够好。此时如果走标量奖励更新,学生只知道“不好”,不知道“怎么改”。OPSA 会转向另一条通道:把教师生成的 y_t 作为完整目标,用交叉熵监督学生。
这个设计可以用一个简化的加权公式理解:
L_OPSA = λ_s * KL(π_θ || π_ref) + λ_t * CE(π_θ, y_t)系数 λ_t 与教师对学生的认可度负相关。认可度越低,λ_t 越大,教师示范的监督权重越强。认可度越高,λ_t 越小,主要保留 On-Policy 的稳定性。
这样做的好处是,既利用了 On-Policy 采样天然贴近当前策略分布的优势,又避免了信息被压缩成标量的问题。教师示范通道保证了学生的输出空间可以扩展,不必永远停留在自己已经探索过的模板里。
从批判角度看,OPSA 最大的贡献是提出了一个显式划分:评价信号只做门控,不做最终的学习目标;最终学习目标保留 token 级监督。
这不是一个绝对完美的方案,但它在概念上解决了“Is On-Policy Distillation Really Distilling?”这个标题问题。
6. 最小示例:从三行代码看三种方法的差异
为了更直观地说明监督信号的区别,我们用一组不依赖特定大模型的 PyTorch 风格代码做对比。这里不追求庞大的训练循环,只展示损失函数层面的差异。
6.1 第一种:真正的 logits 蒸馏
import torch import torch.nn.functional as F def kl_logits_distill(student_logits, teacher_logits, temperature=2.0): """ student_logits: [batch, seq_len, vocab] teacher_logits: [batch, seq_len, vocab] """ p_teacher = F.log_softmax(teacher_logits / temperature, dim=-1) p_student = F.log_softmax(student_logits / temperature, dim=-1) loss = F.kl_div( p_student, p_teacher, reduction="batchmean", log_target=True ) * (temperature ** 2) return loss这里的损失函数比较的是教师在完整词表上的概率分布。学生得到的监督信息不只是一个正确答案,而是每个候选词应分配多少概率。
6.2 第二种:用教师分数做 REINFORCE
def teacher_score_reinforce(student_logits, actions, teacher_scores): """ actions: [batch, seq_len],学生采样得到的 token teacher_scores: [batch, 1],教师返回的标量分数 """ log_probs = F.log_softmax(student_logits, dim=-1) action_log_probs = torch.gather( log_probs, dim=-1, index=actions.unsqueeze(-1) ).squeeze(-1).sum(-1) # 教师分数作为 reward,走策略梯度 loss = -(teacher_scores * action_log_probs).mean() return loss注意,这段代码里教师只给了一个标量分数。即使教师模型内部再强大,更新学生时能够利用的信息也只有这个标量。学生无法根据这个损失学会它没有采样到的新 token 组合。
6.3 第三种:OPSA 的思路示意
def opsa_loss(student_logits, student_answer_ids, teacher_answer_ids, teacher_logprob_for_student, threshold=0.0, kl_weight=0.1): """ 一个简化的 OPSA 概念实现。 teacher_logprob_for_student: 教师模型对学生回答计算的对数概率 """ if teacher_logprob_for_student < threshold: # 学生回答质量不够,教师示范进入 token 监督通道 ce_loss = F.cross_entropy( student_logits.reshape(-1, student_logits.size(-1)), teacher_answer_ids.reshape(-1) ) return ce_loss else: # 学生回答处于安全区域,使用 KL 正则做保守更新 # 这里可替换成与参考策略或旧策略的 KL log_probs = F.log_softmax(student_logits, dim=-1) uniform_logits = torch.zeros_like(student_logits) kl_loss = F.kl_div( log_probs, F.log_softmax(uniform_logits, dim=-1), reduction="batchmean", log_target=True ) return kl_weight * kl_loss这段代码做了明显的逻辑简化,实际的 OPSA 实现会复杂得多,比如门控信号的选择需要平滑、需要避免人工设定的阈值过于武断、需要处理 batch 内的差异化路由。但核心已经展示出来:
- 不把教师分数直接乘到 token log-prob 上。
- 教师分数只用于判断“应该走示范通道”还是“安全保守区域”。
- 一旦走示范通道,监督目标是教师的 token 序列,而不是教师分数本身。
这样,学生既能享受到 On-Policy 采样的稳定性,又能获得真正的 token 级知识输入。
6.4 如何验证这段代码的差异
可以在小型的离散概率分布上做测试。给一个词汇表大小为 5 的随机分布,分别用三种方式计算 loss,观察损失对教师输出信息的敏感程度。
如果教师只改变另一个非采样 token 的概率分布,而完全不影响学生采样到的 token,那么第二种方法(REINFORCE)的损失不会发生任何变化。这是非常直观的证据:教师对非采样 token 的知识在标量奖励通道中被彻底浪费了。
第一种和第三种方法则会在这种情况下产生不同的梯度。因为它们都直接看到了完整分布或教师完整回答。
这个实验不需要很大的算力,适合作为论文复现前的初步验证。
7. 常见的训练问题与排查思路
在实际项目中,实现 On-Policy 蒸馏或 OPSA 类方法时,会经常遇到训练不收敛或收益不明显的现象。这里整理几个高频问题。
| 现象 | 可能原因 | 排查方式 | 解决方案 |
|---|---|---|---|
| 训练 loss 下降但下游评测不提升 | 学生只学会了拟合教师的高频回答,没有学到泛化知识 | 对比教师和学生在新 prompt 上的回答多样性 | 在蒸馏损失中加入多样性正则,增加更多教师示范样本 |
| 学生输出逐渐退化到少数模板 | 教师的 token 监督权重过低,大部分信号来自策略梯度 | 检查学生对同一 prompt 多次采样的输出相似度 | 提高教师示范通道权重,降低标量评分通道影响 |
| 教师越强,学生反而越差 | 教师与学生能力差距过大,学生无法拟合教师回答 | 观察训练集 teacher 回答与 student 回答的 edit distance | 使用中间规模的教师,或加入课程式蒸馏,由浅到深 |
| 门控信号不稳定 | 教师对数概率方差大,阈值设置过死 | 画出 teacher_logprob 的分布图 | 改用百分位数门控或滑动平均门控,避免硬阈值 |
| 显存开销过高 | 同时运行教师模型和学生模型,且保存完整 logits | 观察显存峰值出现阶段 | 使用 batch 内交替前向,或冻结教师采用梯度 checkpointing |
| 蒸馏后模型事实准确性下降 | 教师回答本身存在幻觉,学生逐一记住 | 抽样检查教师生成文本 | 在示范数据中加入事实性过滤或检索校验模块 |
这里最值得强调的问题是第一条。不少团队上线蒸馏模型时只看 loss 曲线,training loss 一路下降,看起来一切正常,一到真实业务场景就发现模型不像教师。
原因通常是训练数据分布太窄,导致教师只展示了自己知识覆盖范围的一小部分。Loss 下降只说明模型在现有数据上记住了模式,并没有任何指标证明模型学到了完整知识。因此蒸馏项目必须预留验证集,并且验证集要覆盖教师可能回答的多样化路径。
8. 工程落地:命名规范、指标设计与风险控制
对于想把类似思路落到生产环境的团队,我的建议不只是“照着论文公式写代码”,而是先建设一套可追踪的指标和配置体系。
第一,要在代码与文档里把“监督类型”写清楚。项目里不能只写 distill_loss 这样一个笼统字段。建议拆成 teacher_logits_kl、teacher_sequence_ce、teacher_reward_reinforce 等不同 loss 名称。这样团队成员在排查问题时,能立刻知道每一个训练步骤消耗的监督信号是什么。
第二,要单独统计监督信号用量。例如记录一个 batch 里有多少样本走了教师示范通道,有多少样本走了保守更新通道。如果 OPSA 的门控判断几乎把所有样本都判到同一个通道,说明门控信号失效,或者教师与学生能力差距不符合预期。
第三,要做 teacher 与 student 的差异指标。不只是看 student 的 benchmark 分数,还要看 student 对同一组 prompt 生成的回答与 teacher 回答之间的语义相似度、答案重合度、错误模式差异。只有差异指标持续收敛,才能说明学生确实在向教师迁移。
第四,要明确安全边界。OPSA 这类无监督方法并不等于无风险。如果教师模型本身存在事实错误,示范通道会把错误一并教给学生。上线前必须做人审抽样和数据过滤,尤其在医疗、法律、金融等领域。
第五,训练过程的回滚机制要预留。无论蒸馏还是强化,模型都可能出现训练到一半突然产生重复输出或安全违规内容的情况。保存每个 checkpoint 对应的超参、数据版本和采样种子,能够帮助你快速定位问题来源。
第六,不需要迷信“教师越大越好”。On-Policy 蒸馏是否真正有效,取决于学生的可学习性,而非教师的绝对能力。能力差距过大时,学生会把教师的复杂推理压缩成表面套话,产生严重的“假蒸馏”现象。
9. 总结与延伸阅读方向
这篇论文真正让人印象深刻的点,不是它提出了一个绝对更优的 SOTA 方案,而是它让研究者重新审视了“蒸馏”一词在强化学习场景里的语义边界。
On-Policy 采样本质上是为了稳定策略更新。蒸馏本质上是为了传递教师的知识分布。这两者并不天然冲突,但很多方案为了稳定而把教师压缩成了打分器,最终牺牲了知识通道。是否真在蒸馏,不取决于教师是否出现在数据流里,而取决于教师知识以什么形式进入了学生。
从实践中看,判断一个方案是否值得深耕,可以问三个问题:
- 教师的输出是否具备 token 级信息?
- 学生的更新梯度是否对非采样 token 的变化敏感?
- 学生的采样空间是否具备被教师示范扩展的可能?
如果三个问题的答案都是否,那么无论方案名叫什么,本质都是带奖励模型的策略优化。
后续想继续深入,可以关注几个方向:如何在 On-Policy 条件下高效保存和复用教师 logits;如何用自适应门控替代人工阈值;如何在多教师场景下分配监督权重;以及如何将这类方法与 DPO、PPO 等偏好优化算法做公平对比。理解框架比追新术语更重要。先把“蒸馏”和“奖励学习”分开,再去看各类模型蒸馏、知识蒸馏、Skill 蒸馏的方法,会发现很多论文的贡献其实比它们声称的小,也会发现一些看起来朴素的工作反而解决的是真问题。