2306.13649-on-policy-distillation-language-models

On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes

GKD 把 LLM 蒸馏从固定参考答案上的 teacher forcing 推进到学生自生成轨迹上的教师分布匹配,用一个可调系数混合离线数据和 on-policy 数据,并允许 forward KL、reverse KL、JSD 等不同散度;它在 T5 系列的摘要、翻译、算术推理和 instruction tuning 上展示了稳定收益,是后续 OPD 系方法的重要早期基线。

已审阅 Archived 2026-07-03 09:16 Updated 2026-07-17 20:32 Reviewed 2026-07-18 17:42 Source

title: "On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes" arxiv: "2306.13649" venue: "ICLR 2024" authors:

  • Rishabh Agarwal
  • Nino Vieillard
  • Yongchao Zhou
  • Piotr Stanczyk
  • Sabela Ramos
  • Matthieu Geist
  • Olivier Bachem tags:
  • on-policy-distillation
  • generalized-knowledge-distillation
  • llm-distillation
  • rl-from-ai-feedback
  • t5
  • xsum
  • wmt14
  • gsm8k

On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes

Source

作者与关系

  • Rishabh Agarwal: Google DeepMind / Mila;历史机构:Google Brain / IIT Bombay。
  • Nino Vieillard: Google DeepMind。
  • Yongchao Zhou: Google DeepMind / University of Toronto。
  • Piotr Stanczyk: Google DeepMind。
  • Sabela Ramos: Google DeepMind;历史机构:Google Research / YouTube / ETH Zurich / University of A Coruna。
  • Matthieu Geist: Google DeepMind;历史机构:Google Brain France / Cohere / University of Lorraine / CentraleSupelec / University Paul Verlaine of Metz / University of Lille。
  • Olivier Bachem: Google DeepMind;历史机构:Google Brain / ETH Zurich。

作者关系上,这篇论文由 Google DeepMind 主导,Rishabh Agarwal 同时关联 Mila,Yongchao Zhou 同时关联 University of Toronto。论文脚注标注 Rishabh Agarwal 与 Nino Vieillard 等贡献并共同通讯,Piotr Stanczyk 与 Sabela Ramos 贡献基础设施。

与本地归档的关系上,这篇论文是 MOPD 之前更基础的 OPD / GKD 节点:它用学生模型自己的生成作为 distillation 状态分布,再让教师在这些学生状态上给 token-level 分布监督;MOPD 将这个思路扩展到多教师、多能力、RL 后教师和能力路由。它也和 Qwen3DeepSeek-V4GLM-5.2 中出现的 on-policy distillation 形成方法谱系。

论文脉络

1. 传统 KD 的核心错位

自回归语言模型的部署成本常通过蒸馏降低。传统 token-level KD 通常在固定训练样本或教师生成样本上做 teacher forcing:给定输入 xx 和参考输出前缀 y<ny_{<n},学生拟合教师的下一个 token 分布。这类训练的状态分布来自数据集或教师输出。

推理时,学生需要沿着自己的输出前缀继续生成。只要学生早期产生了偏离训练轨迹的 token,后续 hidden state 和 token context 就会进入训练中较少覆盖的区域。论文把这个问题表述为 distribution shift:训练时看到的是数据或教师轨迹,测试时展开的是学生轨迹。

这篇论文的直觉很直接:如果学生推理时会犯自己的错误,蒸馏训练就应该让教师在这些学生已经走到的状态上给反馈。

2. Generalized Knowledge Distillation

论文提出 Generalized Knowledge Distillation, GKD。它保留传统 token-level teacher distribution matching,但把训练序列来源推广为两部分:

  • 离线样本 (x,y)(X,Y)(x,y)\sim(X,Y),来自人工数据或教师输出;
  • 学生样本 ypS(x)y\sim p_S(\cdot|x),来自当前学生在输入 xx 上的自回归生成。

