资料摘要:On-Policy Distillation
Thinking Machines Lab 提出 on-policy distillation:让学生模型自己采样轨迹,教师模型对每个 token 打分(dense reward),结合了 RL 的 on-policy 相关性与蒸馏的密集监督信号,在数学推理上达到 9–30× 成本降低,并可恢复领域微调后丢失的指令遵循能力。
核心要点
- On-policy distillation = on-policy 采样 + dense reward:学生自己生成轨迹,教师对每个 token 计算 reverse KL,获得逐 token 的密集梯度信号
- vs RL:RL 每 episode 只教 bits(对/错),蒸馏教 bits(N = token 数),计算效率提升 50–100×
- vs Off-policy SFT:SFT 训练在教师分布上,学生遇到自身错误时会 compounding error;on-policy 始终在学生自己的状态分布上训练
- 实测:Qwen3-8B-Base 做数学推理,on-policy distillation 用 150 步从 60%→70% AIME’24,成本仅为 SFT 的 1/9–1/30
- 持续学习:在领域微调(mid-training)后用 on-policy distillation 恢复指令遵循,IF-eval 从 45% 恢复到 83%,不丢领域知识
详细笔记
背景:三种训练阶段
文章将 LLM 训练分为三个阶段:
- Pre-training:通用能力(语言、推理、世界知识)
- Mid-training:领域知识注入(代码、医疗、公司文档)
- Post-training:目标行为激发(指令遵循、数学推理、对话)
后训练方法分为两类:
- On-policy:采样学生自身 rollout → 分配 reward(如 RL)
- Off-policy:依赖外部源的目标输出 → 学生模仿(如 SFT / distillation)
On-Policy Distillation 核心思想
损失函数:Reverse KL
- Student 采样轨迹 → Teacher 对每个 token 计算 log-prob → 计算 reverse KL → 梯度更新 student
- Discount factor = 0:只优化当前 token,不考虑未来 token
- Reverse KL 是 mode-seeking:学教师的一个具体行为,而非分散到多个次优选项
- Teacher 只需一次 forward pass(无需反向传播),推理成本低
伪代码(基于 Tinker API):
teacher_client = create_sampling_client(teacher_model)
trajectories = do_group_rollout(student_client, env_group_builder)
sampled_logprobs = trajectories.loss_fn_inputs["logprobs"]
teacher_logprobs = teacher_client.compute_logprobs(trajectories)
reverse_kl = sampled_logprobs - teacher_logprobs
trajectories["advantages"] = -reverse_kl
training_client.forward_backward(trajectories, loss_fn="importance_sampling")实验一:数学推理蒸馏
| 方法 | AIME’24 | 教师 FLOPs | 学生 FLOPs | 成本比(vs SFT-2M) |
|---|---|---|---|---|
| SFT-400K (初始化) | 60% | 8.5×10²⁰ | 3.8×10²⁰ | – |
| SFT-2M (外推) | ~70% | 3.4×10²¹ | 1.5×10²¹ | 1× |
| RL (Qwen3 报告) | 68% | – | – | ≈1× |
| On-policy distillation | 70% | 8.4×10¹⁹ | 8.2×10¹⁹ | 9–30× |
- LoRA rank=32 在 SFT 后比全量微调差 13%,但 on-policy distillation 后仅差 6%
- 无需新采集 SFT 数据集时成本降 9×,需采集时降 30×(GPU 小时约 18×)
实验二:个性化 / 内部助手
问题:Mid-training 注入领域知识会破坏指令遵循(IF-eval 从 85% 跌至 45%)
方案:
- Mid-train Qwen3-8B 在内部文档+背景聊天数据上
- 用原始 Qwen3-8B 自身作为教师,在 Tulu3 提示上做 on-policy distillation
- IF-eval 恢复到 83%,内部 QA 从 36% 提升到 41%(正迁移)
关键发现:把”旧版自己”当教师做 distillation,可廉价恢复被微调覆盖的行为——这对持续学习非常有前景。
实验三:Self-Distillation 效率对比
从同一基座出发,RL 训练出的策略作为教师,蒸馏回基座:
- On-policy distillation 用 7–10× 更少梯度步达到教师 AIME 水平
- 累积计算效率提升 50–100×(可短上下文训练 + 小 batch size)
关键洞察
- RL 在搜索语义策略空间,而非参数空间:RL 大部分算力花在 rollout 和 credit assignment,而非梯度更新本身。一旦好策略被发现,蒸馏是学习它的捷径。
- 蒸馏可高效复用数据:同一 prompt 多 epoch 训练不导致死记硬背(vs RL 会),因为学的是完整分布而非单一答案。单 prompt 训练 5120 条轨迹即可匹配教师 AIME 性能。
- On-policy SFT 反会退化:在自己的样本上做 SFT(即使期望 KL=0),因有限 batch 的随机性导致分布漂移,逐步退化为 off-policy。
- On-policy distillation 永不退化:教师固定,学生始终向教师收敛。
引用与数据
- 数学推理数据集:OpenThoughts-3(1.2M 推理 prompts)
- 学生模型:Qwen3-8B-Base / Qwen3-4B-Instruct-2507
- 教师模型:Qwen3-32B / Qwen3-235B-A22B-Instruct-2507
- 实现:Tinker Cookbook — distillation recipes
- 灵感来源:DAGGER(迭代 SFT)、Process Reward Modeling、Agarwal et al.、Gu et al.、Qwen3 技术报告