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

Aug 12, 2026

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

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.

Huipeng Huang, Hongxin Wei

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

Language Model Distillation On-Policy Distillation Knowledge Distillation

Ask PDF

Sep 28, 2026

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

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.

Linjian Meng, Siyuan Gan, YuHan Li +5

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

Language Model Distillation On-Policy Distillation

Ask PDF

Sep 30, 2026

ReTaCo: Residual-Target Control for On-Policy Distillation

On-policy distillation (OPD) trains a student on its own generated prefixes with token-level teacher feedback, but transmitting or storing the teacher's full-vocabulary distribution at every token is costly. Entropy-aware OPD (EOPD) adds forward supervision to reverse KL to help the student recover plausible tokens it underestimates, using only the teacher's top-

k

probabilities to limit cost. Because EOPD renormalizes these probabilities, its target assigns no mass to the omitted vocabulary. We prove that the resulting loss keeps pushing the student's top-

k

mass toward one even after the student matches the teacher's relative probabilities within the top-

k

set, so the teacher itself is not a stationary point whenever the omitted tokens have positive teacher probability. We propose ReTaCo (Residual-Target Control), which keeps the top-

k

tokens individually and groups the remaining tokens into one residual symbol, and pairs this forward target with a single-sample estimator whose expectation equals the full-vocabulary reverse KL. With teacher top-

k

mass

m

, the residual target is

(1-β)(1-m)

for

β\in[0,1]

β=0

preserves the teacher's mass, and larger

β

moves more mass onto the top-

k

tokens without changing their relative probabilities. At a fixed prefix, we prove that the population objective has a unique optimum whose top-

k

mass lies between

m

and

m+β(1-m)

and increases monotonically with

β

; at

β=0

, underestimated top-

k

tokens still receive non-vanishing recovery gradients. Numerical optimization confirms these predictions, and across three teacher-student pairs, ReTaCo outperforms EOPD on most mathematics and code benchmarks.

Zixiang Ni, Zhuo Hu, Renjie Cao +8

Xi’an Jiaotong University · Zhejiang University · Georgia Institute of Technology +6

Language Model Distillation On-Policy Distillation Knowledge Distillation

Ask PDF

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