Tail-Aware Top-$k$ On-Policy Distillation

Huipeng Huang; Hongxin Wei

cs.LGAug 12, 2026

Tail-Aware Top- $k$ On-Policy Distillation

Authors: Huipeng Huang, Hongxin Wei

Organizations: Department of Statistics and Data Science, Southern University of Science and Technology

Abstract

On-policy distillation (OPD) has emerged as an effective paradigm for transferring knowledge between language models, where a student is trained to align its next-token distribution with the teacher's along its own trajectories. To provide dense supervision at tractable cost, many works minimize the reverse Kullback-Leibler (KL) divergence between the student and teacher's normalized distributions over the teacher's top- $k$ tokens. However, this normalized objective discards the information about tail probability: the total probability outside the teacher's top- $k$ tokens. As a result, the optimization can steadily increase the student's tail probability and entropy, empirically degrading downstream accuracy. To address this issue, we propose Tail-Aware Top- $k$ OPD (TA-OPD), a novel distillation method that restores the missing tail probability signal. In particular, TA-OPD minimizes the reverse KL divergence over the top- $k$ tokens plus a tail token that carries the tail probability. In effect, TA-OPD better aligns the student's next-token distribution with the teacher's, preventing the increase in tail probability and entropy caused by top- $k$ normalization. Extensive experiments demonstrate the superiority of TA-OPD, improving Avg@8 by up to 8.05 points on common benchmarks. Our code is available at https://github.com/HuipengHuang/TA-OPD.

Figures & tables

Figure 1: The student’s tail probability, the teacher’s tail probability, and the student’s token-level entropy across training steps. Normalized top-

k

OPD consistently increases the student’s tail probability and entropy across all three model pairs.

Figure 2: Validation accuracy versus student’s tail probability.

Figure 3: Geometric Interpretation of OPD Objectives. (a) Sparse Update (Sampled-Token OPD): The supervision is derived from a single sampled token, making the update direction stochastic and prone to updating some token logits in erroneous directions. (b) Increased Tail Probability (Normalized Top-

k

OPD): The normalized objective discards the tail probability, leading to increased tail probability. (c) Decreased Tail Probability (Tail-Aware Top-

k

OPD, Ours): TA-OPD restores the missing tail probability signal, matching the student’s tail probability to the teacher’s.

Methods	In-Distribution Performance							OOD Performance
Methods	MATH500	Minerva	Olympiad	AMC	AIME 24	AIME 25	Avg.	ARC-c	MMLU-Pro	Avg.
Student: Qwen2.5-7B-Instruct Teacher: OpenThinker3-7B
Sampled-token	74.60	31.25	41.22	45.48	14.58	16.25	37.23	74.99	47.66	61.33
Norm. top- $k$	68.78	25.55	36.02	41.42	13.75	15.00	33.42	23.07	37.38	30.23
TA-OPD (Ours)	77.88	32.58	42.41	45.03	16.67	17.50	38.68	76.30	50.35	63.33
Student: Llama-3.1-8B Teacher: DeepSeek-R1-Distill-Llama-8B

Table 1: Results on math reasoning and out-of-distribution (OOD) benchmarks. We compare TA-OPD with sampled-token OPD (Sampled-token) and normalized top-

k

OPD (Norm. top-

k

). Each Avg. column represents the macro-average. Best results are shown in bold .

Figure 4: Tail probabilities and token-level entropy of TA-OPD and normalized top-

k

OPD over training. The student and teacher models are Qwen2.5-7B-Instruct and OpenThinker3-7B.

Figure 5: Ablation study on

k

Methods	In-Distribution Performance							OOD Performance
Methods	MATH500	Minerva	Olympiad	AMC	AIME 24	AIME 25	Avg.	ARC-c	MMLU-Pro	Avg.
Student: DeepSeek-R1-Distill-Qwen-1.5B Teacher: JustRL-DeepSeek-1.5B
TA-OPD	86.13	30.97	57.04	76.20	41.67	29.73	53.62	35.40	25.98	30.69
SC-TA-OPD	87.00	32.08	57.65	76.51	42.92	30.17	54.39	35.55	25.76	30.66
Student: Qwen3-1.7B Teacher: Qwen3-30B-A3B-Instruct-2507
TA-OPD	83.48	49.82	34.50	52.86	26.25	20.83	44.62	83.59	50.42	67.01

