PaperlayerView exact revision on arXiv

Source-linked AI summary

TIP: Token Importance in On-Policy Distillation

Yuanda Xu, Hejian Sang, Zhengze Zhou, Ran He, Zhipeng Wang, Alborz Geramifard

arXiv:2604.14084v4cs.LGcs.AI

TL;DR

OPD lacks a complete account of which student-generated tokens carry useful teacher-supervised learning signal. The paper introduces TIP, a two-axis entropy–divergence taxonomy and token-selection view, finding that entropy is useful but misses overconfident errors. Across mathematical reasoning and agentic planning, combining uncertainty with disagreement recovers this blind spot, while Q3-only selection can approach or exceed full-token OPD.

  • Problem

    OPD asks which tokens in student-generated rollouts carry the most useful teacher-supervised learning signal, but entropy alone cannot distinguish confident correct from confident wrong tokens.

  • Method

    TIP organizes token importance by student entropy and teacher–student divergence and uses their combination to identify high-entropy and overconfident, high-divergence tokens.

  • Results

    Across mathematical reasoning and DeepPlanning, combined selection improves over entropy-only selection on mathematical reasoning, while Q3-only training with 20% of tokens surpasses full OPD on DeepPlanning.

  • Takeaways & Limitations

    Entropy is an effective first-order proxy, but divergence is needed to recover the low-entropy, high-divergence tokens carrying dense corrective signal.

  • Takeaways & Limitations

    The largest teacher is 70B and rollouts reach 16K tokens, leaving persistence of the quadrant structure and Q3 concentration at larger scales and extreme rollout lengths open.

Abstract

from arXiv · show

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.

1 Introduction

OPD raises the question of which student-generated token positions carry the most useful teacher-supervised learning signal. TIP answers with high-entropy tokens and low-entropy, high-divergence tokens, showing why entropy-only selection is useful but incomplete.

  • OPD transfers teacher corrections token by token on student-generated rollouts, making token importance a property of each student–teacher state.
  • Informative tokens occupy high student-entropy positions and low-entropy positions with high teacher–student divergence.
  • 50% entropy-based retention matches or improves all-token training while substantially reducing end-to-end training memory.
  • Entropy-only selection misses overconfident tokens because low entropy conflates confident and correct with confident and wrong.
  • TIP combines student entropy and teacher–student divergence, and Soft-OR consistently improves over entropy-only selection on mathematical reasoning while remaining competitive on agentic planning.

2 Related Work

Prior work selects examples, sequences, or tokens using difficulty, quality, gradient, entropy, or divergence signals. TIP focuses on token-level importance within student-generated OPD rollouts, where importance must be assessed online.

  • Curriculum learning and importance sampling prioritize training examples by difficulty, gradient norm, or learned weights rather than individual tokens.
  • Off-policy distillation trains on teacher-generated sequences, whereas OPD uses student rollouts with token-level teacher supervision to avoid train–test distribution mismatch.
  • Because token importance depends on the student’s distribution at each position, it cannot be pre-computed from teacher outputs and must be assessed online.
  • Response-level methods select rollouts but treat all tokens within a response uniformly, leaving token-level selection as a complementary problem.
  • Related token-level work identifies high-entropy decision points, overconfident errors, or divergence-based weighting, while TIP specifies Q3 through low entropy plus high divergence.

3 Setup

The setup characterizes each student-generated token by student entropy and teacher–student divergence, two quantities already available during standard OPD training. Their combination defines the token-importance plane and its quadrants.

  • The student samples a rollout from a prompt, and the teacher scores each position in the resulting student-generated sequence.
  • Both entropy and divergence are already computed during training, so the paper studies whether useful signal concentrates in regions of their two-dimensional plane.
  • The taxonomy classifies tokens by student entropy and teacher–student divergence into four quadrants.
  • Student entropy measures uncertainty: high entropy indicates uncertainty, while low entropy indicates confidence.
  • Teacher–student divergence measures disagreement, with high divergence indicating that the teacher disagrees with the student.

