The Gradient That Top-k Distillation Never Computes

Rasul Khanbayov; Hasan Kurban

cs.LGOct 11, 2026

The Gradient That Top-k Distillation Never Computes

Authors: Rasul Khanbayov, Hasan Kurban

Organizations: College of Science and Engineering, Hamad Bin Khalifa University, Doha, Qatar

Abstract

Many deployed distillation pipelines show the student a truncated view of the teacher: keep the top-k probabilities, renormalize, and minimize a divergence on what remains. Renormalization is assumed to account for the missing mass. We prove that it removes the visibility of tail mass, not the freedom to move it. The renormalized objective depends only on log-odds among the retained tokens, so at each position its gradient along the directions that control tail mass is exactly zero; on a real teacher that component measures $2$ - $3\times10^{-15}$ of the gradient. The full-vocabulary gradient equals the truncated gradient plus a residual weighted by the tail masses, and training acts only on the first term. Without the cross-entropy term, a top-k student's tail mass drifts away from the teacher's while a full-vocabulary control moves slightly toward it; with it, the gap shrinks about fourfold but persists. On seven open-weight teachers, the retained set determines the teacher's distribution only from $k=d+1$ ( $d$ the hidden dimension): recovery is exact from there on all seven, and commonly cached sizes sit below it. In our off-policy runs the practical effect is real but smaller than the choice of teacher-student pair. Renormalization is a normalization, not a safeguard. The code is available at https://github.com/KurbanIntelligenceLab/top-k-distillation.

Figures & tables

Figure 1: The claim and its signature. (a) At one position the level sets of

\mathcal{L}_{S}

are parallel to

A^{\perp}

, where

A

is the span of the retained rows’ differences ( Theorem 1 ). So its gradient lies in

A

and is blind to the directions along which tail mass still varies ( Eq. 4 ); at

k=64

and

d^{\prime}=2048

A

has dimension at most

63

2048

. (b) With the cross-entropy term removed (Llama-3.1-8B

\to

Llama-3.2-1B,

k=256

, six seeds, seed spread below

10^{-4}

), the top-

k

student’s tail mass moves away from the teacher’s while the full-vocabulary control’s moves slightly toward it: mean

|m_{p}-m_{q}|

at steps

0,50,250,500

0.0280\to 0.0449\to 0.0544\to 0.0568

against

0.0280\to 0.0260\to 0.0258\to 0.0255

Figure 2: Twelve tokens, $k=4$ : lowering every retained logit of a student $p$ by $\ln 6$ gives $p^{\prime}$ with the same renormalized head as $p$ and the teacher ( $0.50,0.25,0.15,0.10$ ), so $\mathcal{L}_{S}(p)=\mathcal{L}_{S}(p^{\prime})=0$ . But $p^{\prime}$ ’s tail mass has tripled ( $0.20$ to $0.60$ ) and its full-vocabulary KL to the teacher is $0.335$ nats against $p$ ’s $0$ .

Figure 3: A top- $k$ cache crops the teacher’s distribution; a hidden-state cache compresses it. Top: keeping the top- $k$ of the teacher’s logits ( $1{,}536$ bytes at $k=256$ , against $303{,}872$ for all of them) discards the tail mass $m_{q}$ , and renormalizing discards the head mass $1-m_{q}$ , so the loss is computed on $S$ alone and the teacher’s entropy, $m_{q}$ , and top- $k$ sets beyond the cached $k$ cannot be recovered. Bottom: caching the final hidden state $g$ ( $8{,}192$ bytes) and reconstructing $z=Ug+b$ gives the exact full-vocabulary KL and recovers all three, at the cost of access to $U$ and $b$ and $2dV$ FLOP per token. Qwen3-8B, bf16.

Teacher	$d$	Collapse $k$	$\mathbf{1}$ -res. a	$k-(d{+}1)$
Qwen3-1.7B	$2\,048$	2049	$0.071$	$0$
Gemma-3-4B	$2\,560$	2561	$0.084$	$0$
Gemma-3-12B	$3\,840$	3841	$0.130$	$0$
Llama-3.1-8B	$4\,096$	4096	$0.104$	$-1$ b
Qwen3-8B	$4\,096$	4097	$0.067$	$0$
Qwen3-32B	$5\,120$	5121	$0.062$	$0$

Table 1: Threshold: smallest

k

at which the median

\KL(q\|\hat{q})

over

64

positions falls below

10^{-6}

. Theory predicts

k=d+1

. Reconstruction is deterministic. All seven unembeddings have full column rank. Green: collapse at

k=d+1

; for the soft-capped Gemma teachers smaller

k

cannot be solved in closed form and was not measured.

Study	$k$	Teacher $d$	$(d+1)/k$
Liu et al. (2026)	$10$ – $20$	$3584$ – $5120$	$179$ – $512$
Li et al. (2026)	$16$ – $64$	$2560$ – $3584$	$40$ – $224$
Shen et al. (2026)	$32$	$4096$	$128$
Peng et al. (2025)	$100$	$4096$	$41$
Ryskulov et al. (2026)	$100$	$4096$	$41$
Busbridge et al. (2025)	$128$	$1792$ – $4352$	$14$ – $34$

