1% of Tokens Can Be Enough: Gradient Estimation in On-Policy Distillation

An MBZUAI and Ant Group paper (first and corresponding authors at MBZUAI; Ant Group co-authors include Jinjie Gu) that adds a second axis to sparse on-policy distillation. Existing token selectors ask whether a teacher correction is useful; this work asks whether its gradient can be reliably estimated from the single sampled next token that sampled-token OPD actually uses. At a fixed prefix it decomposes the one-sample reverse-KL gradient into signal and sampling noise in Fisher geometry, derives the variance-minimizing scalar baseline, and defines the information-efficiency ratio (IER) as the resulting signal-to-noise ratio. A top-K candidate-set approximation makes IER cheap enough to rank every token, either alone or fused with usefulness scores through soft OR and AND operators, while the training objective stays the ordinary sampled reverse KL.

IER scores are extremely heavy-tailed: fewer than 0.1% of tokens have estimated signal exceeding noise. Across two math pairs (JustRL-Nemotron-1.5B to OpenMath-Nemotron-1.5B, JustRL-Qwen3-4B to Qwen3-1.7B, on DAPO-Math-17k, Bayes@32 on AIME and HMMT 2025/2026) and a medical pair (ClinAlign-4B to Qwen3-4B on HealthBench), IER alone at a 0.1% token budget lands within a point of full OPD (AIME26 58.9 vs 59.9; HealthBench 45.25 vs 45.77), and IER fused with TIP or TA-OPD at 0.1% to 1% matches or beats full OPD. It rescues the weak Prefix selector at 0.1% from 38.30 to 44.98 on HealthBench. Raising the budget toward 80% does not monotonically help, and Entropy-based fusions sometimes hurt. A companion to Amazon's extremely sparse supervision result from earlier in September, and one more data point that OPD's cost lives in rollouts, not in supervised tokens.