Fixing On-Policy Distillation Mismatch

Summary & Methodology Analysis

On-policy distillation, a process where a smaller student model learns from a larger teacher model using live rollouts, frequently suffers from degenerate agreement. In this failure mode, student models exploit repetitive loops to achieve near-local token alignment with the teacher while outputting globally flawed, degenerate responses. Standard on-policy distillation computes trajectory-level reverse KL divergence via detached token-level policy gradients using sampled advantage. To resolve this mismatch, the authors introduce TIDE, which diagnoses teacher-student mismatch at each state from two directional perspectives and computes branch-specific statistics. These statistics flag positions where the student over-estimates probability or under-weights teacher-preferred tokens. TIDE then applies quantile gating to select the most severe mismatch positions for updates while discarding neutral tokens. Furthermore, it applies bounded Hellinger shaping to suppress student-excess tokens, replacing raw unbounded log-ratio penalties, and applies an analytic teacher top-K injection to recover missing probability mass for student-deficit tokens without requiring rare rollouts. Finally, it jointly optimizes the suppression and recovery branches as a weighted loss objective.

In terms of empirical performance evaluated on models including Qwen3, MiMo-V2-Flash, GLM-5, Kimi K3, JustRL-DeepSeek-1.5B, and DeepSeek-R1-Distill-Qwen-1.5B across datasets such as DAPO-Math-17K, AIME'24, AIME'25, AMC23, MATH, Minerva, and Olymp., TIDE demonstrates massive improvements. On the strong-mismatch Qwen3-8B to Qwen3-1.7B-Base pair, TIDE improves Avg@8 from 6.9% to 20.3% and reduces the average response length from 22,395 tokens to 7,294 tokens. Additionally, deficit-only interventions alone improved Avg@8 from 6.87 to 18.32.

Despite these strong results, the approach comes with distinct limitations. TIDE assumes that the teacher model is locally reliable, which may not hold in all scenarios. The authors acknowledge that the effectiveness of TIDE has only been evaluated on two teacher-student pairs and within the domain of mathematical reasoning. Additionally, the theoretical framework provided is based on fixed-state, operator-level analysis and lacks global convergence guarantees. The paper does not specify hardware requirements, exact training time, or dollar costs.