TIP: Token Importance in On-Policy Distillation
About
On-policy knowledge distillation (OPD) trains a student on its own rollouts under token-level supervision from a teacher. Not all token positions matter equally, but existing views of token importance are incomplete. We ask a direct question: which tokens carry the most useful learning signal in OPD? Our answer is that informative tokens come from two regions: positions with high student entropy, and positions with low student entropy plus high teacher--student divergence, where the student is overconfident and wrong. Empirically, student entropy is a strong first-order proxy: retaining $50\%$ of tokens with entropy-based sampling matches or exceeds all-token training while reducing peak memory by up to $47\%$. But entropy alone misses a second important region. When we isolate low-entropy, high-divergence tokens, training on fewer than $10\%$ of all tokens nearly matches full-token baselines, showing that overconfident tokens carry dense corrective signal despite being nearly invisible to entropy-only rules. We organize these findings with TIP (Token Importance in on-Policy distillation), a two-axis taxonomy over student entropy and teacher--student divergence, and give a theoretical explanation for why entropy is useful yet structurally incomplete. This view motivates type-aware token selection rules that combine uncertainty and disagreement. We validate this picture across three teacher--student pairs spanning Qwen3, Llama, and Qwen2.5 on MATH-500 and AIME 2024/2025, and on the DeepPlanning benchmark for long-horizon agentic planning, where Q3-only training on $<$$20\%$ of tokens surpasses full-token OPD. Our experiments are implemented by extending the OPD repository https://github.com/HJSang/OPSD_OnPolicyDistillation, which supports memory-efficient distillation of larger models under limited GPU budgets.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Mathematical Reasoning | MATH 500 | Top-1 Accuracy79.1 | 452 | |
| Scientific Question Answering | GPQA Diamond | Accuracy47.98 | 131 | |
| Function Calling | BFCL V3 | -- | 104 | |
| Logical reasoning | ZebraLogic | Accuracy18.7 | 86 | |
| Mathematical Reasoning | AIME 24 | Accuracy26 | 39 | |
| Mathematical Reasoning | MATH-500 1 (test) | Accuracy86.2 | 38 | |
| Code Generation | HumanEval v1 (test) | Accuracy80.49 | 37 | |
| Logical reasoning | AutoLogi | Accuracy (Avg@8)65.5 | 32 | |
| Mathematical Reasoning | DAPO-Math OOD average (avg of four benchmarks) | Avg@1653.3 | 30 | |
| Mathematical Reasoning | DAPO-Math In-Distribution (test) | Avg@1676.9 | 30 |