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

Sayantan Dasgupta; Trevor Cohn; Timothy Baldwin

cs.CLFeb 24, 2026

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

Authors: Sayantan Dasgupta, Trevor Cohn, Timothy Baldwin

Organizations: School of Computing and Information Systems, University of Melbourne, Melbourne, Australia.

Abstract

The core learning signal used in language model distillation is the standard Kullback-Leibler (KL) divergence between the student and teacher distributions. Traditional KL divergence tends to be dominated by the next tokens with the highest probabilities, i.e., the teacher's modes, thereby diminishing the influence of less probable yet potentially informative components of the output distribution. We propose a new tail-aware divergence that decouples the contribution of the teacher model's top-K predicted probabilities from that of lower-probability predictions, while maintaining the same computational profile as the KL Divergence. Our decoupled approach reduces the impact of the teacher modes and, consequently, increases the contribution of the tail of the distribution. Experimental results demonstrate that our modified distillation method yields competitive performance in both pre-training and supervised distillation of decoder models across various datasets. Furthermore, the distillation process is efficient and can be performed with a modest academic budget for large datasets, eliminating the need for industry-scale computing.

Figures & tables

Panel of Figure 1 — Figure 1: KL divergence on the validation set of Regmix for vanilla KD vs. TAD. The $x$ axis shows training progress in terms of the number of tokens, and the $y$ axis shows held-out KL between the student and teacher, measured on Regmix’s validation set ( Section 3 ).

Panel of Figure 2 — Figure 2: Tail probability mass ( $\alpha^{T}_{K}$ ) against $K$ for different teachers in the first, and the Next Token vs. Mode mismatch rate in percentage in the second plot, measured on the validation set of Regmix (see Section 3.2 )

Teacher/Student		HS	WG	OBQA	ARC-E	ARC-C	PIQA	SIQA	Story	Avg	$\overline{\text{Rel}}$
Qn1.5 1.8B $\downarrow$ 1.2B	CLM (no KD)	39.4	51.8	28.4	46.0	25.7	67.0	39.5	62.2	45.0	$-$ 0.7
	Vanilla KD	40.7	53.2	29.8	46.1	25.5	67.3	39.2	63.5	45.6
	Seq-KD	38.5	51.9	29.2	46.5	25.1	66.3	39.0	61.0	44.7	$-$ 0.9
	MiniLLM	36.1	51.2	28.5	44.1	25.3	65.8	37.9	61.4	43.8	$-$ 1.9
	MiniPLM	42.8	53.3	31.0	46.8	26.9	68.3	39.8	64.0	46.6	$+$ 1.0
	TAD ( $K=1$ )	42.3	53.8	30.5	52.0	27.0	67.3	41.2	63.9	47.2	$+$ 1.2

Table 1: Results for Tail-aware distillation for

\beta=2

over Qwen1.5-1.8B ( “Qn” ), for a 1.2B and 0.5B student model. The best performance for each column, and any value within

0.4

of it, is highlighted. CLM stands for pre-training the model with only the CLM loss, without distillation. The average relative change for the best-case TAD (

K=10

) is 50% to 120% better than MiniPLM.

	$\beta$	0.5	1	2	5	10
1.2B	Avg	47.0	47.6	47.8	47.7	47.6
1.2B	$\overline{\text{Rel}}$	$+$ 1.4	$+$ 2.0	$+$ 2.2	$+$ 2.1	$+$ 2.0
0.5B	Avg	45.0	45.1	45.4	45.1	44.9
0.5B	$\overline{\text{Rel}}$	$+$ 1.0	$+$ 1.2	$+$ 1.5	$+$ 1.2	$+$ 1.0

Table 2: Parameter sensitivity of

\beta

for the distillation of Qwen 1.8B for

K=10

# P(M)	Vanilla	MiniPLM	TAD	MiniLLM	Seq-KD
1.2B	9.2	12.4	9.3	39.0	65.0
0.5B	6.4	9.7	6.5	21.8	43.2

Table 3: PetaFLOPs for the distillation of Qwen-1.5-1.8B ( Section 3.2.1 ) on a subset of 1M tokens from the Regmix dataset. TAD has a similar PFLOP to Vanilla KD, while MiniPLM is higher than both. The methods involving sequence generation (SeqKD or MiniLLM) are too expensive to scale to billions of tokens.

Teacher/Student		HS	WG	OBQA	ARC-E	ARC-C	PIQA	SIQA	Story	Avg	$\overline{\text{Rel}}$	F-ECE $\downarrow$
Phi2 2.8B $\downarrow$ 1.1B	CLM (no KD)	38.2	51.1	27.4	51.2	24.1	66.3	40.8	63.1	45.3	$-$ 4.3	1.57
	CLM (Mat.)	40.2	51.9	28.6	52.3	24.8	67.6	41.7	64.7	46.5	$-$ 2.4	1.50
	Vanilla KD	43.6	53.5	33.0	57.3	30.0	68.0	43.2	64.3	49.1		1.45
	MiniPLM	43.7	52.5	30.6	57.1	29.9	68.1	43.8	64.3	48.8	$-0.4$	1.62
	RKL	42.3	54.1	31.6	58.0	28.7	68.2	43.8	64.9	49.0	$-0.4$	1.77
	TAD ( $K=1$ )	45.2	55.3	34.0	58.0	30.7	68.3	44.4	64.9	50.1	$+$ 0.9	1.19

