Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

arxiv:2603.25562

Original source · Paper JSON

Post-training · 2026-03-26

Cites (5) · Cited by (3)

Complete cleaned Markdown · Source Markdown including bibliography

Read in parts

1. Part 1
2. Part 2
3. Part 3
4. Part 4

Contents

• Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes
• Abstract
• 1 Introduction
• 2 Understanding Sampled-Token OPD: Trade-offs and Failure Modes
• 2.1 From reverse-KL to token-level OPD
• 2.2 Why sampled-token OPD is brittle in practice
• A highly imbalanced sampled-token signal.
• The teacher signal can become unreliable on student-generated prefixes.
• Tokenizer and special-token mismatch.
• 3 Method
• 3.1 Teacher top- K K local support matching
• 3.2 Practical stabilization choices
• Support-set renormalization.
• Top- p p rollout sampling.
• Special-token masking.
• 4 Experiments
• 4.1 Setup
• 4.2 Single-task math reasoning
• 4.3 Multi-task agentic-plus-reasoning training
• 4.4 Ablations
• Top- K K support variants.
• 5 Conclusion
• Appendix A Discussion and Limitations
• Sampled-token augmentation of the teacher top- K K support.
• Support-set KL is a truncated objective.
• Training–inference mismatch.
• Teacher matching remains an imperfect proxy for task success.
• Appendix B Future Directions
• Top- p p truncation as an adaptive support.
• OPD versus RL in multi-task settings.
• Continual learning as a testbed.
• Relation to other stabilization directions.
• Appendix C Related Work
• On-policy distillation.
• KL computation for LLM post-training.
• KL divergence in RL for LLMs.
• Appendix D Bias and variance analysis of token-level versus sequence-level OPD
• D.1 Bias of the token-level estimator
• D.2 Worst-case variance upper bounds
• D.3 Discussion
• Appendix E Toy experiment details
• E.1 Environment
• E.2 Gradient variance estimation
• E.3 Additional Results of Toy Experiments
• Appendix F Experiment Setups
• Appendix G Additional Results
• G.1 Additional Evidence on Single-task WebShop
• G.2 Training Dynamics and Alignment
• Better learning curves.
• More stable optimization.
• Improved teacher–student alignment.
• G.3 Additional Results of Top- K K Variants
• G.3.1 Formulations of Five Top- K K Variants
• Shared notation.
• Renormalized local reverse-KL.
• Variant 1: teacher top- K K renormalization (our method).
• Variant 2: student top- K K + sampled token + renormalization.
• Variant 3: teacher top- K K + sampled token + renormalization.
• EMA-PG-style local loss.
• Variant 4: student top- K K + EMA-PG correction.
• Variant 5: teacher top- K K + EMA-PG correction.
• Summary.
• G.3.2 Results under Single-task and Multi-task Settings
• Appendix H Qualitative OPD reward-hacking case study
• Instructions for reporting errors