给定一条输出序列 yy,教师和学生在每个前缀 y<ny_{<n} 上给出下一个 token 分布,序列级 token divergence 写成:

D(pTpSθ)(yx)=1Lyn=1LyD(pT(y<n,x)pSθ(y<n,x)) \mathcal{D}(p_T\|p_S^\theta)(y|x)= \frac{1}{L_y}\sum_{n=1}^{L_y} \mathcal{D}\left( p_T(\cdot|y_{<n},x)\|p_S^\theta(\cdot|y_{<n},x) \right)

纯 on-policy distillation 的目标是:

LOD(θ)=ExX[EypS(x)[KL(pTpSθ)(yx)]] L_{\mathrm{OD}}(\theta)= \mathbb{E}_{x\sim X} \left[ \mathbb{E}_{y\sim p_S(\cdot|x)} \left[ \mathrm{KL}(p_T\|p_S^\theta)(y|x) \right] \right]

GKD 用 λ\lambda 控制学生样本比例:

LGKD(θ)=(1λ)E(x,y)(X,Y)[D(pTpSθ)(yx)]+λExX[EypS(x)[D(pTpSθ)(yx)]] L_{\mathrm{GKD}}(\theta)= (1-\lambda)\mathbb{E}_{(x,y)\sim(X,Y)} \left[ \mathcal{D}(p_T\|p_S^\theta)(y|x) \right] + \lambda\mathbb{E}_{x\sim X} \left[ \mathbb{E}_{y\sim p_S(\cdot|x)} \left[ \mathcal{D}(p_T\|p_S^\theta)(y|x) \right] \right]

其中 λ=0\lambda=0 退化到监督 KD,λ=1\lambda=1 对应纯 on-policy KD,中间值把固定数据和学生轨迹混合。实际训练中,算法对每个 mini-batch 抽样一个随机数 uu:若 uλu\le\lambda,学生先生成输出,再在这些输出上对齐教师分布;若 u>λu>\lambda,使用原始样本做分布匹配。

一个实现选择很关键:论文不对学生采样过程反向传播。学生生成的序列被当作训练数据,梯度只来自教师和学生在这些 prefix 上的分布散度。这样训练目标接近普通 KD,工程复杂度低于把序列采样过程纳入 policy-gradient 的方案。

3. 散度选择:forward KL、reverse KL 与 JSD

论文把 D\mathcal{D} 设计成可替换项,实验覆盖:

  • forward KL;
  • reverse KL;
  • Jensen-Shannon divergence, JSD,含不同混合系数 0.10.10.50.50.90.9

forward KL 倾向覆盖教师分布中概率质量较大的区域,reverse KL 更偏向 mode-seeking,JSD 在两者之间提供更平滑的折中。论文没有给出一个所有任务通用的最优散度,实验现象显示 on-policy 数据比例和散度选择需要随任务调参;例如 GSM8K 上,学生轨迹比例超过 25%25\% 后更明显地带来提升。

4. 与 RL fine-tuning 的组合

论文还把 GKD 接到 RL fine-tuning 中。以 XSum factual consistency 为例,学生模型通过 reward r(y)r(y) 优化摘要事实一致性,同时保持教师风格和摘要质量。目标写成:

ExX[(1α)EypSθ(x)[r(y)]αEypS(x)[D(pTpSθ)(yx)]] \mathbb{E}_{x\sim X} \left[ (1-\alpha) \mathbb{E}_{y\sim p_S^\theta(\cdot|x)}[r(y)] - \alpha \mathbb{E}_{y\sim p_S(\cdot|x)} \left[ \mathcal{D}(p_T\|p_S^\theta)(y|x) \right] \right]

其中 α\alpha 控制 distillation 约束强度。实验显示,较大的 α\alpha 能提高 ROUGE-2,但事实一致性收益会下降。这说明 GKD 在 RL 中相当于一个 teacher-behavior regularizer:它能减少模型为了 reward 单点优化而偏离摘要质量,但也会限制 reward 目标的提升幅度。

