TIP: Token Importance in On-Policy Distillation

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

cs.LGApr 15, 2026

TIP: Token Importance in On-Policy Distillation

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

Abstract

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.

Figures & tables

Figure 1: Cross-task summary: average accuracy by selection method. Each panel shows one benchmark; bar height is the mean accuracy (mean@16) averaged across three teacher–student pairs for mathematical reasoning (Qwen3-8B

\to

4B, Llama-70B

\to

8B, Qwen2.5-14B

\to

1.5B) and across two teacher sizes (14B, 32B) for DeepPlanning. Methods: Base. = all-token OPD (100%); Ent. 50%/20% = entropy-based token selection at the stated retention ratio; SO 50%/20% = Soft-OR (Eq. 7 ) top-

k

selection. The dashed line marks the all-token baseline. Soft-OR consistently improves over entropy-only selection on the mathematical reasoning benchmarks and remains competitive on DeepPlanning, confirming that augmenting entropy with divergence recovers the Q3 blind spot without sacrificing Q1/Q2 coverage.

Quadrant	$h_{t}$	$\delta_{t}$	Learning role
Q1: High entropy, high divergence	High	High	Correct errors or consolidate fragile knowledge
Q2: High entropy, low divergence	High	Low	Stabilize underconfident predictions
Q3: Overconfident	Low	High	Break systematic confident biases
Q4: Solved	Low	Low	Negligible signal

Table 1: Token taxonomy. Classification by student entropy

h_{t}

and teacher–student divergence

\delta_{t}

Figure 2: TIP taxonomy as a two-axis map. Entropy determines whether the student is uncertain or confident; divergence determines whether the teacher agrees or disagrees. Q1 and Q2 are visible to entropy-based methods, while Q3 is the low-entropy blind spot that requires divergence to detect.

Result	Prediction	Tested in
Proposition 1	Q1/Q2 carry the most signal; removing Q4 improves efficiency	Section 7.2
Proposition 2	Entropy-only selection misses Q3 tokens	Section 7.3
Remark 1	Combined score recovers Q3 and outperforms entropy-only	Section 7.4

Table 2: Theoretical predictions and experimental tests.

Model pair	Benchmark	100%	50%	20%	10%
Qwen3-8B (GRPO) $\to$ 4B	MATH-500	76.7 $\pm$ 0.7	78.6 $\pm$ 0.6	74.1 $\pm$ 0.9	70.8 $\pm$ 1.0
	AIME’24	21.9 $\pm$ 1.2	23.8 $\pm$ 1.3	22.5 $\pm$ 1.1	21.1 $\pm$ 1.2
	AIME’25	19.4 $\pm$ 1.1	20.7 $\pm$ 1.3	21.5 $\pm$ 1.2	19.2 $\pm$ 1.0
Llama-70B $\to$ 8B	MATH-500	71.0 $\pm$ 0.7	74.0 $\pm$ 0.8	73.6 $\pm$ 0.7	72.5 $\pm$ 0.8
	AIME’24	21.5 $\pm$ 1.1	25.3 $\pm$ 1.5	18.8 $\pm$ 1.3	20.6 $\pm$ 1.4
	AIME’25	4.9 $\pm$ 0.9	7.5 $\pm$ 1.1	10.0 $\pm$ 1.2	9.3 $\pm$ 1.0

Table 3: Entropy sampling across model pairs. Accuracy (%, mean@16

\pm

std). Sampling selects tokens with probability

p_{t}\propto h_{t}

. Peak Mem reports end-to-end training peak GPU memory. Bold marks the best per benchmark.

Model pair	Benchmark	Baseline (100%)	Q3 20%	Q3 10%
Qwen3-8B $\to$ 4B	MATH-500	76.7 $\pm$ 0.7	75.7 $\pm$ 0.8	76.1 $\pm$ 0.9
	AIME’24	21.9 $\pm$ 1.2	20.2 $\pm$ 1.1	21.5 $\pm$ 1.3
	AIME’25	19.4 $\pm$ 1.1	19.2 $\pm$ 1.1	17.1 $\pm$ 1.0
Llama-70B $\to$ 8B	MATH-500	71.0 $\pm$ 0.7	71.8 $\pm$ 1.0	70.4 $\pm$ 1.1
	AIME’24	21.5 $\pm$ 1.1	20.2 $\pm$ 1.2	20.8 $\pm$ 1.1
	AIME’25	4.9 $\pm$ 0.9	4.5 $\pm$ 0.7	4.1 $\pm$ 0.7

