资料摘要：On-Policy Distillation

资料摘要更新 2026-07-01置信度 high待阅读原始来源 ↗#蒸馏#后训练#深度#进阶#RLHF

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 只教 $O(1)$ bits（对/错），蒸馏教 $O(N)$ 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

$\text{KL}(\pi_\theta \lvert\rvert \pi_\text{teacher}) = \mathbb{E}_{x \sim \pi_\theta} [\log \pi_\theta(x_{t+1} | x_{1..t}) - \log \pi_\text{teacher}(x_{t+1} | x_{1..t})]$

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 技术报告