Table 2: Comparison between TA-OPD and sample-corrected TA-OPD on math reasoning and out-of-distribution (OOD) benchmarks. Best results are shown in bold .

Figure 6: Performance under different student–teacher capability gaps. The students are Qwen2.5-Math-1.5B and DeepSeek-R1-Distill-Qwen-1.5B. The teacher is JustRL-DeepSeek-1.5B.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: The student’s tail probability, the teacher’s tail probability, and the student’s entropy across training steps under different student–teacher capability gaps. The teacher is Qwen3-8B, and the students are Qwen3-1.7B-Base and Qwen3-8B-Base. The tail probability increase occurs for Qwen3-1.7B-Base but not for Qwen3-8B-Base.

Figure 8: The student’s tail probability, the teacher’s tail probability, and the student’s entropy across training steps. The student and teacher models are Qwen2.5-Math-1.5B and DeepSeek-R1-Distill-Qwen-7B, respectively. The tail probability increase occurs when

k\leq 64

Figure 9: The student’s tail probability, the teacher’s tail probability, and the student’s entropy across training steps under different maximum response lengths. The student and teacher models are Qwen2.5-7B-Instruct and OpenThinker3-7B, respectively. The tail probability increase becomes more pronounced as the maximum response length grows.

Figure 10: The student’s tail probability, the teacher’s tail probability, and the student’s entropy across training steps under the student and teacher top-

k

supports. The student and teacher models are Qwen3-1.7B-Base and Qwen3-8B. The tail probabilities are computed with respect to the support in use. The tail probability increase occurs only under the teacher top-

k

support.

Figure 11: Average response length across training steps under the student and teacher top-

k

supports. The student and teacher models are Qwen3-1.7B-Base and Qwen3-8B, respectively. Under the student top-

k

support, the response length saturates at the maximum response length of

7168

tokens, while the teacher top-

k

variant maintains a healthy response length.

Student: Qwen3-1.7B-Base Teacher: Qwen3-8B
Method	MATH500	Minerva	Olympiad	AMC	AIME 24	AIME 25	Avg.
Student top- $k$	68.58	25.32	33.35	36.00	9.17	6.25	29.78
Teacher top- $k$	69.93	25.46	33.41	36.00	9.17	7.08	30.18

Appendix

Table 3: Performance of normalized top-

k

OPD under the student and teacher top-

k

supports. The Avg. column represents the macro-average across all six math reasoning benchmarks. Best results are shown in bold .

Category	Item	Value
Training	Training temperature	1.0
	Global batch size	72
	Mini batch size	36
	Rollout number	4
	$k$ (Number of Top- $k$ Tokens)	16
	Top- $p$	1.0

Appendix

Table 4: Default training and evaluation settings.

Figure 12: Visualization of the student policies under different OPD objectives. Normalized top-

k

OPD ( Red ) only matches the relative shape of the teacher distribution over the top-

k

tokens, but fails to align the tail probability. In contrast, TA-OPD ( Blue ) and sample-corrected TA-OPD ( Purple ) closely recover the bimodal teacher distribution.

Figure 13: Learning dynamics of different OPD objectives. TA-OPD better approximates the full-vocabulary OPD.

Figure 14: Top-

k

overlap ratio between the student and the teacher across training steps. TA-OPD achieves higher overlap ratios than normalized top-

k

OPD across different model pairs.

Panel of Figure 15 — Figure 15: Computational overhead comparison between sampled-token OPD, normalized Top- $k$ OPD, TA-OPD, and sample-corrected TA-OPD. Left: total training runtime. Right: throughput distribution across training steps. The student and teacher models are Qwen3-1.7B and Qwen3-30B-A3B-Instruct-2507, respectively.

Figure 16: The student’s tail probability, the teacher’s tail probability, and the student’s token-level entropy across training steps. Without the tail token, the student’s tail probability rises above