关键实验/定理

1. 总体设定

论文没有定理主线,贡献主要来自实验验证。模型家族采用 T5:

  • 学生:T5-small 约 77M、T5-base 约 250M、T5-large 约 800M;
  • 教师:T5-XL 约 3B;
  • 压缩比例:约 38×38\times12×12\times3.8×3.8\times

主要基线包括:

  • Supervised KD:在固定数据上拟合教师 token 分布;
  • SeqKD:使用教师生成序列作为训练目标;
  • ImitKD:把训练样本和学生生成样本混合,可视为 GKD 的一个特例风格;
  • f-distill:基于可替换 ff-divergence 的离线蒸馏。

为保证比较公平,论文通常先把学生做 supervised fine-tuning,再从同一学生 checkpoint 启动不同 KD 方法。

2. XSum 摘要:on-policy 数据带来数据效率收益

XSum 是论文最核心的展示场景。所有模型先在 XSum 上 supervised fine-tune,随后比较不同 distillation 方法在 ROUGE-2 上的表现。结果的关键点包括:

  • on-policy GKD 在 T5-small、T5-base、T5-large 上都能改善初始学生;
  • 在数据缩减实验中,T5-small 只使用 5%5\% XSum 数据进行 on-policy GKD,表现能超过使用完整训练集的 Supervised KD 和 ImitKD;
  • 使用温度采样评估时,GKD 对学生自生成分布的覆盖优势更明显;
  • 论文汇总称,on-policy GKD 在摘要任务上相对初始学生的收益约为其他 KD 方法平均收益的 2.1×2.1\times

这个实验支撑了论文的主张:如果蒸馏数据分布更贴近学生推理分布,小模型能用更少数据获得有效教师反馈。

3. WMT14 en-de:翻译任务上的泛化

WMT14 en-de 使用 BLEU 作为指标,教师 T5-XL 的 BLEU 约为 2828,T5-small 初始学生约为 25.5825.58,T5-base 初始学生约为 26.9826.98。GKD 在该任务中同样优于主要基线,附录图注报告 on-policy GKD 的 BLEU 改善平均高于 ImitKD 和 f-distill。

翻译任务的意义在于,它比 XSum 更强调局部 token 选择和较短输出。GKD 仍能带来收益,说明 distribution shift 并不限于长摘要输出。

4. GSM8K:算术推理与 CoT 轨迹

GSM8K 实验使用 few-shot chain-of-thought 格式,并用外部计算器校验答案。教师是 fine-tuned FLAN T5-XL,准确率约为 27.927.9;学生从 FLAN-T5 系列出发,基于 CoT 训练对继续蒸馏。

关键观察是:当 on-policy 样本比例超过 25%25\% 后,GKD 的准确率收益更稳定。推理任务中,学生早期推理步骤的偏差会影响后续步骤,on-policy 蒸馏能让教师在学生真实走到的中间状态上给 token-level feedback。

5. Task-agnostic instruction tuning

论文还做了 task-agnostic 蒸馏:使用 FLAN2021 的 62 个训练任务训练学生 FLAN T5-base,从教师 FLAN T5-XL 获取反馈,然后在 held-out MMLU 57 个任务和 BBH 23 个任务上评估 exact match。

教师约为 MMLU 52.4%52.4\%、BBH 41.0%41.0\%;初始学生约为 MMLU 35.6%35.6\%、BBH 31.25%31.25\%。论文报告 on-policy GKD 尤其是 reverse KL 设定优于 Supervised KD 和 ImitKD,并带来约 1%1\% MMLU 和 2%2\% BBH 绝对准确率提升。

这个实验把 GKD 从单任务蒸馏推到多任务 instruction distillation,也解释了后续大模型训练报告为什么会把 OPD 作为跨任务能力迁移组件。