4 TIP Taxonomy: A Two-Axis View of Token Importance

TIP organizes token importance along student entropy and teacher–student divergence, yielding four quadrants with different prevalence and corrective-signal profiles.

  • TIP uses student entropy and teacher–student divergence as two axes for classifying token importance.
  • Q4 contains roughly 40–47% of tokens, Q1 and Q2 together contain 40–52%, and Q3 contains only 3–15%.
  • Despite its smaller share, Q3 carries disproportionate corrective signal compared with the more prevalent quadrants.

5 Theoretical Analysis

The analysis characterizes useful OPD tokens through oracle descent signal and shows why entropy is effective but structurally incomplete. High-entropy tokens and low-entropy, high-divergence tokens can carry corrective signal, whereas low-signal solved tokens do not.

  • 5.1 Oracle Token Weight: The oracle token weight favors gradients aligned with descent without excessive energy, providing a signal-to-curvature interpretation of token importance.The weight is an oracle because it depends on population-level gradient quantities.
  • 5.1 Oracle Token Weight: Q1, Q2, and Q3 can receive non-negligible oracle weight, while Q4 has negligible weight because its descent signal is near zero.Q3 retains corrective signal despite low student entropy, so entropy-only rules can miss it.
  • 5.2 A Signal-to-Curvature View: Low-entropy Q3 tokens can have large corrective signal because teacher disagreement is high, whereas low-entropy Q4 tokens have near-zero signal because disagreement is small.The distinction between Q3 and Q4 is therefore a signal distinction rather than an entropy distinction.
  • 5.2 A Signal-to-Curvature View: Any non-decreasing entropy-only score with zero weight at zero entropy is structurally blind to Q3 tokens and cannot distinguish confident correctness from confident error.This follows because Q3 and Q4 can both have low student entropy despite differing teacher–student disagreement.
  • 5.2 A Signal-to-Curvature View: The Soft-OR score recovers Q3 coverage while suppressing Q4 and preserving high scores for Q1 without requiring oracle gradient quantities.It is nonzero whenever entropy or divergence is nonzero.

6 Method: Type-Aware Token Selection

The method selects a retained fraction of tokens using a parameter-free Soft-OR score that combines normalized student entropy and teacher–student divergence. It stabilizes entropy normalization by clipping batch outliers before ranking tokens.

  • 6 Method: Type-Aware Token Selection: Given retention ratio ρ, the method retains the top-ρ fraction of tokens ranked by the Soft-OR score st = ˆht + ˆδt − ˆht · ˆδt.The inputs are min-max normalized entropy and divergence.
  • 6 Method: Type-Aware Token Selection: Setting normalized divergence to zero recovers entropy-only selection, while including divergence promotes low-entropy, high-disagreement Q3 tokens.The score is parameter-free and uses quantities already computed during standard distillation.
  • 6 Method: Type-Aware Token Selection: Before scoring, the method clips the top 2% of entropy values within each batch and applies min-max normalization to stabilize token rankings.The clipping suppresses rare entropy outliers.

7 Experiments