Table 4: Pretraining distillation of various teachers to students with

\sim

1B active parameters on 2 billion tokens from Regmix. CLM (no KD) refers to pretraining with only CLM loss, without distillation with the same number of tokens (2B), where CLM (Mat.) refers to computation-matched pretraining, matched to the same FLOPs as training of TAD. The last column “F-ECE” shows the calibration error of the models, measured using Full-ECE, with the lower being better.

Model		Data (#Tkns)	GSM8K	MATH	SVAMP	ASDiv	MAWPS	TAB	MQA	SAT	Avg
TinyLlama(TL) $-$ 1.1B		Web (2.5T)	2.0	2.6	9.5	16.3	20.1	12.7	12.8	15.6	11.4
	CLM (no KD)	+ OWM(2.5B)	3.9	3.8	17.9	29.7	39.5	12.2	10.8	15.6	16.7
Phi3 4B $\downarrow$ TL	Vanilla KD	+ OWM(2.5B)	6.1	4.2	21.1	33.5	41.5	15.5	11.2	16.7	18.7
	MiniPLM	+ OWM(2.5B)	3.3	3.4	13.4	27.3	34.0	10.8	10.5	12.5	14.4
	TAD ( $K=1$ )	+ OWM(2.5B)	6.1	6.2	22.1	33.1	41.5	14.0	11.3	21.9	19.5
	TAD ( $K=5$ )		7.1	4.8	19.2	35.9	46.7	15.9	10.0	22.6	20.3

Table 5: Adaptation to mathematical reasoning via pretraining distillation of Phi-3 into TinyLlama-1B ( “TL” ) on the OpenWebMath (OWM) corpus. The distilled students with TAD outperform pretrained 1B Gemma3 and Llama3.2 models in terms of average score.

Model		Data (#tokens)	GSM8K	MATH	SVAMP	ASDiv	MAWPS	TAB	MQA	SAT	Avg.
TinyLlama(TL) $-$ 1.1B		Web (2.5T)	2.0	2.6	9.5	16.3	20.1	12.7	12.8	15.6	11.4
	CLM $+$ SFT	+OWM(2.5B) ’ +ORCAMEL	19.6	4.0	49.4	58.8	74.3	21.8	18.0	28.1	34.3
Phi3 4B $\downarrow$ TL	Vanilla KD	+OWM(2.5B) ’ +ORCAMEL	30.8	6.8	64.6	62.5	80.7	20.1	16.7	27.5	38.7
	TAD ( $K=1$ )	+OWM (2.5B) ’ +ORCAMEL	36.8	6.8	67.8	67.9	81.7	25.4	16.3	28.1	41.4
	TAD ( $K=5$ )		33.2	7.4	65.4	68.7	85.6	27.6	17.9	34.4	42.5
	TAD ( $K=10$ )		30.1	9.0	65.7	68.4	85.4	24.1	18.2	29.8	41.3

Table 6: Supervised distillation for mathematical reasoning, showing distillation of Phi3-4B into TinyLlama-1.1B ( “TL” ) and Phi3-14B into Llama2-7B on ORCAMEL, alongside GPT4-generated solutions. TAD for TinyLlama is

2.5\times

computationally cheaper than Rho-1 and

9\times

cheaper for Llama2-7B than Llemma-7B (see Section B.1 ), which is the best model created from Llama2-7B.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Teacher	#P(M)	$\|\mathcal{V}\|$	$d_{S}$	$L_{S}$	$n_{H}$	$d_{H}$	$d_{FFN}$
Qwen1.5-1.8B	$1.2$ B	151,936	1,536	24	16	96	4,224
Qwen1.5-1.8B	$0.5$ B	151,936	1,024	24	16	64	2,816
Phi2-2.8B	1.1B	52,000	2,560	16	32	80	5,120
Qwen2.5-3B	1.2B	151,936	2,048	18	16	128	7,680
Gemma2-9B	2B	256,000	3,584	14	16	224	4,096

Appendix

Table 7: The architectures of different students used in distillation for pretraining from scratch

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

Tail-Aware Top- $k$ On-Policy Distillation
Aug 12, 2026Huipeng Huang, Hongxin WeiLanguage Model Distillation On-Policy Distillation
Department of Statistics and Data Science, Southern University of Science and Technology
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
The Gradient That Top-k Distillation Never Computes
Oct 11, 2026Rasul Khanbayov, Hasan KurbanLanguage Model Distillation Knowledge Distillation
College of Science and Engineering, Hamad Bin Khalifa University, Doha, Qatar

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

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

Huipeng Huang, Hongxin Wei

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

Language Model Distillation On-Policy Distillation Knowledge Distillation

Open paper AskPDF arXiv

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.

Abstract

Language Model Distillation

Explore similar work

Tail-Aware Top-kkk On-Policy Distillation

Unbiased Top-kkk Estimation for On-Policy Distillation

The Gradient That Top-k Distillation Never Computes

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

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