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

TL;DR

This paper analyzes the failure modes of on-policy distillation in LLM training and proposes a simple, effective fix using top-K local support matching to improve stability and performance.

Contribution

It introduces a new top-K local support matching objective that addresses key failure modes of sampled-token OPD, enhancing stability and performance in LLM distillation.

Findings

Identified three failure modes of sampled-token OPD: imbalance, unreliability, and mismatch.

Proposed top-K local support matching improves training stability.

Achieved a +19.8% performance gain over standard methods.

Abstract

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.