Experiments test TIP’s predictions across mathematical reasoning and agentic planning: high-entropy selection preserves much of OPD’s benefit, while Q3 selection recovers corrective signal that entropy misses. Combining entropy and divergence is designed to retain both regions under token budgets.

  • 7.2 High-Entropy Tokens (Q1/Q2): Across three model pairs, retaining 50% of tokens with entropy sampling matches or outperforms all-token training on most benchmarks.For Qwen3 Small, MATH improves from 76.7 to 78.6; for Llama, from 71.0 to 74.0.
  • 7.2 High-Entropy Tokens (Q1/Q2): Entropy sampling substantially reduces total peak training memory while preserving benchmark performance.Table 3 reports end-to-end peak GPU memory alongside accuracy.
  • 7.2 High-Entropy Tokens (Q1/Q2): Entropy-only selection becomes incomplete at aggressive retention ratios because discarded low-entropy, high-divergence tokens retain useful learning signal.The experiments therefore isolate Q3 tokens: low student entropy with high teacher–student divergence.
  • 7.3 Overconfident Tokens (Q3): Q3-only training with fewer than 10% of all tokens nearly matches the all-token baseline across model pairs.For Qwen3, 5.7K overconfident tokens reach 76.1 on MATH versus 76.7 for full-token training.
  • 7.3 Overconfident Tokens (Q3): The Q3 selector combines low entropy with high divergence rather than ranking by divergence alone, which underperforms at matched tight budgets.The low-entropy conjunct is reported as the feature that makes selection effective under constrained token budgets.
  • 7.4 Combined Selection: TIP’s Soft-OR score combines normalized entropy and divergence to recover Q3 while retaining high-entropy Q1/Q2 tokens.The main comparison evaluates Baseline, Entropy-only, and Soft-OR with Top-K selection; the complementary bottom-half test treats low-scoring tokens as predominantly Q4.
  • 7.1 Experimental Setup: Teacher entropy provides little discriminative signal; TIP instead uses student entropy and the student–teacher gap.The paper reports teacher entropy as nearly constant across positions.
  • 7.5 Beyond Mathematical Reasoning: Agentic Planning: On DeepPlanning, 20% Q3-only training surpasses full-token OPD for both teachers, while 50% entropy selection matches or exceeds full-token OPD.Reported Avg@16 results are 12.6 versus 11.7 and 13.6 versus 12.8 for Q3-only comparisons, and 12.1 versus 11.7 and 13.1 versus 12.8 for entropy-only comparisons.

8 Discussion and Conclusion

TIP argues that token importance in OPD requires both student entropy and teacher–student divergence, with overconfident errors forming a particularly informative region. Its empirical support spans model families and task domains, while its scale remains bounded by the evaluated teachers and rollout lengths.

  • Discussion and Conclusion: TIP identifies student entropy and teacher–student divergence as jointly necessary axes of token importance in OPD.The taxonomy organizes tokens into four quadrants and reports support across three model families and two task domains.
  • Discussion and Conclusion: Fewer than 10% of overconfident tokens nearly matches full-training performance, showing dense corrective signal in the low-entropy, high-divergence quadrant.This is described as the clearest result supporting the Q3 blind spot.
  • Discussion and Conclusion: 20% Q3-only training surpasses full OPD on DeepPlanning, reaching 12.6 versus 11.7 Avg@16.The paper suggests this value grows when downstream computation depends on a single committed step.
  • Limitations: The study evaluates teachers up to 70B and rollouts up to 16K tokens, leaving larger models and longer agentic rollouts open.The authors specifically identify trillion-parameter scale and extremely long tool-calling rollouts as unresolved settings.
  • Discussion and Conclusion: The local curvature proxy is smaller for near-deterministic students and larger for diffuse student distributions, but it is distinct from Shannon entropy.The appendix frames this quantity as a curvature-side proxy for distributional spread.
  • Discussion and Conclusion: The descent-bound interpretation requires non-negligible gradient magnitude, so low curvature alone cannot justify entropy-only weighting.Q3 and Q4 may both have low entropy and curvature, while Q3 has greater teacher–student disagreement.

A.2 Derivation of the Descent Bound

The derivation approximates token-level descent using local curvature and token-separable interactions. It shows why entropy-weighted sampling preserves coverage but why disagreement is needed to distinguish overconfident errors from correct confident tokens.

  • Derivation: The analysis neglects off-diagonal interactions across token positions to obtain a token-separable quadratic approximation.Mean interaction terms are also treated as lower-order in the approximation.
  • Entropy-Weighted Sampling: Entropy sampling preserves nonzero Q3 coverage, but its variance cost grows as 1/p_t for low-entropy tokens.The sampling probability is proportional to student entropy.
  • Soft-OR Selection: The Soft-OR score separates Q3 from Q4 by adding divergence while preserving the high-entropy ranking without a tuning parameter.The product term prevents double-counting the two normalized axes.