6. RL + GKD:摘要质量和事实性之间的折中

论文在 XSum 上用 T5-XXL NLI 模型构造 factual consistency reward,并比较 RL fine-tuning 与加入 GKD regularization 的版本。结果显示:

  • RL 能提升事实一致性,但可能牺牲 ROUGE;
  • 加入 GKD 后,ROUGE-2 能明显改善;
  • GKD 权重 α\alpha 增大时,摘要质量更接近教师,但 reward 指标改善幅度降低。

这个实验为后续 RL 后模型蒸馏提供了直接线索:OPD 可以把 RL 优化后的行为迁移到小模型,也可以在 RL 过程中约束模型不要过度偏离教师。

方法定位

GKD 可以放在三个坐标上理解。

第一,它是 token-level KD 的 on-policy 化。教师仍然给 soft label,学生仍然通过常规分布散度训练;变化发生在被监督的前缀分布上,从固定答案前缀扩展到学生生成前缀。

第二,它是 RL-style training 和 supervised KD 之间的工程折中。它借用 on-policy 数据分布的思想,但避免对采样动作做 policy-gradient,训练实现更接近标准语言模型 fine-tuning。

第三,它是后续多教师 OPD 的单教师基础版。MOPD 的核心问题是不同教师具备不同能力、不同 token 风格和不同 RL 后行为,需要把教师选择、能力路由、学生自生成轨迹与多教师反馈结合。GKD 先给出了单教师、单学生、单任务或多任务 instruction 设定下的基本形式。

证据链强度评估

强证据

  • 论文在摘要、翻译、算术推理和 task-agnostic instruction tuning 上重复验证 on-policy distillation,相比单任务实验更能支撑方法通用性。
  • XSum 数据缩减实验和 GSM8K on-policy 比例实验直接对应核心主张:学生自生成轨迹上的教师反馈能缓解 train / inference distribution shift。
  • 对 Supervised KD、SeqKD、ImitKD、f-distill 等基线的比较,为 GKD 在当时 KD 方法谱系中的位置提供了清晰参照。

中等强度证据

  • 实验控制了学生初始 checkpoint,并在多个 T5 尺寸上观察到收益,但模型生态集中在 encoder-decoder T5。
  • RL + GKD 只在摘要事实性 reward 上展示,能说明 regularization 价值,不能覆盖复杂 agentic reward。

需要谨慎的推论

  • 将 GKD/OPD 迁移到 frontier decoder-only、长上下文、多轮工具调用和 multi-teacher 设置时,需要重新评估采样成本、teacher serving 成本和行为策略偏移。
  • λ\lambda、divergence 和 on-policy 比例没有统一选择规则,工程使用时仍需任务级调参。

主要启发

  • OPD 的核心是让教师在学生真实访问到的状态分布上给反馈。
  • 稀疏 outcome reward 场景可以把 teacher token distribution 当成 dense learning signal,但 teacher 质量和服务成本会成为系统瓶颈。
  • 对后续 MOPD 或 async OPD 论文,应重点追踪 teacher 来源、trajectory freshness、prefill 开销和 KL 方向。

局限

  • 模型规模停留在 T5-XL 教师和 T5-small/base/large 学生,证据主要来自 2023 年前后的 encoder-decoder 模型生态;对 frontier decoder-only 模型、长上下文推理、多轮 agent 轨迹的适配需要额外验证。
  • 学生必须先有足够可用的初始能力。论文也强调 on-policy 蒸馏需要从 supervised fine-tuned student 启动,否则学生自生成轨迹质量过低,教师反馈覆盖到的状态可能偏离任务目标。
  • on-policy 采样增加训练成本。论文在 GSM8K 上报告,学生采样相对固定输出训练约带来 1.8×1.8\times2.0×2.0\times2.2×2.2\times 的开销,具体取决于学生/教师规模比例。
  • λ\lambda 与 divergence 选择缺少统一理论规则。不同任务上 forward KL、reverse KL、JSD 和 on-policy 比例都有调参空间。
  • 实验的 RL 部分集中在摘要事实性 reward,不能直接推出对代码、数学证明、多工具 agent 等复杂 reward 环境的结论。
  • OpenReview 公开评审细节在当前检索中未能可靠拉取,本文对审稿分歧的判断主要依据论文正文、附录、arXiv/DeepMind 的接收状态以及后续技术报告中的方法引用。