Table 2: Operating points of top-

k

distillation in the literature against the identifiability threshold

d+1

of the teacher used (teachers listed in Appendix C ). Every one sits below it.

Teacher $\to$ student	bytes	cache	top- $1024$	$d_{z}$	feature
Qwen3-1.7B $\to$ 0.6B	4096	18.615	18.336	$40$	$83{,}248$
Qwen3-4B $\to$ 1.7B	5120	13.885	13.747	$25.2$	$117{,}018$
Gemma-3-4B $\to$ 1B	5120	14.275	15.417	$-23.6$	$244{,}352$
Llama-3.2-3B $\to$ 1B	6144	11.680	11.718	$-6.8$	$55{,}517$
SmolLM2-1.7B $\to$ 360M	4096	12.558	12.569	$-2.4$	$684.7$

Table 3: Hidden-state cache against top-

1024

(

6144

bytes): held-out perplexity, six seeds. Every cache-versus-top-

1024

and cache-versus-feature-matching difference has

p=0.03125

;

d_{z}

is for cache minus top-

1024

(negative favors the cache). Green: the better arm in each row.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Teacher	$d$	$V$	Head (GFLOP)
Qwen3-1.7B	$2\,048$	$151\,936$	$0.62$
Gemma-3-4B	$2\,560$	$262\,144$	$1.34$
Gemma-3-12B	$3\,840$	$262\,144$	$2.01$
Llama-3.1-8B	$4\,096$	$128\,256$	$1.05$
Qwen3-8B	$4\,096$	$151\,936$	$1.24$
Qwen3-32B	$5\,120$	$151\,936$	$1.56$

Appendix

Table A1 : Head cost of reconstructing logits from

g

2dV

FLOP per token.

Teacher	$d$	$V$	bf16	int8 a	vs. full	vs. top-64
			Hidden-state cache (B)		Ratio ( $\times$ )
Qwen3-1.7B	$2\,048$	$151\,936$	$4\,096$	$2\,052$	74.2	$5.34$
Gemma-3-4B	$2\,560$	$262\,144$	$5\,120$	$2\,564$	102.4	$6.68$
Gemma-3-12B	$3\,840$	$262\,144$ b	$7\,680$	$3\,844$	68.3	$10.01$
Llama-3.1-8B	$4\,096$	$128\,256$	$8\,192$	$4\,100$	31.3	$10.68$
Qwen3-8B	$4\,096$	$151\,936$	$8\,192$	$4\,100$	37.1	$10.68$

Appendix

Table A2 : Hidden-state cache size per supervised position, and its ratio to the full logit vector and to a top-

64

cache. Green: how much smaller than the full logits the exact cache is.

$k$	Llama-3.1-8B ( $d=4096$ )	Llama-3.1-70B ( $d=8192$ )
$0.90d$	$9.98\times 10^{-5}$	$2.06\times 10^{-5}$
$d$	$1.39\times 10^{-7}$	$3.48\times 10^{-8}$
$d+1$	$2.72\times 10^{-8}$	$2.45\times 10^{-8}$
$d+2$	$-1.48\times 10^{-8}$	$5.75\times 10^{-9}$
$2d$	$-2.59\times 10^{-10}$	$-2.90\times 10^{-9}$

Appendix

Table A3 : Median

\KL(q\|\hat{q})

near the transition for the two teachers marked

-1

in Table 1 : already below

10^{-6}

k=d

but not zero, and floating-point noise from

k=d+1

on (negative entries are rounding). Green: exact recovery.

Teacher	Tail beyond top- $1024$		Tail beyond top- $64$	Entropy	Cache gain
	median	mean	mean	(nats)	(ppl)
Gemma-3-4B	0.0121	0.0201	0.151	3.59	$+1.142$
Llama-3.2-3B	0.0024	0.0097	0.064	2.28	$+0.038$
SmolLM2-1.7B	0.0020	0.0087	0.061	2.20	$+0.011$
Qwen3-1.7B	0.0006	0.0070	0.036	1.68	$-0.279$
Qwen3-4B	0.0004	0.0055	0.028	1.53	$-0.138$

Appendix

Table A4 : Teacher probability outside its own top-

k

, on the matched-storage training stream (

65{,}408

positions), with mean entropy. “Cache gain” is top-

1024

minus cache perplexity from Table 3 .

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
ReTaCo: Residual-Target Control for On-Policy Distillation
Sep 30, 2026Zixiang Ni, Zhuo Hu, Renjie Cao +8Language Model Distillation On-Policy Distillation
Xi’an Jiaotong University · Zhejiang University · Georgia Institute of Technology +6

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

ReTaCo: Residual-Target Control for On-Policy Distillation

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

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