Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes
Yuqian Fu · Haohuan Huang · Kaiwen Jiang · Jiacai Liu · Zhuo Jiang · Yuanheng Zhu · Dongbin Zhao
We identify three failure modes of on-policy distillation in LLM post-training and propose a top-K truncated reverse-KL objective that yields more stable training and improved performance.
摘要
On-policy distillation (OPD) is increasingly used in LLM post-training because it can leverage a teacher model to provide dense supervision on student rollouts. The standard implementation, however, usually reduces distribution matching to a sampled-token log-ratio, which can make the learning signal fragile on long rollouts whose prefixes drift away from the teacher's typical support. We revisit this formulation from both theoretical and implementation perspectives. Theoretically, token-level OPD is biased relative to sequence-level reverse-KL minimization, but admits a substantially tighter worst-case variance bound; a controlled synthetic study further shows that stronger future-reward coupling increases gradient variance and destabilizes training. Empirically, we identify three failure modes of sampled-token OPD: imbalanced token-level supervision, unreliable teacher guidance on student-generated prefixes, and tokenizer or special-token mismatch. These findings motivate teacher top-$K$ local support matching, a truncated reverse-KL objective that compares teacher and student distributions over a teacher-supported token set at each prefix, together with top-$p$ rollout sampling and special-token masking. Across single-task reasoning and multi-task benchmarks spanning agentic and reasoning settings, this objective improves optimization stability and yields a **+19.8\%** performance gain over standard sampled-token OPD baselines, providing a practical recipe for more stable on-policy distillation.