聊聊On-policy Distillation（1）
背景
前段时间，On-policy Distillation这个方法特别火，主要是Thinking Machine Lab这个团队有一篇爆火的文章出圈了，让大家关注到了这个方法，各种新的方法也都雨后春笋般冒出来了。
其中千问AI硬件团队刚中的这篇ICLR的文章引起了我的关注：STAR: Similarity-guided Teacher-Assisted Refinement for Super-Tiny Function Calling Models
AI硬件这个领域现在逐渐步入大众视野，如果所有需求都调用云端大模型，势必会有token吃不消的情况，小模型在端侧的需求是非常强劲的，因此，On-policy Distillation的价值就体现出来了。
但这篇文章没有简单使用On-policy Distillation，而是进行了效果优化，简单来说
为了让小模型从大模型 “学到更多有用的概率输出”，并避免训练崩溃，错误输出过于自信。
-
传统 KD 会用 KL 散度衡量学生和教师之间的输出分布差异，但：
-
直接用传统 KL 在 “top-k” 截断时不稳定
-
小模型容易过信错误的输出，导致学习效果差
-
文中提出的CKD设计思路：
- 重点保留教师模型 top-k 分布（概率最高的前 k 个 token）
- 对学生模型在 top-m（比 k 大）但不在 top-k 的错误概率进行惩罚
实现方式：
- 用前向 KL（Forward KL）保留教师 top-k 的概率分布
- 加了一个额外的正则项来压制学生错误输出概率（L1 抑制）
这篇文章还针对这个端侧小模型进行了进一步的强化学习，也没有基于传统的verifiable的结果来直接做，而是做了基于相似度的优化。这里没有非常强硬使用Verifiable的结果，相似度计算不仅能带来结果上一定的鲁棒性，同时也间接提升了Reward的density，对于训练稳定有至关重要的作用。
该方法在Qwen-0.6B上验证的效果，在某些领域甚至超过了3-4B的大模型。
总结
相比于基于SFT的蒸馏来说，On-policy Distillation的想法非常直白，就是用RL的方法，通过token Level来约束两个大小不同的概率模型的分布，提升小模型的效果。相比于SFT蒸馏而言，RL蒸馏的最大区别在自上下文学习。
传统意义上讲，如果我们直接用更大模型产生的数据来蒸馏小模型，是在更大模型的上下文里学习，假设这个更大模型和小模型的分布不同，比如在某个token上的entropy特别高，就会导致小模型的更新梯度更大，从而造成灾难性遗忘等问题。因此简单来说，On-policy Distillation的两个重大优势就是
- 能显著避免小模型的灾难性遗忘问题；
- 能显著提升训练效率，用极少的数据和业务侧的数据混合就能训练出和混大量数据的SFT效果。
更准确点说，我觉得更吸引人的是这个自上下文学习的概念，用这个比喻更好些：老师告诉你的并非题的解法，而是告诉你你错在哪里了，你本身就大概知道怎么改，适合模型能力本身就还可以，但是需要引入额外能力的场景。
之前听到蒸馏，总有种拉不高模型能力上限的感觉；但基于应用场景看，才发现蒸馏是一种智慧：在资源有限场景下，如何提升小模型的专项能力，如何让小模型保持更好的通用能力。千问的这篇文章，就是未来AI硬件卷模型的一个很重要的参考，什么时候才用到云端的大模型，什么时候端侧的小模型就能解决。我相信在未来，谁能更好的利用好用户手里的设备的算力，谁就能在运行成本上和别人拉开更大的差距，从而就能更多投入到研发上，进而拉大和别家的差距。