1. 为什么值得花时间啃 On-Policy Distillation第一次看到 On-Policy Distillation 这个词是在翻一份模型压缩的综述时。当时我的第一反应是蒸馏就蒸馏为什么还要强调 On-Policy后来把几篇相关论文和开源实现翻了一遍才意识到这个方向解决的是一个非常具体的痛点——传统知识蒸馏里学生模型是在老师生成的静态数据分布上学习的而推理时它面对的却是自己生成的分布两者之间的错位会随着生成步数累积越走越偏。On-Policy Distillation 的核心思路就是让学生模型在自己生成的轨迹上接受老师的纠正信号把训练分布和推理分布对齐。它最早在序列生成任务尤其是语言模型和序列决策里被系统讨论后来逐渐扩展到对话模型、代码生成、多步推理等场景。如果你正在做模型压缩、小模型能力对齐、或者想让一个便宜的小模型在特定任务上逼近大模型的表现这个方向值得认真看。这篇文章适合三类人一是刚接触知识蒸馏、想搞清楚 On-Policy 到底改变了什么的新手二是已经在用传统蒸馏、但发现学生模型“训练时好好的、一推理就崩”的从业者三是想快速判断这个方向值不值得投入的工程负责人。我会从设计动机、核心机制、实操流程、踩坑经验四个层面展开尽量把论文里那些公式背后的直觉讲清楚同时给出可以直接参考的实现思路。需要提前说明的是On-Policy Distillation 目前还没有一个像传统 KD 那样统一的教科书式定义不同论文的术语和实现细节有差异。我下面讲的是基于常见实践归纳出来的主流框架具体到某个任务时你需要根据实际情况调整。2. 核心概念拆解On-Policy 到底改变了什么2.1 传统知识蒸馏的分布错位问题传统知识蒸馏的流程很直观拿一个训练好的大模型老师在一个固定数据集上跑一遍得到每个位置的输出分布软标签然后让学生模型去拟合这些分布。损失函数通常是 KL 散度或者交叉熵学生尽量让自己的输出分布靠近老师的。这个流程在分类任务上工作得很好因为分类任务的输入是固定的学生和老师看到的是同一批样本。但到了序列生成任务问题就来了。序列生成是一个自回归过程学生每一步的输出会作为下一步的输入。训练时学生看到的是真实数据里的前缀teacher forcing推理时学生看到的是自己生成的前缀。一旦某一步生成错了后面的输入分布就完全偏离训练分布错误会累积放大。这就是经典的 exposure bias 问题。传统蒸馏虽然把老师的软标签引入了但并没有解决这个分布错位——因为学生仍然是在真实前缀上学习老师的分布而不是在自己生成的前缀上。2.2 On-Policy 的核心思想在自己的轨迹上学习On-Policy Distillation 的关键改动是让学生模型先自己生成一段轨迹然后在这些自己生成的轨迹上让老师给出纠正信号。具体来说流程大致是这样的学生模型对当前输入进行自回归生成得到一条完整的输出序列。把学生生成的序列而不是真实标签序列作为输入喂给老师模型。老师模型在每个位置输出它的分布作为学生学习的监督信号。学生根据老师的分布调整自己的参数。这样做的直接好处是学生学习的分布就是它自己推理时会遇到的分布。训练和推理的分布对齐了exposure bias 被大幅缓解。用一个生活化的类比传统蒸馏像是老师把标准答案写在黑板上让学生照着抄On-Policy 蒸馏像是让学生先自己做一遍题老师看着学生做的过程在学生出错的地方指出“这里应该这样想”。后者显然更贴近学生真实的学习状态。2.3 和强化学习、DAgger 的关系如果你熟悉强化学习会发现 On-Policy Distillation 和 DAggerDataset Aggregation的思路非常像。DAgger 也是让策略自己跑然后让专家在策略访问到的状态上给出动作标签再把这些数据聚合起来训练。区别在于DAgger 通常聚合多轮数据而 On-Policy Distillation 更强调在每一步都用老师的分布做监督且监督信号是完整的分布而不是单个动作。另一个相关概念是 RLHF 里的 KL 约束。RLHF 也强调让学生策略不要偏离参考模型太远但它是通过奖励信号和 KL 惩罚来间接约束的。On-Policy Distillation 则是直接用老师的分布做监督信号更密集、更直接。理解这些关系有助于你在实际选型时判断如果你的任务有明确的奖励函数RLHF 可能更合适如果你有一个强老师模型、且希望学生快速逼近它On-Policy Distillation 更直接。3. 核心机制与关键技术点3.1 损失函数的设计与选择On-Policy Distillation 的损失函数核心还是让学生分布逼近老师分布。常见的选择有几种前向 KL 散度KL(p_student || p_teacher)。这个方向会让学生尽量覆盖老师分布的所有模式但容易在老师分布平坦的地方产生较大惩罚。反向 KL 散度KL(p_teacher || p_student)。这个方向会让学生集中在老师分布的高概率区域生成更保守、更确定但可能丢失多样性。Jensen-Shannon 散度对称版本折中两者。交叉熵如果老师分布是 one-hot硬标签就退化成普通的监督学习。实际用的时候前向 KL 和反向 KL 的选择会显著影响生成质量。我试过在对话任务上用反向 KL学生生成的回复更稳、更少胡言乱语但多样性下降明显换成前向 KL 后多样性回来了但偶尔会出现一些奇怪的重复。一个常见的折中方案是混合两种 KL或者用温度参数调节老师分布的平滑程度。注意温度参数 T 在蒸馏里非常关键。T 越大老师分布越平滑学生能学到的“暗知识”越多T 太小老师分布接近 one-hot蒸馏退化成普通监督。常见取值在 2 到 10 之间需要根据任务调。3.2 老师信号的获取方式On-Policy Distillation 里老师信号的获取有两种主流方式第一种是离线老师老师模型固定不动学生生成轨迹后把轨迹喂给老师老师前向计算一次得到分布。这种方式实现简单老师可以是任意一个已经训练好的大模型甚至可以是多个老师的集成。第二种是在线老师老师模型也参与训练或者老师本身是一个不断更新的模型。这种方式更复杂但理论上可以让学生学到老师“进化中”的知识。不过在线老师容易导致训练不稳定实际用得比较少。还有一种变体是自蒸馏老师就是学生自己的一个旧版本比如 EMA 版本。这种方式不需要额外的大模型适合没有强老师可用的场景。自蒸馏在 On-Policy 框架下效果也不错因为学生自己生成的轨迹天然就是 on-policy 的。3.3 轨迹采样与批次构造On-Policy 的核心是“用自己的轨迹”所以轨迹采样策略直接影响训练效果。几个关键点采样温度学生生成轨迹时用的温度最好和推理时一致。如果训练时用低温、推理时用高温分布又会错位。轨迹长度长轨迹能暴露更多错误累积问题但计算成本高。常见做法是截断到固定长度或者按任务需要动态调整。批次内多样性如果一批轨迹都差不多老师信号就缺乏多样性。可以通过调整采样温度、加噪声、或者用不同的 prompt 来增加多样性。我踩过的一个坑是一开始为了省算力只采样了很短的轨迹结果学生只在短序列上对齐了长序列生成时还是崩。后来把轨迹长度拉到和推理时一致问题才解决。所以采样配置一定要贴近推理配置这是 On-Policy 的前提。3.4 训练稳定性与收敛判断On-Policy Distillation 的训练比传统蒸馏更容易不稳定因为学生生成的轨迹质量会波动老师信号也跟着波动。几个稳定训练的技巧学习率要小比传统蒸馏小一个量级左右因为 on-policy 信号的方差更大。梯度裁剪防止个别坏轨迹导致梯度爆炸。老师分布的温度平滑适当提高温度让老师信号更柔和。定期评估不要只看训练 loss要在验证集上实际生成看效果。训练 loss 下降但生成质量变差的情况很常见。收敛判断上我一般看两个指标一是学生和老师在验证集上的生成分布距离比如用 KL 或者 perplexity 差二是人工或自动评估的生成质量。两者都稳定了才算收敛。4. 完整实操流程从零搭一个 On-Policy Distillation 实验4.1 环境准备与依赖选择假设你要在一个序列生成任务上做 On-Policy Distillation基础环境大概需要PyTorch 或 JAX看你的模型框架一个预训练的老师模型和一个学生模型数据加载和生成工具比如 HuggingFace 的 transformers 库评估工具BLEU、ROUGE、或者任务特定的指标如果你用的是 HuggingFace 生态老师模型可以直接加载一个大的预训练模型学生模型可以用小一号的架构。注意老师和学生的 tokenizer 最好一致否则分布对齐会很麻烦。4.2 数据准备与预处理On-Policy Distillation 的数据准备和传统蒸馏略有不同。传统蒸馏需要 (输入, 老师输出) 对On-Policy 只需要输入老师输出是训练时动态生成的。所以数据准备的重点是准备一批高质量的输入 prompt覆盖任务的主要场景。对输入做和推理时一致的预处理tokenization、截断、padding。如果任务有参考答案可以留着做评估但训练时不直接用。数据量上On-Policy 通常比传统蒸馏需要更少的样本因为每条样本会被多次利用每轮学生生成不同轨迹。但样本质量要求更高因为坏样本会让学生学到坏分布。4.3 学生模型生成轨迹的实现这是整个流程的核心步骤。伪代码大概长这样# 学生生成轨迹 student.eval() with torch.no_grad(): generated student.generate( input_ids, max_lengthmax_len, temperaturesample_temp, do_sampleTrue, return_dict_in_generateTrue, output_scoresTrue )关键参数说明temperature要和推理时一致常见 0.7 到 1.0。do_sampleTrue必须采样不能贪心否则轨迹缺乏多样性。max_length和推理时一致不要为了省算力截短。生成完后把generated.sequences作为老师模型的输入。4.4 老师分布的计算与对齐老师模型对学生生成的轨迹做前向计算得到每个位置的 logits然后 softmax 得到分布teacher.eval() with torch.no_grad(): teacher_outputs teacher(generated.sequences) teacher_logits teacher_outputs.logits / teacher_temp teacher_probs softmax(teacher_logits, dim-1)注意这里用了teacher_temp做温度缩放。温度的选择会影响蒸馏效果建议从 2 开始试。然后学生模型也对同一轨迹做前向计算得到学生分布计算 KL 散度student_outputs student(generated.sequences) student_logits student_outputs.logits / student_temp student_log_probs log_softmax(student_logits, dim-1) # 前向 KL: KL(student || teacher) loss (student_probs * (student_log_probs - teacher_log_probs)).sum(-1) loss loss.mean()这里有个细节学生和老师的温度最好一致或者至少保持固定比例否则分布对齐会有偏差。4.5 参数更新与迭代训练拿到 loss 后正常反向传播更新学生参数。然后进入下一轮学生重新生成轨迹老师重新给信号。整个训练循环大概是采样一批输入。学生生成轨迹。老师计算分布。计算 KL loss更新学生。重复 1-4 直到收敛。每轮之间可以清空生成缓存确保轨迹是新鲜的。如果算力有限可以每隔几轮才重新生成轨迹中间复用旧轨迹——但这会削弱 On-Policy 的效果属于折中方案。4.6 评估与对比实验设计评估 On-Policy Distillation 的效果建议至少做三组对比实验组训练方式预期效果Baseline学生直接微调基准传统 KD在固定数据上蒸馏比 baseline 好但推理时可能退化On-Policy KD在自己轨迹上蒸馏推理时表现最好评估指标除了任务指标BLEU、准确率等还建议看生成多样性distinct-n、重复率、以及长序列上的表现。On-Policy 的优势通常在长序列和复杂任务上更明显。5. 常见问题与排查技巧实录5.1 训练 loss 下降但生成质量变差这是最常见的问题。原因通常是学生过拟合了老师的某些模式或者 KL 方向选错了。排查思路检查 KL 方向如果用的是反向 KL学生容易变得过于保守生成重复内容。换成前向 KL 或混合 KL 试试。检查温度温度太低会导致学生只学老师的 argmax丢失多样性。提高温度。检查轨迹质量如果学生生成的轨迹本身就很差老师信号也会被带偏。可以先让学生做一轮监督微调再开始 On-Policy。5.2 训练不稳定、loss 震荡On-Policy 的方差天然比传统蒸馏大。缓解方法降低学习率通常比传统蒸馏小 5 到 10 倍。增大 batch size或者用梯度累积。对老师分布做温度平滑。用 EMA指数移动平均稳定学生参数。我自己的经验是学习率从 1e-5 降到 2e-6 后震荡明显减轻。5.3 学生生成轨迹质量太差如果学生初始能力太弱生成的轨迹全是垃圾老师信号也没法救。这时候需要先做一轮 warm-up用传统蒸馏或监督学习先让学生达到一个基本水平。或者用老师生成的数据做一轮 SFT再切换到 On-Policy。不要指望 On-Policy 能从零教会一个学生它更适合“锦上添花”而不是“雪中送炭”。5.4 计算成本太高On-Policy 需要学生生成 老师前向计算量比传统蒸馏大不少。优化方向用更小的老师模型比如用中等模型当老师而不是最大的。减少轨迹长度但要保证和推理一致。用梯度检查点、混合精度等常规优化。如果老师是 API 模型考虑缓存老师输出避免重复调用。5.5 常见问题速查表问题可能原因解决方向loss 降但质量差KL 方向错、温度低换 KL、提温度训练震荡学习率大、方差大降学习率、增 batch轨迹质量差学生初始弱先做 warm-up成本高生成老师前向小老师、短轨迹、缓存长序列崩轨迹长度不匹配拉长训练轨迹6. 几个容易忽略的实操心得第一个心得是关于老师模型的选择。很多人默认老师越大越好但实际上老师和学生的能力差距太大时学生反而学不动。我试过用一个大模型当老师学生怎么训都追不上后来换成一个中等模型当老师学生很快就逼近了。所以老师的选择要匹配学生的容量差距在一个合理范围内效果最好。第二个心得是关于轨迹的多样性。如果一批轨迹都差不多老师信号就缺乏区分度。我通常会在采样时用稍高的温度或者在 prompt 层面做多样化确保学生看到不同的生成路径。多样性够了蒸馏信号才丰富。第三个心得是关于评估的频率。On-Policy 训练中生成质量可能在某个点突然变好也可能突然崩掉。我一般每几百步就做一次小规模生成评估而不是等到训练结束才看。早发现早调整能省很多时间。第四个心得是关于和 RLHF 的结合。如果你的任务有奖励信号可以考虑把 On-Policy Distillation 和 RLHF 结合用蒸馏 loss 做主要监督用奖励做辅助信号。这样既能快速逼近老师又能优化任务特定目标。不过这种组合调参比较复杂建议先把纯蒸馏跑通再尝试。最后再分享一个小技巧在计算 KL loss 时可以对不同位置加权。比如序列开头的 token 权重高一些因为开头错了后面全错或者对老师置信度高的位置权重高一些因为这些位置的信号更可靠。这个加权策略在长序列任务上效果很明显我实测能提升几个点的生成质量。这个方向后续还可以往多老师蒸馏、在线蒸馏、以及和推理加速结合的方向扩展。如果你已经在做相关实验欢迎交流踩坑经验。