關於 On-Policy 的相關工作,我感覺大部分是在賣概念,沒有太本質上的創新。例如一會兒說 On-Policy Distillation 是 SFT,一會兒說是 RL,但感覺不如就直接叫它 Distillation 好了。
比較公認的觀點是:
- On-Policy Distillation 相較於 GRPO 提升性能的同時,避免了過多的 Aha Moment(頓悟時刻)。
- On-Policy Distillation 可以緩解災難性遺忘(Catastrophic Forgetting)。
- On-Policy Distillation 天然適合與 GRPO 結合,提供 token-level 的密集獎勵信號。
- On-Policy Distillation 的難點是獲取 Teacher 模型,而 Self-Distillation(即使用 Policy 模型當作 Teacher 模型)已被驗證為可行。
- On-Policy Distillation 的實現兼容 RL,所以直接在 RL 框架上比較容易開發。
本文重點介紹 On-Policy/Self-Distillation (OPSD),即使用 Policy 本身當作 Teacher 模型。
1. On-Policy Distillation 的目標及梯度
On-Policy Distillation 旨在最小化學生策略與教師策略在學生策略自身生成的軌跡分佈上的 KL 散度:
其中 KL 可以是 Reverse KL,也可以是 Forward KL。文獻 [1-2] 中使用的是 Reverse KL,而文獻 [3] 中使用的是 Forward KL。
對於 Forward KL,其梯度可以推導為:
對於 Reverse KL,其梯度可以推導為:
可以發現其實和 RL 的目標很相似,都有 ,只是前面所有的加權不一樣。RL 加權的是 Reward 或者 Advantage。
2. On-Policy Self-Distillation
Self-Distillation 旨在用 Policy 充當 Teacher 模型:
其中 代表 stop gradient,
代表額外知識,
可以通過
獲得。
On-Policy Self-Distillation 的關鍵在於 Teacher 構造。OPSD 中 Teacher 構造方式和 MoCo、SimCLR、DINO、SigLIP 等有相似之處。
- 需要引入額外的知識,以來 In-Context Learning 能力。
- Policy model 在不斷更新,Teacher 模型的參數要相對穩定一點。
這兩點就是後面實現細節要討論的。
3. 實現細節
3.1 如何引入額外的知識
目前看到的有下面兩種方式:
方法一:直接把 Ground-Truth 透露給 Policy 模型參考
方法二:來源於環境的反饋
3.2 Teacher 模型參數如何確定
- 直接使用原本的 frozen 模型前期可以,後面就崩潰了。
- Teacher 模型一直使用 Policy 模型也是可以的,但效果不如 EMA。
- Trust-region 和 EMA 效果差不多,都是為了獲得更穩定的 Teacher,避免隨著優化過程劇烈變化。
Trust-region 的更新策略為:
4. On-Policy Self-Distillation 的優點
4.1 緩解災難遺忘
「Self-Distillation Bridges Distribution Gap in Language Model Fine-Tuning」這篇工作比較早研究了,但是當時 On-Policy 概念還沒有火。
4.2 減少 Aha Moment
4.3 減少 Train-Test Gap
提前在 train-time 暴露學生模型 test-time 的分佈,緩解暴露偏差。
5. Scaling On-Policy Self-Distillation
正常來說,模型尺寸越大,超過 GRPO 越多。因為模型越大,in-context learning 能力通常越強。
Reference
[1] Reinforcement Learning via Self-Distillation
https://arxiv.org/html/2601.20802
[2] Self-Distillation Enables Continual Learning
https://arxiv.org/html/2601.19897
[3] Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models
https://arxiv.org/html/2601.18734延伸閱讀
通往 AGI 的必經之路:Agent 自進化到底是在「進化」什麼?