Table 4: Training on Q3 (overconfident) tokens only. Accuracy (%, mean@16

\pm

std) when training exclusively on low-entropy, high-divergence tokens. Q3-only training with

<

10% of all tokens can nearly match the all-token baseline across model pairs. Peak Mem reports end-to-end training peak GPU memory.

		Baseline	Entropy-only		Soft-OR
Model pair	Benchmark	100%	50%	20%	50%	20%
Qwen3-8B $\to$ 4B	MATH-500	76.7 $\pm$ 0.7	78.6 $\pm$ 0.6	74.1 $\pm$ 0.9	79.1 $\pm$ 0.8	77.6 $\pm$ 0.7
	AIME’24	21.9 $\pm$ 1.2	23.8 $\pm$ 1.3	22.5 $\pm$ 1.1	25.7 $\pm$ 1.4	24.5 $\pm$ 1.2
	AIME’25	19.4 $\pm$ 1.1	20.7 $\pm$ 1.3	21.5 $\pm$ 1.2	21.9 $\pm$ 1.2	23.2 $\pm$ 1.2
Llama-70B $\to$ 8B	MATH-500	71.0 $\pm$ 0.7	74.0 $\pm$ 0.8	73.6 $\pm$ 0.7	74.7 $\pm$ 1.0	74.2 $\pm$ 0.7
	AIME’24	21.5 $\pm$ 1.1	25.3 $\pm$ 1.5	18.8 $\pm$ 1.3	26.0 $\pm$ 1.4	21.0 $\pm$ 1.5

Table 5: Main results: Baseline vs. Entropy-only vs. Soft-OR. Accuracy (%, mean@16

\pm

std). Soft-OR uses

s_{t}=\hat{h}_{t}+\hat{\delta}_{t}-\hat{h}_{t}\cdot\hat{\delta}_{t}

(Eq. 7 ) with Top-K selection. Bold marks the best per benchmark.

	Qwen3-8B $\to$ 4B		Llama-70B $\to$ 8B		Qwen2.5-14B $\to$ 1.5B
Benchmark	Top 50%	Bot. 50%	Top 50%	Bot. 50%	Top 50%	Bot. 50%
MATH-500	79.1 $\pm$ 0.8	72.3 $\pm$ 0.9	74.7 $\pm$ 1.0	67.4 $\pm$ 1.1	56.2 $\pm$ 1.2	50.3 $\pm$ 1.0
AIME’24	25.7 $\pm$ 1.4	15.5 $\pm$ 1.3	26.0 $\pm$ 1.4	17.2 $\pm$ 1.3	3.8 $\pm$ 1.2	1.5 $\pm$ 0.7
AIME’25	21.9 $\pm$ 1.2	12.8 $\pm$ 1.2	11.5 $\pm$ 1.1	2.9 $\pm$ 0.8	1.5 $\pm$ 0.7	0.8 $\pm$ 0.5

Table 6: Top 50% vs. bottom 50% by Soft-OR score. Accuracy (%, mean@16

\pm

std). “Top” trains on the highest-scoring half by

s_{t}

; “Bot.” trains on the lowest-scoring half. The bottom tokens carry substantially less signal.

Method	Teacher 14B	Teacher 32B
OPD, all tokens (100%)	$11.7{\pm}0.07$	$12.8{\pm}0.07$
+ Entropy-only 50%	$12.1{\pm}0.06$	$13.1{\pm}0.07$
+ Entropy-only 20%	$11.6{\pm}0.07$	$12.7{\pm}0.06$

Table 7: Agentic planning on DeepPlanning (Qwen3-1.7B student, thinking-enabled, Avg@16 %). Left : reference and entropy-only methods. Right : Q3-only and Soft-OR. Q3-only 20% surpasses full-token OPD; Soft-OR matches or exceeds entropy-only.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Entropy sampling across retention ratios. Accuracy (mean@16) on three benchmarks as a function of retention ratio. Retaining 50% of tokens with entropy-based sampling matches or outperforms the all-token baseline across model pairs. At very low retention, entropy-only selection begins to plateau or degrade.

		Baseline	Div-only	Div-only	Q3-only
	Benchmark	100%	50%	10%	10%
Accuracy	MATH-500	76.7 $\pm$ 0.7	76.3 $\pm$ 0.8	74.3 $\pm$ 0.8	76.1 $\pm$ 0.9
	AIME’24	21.9 $\pm$ 1.2	21.5 $\pm$ 1.1	19.2 $\pm$ 1.2	21.5 $\pm$ 1.3
	AIME’25	19.4 $\pm$ 1.1	21.1 $\pm$ 1.2	15.3 $\pm$ 1.1	17.1 $\pm$ 1.0

