1. 从老师教学生说起On-Policy Distillation到底在解决什么问题如果你接触过大模型训练一定对知识蒸馏不陌生。传统做法很简单先训练一个大的教师模型然后让小模型去拟合教师模型输出的概率分布用KL散度做损失让小模型的输出分布尽量靠近教师。这套流程在工程上跑了很多年效果也还算稳定。但真正上手做过大规模蒸馏的人会发现一个很别扭的现象——学生模型在训练集上拟合得挺好一到自己生成文本的时候就开始跑偏错误会像滚雪球一样越滚越大。这个问题的根源在于训练和推理的分布不一致。传统蒸馏里学生模型是在教师模型生成的序列或者真实数据上做拟合它见到的输入分布是教师走过的路。可一旦部署上线学生模型要自己一步步生成它见到的输入分布变成了自己走过的路。这两条路一旦出现偏差后面每一步都会在错误的前提下继续决策最终输出质量断崖式下跌。这在序列生成任务里尤其明显因为一个token的偏差会通过自回归机制被无限放大。On-Policy Distillation在线策略蒸馏也有文献叫On-Policy Knowledge Distillation就是冲着这个痛点来的。它的核心思想用一句话概括让学生模型用自己当前策略生成的样本来接受教师的指导。也就是说训练数据不再完全来自教师或固定数据集而是学生模型自己走出来的轨迹教师在这些轨迹的每一个状态上给出监督信号。这样一来训练分布和推理分布就对齐了学生见到的问题正是它将来会遇到的。我第一次接触这个概念的时候直觉上觉得它和强化学习里的on-policy思路是一脉相承的。事实上确实如此On-Policy Distillation在方法论上和策略梯度、DAggerDataset Aggregation这些经典工作有很深的渊源。DAgger当年解决的就是模仿学习中分布偏移的问题让学习者自己跑然后让专家对学习者访问到的状态打标签。On-Policy Distillation把这个思路搬到了知识蒸馏场景只不过专家从人类标注者换成了教师模型。这篇文章我打算从几个层面把这件事讲透先拆清楚它和off-policy蒸馏的本质区别再讲清楚几种主流实现路线的技术细节然后聊聊实际训练中那些文档里不会写的坑最后给出一个可以照着复现的最小实验框架。适合已经了解基础蒸馏、想进一步搞清楚在线蒸馏怎么落地的人也适合正在做模型压缩、想让小模型在生成任务上不掉点的工程师。2. Off-Policy与On-Policy两种蒸馏范式的本质分歧2.1 分布偏移是怎么一步步毁掉蒸馏效果的要理解On-Policy Distillation的价值得先把Off-Policy蒸馏的问题看清楚。假设我们有一个教师模型T和一个学生模型S标准蒸馏的损失函数长这样L KL(T(x) || S(x))这里的x来自哪里通常有两个来源一是真实数据集里的输入二是教师模型自己生成的序列。无论哪种x的分布都是由教师或者数据决定的和学生模型当前的生成策略无关。训练的时候学生模型在这些x上被反复纠正loss降得很漂亮。但推理的时候学生模型自己生成第一个token这个token可能和教师的选择略有不同于是第二个token的输入就变成了一个教师从没见过的状态。学生在这个陌生状态上的表现没有任何保证错误就这样累积下去。这个问题在分类任务里不明显因为分类任务的输入是固定的不存在自己生成输入这回事。但在文本生成、语音识别、机器翻译、对话系统这些自回归任务里分布偏移是致命的。有研究做过统计在长序列生成任务中off-policy蒸馏的学生模型在生成到第20个token之后输出分布和教师分布的KL散度会比训练时高出好几倍。2.2 On-Policy的核心机制让学生自己踩坑教师现场纠正On-Policy Distillation的做法是训练时让学生模型自己生成完整序列然后在生成的每一步用教师模型对该状态下的输出分布做监督。形式化地如果学生模型在状态s_t即前面已经生成的token序列下采样出了动作a_t下一个token那么损失可以写成L E_{s_t ~ π_S}[ KL(T(·|s_t) || S(·|s_t)) ]注意这里的期望是对学生策略π_S采样的状态求的而不是对教师策略或数据分布求的。这就是on-policy的含义——训练数据分布跟随学生当前策略。这个改动看起来很小但带来的效果是根本性的。学生模型在训练时见到的每一个状态都是它自己真实会走到的状态。教师的监督信号直接作用在这些状态上告诉学生在你现在这个处境下应该怎么选。这相当于一个老师跟着学生做题学生做到哪一步卡住了老师就在那一步讲解而不是事先给学生一堆标准答案让他背。2.3 两种范式的对比不只是数据来源的差别很多人以为On-Policy Distillation只是把训练数据换成了学生自己生成的其他都一样。实际上两者的差别远不止于此我整理了一个对比表维度Off-Policy蒸馏On-Policy蒸馏训练数据来源固定数据集或教师生成学生模型实时生成状态分布教师/数据分布学生当前策略分布训练与推理一致性低存在分布偏移高分布对齐每步计算开销低可离线准备数据高需在线生成教师前向训练稳定性较稳定loss单调下降波动大需要技巧对教师依赖可离线蒸馏训练全程需要教师适合任务分类、短序列长序列生成、对话、翻译从工程角度看On-Policy蒸馏的代价是明显的每一步训练都要学生生成、教师前向计算量比离线蒸馏大得多。但它换来的是生成质量的显著提升尤其是在长序列任务上。这个trade-off是否值得取决于你的任务对生成质量的要求和可用的算力预算。还有一个容易被忽略的点On-Policy蒸馏天然适合和强化学习结合。因为学生自己生成序列这个过程本身就可以同时收集reward信号做RLHF或者PPO。事实上现在很多大模型的训练管线里蒸馏和RL是交织在一起的On-Policy的思路让这两者可以共享同一套采样流程。3. 三条主流技术路线从序列级到token级的实现拆解3.1 序列级On-Policy蒸馏最直接但方差最大最朴素的实现方式是序列级蒸馏。学生模型生成一整条序列然后教师模型对这条序列的每个位置给出logits计算逐token的KL散度并求和作为损失。伪代码大概是这样for batch in dataloader: # 学生生成序列 with torch.no_grad(): student_outputs student.generate(input_ids, max_length512) # 教师对同一序列打分 with torch.no_grad(): teacher_logits teacher(student_outputs).logits # 学生重新前向计算logits student_logits student(student_outputs).logits # 逐token KL loss kl_divergence( F.log_softmax(student_logits, dim-1), F.softmax(teacher_logits, dim-1) ) loss.backward()这种方式的优点是实现简单直接复用现有的蒸馏损失函数。但问题也很明显整条序列共享一个loss序列里每个token的贡献被平均了而实际上不同位置的token重要性差异很大。更麻烦的是学生生成序列用的是采样temperature0同一个输入每次生成的序列都不一样导致loss方差很大训练曲线会非常抖。我实测下来的经验是序列级蒸馏适合作为warm-up阶段先让学生模型大致对齐教师的输出风格然后再切换到更细粒度的方案。如果一上来就用序列级做主力训练收敛会非常慢。3.2 Token级On-Policy蒸馏细粒度监督的代价与收益Token级蒸馏的思路是不等到整条序列生成完而是在每个生成步骤都计算损失。学生生成一个token教师立刻对这个状态下的分布做监督然后学生继续生成下一个。这种方式的好处是监督信号更密集每个决策点都能得到反馈。实现上通常用teacher-forcing和free-running的混合模式。具体来说前向生成时用学生自己的输出作为下一步输入free-running但在计算loss时对每个位置都算KL。代码结构大致如下generated input_ids losses [] for step in range(max_length): student_logits student(generated).logits[:, -1, :] with torch.no_grad(): teacher_logits teacher(generated).logits[:, -1, :] step_loss kl_divergence( F.log_softmax(student_logits, dim-1), F.softmax(teacher_logits, dim-1) ) losses.append(step_loss) # 学生采样下一个token next_token sample_from(student_logits) generated torch.cat([generated, next_token], dim-1) total_loss sum(losses) / len(losses)Token级蒸馏的收益是监督更精准但代价是计算量成倍增加。因为每个step都要教师前向一次而教师模型通常比学生大好几倍这个开销在长序列上非常可观。实践中常用的优化是每隔k步才做一次教师前向或者只对关键位置比如实体、关键词做细粒度监督。3.3 混合策略与重要性采样降低方差的关键技巧纯On-Policy蒸馏有个统计上的问题学生自己生成的样本分布可能很窄导致训练数据多样性不足。比如学生如果总是生成相似的句子教师在这些句子上的监督就无法覆盖其他可能的输出空间。为了解决这个问题实践中常用混合策略。一种做法是混合学生和教师的生成结果。比如每个batch里70%的序列由学生生成30%由教师生成。这样既保证了训练分布和学生推理分布对齐又引入了教师的高质量样本避免学生陷入局部模式。另一种做法是引入重要性采样权重。因为On-Policy的样本是用学生策略采样的而我们的目标分布可能更接近教师分布所以可以用重要性权重来修正w π_T(a|s) / π_S(a|s) L E_{s~π_S}[ w * KL(T(·|s) || S(·|s)) ]这个权重的作用是给那些教师更可能选择但学生不太会选择的样本更高权重从而让训练更关注学生薄弱的地方。不过重要性权重的方差通常很大需要做截断clipping比如把w限制在[0.1, 10]之间。还有一种更工程化的做法是课程学习训练初期用较高比例的教师数据随着训练进行逐步提高学生数据的比例。这样学生先学会基本能力再逐步适应自己的分布。我在几个项目里用这个策略收敛速度和最终效果都比固定比例好。4. 训练稳定性On-Policy蒸馏最容易被低估的工程难题4.1 为什么loss会剧烈震荡如果你第一次跑On-Policy蒸馏大概率会被loss曲线吓到。它不像传统蒸馏那样平滑下降而是上下大幅震荡有时候甚至发散。这不是代码写错了而是On-Policy方法的内在特性。原因有几个。第一学生策略在训练过程中不断变化每一轮生成的样本分布都不一样导致loss的基准线在移动。第二采样本身有随机性同一个状态下的不同采样会带来不同的loss。第三教师和学生的能力差距如果太大教师给出的分布在学生看来可能非常尖锐KL散度对参数非常敏感。我踩过的一个坑是一开始用和off-policy一样的learning rate结果loss直接炸了。后来把学习率降到原来的1/5到1/10配合梯度裁剪才稳定下来。这个经验后来在好几个项目里都验证过On-Policy蒸馏对学习率比off-policy敏感得多。4.2 梯度裁剪、学习率与batch size的联动调参稳定On-Policy蒸馏训练核心是控制梯度的方差。我总结了一套调参顺序供参考先固定一个较小的学习率比如1e-5到5e-5之间比常规蒸馏小一个量级。开启梯度裁剪max_norm设在0.5到1.0之间。这个值比常规训练小因为On-Policy的梯度本身方差大需要更严格的裁剪。增大batch size。大batch能平均掉采样噪声对On-Policy尤其重要。如果显存不够用梯度累积模拟大batch。调整采样temperature。训练时学生的采样温度不要太高0.7到1.0之间比较合适。温度太高会导致生成质量差教师监督信号失去意义温度太低则样本多样性不足。还有一个技巧是用**EMA指数移动平均**维护一个学生模型的影子副本用影子模型来生成训练数据。这样生成策略的变化更平滑loss震荡会明显减小。这个技巧在RL里很常见搬到蒸馏里同样有效。4.3 教师信号的可信度问题什么时候该信老师On-Policy蒸馏有一个隐含假设教师在任何状态下都能给出正确的指导。但实际中这个假设不一定成立。教师模型在学生走到的某些奇怪状态下可能自己也不确定该输出什么给出的分布可能是高熵的、模糊的。如果学生盲目拟合这种分布反而会学到噪声。我在实际项目里遇到过这种情况学生生成了一些语法不通的序列教师在这些序列上的输出分布非常平坦几乎接近均匀分布。这时候如果强行做KL等于让学生去拟合一个随机分布没有任何信息量。解决办法是引入置信度过滤。具体来说对教师输出分布计算熵如果熵超过某个阈值就降低这个位置的loss权重或者直接跳过。另一种做法是用教师输出的top-1概率作为置信度低于阈值的样本降权。这个技巧看起来简单但对最终效果影响很大尤其是在训练后期。5. 一个可复现的最小实验框架5.1 环境与模型准备说了这么多原理还是得落到能跑的代码上。下面给一个最小可复现的On-Policy蒸馏框架基于PyTorch和HuggingFace Transformers。假设你已经有一个教师模型和一个学生模型两者都是因果语言模型。import torch import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoTokenizer teacher AutoModelForCausalLM.from_pretrained(teacher_path).eval() student AutoModelForCausalLM.from_pretrained(student_path) tokenizer AutoTokenizer.from_pretrained(student_path) for p in teacher.parameters(): p.requires_grad False optimizer torch.optim.AdamW(student.parameters(), lr2e-5)关键配置学习率2e-5梯度裁剪1.0batch size尽量大。教师模型全程eval模式不计算梯度。5.2 核心训练循环的逐行注释def on_policy_distill_step(prompts, max_new_tokens128, temperature0.8): student.train() input_ids tokenizer(prompts, return_tensorspt, paddingTrue).input_ids input_ids input_ids.to(student.device) generated input_ids total_loss 0.0 num_steps 0 for step in range(max_new_tokens): # 学生前向 student_logits student(generated).logits[:, -1, :] # 教师前向无梯度 with torch.no_grad(): teacher_logits teacher(generated).logits[:, -1, :] # 计算教师分布熵做置信度过滤 teacher_probs F.softmax(teacher_logits, dim-1) entropy -(teacher_probs * torch.log(teacher_probs 1e-8)).sum(dim-1) confidence 1.0 - entropy / torch.log(torch.tensor(teacher_probs.size(-1), dtypetorch.float)) confidence confidence.clamp(min0.0, max1.0) # KL损失按置信度加权 student_log_probs F.log_softmax(student_logits, dim-1) kl F.kl_div(student_log_probs, teacher_probs, reductionnone).sum(dim-1) weighted_kl (kl * confidence).mean() total_loss total_loss weighted_kl num_steps 1 # 学生采样下一个token with torch.no_grad(): probs F.softmax(student_logits / temperature, dim-1) next_token torch.multinomial(probs, num_samples1) generated torch.cat([generated, next_token], dim-1) loss total_loss / num_steps loss.backward() torch.nn.utils.clip_grad_norm_(student.parameters(), max_norm1.0) optimizer.step() optimizer.zero_grad() return loss.item()这段代码有几个关键点值得说明。第一教师前向放在torch.no_grad()里避免不必要的显存占用。第二置信度过滤用熵的归一化值把熵映射到[0,1]区间熵越低置信度越高。第三采样时用temperature控制多样性但计算loss时用的是原始logits不受temperature影响。5.3 训练过程中的监控指标跑起来之后光看loss是不够的。我建议至少监控这几个指标指标含义健康范围KL loss学生与教师的分布差距逐步下降波动但趋势向下教师置信度均值教师对当前状态的确定程度0.5以上太低说明学生走偏了学生生成熵学生输出的多样性与教师熵接近过高过低都不好梯度范数梯度大小裁剪前不超过5裁剪后稳定生成样本质量人工或自动评估定期抽样检查其中教师置信度均值是我特别推荐关注的。如果这个值持续走低说明学生生成的状态越来越离谱教师都看不懂了。这时候需要回退到更保守的训练策略比如提高教师数据的混合比例。6. 那些文档里不会写的实操心得6.1 教师模型不是越大越好一个反直觉的经验在On-Policy蒸馏里教师模型太大反而可能拖累效果。原因是教师和学生能力差距过大时教师给出的分布在学生看来过于尖锐学生很难拟合。而且大教师意味着每一步前向的开销都很大训练速度慢迭代次数少最终效果未必好。我试过用不同规模的教师做对比发现教师规模是学生的3到5倍时效果最好。超过10倍之后收益递减甚至变负。这个比例不是绝对的和任务难度、数据量都有关系但可以作为初始参考。6.2 训练数据的prompt选择很关键On-Policy蒸馏虽然用学生自己生成的序列但初始prompt还是需要从数据集里选。prompt的质量直接影响学生探索的状态空间。如果prompt太简单学生生成的都是简单样本学不到复杂能力如果prompt太难学生一开始就生成一堆垃圾教师也救不回来。我的做法是做一个难度分层先用中等难度的prompt训练等学生稳定后再逐步加入难样本。这个课程学习的思路在On-Policy蒸馏里特别有效因为学生的能力是动态变化的训练数据也应该跟着变。6.3 什么时候该停止On-Policy训练On-Policy蒸馏不是训练越久越好。训练到后期学生和教师的分布已经很接近继续做KL的收益很小反而可能因为过拟合教师的某些偏见而损害学生自身的泛化能力。判断停止时机的实用信号教师置信度均值稳定在高位比如0.8以上KL loss下降趋缓生成样本的人工评估不再提升。这时候可以考虑切换到off-policy做微调或者直接停止。我一般会保留几个checkpoint用验证集上的生成质量做最终选择而不是只看训练loss。6.4 和RLHF的衔接如果你的最终目标是做RLHFOn-Policy蒸馏可以作为一个很好的初始化阶段。先用蒸馏把学生的生成质量拉到接近教师的水平再用RL做偏好对齐。这样RL阶段不需要从零探索收敛更快也更稳定。衔接的时候要注意蒸馏阶段用的采样temperature和RL阶段要匹配。如果蒸馏用0.8RL突然改成1.0策略分布会突变训练不稳定。最好是逐步调整让策略平滑过渡。7. 写在最后一些个人体会On-Policy Distillation这个概念听起来学术但落地的时候全是工程细节。我最大的体会是不要指望一次跑通。分布偏移这个问题本身就很微妙训练过程中会出现各种意料之外的现象loss震荡、生成退化、教师置信度崩塌每一个都需要具体分析。另一个体会是On-Policy蒸馏和off-policy蒸馏不是对立的而是互补的。实际项目里我通常先用off-policy做warm-up让学生有个基本能力再切到on-policy做精细对齐。这个组合比单用任何一种都稳。最后分享一个小技巧训练的时候定期保存学生生成的样本人工看几条。不要只看指标指标好看但生成质量差的情况太常见了。有时候一条样本就能告诉你问题出在哪比盯着loss曲线猜半天有用得多。