0.9

and training collapses.

Student: Qwen2.5-7B-Instruct Teacher: OpenThinker3-7B
Methods	MATH500	Minerva	Olympiad	AMC	AIME 24	AIME 25	Avg.
UnNorm. top- $k$	3.35	1.70	1.42	3.00	0.42	0.42	1.72
TA-OPD (Ours)	77.88	32.58	42.41	45.03	16.67	17.50	38.68

Appendix

Table 5: Avg@8 of unnormalized top-

k

OPD and TA-OPD. The Avg. column represents the macro-average across all six benchmarks. Best results are shown in bold .

Panel of Figure 17 — Figure 17: Tail probability and token-level entropy over training for TA-OPD and Normalized top- $k$ OPD. Each row corresponds to a different student-teacher pair.

Panel of Figure 18 — Figure 18: Tail probability and token-level entropy over training for TA-OPD and Normalized top- $k$ OPD under different student-teacher capability gaps. The tail probability increase only occurs for Qwen2.5-Math-1.5B but not for DeepSeek-R1-Distill-Qwen-1.5B.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 508 papers

Unbiased Top- $k$ Estimation for On-Policy Distillation
Sep 28, 2026Linjian Meng, Siyuan Gan, YuHan Li +5Language Model Distillation On-Policy Distillation
Shanghai Artificial Intelligence Laboratory, Shanghai, China · State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing, China
Don't Ignore the Tail: Decoupling top-K Probabilities for Efficient Language Model Distillation
Feb 24, 2026Sayantan Dasgupta, Trevor Cohn, Timothy BaldwinLanguage Model Distillation Knowledge Distillation
School of Computing and Information Systems, University of Melbourne, Melbourne, Australia.
Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes
Mar 26, 2026Yuqian Fu, Haohuan Huang, Kaiwen Jiang +4Language Model Distillation On-Policy Distillation
State Key Laboratory of Multimodal Artificial Intelligence Systems, CASIA · School of Artificial Intelligence, UCAS · Fudan University +1

Sep 28, 2026J/K move · Enter open · S save

Unbiased Top- $k$ Estimation for On-Policy Distillation

Linjian Meng, Siyuan Gan, YuHan Li, Xiran Wang+4

Shanghai Artificial Intelligence Laboratory, Shanghai, China · State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing, China

Language Model Distillation On-Policy Distillation

Open paper AskPDF arXiv

On-policy distillation (OPD) is becoming an important component of large language model (LLM) post-training for transferring the reasoning capability of a strong teacher LLM to a weaker student LLM. OPD trains the student by minimizing the reverse KL divergence between the teacher and the student via rollouts generated by the student's policy. However, estimating the gradient of the reverse KL divergence in OPD remains a challenge. Using only the sampled token from the student-generated rollout is computationally cheap but provides limited distributional supervision, which will degrade accuracy. In addition, using the full vocabulary provides complete distributional supervision but is computationally expensive. Therefore, recent works propose Top-

k

OPD (TK-OPD) that use selected top-

k

tokens, which provides richer distributional supervision than sampled-token estimation at substantially lower computational cost than full-vocabulary estimation. Unfortunately, using only the selected top-

k

tokens induces bias, leading to accuracy degradation, as the probability mass outside the selected top-

k

tokens is discarded. To address the bias of TK-OPD, we propose Tail-Corrected Top-

k

On-Policy Distillation (TT-OPD). It preserves the advantages of TK-OPD, including rich distributional supervision and low computational cost, while providing an unbiased estimator of the gradient of the reverse KL divergence. The key insight of TT-OPD is to use not only the selected top-

k

tokens, but also the sampled token from the student-generated rollout, thereby recovering the discarded probability mass in expectation, avoiding the bias. Experimental results demonstrate that TT-OPD significantly outperforms other tested OPD variants.

Abstract

Language Model Distillation

Explore similar work

Unbiased Top-kkk Estimation for On-Policy Distillation

Don't Ignore the Tail: Decoupling top-K Probabilities for Efficient Language Model Distillation

Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

Unbiased Top- $k$ Estimation for On-Policy Distillation