Research map / Label-efficient and robust learning / Knowledge distillation
On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Generalized Knowledge Distillation (GKD), which trains an auto-regressive student on its own self-generated (on-policy) outputs with teacher token-level feedback, using flexible divergences (forward/reverse KL, generalized JSD) and optionally mixing with RL fine-tuning.
Problem
Knowledge distillation for auto-regressive sequence models suffers from train-inference distribution mismatch, since students are trained on fixed teacher or ground-truth sequences but generate their own outputs at inference, and forward KL can force a low-capacity student to cover teacher modes.
Approach
Frame distillation as imitation learning with an interactive expert: the student samples its own outputs and minimizes a token-level divergence to the teacher on those outputs, generalized by a student data fraction (mixing on-policy, supervised, and teacher-generated data) and a choice of divergence, and combined with RL fine-tuning by regularizing toward the teacher.
Key results
- On XSum, on-policy GKD with JSD(0.9) outperforms baselines and, with T5-small on 5% of data without ground-truth summaries, beats supervised KD and ImitKD trained on full data.
- Averaged across T5 student sizes, on-policy GKD gives relative gains of 2.1x on summarization, 1.7x on machine translation, and 1.9x on arithmetic reasoning over baseline KD methods.
- On instruction tuning with FLAN-T5-base student, on-policy GKD with reverse KL yields about 2% and 1% absolute accuracy gains on held-out BBH and MMLU over supervised KD and ImitKD.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use on-policy GKD instead of baseline KD | T5 students · Summarization | performance improvement over the initial student: 2.1× | Helped |
| Use on-policy GKD instead of baseline KD | T5 students · Machine translation | performance improvement over the initial student: 1.7× | Helped |
| Use on-policy GKD instead of baseline KD | T5 students · Arithmetic reasoning | performance improvement over the initial student: 1.9× | Helped |
| Distill with on-policy GKD | FLAN-T5-Base · BBH | accuracy: 2% | Helped |
| Distill with on-policy GKD | FLAN-T5-Base · MMLU | accuracy: 1% | Helped |
| Use on-policy student-generated sequences instead of fixed supervised sequences (ablation) | GKD · XSum | ROUGE-2 | Helped |
| Use mixed student-generated and fixed sequences instead of fixed supervised sequences (ablation) | GKD · XSum | ROUGE-2 | Helped |
| Use reverse KL or JSD (0.9) instead of forward KL (ablation) | On-policy GKD · XSum, temperature sampling | ROUGE-2 | Helped |
| Change the divergence among forward KL, reverse KL, and JSD variants (ablation) | On-policy GKD · XSum, greedy sampling | ROUGE-2 | No clear effect |
| Use on-policy or mixed sequences instead of only fixed supervised sequences (ablation) | GKD · WMT14 en-de | BLEU | Helped |
| Use generalized JSD instead of forward or reverse KL (ablation) | GKD · WMT14 en-de | BLEU | Helped |
| Use solely student-generated CoTs instead of fixed or mixed CoT data (ablation) | GKD · GSM8K | accuracy | Helped |
Surprising findings
- Forward KL performs poorly with temperature sampling on XSum, while reverse KL and JSD(0.9) perform best.
- Reverse KL performs well even when training only on a fixed dataset for GSM8K.
Limitations
- Requires a student that already produces reasonable-quality samples (starts from supervised fine-tuned checkpoint).
- Student sampling adds 1.8x-2.2x training compute relative to fixed-dataset distillation.
- Optimal divergence and student data fraction are task-dependent and must be tuned.
Experimental setup
Datasets: XSum WMT14 en-de GSM8K FLAN2021 PaLM-540B generated CoT data
Benchmarks: XSum WMT14 en-de GSM8K MMLU BBH
Metrics: ROUGE-2 BLEU accuracy Self-BLEU factual consistency (entailment)
Baselines: SeqKD Supervised KD ImitKD f-distill MiniLLM PaLM 540B
Builds on: ImitKD SeqKD MiniLLM RLEF
Closest papers by meaning
- DistiLLM: Towards Streamlined Distillation for Large Language Models (ICML 2024 · similarity 0.94)
- KL for a KL: On-Policy Distillation with Control Variate Baseline (NeurIPS 2026 · similarity 0.90)
- DDK: Distilling Domain Knowledge for Efficient Large Language Models (NeurIPS 2024 · similarity 0.90)
- Hybrid Policy Distillation for LLMs (ICML 2026 · similarity 0.90)
- DA-KD: Difficulty-Aware Knowledge Distillation for Efficient Large Language Models (ICML 2025 · similarity 0.90)
- Why Knowledge Distillation Works in Generative Models: A Minimal Working Explanation (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.