Appendix

Table 8: Isolating the entropy axis: Div-only vs Q3-only on Qwen3-8B (GRPO)

\to

4B (mean@16 %,

\pm

std). Div-only is a hard top-

k

instantiation of LATF-style divergence selection [ Xie et al., 2026 ] . We compare Q3-only at 10% against Div-only at both

5\times

the budget (50%) and matched budget (10%). Bold marks the best per benchmark.

Figure 4: Token selection for agentic OPD on 20% held-out travel-planning queries. Top row : Avg@16; Bottom row : Best@16 (Pass@16). Within each row the left panel uses the 14B teacher and the right panel uses the 32B teacher. Q3-only 20% matches or exceeds the full-token baseline in every setting, consistent with Table 7 . Best@16 results show the same pattern: overconfident-token training improves the upper tail of performance, not just the mean.

	Qwen3 (8B $\to$ 4B)	Llama (70B $\to$ 8B)	Qwen2.5 (14B $\to$ 1.5B)
Optimizer	AdamW	AdamW	AdamW
Learning rate	$1\times 10^{-6}$	$3\times 10^{-7}$	$1\times 10^{-6}$
Batch size (rollouts)	8	8	8
Rollouts per prompt	16	16	16
Max response length	8192	8192	8192
Max prompt length	2048	2048	2048

Appendix

Table 9: Training hyperparameters across model pairs.

Knowledge Distillation

Knowledge distillation transfers useful behavior or representations from larger or specialized models into more compact or adaptable learners.

See all 811 papers

Influence-Directed Distillation: Solving the Diversity Bottleneck in Sampled-Token On-Policy Distillation
Aug 30, 2026Run Yang, Runpeng Dai, Jie Sun +5Language Model Distillation On-Policy Distillation
University of North Carolina at Chapel Hill · University of Science and Technology of China · Shanghai University of Finance and Economics
Not All Disagreement Is Learnable: Token Teachability in On-Policy Distillation
May 26, 2026Yuanyi Wang, Su Lu, Yanggan Gu +6Selective Knowledge Distillation Language Model Distillation
The Hong Kong Polytechnic University, PolyU · Hong Kong Polytechnic University Daya Bay Technology and Innovation Research Institute
Compatibility Matters: A Framework for Token-Level Teacher Selection in On-Policy Distillation
Oct 15, 2026Stanley Wei, Jiin Woo, Kaan Ozkara +1On-Policy Distillation Multi-Teacher Knowledge Distillation
Princeton University · Amazon Annapurna Labs

Aug 30, 2026J/K move · Enter open · S save

Influence-Directed Distillation: Solving the Diversity Bottleneck in Sampled-Token On-Policy Distillation

Run Yang, Runpeng Dai, Jie Sun, Jielei Zhang+4

University of North Carolina at Chapel Hill · University of Science and Technology of China · Shanghai University of Finance and Economics

Language Model Distillation On-Policy Distillation Knowledge Distillation

Open paper AskPDF arXiv

Sampled-token on-policy distillation (OPD) efficiently transfers capabilities from teacher to student using student-generated tokens, requiring teacher probabilities only for sampled tokens. Yet it frequently suffers from diversity distillation failure: the student's pass@1 improves while its pass@

k

plateaus, failing to inherit the teacher's diversity. To explain this, we introduce First-Order Local Entropy Influence, a signed first-order proxy that decouples each update's entropy effect into the teacher--student log-probability gap and the student's local probability structure, and empirically links entropy contraction to negative-influence positions. Motivated by this, we propose Influence-Directed Adaptive On-Policy Distillation (IDA-OPD): rather than relying on costly full-vocabulary Forward-KL objectives, it preserves entropy-expanding updates while replacing entropy-contracting ones with divergence-adaptive advantage shrinkage, using only the teacher's sampled-token log-probability. Experiments on reasoning-oriented distillation show IDA-OPD consistently improves pass@

k

, inheriting the teacher's diversity through distillation, matches the strongest teacher-informed methods at strictly lower cost, and broadly maintains vanilla OPD's pass@1, all without full-vocabulary teacher information.

Abstract

Knowledge Distillation

Explore similar work

Influence-Directed Distillation: Solving the Diversity Bottleneck in Sampled-Token On-Policy Distillation

Not All Disagreement Is Learnable: Token Teachability in On-Policy Distillation

Compatibility Matters: A Framework for Token-Level Teacher Selection in On-Policy Distillation