B Supplementary Experiments

Entropy-based sampling retains strong performance at moderate retention, but teacher entropy is nearly constant and therefore offers little discriminative information for token selection.

  • Entropy Sampling: 50% entropy-based retention matches or outperforms all-token training across model pairs and benchmarks.At very low retention, entropy-only selection begins to plateau or degrade.
  • Entropy Sampling: Teacher entropy is near-zero and nearly constant across evaluated models, so the useful selection axes are student entropy and the student–teacher gap.Reported teacher-entropy statistics include mean 0.031 for Qwen3 and 0.067 for Llama.

B.2 Isolating the Entropy Axis: Comparison with Divergence-Only Selection

The entropy-axis comparison shows that low entropy is not redundant when selecting divergent tokens: Q3-only selection is more efficient at matched budgets, while divergence-only selection catches up only with a larger budget. In agentic planning, Q3-only or Soft-OR selection also improves both average and upper-tail outcomes.

  • Equal-Budget Comparison: At 10% retention, Q3-only beats Div-only on MATH-500, AIME’24, and AIME’25 by +1.8, +2.3, and +1.8, respectively.The corresponding scores are 76.1 vs 74.3, 21.5 vs 19.2, and 17.1 vs 15.3.
  • Interpretation: Divergence-only ranking is not sufficient at a matched budget because it selects the high-divergence tail without isolating the low-entropy conjunct.The paper frames the Q3 specification as a distinct, structurally identifiable region rather than a restatement of either axis.
  • Interpretation: Div-only 10% underperforms the all-token baseline, so a larger individual gradient does not establish greater usefulness under a fixed token budget.The comparison distinguishes pointwise gradient size from informative signal retained for optimization.
  • Agentic Planning: Q3-only 20% leads Avg@16 with both teacher sizes, while Best@16 is led by Soft-OR with the 14B teacher and Q3-only with the 32B teacher.The reported Avg@16 values are 12.6 and 13.6 versus baselines of 11.7 and 12.8; Best@16 values are 20.3 versus 18.9 and 20.1 versus 19.7.
  • Experimental Setup: The experiments use mean@16, defined as average accuracy over 16 independent samples per problem, with temperature 1.0.Top-K experiments select the top-ρ fraction deterministically by score.

B.5 Qualitative Examples Across Quadrants

The examples span Q1 and Q3 tokens, showing that entropy captures uncertain disagreements but misses confident errors that require divergence. Across mathematical reasoning cases, low-entropy, high-divergence tokens encode dense corrective signals.

  • Q3: low entropy, high divergence: 0.02 entropy and 5.27 forward KL identify a confident generic-variable choice that should instead advance to a concrete substitution.The teacher places 49.9% probability on the concrete value 2, while the student assigns low entropy to the incorrect continuation.
  • Q3: low entropy, high divergence: 0.40 entropy and 3.54 forward KL mark a confident arithmetic error where the student favors 190 but the teacher prefers 6 or 2.The student places 91.2% on digit 0, whereas the teacher distributes probability across the mathematically relevant alternatives.
  • Q3: low entropy, high divergence: 0.12 entropy and 5.58 forward KL identify a very confident wrong-variable continuation that entropy-only selection would miss.The student assigns 98.2% to B, while the teacher prefers A and the student’s cosine-based continuation is mathematically incorrect.
  • Q1: high entropy, high divergence: 1.38 entropy and 4.27 forward KL describe a Q1 symbol error that entropy detects while the teacher strongly favors the relevant function cot.The student spreads probability across irrelevant symbols, whereas the teacher assigns 91.7% to cot.
  • Q3: low entropy, high divergence: Q3 tokens combine low student entropy with high teacher–student divergence, so entropy-only selection under-weights confident mathematical errors.Examples include generic-variable misuse, arithmetic mistakes, and an incorrect variable-level derivation.
Loading 2604.14084v4…