OpenReview / 审稿意见吸收

  • Venue status: ICLR 2024 accepted,arXiv comment 和 Google DeepMind publication page 均给出该状态。
  • Public reviews: OpenReview forum 可定位到 3zKtaqxLhW,当前浏览/API 会话未能可靠取得 review note payload。
  • Reviewer consensus: 当前不做评分级复述。
  • Credibility impact: 这篇论文的可信度主要来自完整实验覆盖、附录超参、公开 ICLR 接收状态,以及后续 Qwen/DeepSeek/GLM/MOPD 等技术报告对 OPD/GKD 风格训练的延续。

本地讨论补充

  • 用户曾询问 OPD 的经典论文。就公开论文谱系看,这篇 GKD 是 LLM 蒸馏语境下非常核心的早期 OPD 论文:它明确把学生自生成输出纳入 distillation objective,并把 supervised KD、on-policy KD、不同 divergence 和 RL regularization 放进同一框架。
  • 与 MOPD 的关系可以概括为:GKD 解决“学生在哪些状态上接受教师反馈”,MOPD 进一步解决“多个教师分别在哪些能力区域给学生反馈”。
  • 对后续系统训练的启发是:OPD 的收益取决于学生轨迹质量、教师反馈质量、采样成本和 divergence 选择。把 OPD 放进大规模 RL 或 agent rollout 系统时,核心工程问题会转为调度、缓存复用、teacher routing 和在线样本质量控制。

跨论文关系

  • SDFT:两者都让 teacher 在 student 自生成 prefix 上提供 token distribution;SDFT 用 query-specific demonstration 构造同尺寸 self-teacher。SDFT 项目 README 进一步披露其全部论文结果实际使用 per-token forward KL,并明确类比 GKD,因此 GKD 是 SDFT 当前实验实现最直接的前序方法。
  • MOPD:GKD 解决学生在哪些自生成状态上接受教师反馈;MOPD 继续处理多教师的能力路由、teacher mismatch 和多领域能力整合,是最直接的方法延展。
  • Qwen3DeepSeek-V4GLM-5.2:这些模型报告分别用 strong-to-weak distillation、OPD 或 parallel OPD 将 on-policy teacher signal 推进到 decoder-only frontier model 的后训练流程。
  • verl 官方仓库:verl 的 async on-policy distill 将教师查询、学生 rollout 和 policy update 放入异步流水线,补充了 GKD 在大规模训练中的调度与吞吐问题。

Reference Intake Brief

Source Used for Reliability Notes
arXiv abs 2306.13649 Title, authors, dates, abstract, ICLR 2024 comment High arXiv canonical metadata.
arXiv HTML 2306.13649v3 Abstract, method, contributions, experiment descriptions High Used for fast line-level reading.
arXiv TeX source Equations, algorithms, appendices, author footnotes, hyperparameters High Primary source; resolved details missing from abstract.
Google DeepMind publication page Venue/date confirmation and institutional source High Official publication entry.
OpenReview forum Venue page pointer Medium Forum located, but public review payload was not reliably available in current session.
Author homepages / Scholar / institutional pages Author profile and historical affiliation update Medium-High Used only when identity matched name, publication topic, and institution.
Existing local paper notes Cross-paper relationship mapping High Used to connect GKD to MOPD, Qwen3, DeepSeek-V4, GLM-5.2, and verl-related OPD discussion.