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

# Tail-Aware Top- k On-Policy Distillation
Source: [https://arxiv.org/html/2608.14728](https://arxiv.org/html/2608.14728)
Huipeng HuangAffiliation:Department of Statistics and Data Science, Southern University of Science and TechnologyHongxin WeiThanks:Corresponding author \(weihx@sustech\.edu\.cn\)Affiliation: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\-kktokens\. However, this normalized objective discards the information about tail probability: the total probability outside the teacher’s top\-kktokens\. 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\-kkOPD \(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\-kktokens 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\-kknormalization\. 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](https://github.com/HuipengHuang/TA-OPD)\.

## 1Introduction

Knowledge distillation\([15](https://arxiv.org/html/2608.14728#bib.bib2)\)is a promising approach for transferring the capabilities of large language models \(LLMs\) to smaller models\. However, standard distillation is performed off\-policy: the student is trained on teacher\-generated sequences\([17](https://arxiv.org/html/2608.14728#bib.bib13);[9](https://arxiv.org/html/2608.14728#bib.bib40)\), yet at inference time it must condition on its own generations\. This discrepancy between training and inference yields errors that accumulate quickly along the generated sequence\([3](https://arxiv.org/html/2608.14728#bib.bib38);[2](https://arxiv.org/html/2608.14728#bib.bib39)\)and hinders effective learning\([41](https://arxiv.org/html/2608.14728#bib.bib41)\)\. This motivates on\-policy distillation \(OPD\)\([1](https://arxiv.org/html/2608.14728#bib.bib4);[22](https://arxiv.org/html/2608.14728#bib.bib1);[40](https://arxiv.org/html/2608.14728#bib.bib5)\), where the student is trained to align its next\-token distribution with the teacher’s along its own trajectories\.

A popular idea, normalized top\-kkOPD\([7](https://arxiv.org/html/2608.14728#bib.bib3);[38](https://arxiv.org/html/2608.14728#bib.bib7)\), minimizes the reverse Kullback\-Leibler \(KL\) divergence between the student and teacher’s normalized distributions over the teacher’s top\-kktokens\. However, it discards the information about tail probability: the total probability mass outside the teacher’s top\-kktokens\. Theoretically, the optimization increases the student’s tail probability whenever the student’s distribution is not well aligned to the teacher’s, which commonly happens in distillation\. Empirically, as the student’s tail probability steadily increases, it more frequently samples tokens outside the teacher’s top\-kktokens, drifting toward prefixes where the teacher’s supervision is unreliable\. Consequently, the student fails to imitate the teacher’s next\-token distribution, degrading downstream performance\.

To address this, we propose Tail\-Aware Top\-kkOPD \(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\-kktokens plus a tail token that carries the tail probability\. In effect, TA\-OPD aligns the student’s tail probability with the teacher’s, preventing the tail probability and entropy increase caused by top\-kknormalization\. Theoretically, we show that TA\-OPD’s objective is a tight lower bound of the full\-vocabulary reverse KL divergence\. In addition, when the sampled token’s probability is available, we derive a sampled variant of TA\-OPD that debiases the TA\-OPD objective with the sampled token and yields an unbiased estimate of the full\-vocabulary reverse KL divergence\.

Extensive experiments on mathematical benchmarks demonstrate the superiority of TA\-OPD over baselines\. Notably, when distilling OpenThinker3\-7B\([10](https://arxiv.org/html/2608.14728#bib.bib37)\)into Qwen2\.5\-7B\-Instruct\([26](https://arxiv.org/html/2608.14728#bib.bib36)\), normalized top\-kkOPD’s training collapses: its training entropy increases to around 6, and its Avg@8 accuracy on MATH500\([14](https://arxiv.org/html/2608.14728#bib.bib10)\)drops to 68\.78%\. In contrast, TA\-OPD stabilizes distillation, keeps the training entropy below 1\.5, and improves Avg@8 accuracy on MATH500 to 77\.88%\. Through additional analyses, we show that TA\-OPD is most effective when the student\-teacher capability gap is large, and TA\-OPD can be applied with a smallkk\.

We summarize our contributions as follows:

- •We show that normalized top\-kkOPD discards the tail probability\. As a result, minimizing the normalized objective can steadily increase the student’s tail probability and entropy, empirically degrading downstream accuracy\.
- •We propose TA\-OPD, a novel distillation method that restores the tail probability signal\. We show that TA\-OPD addresses the increase in tail probability and entropy caused by top\-kknormalization\. When the sampled token’s probability is available, we derive a sampled variant of TA\-OPD, an unbiased estimate of the full\-vocabulary reverse KL divergence\.
- •We conduct extensive experiments to show that TA\-OPD achieves superior performance compared with normalized top\-kkOPD, improving Avg@8 accuracy by up to 8\.05 points on common benchmarks\.

## 2Preliminaries

On\-policy distillation \(OPD\) aims to distill knowledge from a teacher LLMπte\\pi\_\{\\text\{te\}\}to a student LLMπθ\\pi\_\{\\theta\}on trajectories sampled from the student\. Given a promptx∼𝒟xx\\sim\\mathcal\{D\}\_\{x\}, the student samples a responsey^=\(y^1,…,y^T\)∼πθ\(⋅∣x\)\\hat\{y\}=\(\\hat\{y\}\_\{1\},\\ldots,\\hat\{y\}\_\{T\}\)\\sim\\pi\_\{\\theta\}\(\\cdot\\mid x\)\. At each steptt, on the prefixy^<t=\(y^1,…,y^t−1\)\\hat\{y\}\_\{<t\}=\(\\hat\{y\}\_\{1\},\\ldots,\\hat\{y\}\_\{t\-1\}\), we define the student’s and teacher’s next\-token distributions over a vocabulary𝒱\\mathcal\{V\}aspt=πθ\(⋅∣x,y^<t\)p\_\{t\}=\\pi\_\{\\theta\}\(\\cdot\\mid x,\\hat\{y\}\_\{<t\}\)andqt=πte\(⋅∣x,y^<t\)q\_\{t\}=\\pi\_\{\\text\{te\}\}\(\\cdot\\mid x,\\hat\{y\}\_\{<t\}\)\. OPD minimizes reverse KL divergence over student\-sampled trajectories:

ℒOPD\(θ\)=𝔼x∼𝒟x,y^∼πθ\(⋅∣x\)\[∑t=1TDKL\(pt∥qt\)\],\\mathcal\{L\}\_\{\\mathrm\{OPD\}\}\(\\theta\)=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\_\{x\},\\;\\hat\{y\}\\sim\\pi\_\{\\theta\}\(\\cdot\\mid x\)\}\\left\[\\sum\_\{t=1\}^\{T\}D\_\{\\mathrm\{KL\}\}\(p\_\{t\}\\\|q\_\{t\}\)\\right\],\(1\)whereDKL\(pt∥qt\)=∑v∈𝒱pt\(v\)logpt​\(v\)qt​\(v\)D\_\{\\mathrm\{KL\}\}\(p\_\{t\}\\\|q\_\{t\}\)=\\sum\_\{v\\in\\mathcal\{V\}\}p\_\{t\}\(v\)\\log\\tfrac\{p\_\{t\}\(v\)\}\{q\_\{t\}\(v\)\}\. Different OPD methods vary in how they approximateDKL\(pt∥qt\)D\_\{\\mathrm\{KL\}\}\(p\_\{t\}\\\|q\_\{t\}\)with their per\-token lossesℓt\\ell\_\{t\}\. Given a student\-sampled tokeny^t∼pt\\hat\{y\}\_\{t\}\\sim p\_\{t\}, sampled\-token OPD\([22](https://arxiv.org/html/2608.14728#bib.bib1);[36](https://arxiv.org/html/2608.14728#bib.bib11)\)computesℓtsample=log⁡pt​\(y^t\)−log⁡qt​\(y^t\)\\ell\_\{t\}^\{\\mathrm\{sample\}\}=\\log p\_\{t\}\(\\hat\{y\}\_\{t\}\)\-\\log q\_\{t\}\(\\hat\{y\}\_\{t\}\), which is cheap to compute but discards the dense information over the remaining vocabulary\. Full\-vocabulary OPD\([39](https://arxiv.org/html/2608.14728#bib.bib14)\)computesℓtfull=DKL\(pt∥qt\)\\ell\_\{t\}^\{\\mathrm\{full\}\}=D\_\{\\mathrm\{KL\}\}\(p\_\{t\}\\\|q\_\{t\}\)exactly over𝒱\\mathcal\{V\}\. It provides dense supervision but is prohibitively expensive for LLMs\.

Normalized top\-kkOPD provides an intermediate design between sampled\-token and full\-vocabulary OPD by restricting the divergence computation to a subset\. Prior works adopt two choices of the subset: the student’s top\-kktokens\([21](https://arxiv.org/html/2608.14728#bib.bib9)\)or the teacher’s top\-kktokens\([16](https://arxiv.org/html/2608.14728#bib.bib8);[7](https://arxiv.org/html/2608.14728#bib.bib3)\)\. In this work, we primarily focus on the teacher top\-kkvariant\. Formally, we define the teacher’s top\-kktokens asStk=TopK⁡\(qt,k\)S\_\{t\}^\{k\}=\\operatorname\{TopK\}\(q\_\{t\},k\)\. The normalized distributions onStkS\_\{t\}^\{k\}are given by:

p¯t\(Stk\)​\(v\)=pt\(v\)1\[v∈Stk\]∑u∈Stkpt​\(u\),q¯t\(Stk\)​\(v\)=qt\(v\)1\[v∈Stk\]∑u∈Stkqt​\(u\),\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\(v\)=\\frac\{p\_\{t\}\(v\)\\,\\mathbf\{1\}\[v\\in S\_\{t\}^\{k\}\]\}\{\\sum\_\{u\\in S\_\{t\}^\{k\}\}p\_\{t\}\(u\)\},\\qquad\\bar\{q\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\(v\)=\\frac\{q\_\{t\}\(v\)\\,\\mathbf\{1\}\[v\\in S\_\{t\}^\{k\}\]\}\{\\sum\_\{u\\in S\_\{t\}^\{k\}\}q\_\{t\}\(u\)\},where𝟏​\[⋅\]\\mathbf\{1\}\[\\cdot\]denotes the indicator function\. Distillation is then performed by minimizing the subset reverse KL divergenceℓtnorm=DKL\(p¯t\(Stk\)∥q¯t\(Stk\)\)\\ell\_\{t\}^\{\\mathrm\{norm\}\}=D\_\{\\mathrm\{KL\}\}\\\!\\bigl\(\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\\,\\\|\\,\\bar\{q\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\\bigr\), yielding the trajectory\-level objective:

ℒOPDnorm\(θ\)=𝔼x∼𝒟x,y^∼πθ\(⋅∣x\)\[∑t=1TDKL\(p¯t\(Stk\)∥q¯t\(Stk\)\)\]\.\\mathcal\{L\}\_\{\\mathrm\{OPD\}\}^\{\\mathrm\{norm\}\}\(\\theta\)=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\_\{x\},\\;\\hat\{y\}\\sim\\pi\_\{\\theta\}\(\\cdot\\mid x\)\}\\left\[\\sum\_\{t=1\}^\{T\}D\_\{\\mathrm\{KL\}\}\\\!\\bigl\(\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\\,\\\|\\,\\bar\{q\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\\bigr\)\\right\]\.\(2\)
Normalized top\-kkOPD reduces the teacher\-query cost while retaining dense, multi\-token supervision overStkS\_\{t\}^\{k\}, making it a practical approximation to full\-vocabulary OPD\. However, the objective discards the tail probability, i\.e\., the total probability mass outside the teacher’s top\-kktokens:

pttail=∑v∉Stkpt​\(v\),qttail=∑v∉Stkqt​\(v\)\.p\_\{t\}^\{\\mathrm\{tail\}\}=\\sum\_\{v\\notin S\_\{t\}^\{k\}\}p\_\{t\}\(v\),\\qquad q\_\{t\}^\{\\mathrm\{tail\}\}=\\sum\_\{v\\notin S\_\{t\}^\{k\}\}q\_\{t\}\(v\)\.\(3\)Concerningly, the objective can be minimized even when the student’s tail probability is much higher than the teacher’s, making the student’s next\-token distribution diverge substantially from the teacher’s\. We proceed by analyzing how the method affects the student’s tail probability\.

## 3Motivation

In this section, we investigate how the normalized objective in Eq\. \([2](https://arxiv.org/html/2608.14728#S2.E2)\) changes the student’s tail probability\. We find that minimizing the objective can steadily increase the student’s tail probability\.

### 3\.1Theoretical Analysis on the Tail Probability

Fix a stepttand letzt,vz\_\{t,v\}denote the student logit for tokenvvunder the prefix, so thatpt​\(v\)=ezt,v∑u∈𝒱ezt,up\_\{t\}\(v\)=\\frac\{e^\{z\_\{t,v\}\}\}\{\\sum\_\{u\\in\\mathcal\{V\}\}e^\{z\_\{t,u\}\}\}\. We have the following propositions for the optimization of the normalized objective\.

###### Proposition 1\.

For every steptt, the gradient ofℓtnorm\\ell\_\{t\}^\{\\mathrm\{norm\}\}with respect to the student logitzt,vz\_\{t,v\}is

∂ℓt∂zt,v=\{p¯t\(Stk\)​\(v\)​\[log⁡p¯t\(Stk\)​\(v\)q¯t\(Stk\)​\(v\)−ℓt\],v∈Stk,0,v∉Stk\.\\frac\{\\partial\\ell\_\{t\}\}\{\\partial z\_\{t,v\}\}=\\begin\{cases\}\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\(v\)\\\!\\left\[\\log\\dfrac\{\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\(v\)\}\{\\bar\{q\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\(v\)\}\-\\ell\_\{t\}\\right\],&v\\in S\_\{t\}^\{k\},\\\\\[10\.00002pt\] 0,&v\\notin S\_\{t\}^\{k\}\.\\end\{cases\}Moreover, the gradients on the logits of the top\-kktokens sum to zero:∑v∈Stk∂ℓt∂zt,v=0\\sum\_\{v\\in S\_\{t\}^\{k\}\}\\frac\{\\partial\\ell\_\{t\}\}\{\\partial z\_\{t,v\}\}=0\.

The proof is provided in Appendix[A\.1](https://arxiv.org/html/2608.14728#A1.SS1)\. The proposition shows that the normalized objective provides no explicit mechanism for decreasing the student’s tail probability\. Decreasing the tail probability requires raising the top\-kklogits relative to the tail logits\. However, the normalized objective can do neither: the tail logits receive zero gradient, and the gradients on the top\-kklogits sum to zero\. In the following, we further explore how the zero\-sum gradient within the top\-kktokens affects the student’s tail probability\.

###### Proposition 2\.

Assume that the studentπθ\\pi\_\{\\theta\}is a tabular softmax policy, where each tokenvvat each stepttis associated with an independent logit parameterzt,v=θt,vz\_\{t,v\}=\\theta\_\{t,v\}\. Writep¯:=p¯t\(Stk\)\\bar\{p\}:=\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}andq¯:=q¯t\(Stk\)\\bar\{q\}:=\\bar\{q\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}for the normalized distributions onStkS\_\{t\}^\{k\}, and letCovv∼r\(f,g\)\\operatorname\*\{Cov\}\_\{v\\sim r\}\(f,g\)denote the covariance off⁡\(v\)f\(v\)andg⁡\(v\)g\(v\)underv∼rv\\sim r\. After one gradient descent stepθ←θ−η​∂ℓtnorm∂θ\\theta\\leftarrow\\theta\-\\eta\\,\\frac\{\\partial\\ell\_\{t\}^\{\\mathrm\{norm\}\}\}\{\\partial\\theta\}with learning rateη\\eta, the student’s tail probability changeΔ​pttail\\Delta p\_\{t\}^\{\\mathrm\{tail\}\}is given by:

Δ​pttail=η​pttail​\(1−pttail\)​\[Covv∼p¯\(p¯​\(v\),log⁡p¯​\(v\)\)⏟student self\-covariance−Covv∼p¯\(p¯​\(v\),log⁡q¯​\(v\)\)⏟student\-teacher covariance\]\+O⁡\(η2\)\.\\Delta p\_\{t\}^\{\\mathrm\{tail\}\}=\\eta\\,p\_\{t\}^\{\\mathrm\{tail\}\}\\bigl\(1\-p\_\{t\}^\{\\mathrm\{tail\}\}\\bigr\)\\Bigl\[\\underbrace\{\\operatorname\*\{Cov\}\_\{v\\sim\\bar\{p\}\}\\bigl\(\\bar\{p\}\(v\),\\log\\bar\{p\}\(v\)\\bigr\)\}\_\{\\text\{ student self\-covariance\}\}\-\\underbrace\{\\operatorname\*\{Cov\}\_\{v\\sim\\bar\{p\}\}\\bigl\(\\bar\{p\}\(v\),\\log\\bar\{q\}\(v\)\\bigr\)\}\_\{\\text\{student\-teacher covariance\}\}\\Bigr\]\+O\(\\eta^\{2\}\)\.Consequently, the student’s tail probability increases under first\-order approximation if and only if the student self\-covariance exceeds the student–teacher covariance\.

The proof is provided in Appendix[A\.2](https://arxiv.org/html/2608.14728#A1.SS2)\. The proposition111We use the proposition as a motivation rather than a formal theoretical guarantee\.shows that minimizing the normalized objective strictly increases the student’s tail probability if the student self\-covariance exceeds the student–teacher covariance\. The student self\-covariance is nonnegative, sincelog⁡p¯\\log\\bar\{p\}is monotone withp¯\\bar\{p\}\. In contrast, the student\-teacher covariance is positive only when the student concentrates its probability on the teacher’s high\-probability tokens withinStkS\_\{t\}^\{k\}, and can be negative when the two distributions are poorly aligned\.

In practice, when the student\-teacher capability gap is large, the student cannot imitate the teacher’s next\-token distribution well, and its probability mass cannot concentrate on the teacher’s high\-probability tokens\. As a result, the student–teacher covariance tends to be smaller than the student self\-covariance, making the normalized objective exhibit a systematic bias toward increasing the student’s tail probability\. This analysis motivates us to empirically examine the existence of the tail probability increase and its influence on the training dynamics of OPD\.

### 3\.2Empirical Study on the Tail Probability Increase

#### Setup\.

We conduct experiments with three student\-teacher model pairs: Qwen3\-1\.7B\([43](https://arxiv.org/html/2608.14728#bib.bib15)\)with Qwen3\-30B\-A3B\-Instruct\-2507, Qwen2\.5\-7B\-Instruct with OpenThinker3\-7B, and Llama\-3\.1\-8B\([8](https://arxiv.org/html/2608.14728#bib.bib42)\)with DeepSeek\-R1\-Distill\-Llama\-8B\([11](https://arxiv.org/html/2608.14728#bib.bib6)\)\. All models are trained on DAPO\-MATH\-17K\([44](https://arxiv.org/html/2608.14728#bib.bib16)\)withk=16k=16\.

Figure 1:The student’s tail probability, the teacher’s tail probability, and the student’s token\-level entropy across training steps\.Normalized top\-kkOPD consistently increases the student’s tail probability and entropy across all three model pairs\.
#### Normalized top\-kkOPD steadily increases the student’s tail probability\.

Figure[1](https://arxiv.org/html/2608.14728#S3.F1)visualizes the students’ and teachers’ tail probabilities on student\-generated prefixes, together with the students’ entropy across training steps\. The results show that across three model pairs, the student’s tail probability and entropy steadily increase during on\-policy training, validating our theoretical analysis\. Furthermore, as the student’s tail probability increases, it more frequently samples tokens outside the teacher’s top\-kktokens\. This drives the student toward reaching prefixes where the teacher is uncertain, as reflected by the increase in the teacher’s tail probability\. Prior works show that the teacher’s supervision is unreliable on such uncertain prefixes\([7](https://arxiv.org/html/2608.14728#bib.bib3);[37](https://arxiv.org/html/2608.14728#bib.bib34)\), as the teacher itself exhibits significantly reduced accuracy given these prefixes\. Consequently, normalized top\-kkOPD progressively shifts probability mass away from the teacher’s top\-kktokens, making the student fail to imitate the teacher’s next\-token distribution\.

Figure 2:Validation accuracy versus student’s tail probability\.
#### Tail probability increase hurts downstream performance\.

Figure[2](https://arxiv.org/html/2608.14728#S3.F2)plots validation accuracy against the student’s tail probability\. Each point is a checkpoint from a normalized top\-kkOPD run with Qwen2\.5\-7B\-Instruct as the student and OpenThinker3\-7B as the teacher\. The validation set is MATH\-500, and the validation sampling temperature is11\. The validation accuracy degrades significantly from above0\.70\.7to below0\.20\.2as the tail probability grows from near00to around0\.70\.7, with a Pearson correlation coefficient of−0\.91\-0\.91\.

In Appendix[B](https://arxiv.org/html/2608.14728#A2), we conduct a detailed analysis on the tail probability increase phenomenon\. We find that it occurs when there is a large capability gap between the student and teacher and whenkkis not sufficiently large \(e\.g\.,k≤64k\\leq 64\)\. Moreover, the phenomenon is more pronounced when the maximum training response length becomes longer\.

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\-kkOPD\):The normalized objective discards the tail probability, leading to increased tail probability\.\(c\) Decreased Tail Probability \(Tail\-Aware Top\-kkOPD, Ours\):TA\-OPD restores the missing tail probability signal, matching the student’s tail probability to the teacher’s\.

## 4Method

In our previous analysis, we show that the normalized objective discards information about the tail probability, leading to the tail probability increase issue \(Figure[3](https://arxiv.org/html/2608.14728#S3.F3)\(b\)\)\. To address this problem, our key idea is to restore the missing tail probability signal\.

#### Tail\-Aware Top\-kkOPD\.

We propose Tail\-Aware Top\-kkOPD \(TA\-OPD\), which explicitly aligns the student’s tail probability with the teacher’s \(Figure[3](https://arxiv.org/html/2608.14728#S3.F3)\(c\)\)\. In particular, we introduce a tail tokenvtailv\_\{\\mathrm\{tail\}\}that carries the tail probability\. We then minimize the reverse KL divergence over the top\-kktokens plus the tail token\. Formally, we define the augmented token setSt\+=Stk∪\{vtail\}S\_\{t\}^\{\+\}=S\_\{t\}^\{k\}\\cup\\\{v\_\{\\mathrm\{tail\}\}\\\}\. OnSt\+S\_\{t\}^\{\+\}, each top\-kktoken retains its original probability,pt​\(v\)=πθ​\(v∣x,y^<t\)p\_\{t\}\(v\)=\\pi\_\{\\theta\}\(v\\mid x,\\hat\{y\}\_\{<t\}\), while the tail token carries the tail probability,pt​\(vtail\)=pttailp\_\{t\}\(v\_\{\\mathrm\{tail\}\}\)=p\_\{t\}^\{\\mathrm\{tail\}\}\. The per\-token loss function of TA\-OPD is then given by:

ℓtTA=∑v∈St\+pt​\(v\)​log⁡pt​\(v\)qt​\(v\)=∑v∈Stkpt​\(v\)​log⁡pt​\(v\)qt​\(v\)\+pttail​log⁡pttailqttail\.\\ell\_\{t\}^\{\\mathrm\{TA\}\}=\\sum\_\{v\\in S\_\{t\}^\{\+\}\}p\_\{t\}\(v\)\\log\\frac\{p\_\{t\}\(v\)\}\{q\_\{t\}\(v\)\}=\\sum\_\{v\\in S\_\{t\}^\{k\}\}p\_\{t\}\(v\)\\log\\frac\{p\_\{t\}\(v\)\}\{q\_\{t\}\(v\)\}\+p\_\{t\}^\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{t\}^\{\\mathrm\{tail\}\}\}\{q\_\{t\}^\{\\mathrm\{tail\}\}\}\.\(4\)By introducing the tail tokenvtailv\_\{\\mathrm\{tail\}\}, the termpttail​log⁡pttailqttailp\_\{t\}^\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{t\}^\{\\mathrm\{tail\}\}\}\{q\_\{t\}^\{\\mathrm\{tail\}\}\}explicitly compares the student’s and teacher’s tail probabilities, thereby driving the student’s tail probability toward the teacher’s\. We next characterize the gradient signal ofℓtTA\\ell\_\{t\}^\{\\mathrm\{TA\}\}\.

###### Proposition 3\.

For every steptt, the gradient ofℓtTA\\ell\_\{t\}^\{\\mathrm\{TA\}\}with respect to the student logitzt,vz\_\{t,v\}is

∂ℓtTA∂zt,v=\{pt​\(v\)​\(log⁡pt​\(v\)qt​\(v\)−ℓtTA\),v∈Stk,pt​\(v\)​\(log⁡pttailqttail−ℓtTA\),v∉Stk\.\\frac\{\\partial\\ell\_\{t\}^\{\\mathrm\{TA\}\}\}\{\\partial z\_\{t,v\}\}=\\begin\{cases\}p\_\{t\}\(v\)\\Big\(\\log\\dfrac\{p\_\{t\}\(v\)\}\{q\_\{t\}\(v\)\}\-\\ell\_\{t\}^\{\\mathrm\{TA\}\}\\Big\),&v\\in S\_\{t\}^\{k\},\\\\\[11\.99998pt\] p\_\{t\}\(v\)\\Big\(\\log\\dfrac\{p\_\{t\}^\{\\mathrm\{tail\}\}\}\{q\_\{t\}^\{\\mathrm\{tail\}\}\}\-\\ell\_\{t\}^\{\\mathrm\{TA\}\}\\Big\),&v\\notin S\_\{t\}^\{k\}\.\\end\{cases\}\(5\)

The proposition reveals two key properties of TA\-OPD\. First, while the normalized objective gives no gradient to logits outside the top\-kktokens, TA\-OPD updates them through the tail log\-ratiolog⁡pttailqttail\\log\\tfrac\{p\_\{t\}^\{\\mathrm\{tail\}\}\}\{q\_\{t\}^\{\\mathrm\{tail\}\}\}, which compares the student’s tail probability with the teacher’s\. Second, the gradients on the top\-kklogits no longer sum to zero\. The top\-kkgradient sum∑v∈Stk∂ℓtTA∂zt,v=pttail​\(ℓtTA−log⁡pttailqttail\)\\sum\_\{v\\in S\_\{t\}^\{k\}\}\\frac\{\\partial\\ell\_\{t\}^\{\\mathrm\{TA\}\}\}\{\\partial z\_\{t,v\}\}=p\_\{t\}^\{\\mathrm\{tail\}\}\\Bigl\(\\ell\_\{t\}^\{\\mathrm\{TA\}\}\-\\log\\tfrac\{p\_\{t\}^\{\\mathrm\{tail\}\}\}\{q\_\{t\}^\{\\mathrm\{tail\}\}\}\\Bigr\)is negative whenever the tail log\-ratio is larger than the average log\-ratio overStkS\_\{t\}^\{k\}\. It therefore can explicitly raise the top\-kklogits relative to the tail whenever the student over\-weights the tail relative to the teacher\. In the following, we show thatℓtTA\\ell\_\{t\}^\{\\mathrm\{TA\}\}is not merely a heuristic modification but a principled approximation of the full\-vocabulary objective\.

###### Proposition 4\(Lower bound of the full\-vocabulary reverse KL\)\.

Letp~t​\(v\)=pt​\(v\)/pttail\\tilde\{p\}\_\{t\}\(v\)=p\_\{t\}\(v\)/p\_\{t\}^\{\\mathrm\{tail\}\}andq~t​\(v\)=qt​\(v\)/qttail\\tilde\{q\}\_\{t\}\(v\)=q\_\{t\}\(v\)/q\_\{t\}^\{\\mathrm\{tail\}\}forv∉Stkv\\notin S\_\{t\}^\{k\}be the normalized tail distributions\. Then

ℓtTA=ℓtfull−pttailDKL\(p~t∥q~t\)≤ℓtfull\.\\ell\_\{t\}^\{\\mathrm\{TA\}\}=\\ell\_\{t\}^\{\\mathrm\{full\}\}\-p\_\{t\}^\{\\mathrm\{tail\}\}\\,D\_\{\\mathrm\{KL\}\}\(\\tilde\{p\}\_\{t\}\\\|\\tilde\{q\}\_\{t\}\)\\leq\\ell\_\{t\}^\{\\mathrm\{full\}\}\.\(6\)

The proofs of the above propositions are presented in Appendix[A\.3](https://arxiv.org/html/2608.14728#A1.SS3)and[A\.4](https://arxiv.org/html/2608.14728#A1.SS4)\. This proposition shows that the loss functions of TA\-OPD and full\-vocabulary OPD differ only by a non\-negative residual termpttailDKL\(p~t∥q~t\)p\_\{t\}^\{\\mathrm\{tail\}\}\\,D\_\{\\mathrm\{KL\}\}\(\\tilde\{p\}\_\{t\}\\\|\\tilde\{q\}\_\{t\}\)\. In practice, TA\-OPD drives both the student’s and teacher’s tail probabilities close to zero \(see Figure[4](https://arxiv.org/html/2608.14728#S5.F4)\(a\) and \(b\)\)\. As a result, the residual term becomes small, makingℓtTA\\ell\_\{t\}^\{\\mathrm\{TA\}\}a tight lower bound to the ideal full\-vocabulary reverse KL objective\.

## 5Experiments

### 5\.1Experimental Setup

#### Models and Training Dataset\.

For main experiments, we use the three student\-teacher model pairs as in section[3\.2](https://arxiv.org/html/2608.14728#S3.SS2): Qwen3\-1\.7B with Qwen3\-30B\-A3B\-Instruct\-2507, Qwen2\.5\-7B\-Instruct with OpenThinker3\-7B, and Llama\-3\.1\-8B with DeepSeek\-R1\-Distill\-Llama\-8B\. We disable thinking modes for Qwen3 models\. The training set is DAPO\-MATH\-17K\.

#### Evaluation\.

We evaluate on six math reasoning benchmarks—MATH500, Minerva\([19](https://arxiv.org/html/2608.14728#bib.bib20)\), OlympiadBench\([13](https://arxiv.org/html/2608.14728#bib.bib21)\), AMC\([20](https://arxiv.org/html/2608.14728#bib.bib44)\)and AIME24/25\([20](https://arxiv.org/html/2608.14728#bib.bib44)\)—and two out\-of\-distribution benchmarks, ARC\-c\([4](https://arxiv.org/html/2608.14728#bib.bib43)\)and MMLU\-Pro\([35](https://arxiv.org/html/2608.14728#bib.bib47)\)\. We use a rollout temperature of 0\.7, top\-p sampling with p = 0\.95, and a maximum response length of 8192 tokens\. For MMLU\-Pro, we report Pass@1\. For other datasets, we sample 8 responses per question and report the average accuracy \(Avg@8\)\.

#### Compared Methods\.

We compare our method with sampled\-token OPD and normalized top\-kkOPD\. In Appendix[E\.4](https://arxiv.org/html/2608.14728#A5.SS4), we conduct an additional comparison with unnormalized top\-kkOPD\. We provide a detailed introduction to these methods in Appendix[C](https://arxiv.org/html/2608.14728#A3)\.

#### Implementation\.

All the OPD experiments use the same training setup: 300 training steps with a learning rate of1×10−61\\times 10^\{\-6\}\. Unless otherwise specified, we setk=16k=16for top\-kkOPD\. We use a prompt batch size of 72 and sample 4 rollouts per prompt, with a maximum generation length of 7168 tokens\. More details of implementation are provided in Appendix[D](https://arxiv.org/html/2608.14728#A4)\.

### 5\.2Results

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\-kkOPD \(Norm\. top\-kk\)\. EachAvg\.column represents the macro\-average\. Best results are shown inbold\.MethodsIn\-Distribution PerformanceOOD PerformanceMATH500MinervaOlympiadAMCAIME 24AIME 25Avg\.ARC\-cMMLU\-ProAvg\.Student:Qwen2\.5\-7B\-InstructTeacher:OpenThinker3\-7BSampled\-token74\.6031\.2541\.2245\.4814\.5816\.2537\.2374\.9947\.6661\.33Norm\. top\-kk68\.7825\.5536\.0241\.4213\.7515\.0033\.4223\.0737\.3830\.23TA\-OPD \(Ours\)77\.8832\.5842\.4145\.0316\.6717\.5038\.6876\.3050\.3563\.33Student:Llama\-3\.1\-8BTeacher:DeepSeek\-R1\-Distill\-Llama\-8BSampled\-token43\.4512\.7816\.5617\.622\.081\.2515\.6241\.8828\.9935\.44Norm\. top\-kk31\.309\.249\.9611\.451\.250\.4210\.6010\.0118\.3314\.17TA\-OPD \(Ours\)50\.8314\.9420\.6520\.482\.502\.5018\.6565\.7034\.5750\.14Student:Qwen3\-1\.7BTeacher:Qwen3\-30B\-A3B\-Instruct\-2507Sampled\-token82\.9849\.4934\.9352\.4126\.2519\.5844\.2783\.9150\.0867\.00Norm\. top\-kk81\.8047\.2033\.2850\.6023\.3317\.0842\.2282\.4849\.6666\.07TA\-OPD \(Ours\)83\.4849\.8234\.5052\.8626\.2520\.8344\.6283\.5950\.4267\.01

#### TA\-OPD achieves the best performance across student–teacher model pairs\.

Table[1](https://arxiv.org/html/2608.14728#S5.T1)compares TA\-OPD with sampled\-token OPD and normalized top\-kkOPD\. Evaluated on six math reasoning benchmarks, TA\-OPD attains the best average accuracy on every student–teacher pair, outperforming normalized top\-kkOPD by\+5\.26points on Qwen2\.5\-7B\-Instruct and\+8\.05points on Llama\-3\.1\-8B\. Notably, TA\-OPD demonstrates a significantly greater advantage on MATH500 with Llama\-3\.1\-8B over normalized top\-kkOPD, improving Avg@8 from 31\.30 to 50\.83 \(\+19\.53points\)\. Regarding out\-of\-distribution performance, TA\-OPD also demonstrates strong performance gain: on ARC\-c, it achieves 76\.30 with Qwen2\.5\-7B\-Instruct and 65\.70 with Llama\-3\.1\-8B, whereas normalized top\-kkOPD degrades to 23\.07 and 10\.01, indicating that the tail probability increase not only hurts mathematical reasoning but also degrades the student’s general capabilities\. Overall, these results demonstrate that TA\-OPD consistently achieves the best in\-distribution and out\-of\-distribution performance across different model pairs\.

Figure 4:Tail probabilities and token\-level entropy of TA\-OPD and normalized top\-kkOPD over training\.The student and teacher models are Qwen2\.5\-7B\-Instruct and OpenThinker3\-7B\.
#### TA\-OPD prevents the increase in tail probability and entropy\.

In Figure[4](https://arxiv.org/html/2608.14728#S5.F4), we compare the training dynamics of TA\-OPD and normalized top\-kkOPD on the Qwen2\.5\-7B\-Instruct and OpenThinker3\-7B model pair\. With normalized top\-kkOPD, the student’s tail probability, the teacher’s tail probability, and the student’s entropy increase to approximately0\.60\.6,0\.50\.5, and66over the course of training\. In contrast, TA\-OPD keeps the two tail probabilities close to00and the student’s entropy below1\.51\.5\. The same trends are consistently observed for three alternative model pairs in Figure[17](https://arxiv.org/html/2608.14728#A6.F17)\. In short, these results show that TA\-OPD addresses the tail probability and entropy increase caused by top\-kknormalization\.

Figure 5:Ablation study onkk\.
#### TA\-OPD is effective with a smallkk\.

In Figure[5](https://arxiv.org/html/2608.14728#S5.F5), we ablate how the number of top\-kktokens affects the performance of TA\-OPD\. In particular, we fix Qwen3\-1\.7B\-Base as the student and Qwen3\-8B as the teacher, and only varyk∈\{1,2,4,8,16,32\}k\\in\\\{1,2,4,8,16,32\\\}\. For eachkk, we report the average accuracy macro\-averaged across the six math reasoning benchmarks\. The results show that the accuracy is insensitive tokkoncek≥2k\\geq 2:k=2k=2attains an average accuracy of30\.4330\.43, only0\.430\.43points below the best result of30\.8630\.86atk=32k=32\. Overall, TA\-OPD can be applied with a smallkk\.

#### Additional results\.

Due to space constraints, we defer additional analyses to the appendix\. Appendix[E\.1](https://arxiv.org/html/2608.14728#A5.SS1)conducts a synthetic experiment to provide an intuitive understanding of TA\-OPD and normalized top\-kkOPD’s optimization\. Appendix[E\.2](https://arxiv.org/html/2608.14728#A5.SS2)shows that TA\-OPD consistently achieves higher student\-teacher top\-kkoverlap ratios compared to normalized top\-kkOPD\. Appendix[E\.3](https://arxiv.org/html/2608.14728#A5.SS3)shows that TA\-OPD incurs negligible computational overhead over normalized top\-kkOPD\.

## 6Discussion

Table 2:Comparison between TA\-OPD and sample\-corrected TA\-OPD on math reasoning and out\-of\-distribution \(OOD\) benchmarks\.Best results are shown inbold\.MethodsIn\-Distribution PerformanceOOD PerformanceMATH500MinervaOlympiadAMCAIME 24AIME 25Avg\.ARC\-cMMLU\-ProAvg\.Student:DeepSeek\-R1\-Distill\-Qwen\-1\.5BTeacher:JustRL\-DeepSeek\-1\.5BTA\-OPD86\.1330\.9757\.0476\.2041\.6729\.7353\.6235\.4025\.9830\.69SC\-TA\-OPD87\.0032\.0857\.6576\.5142\.9230\.1754\.3935\.5525\.7630\.66Student:Qwen3\-1\.7BTeacher:Qwen3\-30B\-A3B\-Instruct\-2507TA\-OPD83\.4849\.8234\.5052\.8626\.2520\.8344\.6283\.5950\.4267\.01SC\-TA\-OPD84\.0049\.5434\.9353\.1626\.6719\.5844\.6583\.9150\.2867\.10

#### Sample\-corrected TA\-OPD\.

While TA\-OPD’s loss function is a tight lower bound of the full\-vocabulary reverse KL divergence, it remains a biased estimate of the latter\. By Proposition[4](https://arxiv.org/html/2608.14728#Thmproposition4), the bias ispttailDKL\(p~t∥q~t\)p\_\{t\}^\{\\mathrm\{tail\}\}\\,D\_\{\\mathrm\{KL\}\}\(\\tilde\{p\}\_\{t\}\\,\\\|\\,\\tilde\{q\}\_\{t\}\)\. When the sampled token’s probability is additionally available, we can estimate the bias with the sampled tokeny^t\\hat\{y\}\_\{t\}, and add this estimate back toℓtTA\\ell\_\{t\}^\{\\mathrm\{TA\}\}to yield an unbiased estimate of the full\-vocabulary reverse KL divergence\. Formally, the loss function of sample\-corrected TA\-OPD is given by:

ℓtSC​\-​TA=\\displaystyle\\ell\_\{t\}^\{\\mathrm\{SC\\text\{\-\}TA\}\}=∑v∈St\+pt\(v\)sg\(logpt​\(v\)qt​\(v\)\)\+𝟏\[y^t∉Stk\]pt​\(y^t\)sg⁡\(pt​\(y^t\)\)sg\(logpt​\(y^t\)/pttailqt​\(y^t\)/qttail\),\\displaystyle\\sum\_\{v\\in S\_\{t\}^\{\+\}\}p\_\{t\}\(v\)\\operatorname\{sg\}\\\!\\left\(\\log\\frac\{p\_\{t\}\(v\)\}\{q\_\{t\}\(v\)\}\\right\)\+\\mathbf\{1\}\[\\hat\{y\}\_\{t\}\\notin S\_\{t\}^\{k\}\]\\,\\frac\{p\_\{t\}\(\\hat\{y\}\_\{t\}\)\}\{\\operatorname\{sg\}\(p\_\{t\}\(\\hat\{y\}\_\{t\}\)\)\}\\operatorname\{sg\}\\\!\\left\(\\log\\frac\{p\_\{t\}\(\\hat\{y\}\_\{t\}\)/p\_\{t\}^\{\\mathrm\{tail\}\}\}\{q\_\{t\}\(\\hat\{y\}\_\{t\}\)/q\_\{t\}^\{\\mathrm\{tail\}\}\}\\right\),\(7\)wheresg⁡\(⋅\)\\operatorname\{sg\}\(\\cdot\)denotes the stop\-gradient operator\. We next characterize its theoretical property\.

###### Proposition 5\(Unbiasedness\)\.

For every steptt, taking the randomness overy^t∼pt\\hat\{y\}\_\{t\}\\sim p\_\{t\}, we have

𝔼y^t∼pt\[ℓtSC​\-​TA\]=DKL\(pt∥qt\),𝔼y^t∼pt\[∇θℓtSC​\-​TA\]=∇θDKL\(pt∥qt\)\.\\mathbb\{E\}\_\{\\hat\{y\}\_\{t\}\\sim p\_\{t\}\}\\bigl\[\\ell\_\{t\}^\{\\mathrm\{SC\\text\{\-\}TA\}\}\\bigr\]=D\_\{\\mathrm\{KL\}\}\(p\_\{t\}\\,\\\|\\,q\_\{t\}\),\\quad\\mathbb\{E\}\_\{\\hat\{y\}\_\{t\}\\sim p\_\{t\}\}\\bigl\[\\nabla\_\{\\theta\}\\,\\ell\_\{t\}^\{\\mathrm\{SC\\text\{\-\}TA\}\}\\bigr\]=\\nabla\_\{\\theta\}\\,D\_\{\\mathrm\{KL\}\}\(p\_\{t\}\\,\\\|\\,q\_\{t\}\)\.

The proof is provided in Appendix[A\.5](https://arxiv.org/html/2608.14728#A1.SS5)\. Proposition[5](https://arxiv.org/html/2608.14728#Thmproposition5)shows thatℓtSC​\-​TA\\ell\_\{t\}^\{\\mathrm\{SC\\text\{\-\}TA\}\}is an unbiased estimate of the full\-vocabulary reverse KL in both value and gradient, thereby combining the advantages of sampled\-token OPD’s unbiasedness and top\-kkOPD’s dense supervision\.

Table[2](https://arxiv.org/html/2608.14728#S6.T2)compares TA\-OPD with its sample\-corrected variant on two student–teacher pairs\. Removing the bias term yields only marginal gains: the macro\-average on mathematical reasoning improves from53\.6253\.62to54\.3954\.39for DeepSeek\-R1\-Distill\-Qwen\-1\.5B and from44\.6244\.62to44\.6544\.65for Qwen3\-1\.7B, while the OOD averages remain essentially unchanged in both settings \(30\.6930\.69vs\.30\.6630\.66and67\.0167\.01vs\.67\.1067\.10\)\. We attribute this to the tightness of TA\-OPD: its bias to full\-vocabulary OPDpttailDKL\(p~t∥q~t\)p\_\{t\}^\{\\mathrm\{tail\}\}\\,D\_\{\\mathrm\{KL\}\}\(\\tilde\{p\}\_\{t\}\\,\\\|\\,\\tilde\{q\}\_\{t\}\)is negligible in practice\. Overall, SC\-TA\-OPD provides rigorous theoretical guarantees of unbiasedness, while TA\-OPD attains comparable performance to SC\-TA\-OPD\.

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\.
#### TA\-OPD is most effective when the student–teacher capability gap is large\.

We fix JustRL\-DeepSeek\-1\.5B\([12](https://arxiv.org/html/2608.14728#bib.bib48)\)as the teacher and compare two students, Qwen2\.5\-Math\-1\.5B and DeepSeek\-R1\-Distill\-Qwen\-1\.5B\. DeepSeek\-R1\-Distill\-Qwen\-1\.5B is obtained by applying supervised fine\-tuning to Qwen2\.5\-Math\-1\.5B, and the teacher is obtained by applying reinforcement learning to DeepSeek\-R1\-Distill\-Qwen\-1\.5B\. The capability gap between Qwen2\.5\-Math\-1\.5B and the teacher is therefore larger than that between DeepSeek\-R1\-Distill\-Qwen\-1\.5B and the teacher\.

Figure[6](https://arxiv.org/html/2608.14728#S6.F6)presents the Avg@8 accuracy of four OPD objectives on MATH500 for the two students\. On Qwen2\.5\-Math\-1\.5B, TA\-OPD attains an Avg@8 of 75\.20 on MATH500, outperforming normalized top\-kkOPD by13\.50points\. On DeepSeek\-R1\-Distill\-Qwen\-1\.5B, the four objectives achieve similar performance, with sample\-corrected TA\-OPD attaining the best accuracy of 87\.00\. This is because, as shown in Figure[18](https://arxiv.org/html/2608.14728#A6.F18), the tail probability increase is more pronounced under the larger capability gap, where TA\-OPD therefore brings a larger improvement\. In summary, TA\-OPD brings substantial improvements over baselines when the student–teacher capability gap is large\.

## 7Related Work

#### On\-policy Distillation\.

On\-policy distillation \(OPD\) is an effective post\-training paradigm that has attracted a surge of interest\([33](https://arxiv.org/html/2608.14728#bib.bib31)\)\. To provide dense supervision at tractable cost, top\-kkOPD has become a popular research direction\([47](https://arxiv.org/html/2608.14728#bib.bib12);[38](https://arxiv.org/html/2608.14728#bib.bib7)\)\. Normalized top\-kkOPD\([21](https://arxiv.org/html/2608.14728#bib.bib9);[7](https://arxiv.org/html/2608.14728#bib.bib3)\)directly uses normalized top\-kkreverse KL as the loss function\. Entropy\-aware OPD\([16](https://arxiv.org/html/2608.14728#bib.bib8)\)improves sampled\-token OPD by applying normalized top\-kkforward KL at high\-entropy positions\. vOPD\([23](https://arxiv.org/html/2608.14728#bib.bib33)\)uses normalized top\-kkreverse KL as a reward baseline to reduce variance of sampled\-token OPD\. While these methods use normalized top\-kkKL for different purposes, they share a common limitation: the normalized objective does not faithfully approximate the full\-vocabulary KL\. We note that[6](https://arxiv.org/html/2608.14728#bib.bib32)study a related tail probability issue, but it considers off\-policy full\-vocabulary distillation and is not applicable to top\-kkOPD\. Moreover, the failure mode it studies differs from ours: it studies how full\-vocabulary forward KL in the off\-policy setting makes the student’s tail probability too small, while we study how normalized top\-kkobjectives in the on\-policy setting steadily increase it\.

#### KL computation for LLM post\-training\.

KL divergence is a crucial component in post\-training of LLMs\([34](https://arxiv.org/html/2608.14728#bib.bib25);[24](https://arxiv.org/html/2608.14728#bib.bib23)\)\. In OPD, the student is optimized by minimizing the reverse KL divergence to the teacher\([22](https://arxiv.org/html/2608.14728#bib.bib1)\), while in reinforcement learning \(RL\), a KL constraint against the base model is commonly imposed to prevent the policy from drifting away\([48](https://arxiv.org/html/2608.14728#bib.bib24);[31](https://arxiv.org/html/2608.14728#bib.bib26)\)\. For efficiency, most existing methods estimate the divergence using only the sampled tokens\([28](https://arxiv.org/html/2608.14728#bib.bib28);[30](https://arxiv.org/html/2608.14728#bib.bib27)\)\. These estimators suffer from high variance and ignore the dense information in the logit space\. To provide dense information at tractable cost, top\-k approximations of the KL divergence have become an appealing alternative\([46](https://arxiv.org/html/2608.14728#bib.bib29)\)\. Existing top\-kkOPD methods predominantly adopt a normalized top\-k formulation\([21](https://arxiv.org/html/2608.14728#bib.bib9);[7](https://arxiv.org/html/2608.14728#bib.bib3)\), which, as we reveal in this work, leads to a steady increase in the student’s tail probability\. In the context of RL, several top\-k KL approximations have also been explored, yet they are devised for purposes different from ours\. For instance, DPPO\([25](https://arxiv.org/html/2608.14728#bib.bib30)\)derives a KL approximation similar to TA\-OPD’s loss objective, but uses it as a quantity to decide whether to clip the policy update on a token, thereby substituting heuristic PPO clipping\([29](https://arxiv.org/html/2608.14728#bib.bib22)\)with a principled constraint\. In contrast, we directly use the top\-kkKL approximation as the optimization objective and establish its theoretical properties\. Besides, EMA\-PG\([46](https://arxiv.org/html/2608.14728#bib.bib29)\)proposes a KL approximation similar to our sample\-corrected TA\-OPD\. We provide a discussion of sample\-corrected TA\-OPD and EMA\-PG in Appendix[E\.5](https://arxiv.org/html/2608.14728#A5.SS5)\.

## 8Conclusion

In this paper, we introduce Tail\-Aware Top\-k On\-Policy Distillation \(TA\-OPD\), a novel distillation method that restores the missing tail probability signal\. In particular, TA\-OPD minimizes the reverse KL divergence over the teacher’s top\-kktokens plus a tail token that carries the tail probability\. In effect, TA\-OPD explicitly aligns the student’s tail probability with the teacher’s, addressing the tail probability and entropy increase caused by top\-kknormalization\. We further derive a sampled variant that yields an unbiased estimate of the full\-vocabulary reverse KL when the sampled token’s probability is available\. Extensive experiments show that TA\-OPD consistently improves accuracy across benchmarks\. Our method can be easily adopted in practice: it is straightforward to implement with existing OPD frameworks and requires no additional teacher queries beyond the top\-kkprobabilities\. We hope that our insights inspire future research to further explore loss function designs for OPD\.

#### Limitations\.

The performance gain of TA\-OPD over normalized top\-kkOPD diminishes askkincreases or the capability gap between student and teacher decreases\. Additionally, due to limited computational resources, our experiments are restricted to models with up to 8B parameters\. While these sizes are standard for research\-stage OPD studies, we do not presume automatic transfer to 30B\+ models\.

## References

- R\. Agarwal, N\. Vieillard, Y\. Zhou, P\. Stanczyk, S\. Ramos Garea, M\. Geist, and O\. BachemOn\-policy distillation of language models: learning from self\-generated mistakes\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 21246–21263\.Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p1.1)\.
- Aroraet al\.\(2022\)K\. Arora, L\. El Asri, H\. Bahuleyan, and J\. C\. K\. CheungWhy exposure bias matters: an imitation learning perspective of error accumulation in language generation\.InFindings of the Association for Computational Linguistics: ACL 2022,pp\. 700–710\.Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p1.1)\.
- Bengioet al\.\(2015\)S\. Bengio, O\. Vinyals, N\. Jaitly, and N\. ShazeerScheduled sampling for sequence prediction with recurrent neural networks\.Advances in Neural Information Processing Systems28\.Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p1.1)\.
- Clarket al\.\(2018\)P\. Clark, I\. Cowhey, O\. Etzioni, T\. Khot, A\. Sabharwal, C\. Schoenick, and O\. TafjordThink you have solved question answering? try arc, the ai2 reasoning challenge\.arXiv preprint arXiv:1803\.05457\.Cited by:[§5\.1](https://arxiv.org/html/2608.14728#S5.SS1.SSS0.Px2.p1.1)\.
- Cuiet al\.\(2025\)G\. Cui, Y\. Zhang, J\. Chen, L\. Yuan, Z\. Wang, Y\. Zuo, H\. Li, Y\. Fan, H\. Chen, W\. Chen,et al\.The entropy mechanism of reinforcement learning for reasoning language models\.arXiv preprint arXiv:2505\.22617\.Cited by:[§A\.2](https://arxiv.org/html/2608.14728#A1.SS2.p1.1.1)\.
- Dasguptaet al\.\(2026\)S\. Dasgupta, T\. Cohn, and T\. BaldwinDon’t ignore the tail: decoupling top\-k probabilities for efficient language model distillation\.arXiv preprint arXiv:2602\.20816\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px1.p1.1)\.
- Fuet al\.\(2026\)Y\. Fu, H\. Huang, K\. Jiang, J\. Liu, Z\. Jiang, Y\. Zhu, and D\. ZhaoRevisiting on\-policy distillation: empirical failure modes and simple fixes\.arXiv preprint arXiv:2603\.25562\.Cited by:[Appendix C](https://arxiv.org/html/2608.14728#A3.SS0.SSS0.Px3.p1.1),[§1](https://arxiv.org/html/2608.14728#S1.p2.1),[§2](https://arxiv.org/html/2608.14728#S2.p2.1),[§3\.2](https://arxiv.org/html/2608.14728#S3.SS2.SSS0.Px2.p1.1),[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px1.p1.1),[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Grattafioriet al\.\(2024\)A\. Grattafiori, A\. Dubey, A\. Jauhri, A\. Pandey, A\. Kadian, A\. Al\-Dahle, A\. Letman, A\. Mathur, A\. Schelten, A\. Vaughan,et al\.The llama 3 herd of models\.arXiv preprint arXiv:2407\.21783\.Cited by:[§3\.2](https://arxiv.org/html/2608.14728#S3.SS2.SSS0.Px1.p1.1)\.
- Guet al\.\(2024\)Y\. Gu, L\. Dong, F\. Wei, and M\. HuangMinillm: knowledge distillation of large language models\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 32694–32717\.Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p1.1)\.
- Guhaet al\.\(2025\)E\. Guha, R\. Marten, S\. Keh, N\. Raoof, G\. Smyrnis, H\. Bansal, M\. Nezhurina, J\. Mercat, T\. Vu, Z\. Sprague, A\. Suvarna, B\. Feuer, L\. Chen, Z\. Khan, E\. Frankel, S\. Grover, C\. Choi, N\. Muennighoff, S\. Su, W\. Zhao, J\. Yang, S\. Pimpalgaonkar, K\. Sharma, C\. C\. Ji, Y\. Deng, S\. Pratt, V\. Ramanujan, J\. Saad\-Falcon, J\. Li, A\. Dave, A\. Albalak, K\. Arora, B\. Wulfe, C\. Hegde, G\. Durrett, S\. Oh, M\. Bansal, S\. Gabriel, A\. Grover, K\. Chang, V\. Shankar, A\. Gokaslan, M\. A\. Merrill, T\. Hashimoto, Y\. Choi, J\. Jitsev, R\. Heckel, M\. Sathiamoorthy, A\. G\. Dimakis, and L\. SchmidtOpenThoughts: data recipes for reasoning models\.External Links:2506\.04178,[Link](https://arxiv.org/abs/2506.04178)Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p4.1)\.
- Guoet al\.\(2025\)D\. Guo, D\. Yang, H\. Zhang, J\. Song, P\. Wang, Q\. Zhu, R\. Xu, R\. Zhang, S\. Ma, X\. Bi,et al\.Deepseek\-r1: incentivizing reasoning capability in llms via reinforcement learning\.arXiv preprint arXiv:2501\.12948\.Cited by:[§3\.2](https://arxiv.org/html/2608.14728#S3.SS2.SSS0.Px1.p1.1)\.
- Heet al\.\(2025\)B\. He, Z\. Qu, Z\. Liu, Y\. Chen, Y\. Zuo, C\. Qian, K\. Zhang, W\. Chen, C\. Xiao, G\. Cui,et al\.Justrl: scaling a 1\.5 b llm with a simple rl recipe\.arXiv preprint arXiv:2512\.16649\.Cited by:[§6](https://arxiv.org/html/2608.14728#S6.SS0.SSS0.Px2.p1.1)\.
- Heet al\.\(2024\)C\. He, R\. Luo, Y\. Bai, S\. Hu, Z\. Thai, J\. Shen, J\. Hu, X\. Han, Y\. Huang, Y\. Zhang,et al\.Olympiadbench: a challenging benchmark for promoting agi with olympiad\-level bilingual multimodal scientific problems\.InProceedings of the 62nd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),pp\. 3828–3850\.Cited by:[§5\.1](https://arxiv.org/html/2608.14728#S5.SS1.SSS0.Px2.p1.1)\.
- Hendryckset al\.\(2021\)D\. Hendrycks, C\. Burns, S\. Kadavath, A\. Arora, S\. Basart, E\. Tang, D\. Song, and J\. SteinhardtMeasuring mathematical problem solving with the math dataset\.Advances in Neural Information Processing Systems\.Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p4.1)\.
- Hintonet al\.\(2015\)G\. Hinton, O\. Vinyals, and J\. DeanDistilling the knowledge in a neural network\.arXiv preprint arXiv:1503\.02531\.Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p1.1)\.
- Jinet al\.\(2026\)W\. Jin, T\. Min, Y\. Yang, S\. R\. Kadhe, Y\. Zhou, D\. Wei, N\. Baracaldo, and K\. LeeEntropy\-aware on\-policy distillation of language models\.arXiv preprint arXiv:2603\.07079\.Cited by:[§2](https://arxiv.org/html/2608.14728#S2.p2.1),[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px1.p1.1)\.
- Kim and Rush \(2016\)Y\. Kim and A\. M\. RushSequence\-level knowledge distillation\.InProceedings of the 2016 conference on empirical methods in natural language processing,pp\. 1317–1327\.Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p1.1)\.
- Kwonet al\.\(2023\)W\. Kwon, Z\. Li, S\. Zhuang, Y\. Sheng, L\. Zheng, C\. H\. Yu, J\. Gonzalez, H\. Zhang, and I\. StoicaEfficient memory management for large language model serving with pagedattention\.InProceedings of the 29th symposium on operating systems principles,pp\. 611–626\.Cited by:[§E\.3](https://arxiv.org/html/2608.14728#A5.SS3.p1.1)\.
- Lewkowyczet al\.\(2022\)A\. Lewkowycz, A\. Andreassen, D\. Dohan, E\. Dyer, H\. Michalewski, V\. Ramasesh, A\. Slone, C\. Anil, I\. Schlag, T\. Gutman\-Solo,et al\.Solving quantitative reasoning problems with language models\.Advances in Neural Information Processing Systems35,pp\. 3843–3857\.Cited by:[§5\.1](https://arxiv.org/html/2608.14728#S5.SS1.SSS0.Px2.p1.1)\.
- Liet al\.\(2024\)J\. Li, E\. Beeching, L\. Tunstall, B\. Lipkin, R\. Soletskyi, S\. Huang, K\. Rasul, L\. Yu, A\. Q\. Jiang, Z\. Shen,et al\.Numinamath: the largest public dataset in ai4maths with 860k pairs of competition math problems and solutions\.Hugging Face repository13\(9\),pp\. 9\.Cited by:[§5\.1](https://arxiv.org/html/2608.14728#S5.SS1.SSS0.Px2.p1.1)\.
- Liet al\.\(2026\)Y\. Li, Y\. Zuo, B\. He, J\. Zhang, C\. Xiao, C\. Qian, T\. Yu, H\. Gao, W\. Yang, Z\. Liu,et al\.Rethinking on\-policy distillation of large language models: phenomenology, mechanism, and recipe\.arXiv preprint arXiv:2604\.13016\.Cited by:[Appendix B](https://arxiv.org/html/2608.14728#A2.SS0.SSS0.Px4.p1.1),[Appendix C](https://arxiv.org/html/2608.14728#A3.SS0.SSS0.Px3.p1.1),[§E\.2](https://arxiv.org/html/2608.14728#A5.SS2.p1.1),[§2](https://arxiv.org/html/2608.14728#S2.p2.1),[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px1.p1.1),[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Lu and Lab \(2025\)K\. Lu and T\. M\. LabOn\-policy distillation\.Thinking Machines Lab: Connectionism\.Note:https://thinkingmachines\.ai/blog/on\-policy\-distillationExternal Links:[Document](https://dx.doi.org/10.64434/tml.20251026)Cited by:[Appendix C](https://arxiv.org/html/2608.14728#A3.SS0.SSS0.Px1.p1.1),[§1](https://arxiv.org/html/2608.14728#S1.p1.1),[§2](https://arxiv.org/html/2608.14728#S2.p1.2),[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Ohet al\.\(2026\)M\. Oh, S\. Song, G\. Choi, Y\. Choi, and Y\. JoKL for a kl: on\-policy distillation with control variate baseline\.arXiv preprint arXiv:2605\.07865\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px1.p1.1)\.
- Ouyanget al\.\(2022\)L\. Ouyang, J\. Wu, X\. Jiang, D\. Almeida, C\. Wainwright, P\. Mishkin, C\. Zhang, S\. Agarwal, K\. Slama, A\. Ray,et al\.Training language models to follow instructions with human feedback\.Advances in Neural Information Processing Systems35,pp\. 27730–27744\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Qiet al\.\(2026\)P\. Qi, X\. Zhou, Z\. Liu, T\. Pang, C\. Du, M\. Lin, and W\. S\. LeeRethinking the trust region in llm reinforcement learning\.arXiv preprint arXiv:2602\.04879\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Qwenet al\.\(2025\)Qwen, :, A\. Yang, B\. Yang, B\. Zhang, B\. Hui, B\. Zheng, B\. Yu, C\. Li, D\. Liu, F\. Huang, H\. Wei, H\. Lin, J\. Yang, J\. Tu, J\. Zhang, J\. Yang, J\. Yang, J\. Zhou, J\. Lin, K\. Dang, K\. Lu, K\. Bao, K\. Yang, L\. Yu, M\. Li, M\. Xue, P\. Zhang, Q\. Zhu, R\. Men, R\. Lin, T\. Li, T\. Tang, T\. Xia, X\. Ren, X\. Ren, Y\. Fan, Y\. Su, Y\. Zhang, Y\. Wan, Y\. Liu, Z\. Cui, Z\. Zhang, and Z\. QiuQwen2\.5 technical report\.External Links:2412\.15115,[Link](https://arxiv.org/abs/2412.15115)Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p4.1)\.
- Ren and Sutherland \(2025\)Y\. Ren and D\. SutherlandLearning dynamics of llm finetuning\.InInternational Conference on Learning Representations,Vol\.2025,pp\. 70523–70563\.Cited by:[§A\.2](https://arxiv.org/html/2608.14728#A1.SS2.p1.1.1)\.
- Schulmanet al\.\(2015\)J\. Schulman, S\. Levine, P\. Abbeel, M\. Jordan, and P\. MoritzTrust region policy optimization\.InInternational Conference on Machine Learning,pp\. 1889–1897\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Schulmanet al\.\(2017\)J\. Schulman, F\. Wolski, P\. Dhariwal, A\. Radford, and O\. KlimovProximal policy optimization algorithms\.arXiv preprint arXiv:1707\.06347\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Schulman \(2020\)J\. SchulmanApproximating kl divergence\.John Schulman’s Homepage5\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Shahet al\.\(2025\)V\. Shah, J\. Obando\-Ceron, V\. Jain, B\. Bartoldson, B\. Kailkhura, S\. Mittal, G\. Berseth, P\. S\. Castro, Y\. Bengio, N\. Malkin,et al\.A comedy of estimators: on kl regularization in rl training of llms\.arXiv preprint arXiv:2512\.21852\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Shenget al\.\(2025\)G\. Sheng, C\. Zhang, Z\. Ye, X\. Wu, W\. Zhang, R\. Zhang, Y\. Peng, H\. Lin, and C\. WuHybridflow: a flexible and efficient rlhf framework\.InProceedings of the Twentieth European Conference on Computer Systems,pp\. 1279–1297\.Cited by:[§D\.1](https://arxiv.org/html/2608.14728#A4.SS1.p1.1)\.
- Song and Zheng \(2026\)M\. Song and M\. ZhengA survey of on\-policy distillation for large language models\.arXiv preprint arXiv:2604\.00626\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px1.p1.1)\.
- Vieillardet al\.\(2020\)N\. Vieillard, T\. Kozuno, B\. Scherrer, O\. Pietquin, R\. Munos, and M\. GeistLeverage the average: an analysis of kl regularization in reinforcement learning\.Advances in Neural Information Processing Systems33,pp\. 12163–12174\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Wanget al\.\(2024\)Y\. Wang, X\. Ma, G\. Zhang, Y\. Ni, A\. Chandra, S\. Guo, W\. Ren, A\. Arulraj, X\. He, Z\. Jiang,et al\.Mmlu\-pro: a more robust and challenging multi\-task language understanding benchmark\.Advances in Neural Information Processing Systems37,pp\. 95266–95290\.Cited by:[§5\.1](https://arxiv.org/html/2608.14728#S5.SS1.SSS0.Px2.p1.1)\.
- Xiaoet al\.\(2026\)B\. Xiao, B\. Xia, B\. Yang, B\. Gao, B\. Shen, C\. Zhang, C\. He, C\. Lou, F\. Luo, G\. Wang,et al\.Mimo\-v2\-flash technical report\.arXiv preprint arXiv:2601\.02780\.Cited by:[Appendix C](https://arxiv.org/html/2608.14728#A3.SS0.SSS0.Px1.p1.1),[§2](https://arxiv.org/html/2608.14728#S2.p1.2)\.
- Xieet al\.\(2026\)Y\. Xie, S\. Zhu, T\. Wen, B\. Chen, and Y\. WangOn the position bias of on\-policy distillation\.arXiv preprint arXiv:2606\.22600\.Cited by:[§3\.2](https://arxiv.org/html/2608.14728#S3.SS2.SSS0.Px2.p1.1)\.
- Xinget al\.\(2026\)X\. Xing, H\. Wang, B\. Gao, Z\. Li, and Y\. TangTrust region on\-policy distillation\.arXiv preprint arXiv:2606\.01249\.Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p2.1),[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px1.p1.1)\.
- Xuet al\.\(2026\)A\. Xu, B\. Lin, B\. Xue, B\. Wang, B\. Xu, B\. Wu, B\. Zhang, C\. Lin, C\. Dong, C\. Ling,et al\.DeepSeek\-v4: towards highly efficient million\-token context intelligence\.arXiv preprint arXiv:2606\.19348\.Cited by:[§2](https://arxiv.org/html/2608.14728#S2.p1.2)\.
- Xuet al\.\(2025a\)H\. Xu, Q\. Zhu, H\. Deng, J\. Li, L\. Hou, Y\. Wang, L\. Shang, R\. Xu, and F\. MiKdrl: post\-training reasoning llms via unified knowledge distillation and reinforcement learning\.arXiv preprint arXiv:2506\.02208\.Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p1.1)\.
- Xuet al\.\(2025b\)W\. Xu, R\. Han, Z\. Wang, L\. Le, D\. Madeka, L\. Li, W\. Wang, R\. Agarwal, C\. Lee, and T\. PfisterSpeculative knowledge distillation: bridging the teacher\-student gap through interleaved sampling\.InInternational Conference on Learning Representations,Vol\.2025,pp\. 64616–64646\.Cited by:[§1](https://arxiv.org/html/2608.14728#S1.p1.1)\.
- Yanet al\.\(2026\)J\. Yan, Y\. Li, Z\. Hu, Z\. Wang, G\. Cui, X\. Qu, Y\. Cheng, and Y\. ZhangLearning to reason under off\-policy guidance\.Advances in Neural Information Processing Systems38,pp\. 117157–117186\.Cited by:[§D\.1](https://arxiv.org/html/2608.14728#A4.SS1.p1.1)\.
- Yanget al\.\(2025\)A\. Yang, A\. Li, B\. Yang, B\. Zhang, B\. Hui, B\. Zheng, B\. Yu, C\. Gao, C\. Huang, C\. Lv,et al\.Qwen3 technical report\.arXiv preprint arXiv:2505\.09388\.Cited by:[§3\.2](https://arxiv.org/html/2608.14728#S3.SS2.SSS0.Px1.p1.1)\.
- Yuet al\.\(2026\)Q\. Yu, Z\. Zhang, R\. Zhu, Y\. Yuan, X\. Zuo, Y\. Yue, W\. Dai, T\. Fan, G\. Liu, L\. Liu,et al\.Dapo: an open\-source llm reinforcement learning system at scale\.Advances in Neural Information Processing Systems38,pp\. 113222–113244\.Cited by:[§3\.2](https://arxiv.org/html/2608.14728#S3.SS2.SSS0.Px1.p1.1)\.
- Zhanget al\.\(2026\)H\. Zhang, Y\. Li, Z\. Wang, Z\. Wang, S\. Zhang, X\. Qu, and Y\. ChengCharacterizing, evaluating, and optimizing complex reasoning\.InForty\-third International Conference on Machine Learning,External Links:[Link](https://openreview.net/forum?id=IMFgiWw4jd)Cited by:[§A\.2](https://arxiv.org/html/2608.14728#A1.SS2.p1.1.1)\.
- Zhang and Ba \(2026\)L\. Zhang and J\. BaEMA policy gradient: taming reinforcement learning for llms with ema anchor and top\-k kl\.arXiv preprint arXiv:2602\.04417\.Cited by:[§E\.5](https://arxiv.org/html/2608.14728#A5.SS5.p1.1),[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.
- Zhuet al\.\(2026\)S\. Zhu, X\. Ye, H\. Lu, W\. Shi, and G\. LiuThe many faces of on\-policy distillation: pitfalls, mechanisms, and fixes\.arXiv preprint arXiv:2605\.11182\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px1.p1.1)\.
- Ziegleret al\.\(2019\)D\. M\. Ziegler, N\. Stiennon, J\. Wu, T\. B\. Brown, A\. Radford, D\. Amodei, P\. Christiano, and G\. IrvingFine\-tuning language models from human preferences\.arXiv preprint arXiv:1909\.08593\.Cited by:[§7](https://arxiv.org/html/2608.14728#S7.SS0.SSS0.Px2.p1.1)\.

## Appendix AProof

### A\.1Proof of Proposition[1](https://arxiv.org/html/2608.14728#Thmproposition1)

###### Proof\.

Fix a stepttand a prefixy^<t\\hat\{y\}\_\{<t\}, and omit the subscriptttthroughout: writeS=Stk=TopK⁡\(qt,k\)S=S\_\{t\}^\{k\}=\\operatorname\{TopK\}\(q\_\{t\},k\)for the teacher’s top\-kktokens,zv=zt,vz\_\{v\}=z\_\{t,v\}for the student logit of tokenvv, andp¯=p¯t\(Stk\)\\bar\{p\}=\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\},q¯=q¯t\(Stk\)\\bar\{q\}=\\bar\{q\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}for the normalized distributions onSS, so that

ℓnorm=∑u∈Sp¯​\(u\)​log⁡p¯​\(u\)q¯​\(u\)\.\\ell^\{\\mathrm\{norm\}\}=\\sum\_\{u\\in S\}\\bar\{p\}\(u\)\\log\\frac\{\\bar\{p\}\(u\)\}\{\\bar\{q\}\(u\)\}\.Since the token setSSis selected from the teacher distribution, bothSSandq¯\\bar\{q\}are constants with respect to the student logits\. We proceed in three steps: we first show thatp¯\\bar\{p\}is a softmax over the logits restricted to top\-kktokensSS, then differentiate the loss through it, and finally sum the resulting gradients overSS\.

#### Step 1:p¯\\bar\{p\}can be expressed as a softmax restricted to top\-kktokens\.

Writing the student distribution asp⁡\(v\)=ezv/∑u∈𝒱ezup\(v\)=e^\{z\_\{v\}\}/\\sum\_\{u\\in\\mathcal\{V\}\}e^\{z\_\{u\}\}, the vocabulary\-level normalizer appears in both the numerator and the denominator ofp¯\\bar\{p\}and therefore cancels,

p¯​\(v\)=p⁡\(v\)∑w∈Sp⁡\(w\)=ezv∑w∈Sezw,v∈S\.\\bar\{p\}\(v\)=\\frac\{p\(v\)\}\{\\sum\_\{w\\in S\}p\(w\)\}=\\frac\{e^\{z\_\{v\}\}\}\{\\sum\_\{w\\in S\}e^\{z\_\{w\}\}\},\\qquad v\\in S\.\(8\)Two consequences of Eq\. \([8](https://arxiv.org/html/2608.14728#A1.E8)\) drive the whole proof:p¯\\bar\{p\}is a softmax over the sub\-vector\(zw\)w∈S\(z\_\{w\}\)\_\{w\\in S\}, and it does not depend on any logit outsideSS\. Its Jacobian is therefore the standard softmax Jacobian onSSand vanishes elsewhere,

∂p¯​\(u\)∂zv=\{p¯\(u\)\(𝟏\[u=v\]−p¯\(v\)\),u,v∈S,0,v∉S\.\\frac\{\\partial\\bar\{p\}\(u\)\}\{\\partial z\_\{v\}\}=\\begin\{cases\}\\bar\{p\}\(u\)\\bigl\(\\mathbf\{1\}\[u=v\]\-\\bar\{p\}\(v\)\\bigr\),&u,v\\in S,\\\\\[3\.99994pt\] 0,&v\\notin S\.\\end\{cases\}\(9\)

#### Step 2: the gradient with respect to logit\.

Becauseq¯\\bar\{q\}is fixed with respect to the student logits, the loss depends on the student only through\{p¯​\(u\)\}u∈S\\\{\\bar\{p\}\(u\)\\\}\_\{u\\in S\}, with

∂ℓnorm∂p¯​\(u\)=log⁡p¯​\(u\)q¯​\(u\)\+1,u∈S\.\\frac\{\\partial\\ell^\{\\mathrm\{norm\}\}\}\{\\partial\\bar\{p\}\(u\)\}=\\log\\frac\{\\bar\{p\}\(u\)\}\{\\bar\{q\}\(u\)\}\+1,\\qquad u\\in S\.The chain rule over the vocabulary thus gives

∂ℓnorm∂zv=∑u∈S∂ℓnorm∂p¯​\(u\)×∂p¯​\(u\)∂zv=∑u∈S\(log⁡p¯​\(u\)q¯​\(u\)\+1\)​∂p¯​\(u\)∂zv,\\frac\{\\partial\\ell^\{\\mathrm\{norm\}\}\}\{\\partial z\_\{v\}\}=\\sum\_\{u\\in S\}\\frac\{\\partial\\ell^\{\\mathrm\{norm\}\}\}\{\\partial\\bar\{p\}\(u\)\}\\times\\frac\{\\partial\\bar\{p\}\(u\)\}\{\\partial z\_\{v\}\}=\\sum\_\{u\\in S\}\\left\(\\log\\frac\{\\bar\{p\}\(u\)\}\{\\bar\{q\}\(u\)\}\+1\\right\)\\frac\{\\partial\\bar\{p\}\(u\)\}\{\\partial z\_\{v\}\},\(10\)where the terms withu∉Su\\notin Sare already dropped becauseℓnorm\\ell^\{\\mathrm\{norm\}\}does not involve them\. Forv∉Sv\\notin S, every Jacobian entry in Eq\. \([9](https://arxiv.org/html/2608.14728#A1.E9)\) is zero, hence∂ℓnorm/∂zv=0\\partial\\ell^\{\\mathrm\{norm\}\}/\\partial z\_\{v\}=0, which is the second case of the claim: the loss is blind to the logits outside the teacher’s top\-kktokens\. Forv∈Sv\\in S, substituting Eq\. \([9](https://arxiv.org/html/2608.14728#A1.E9)\) into Eq\. \([10](https://arxiv.org/html/2608.14728#A1.E10)\) and separating the diagonal termu=vu=vgives

∂ℓnorm∂zv=p¯​\(v\)​\(log⁡p¯​\(v\)q¯​\(v\)\+1\)−p¯​\(v\)​∑u∈Sp¯​\(u\)​\(log⁡p¯​\(u\)q¯​\(u\)\+1\)\.\\frac\{\\partial\\ell^\{\\mathrm\{norm\}\}\}\{\\partial z\_\{v\}\}=\\bar\{p\}\(v\)\\left\(\\log\\frac\{\\bar\{p\}\(v\)\}\{\\bar\{q\}\(v\)\}\+1\\right\)\-\\bar\{p\}\(v\)\\sum\_\{u\\in S\}\\bar\{p\}\(u\)\\left\(\\log\\frac\{\\bar\{p\}\(u\)\}\{\\bar\{q\}\(u\)\}\+1\\right\)\.\(11\)Note that

∑u∈Sp¯​\(u\)​\(log⁡p¯​\(u\)q¯​\(u\)\+1\)=∑u∈Sp¯​\(u\)​log⁡p¯​\(u\)q¯​\(u\)\+∑u∈Sp¯​\(u\)=ℓnorm\+1\.\\sum\_\{u\\in S\}\\bar\{p\}\(u\)\\left\(\\log\\frac\{\\bar\{p\}\(u\)\}\{\\bar\{q\}\(u\)\}\+1\\right\)=\\sum\_\{u\\in S\}\\bar\{p\}\(u\)\\log\\frac\{\\bar\{p\}\(u\)\}\{\\bar\{q\}\(u\)\}\+\\sum\_\{u\\in S\}\\bar\{p\}\(u\)=\\ell^\{\\mathrm\{norm\}\}\+1\.By substituding it into Eq[11](https://arxiv.org/html/2608.14728#A1.E11), we obtain

∂ℓnorm∂zv=p¯​\(v\)​\[log⁡p¯​\(v\)q¯​\(v\)−ℓnorm\],v∈S,\\frac\{\\partial\\ell^\{\\mathrm\{norm\}\}\}\{\\partial z\_\{v\}\}=\\bar\{p\}\(v\)\\left\[\\log\\frac\{\\bar\{p\}\(v\)\}\{\\bar\{q\}\(v\)\}\-\\ell^\{\\mathrm\{norm\}\}\\right\],\\qquad v\\in S,\(12\)which is the first case of the claim\. Sinceℓnorm=𝔼v∼p¯​\[log⁡\(p¯​\(v\)/q¯​\(v\)\)\]\\ell^\{\\mathrm\{norm\}\}=\\mathbb\{E\}\_\{v\\sim\\bar\{p\}\}\\bigl\[\\log\(\\bar\{p\}\(v\)/\\bar\{q\}\(v\)\)\\bigr\]by definition, Eq\. \([12](https://arxiv.org/html/2608.14728#A1.E12)\) says that the gradient on a top\-kklogit is its log\-ratio centered at thep¯\\bar\{p\}\-weighted mean log\-ratio\.

#### Step 3: the top\-kkgradients sum to zero\.

Summing Eq\. \([12](https://arxiv.org/html/2608.14728#A1.E12)\) overv∈Sv\\in Sand using∑v∈Sp¯​\(v\)​log⁡\(p¯​\(v\)/q¯​\(v\)\)=ℓnorm\\sum\_\{v\\in S\}\\bar\{p\}\(v\)\\log\\bigl\(\\bar\{p\}\(v\)/\\bar\{q\}\(v\)\\bigr\)=\\ell^\{\\mathrm\{norm\}\}together with∑v∈Sp¯​\(v\)=1\\sum\_\{v\\in S\}\\bar\{p\}\(v\)=1yields

∑v∈S∂ℓnorm∂zv=ℓnorm−ℓnorm=0\.\\sum\_\{v\\in S\}\\frac\{\\partial\\ell^\{\\mathrm\{norm\}\}\}\{\\partial z\_\{v\}\}=\\ell^\{\\mathrm\{norm\}\}\-\\ell^\{\\mathrm\{norm\}\}=0\.This is precisely the mean\-centering exhibited in Eq\. \([12](https://arxiv.org/html/2608.14728#A1.E12)\): the update redistributes logit mass among the top\-kktokens without changing their total, while leaving the tail logits untouched by Step 2\. The normalized objective therefore acts only on the shape of the student distribution withinSS, and never on how much probability the student assigns toSSas a whole\. ∎

### A\.2Proof of Proposition[2](https://arxiv.org/html/2608.14728#Thmproposition2)

###### Proof\.

Fix a stepttand a prefixy^<t\\hat\{y\}\_\{<t\}, and omit the subscriptttthroughout, following the convention of Appendix[A\.1](https://arxiv.org/html/2608.14728#A1.SS1): writeS=StkS=S\_\{t\}^\{k\},zv=zt,vz\_\{v\}=z\_\{t,v\},p¯=p¯t\(Stk\)\\bar\{p\}=\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\},q¯=q¯t\(Stk\)\\bar\{q\}=\\bar\{q\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}, andptail=pttailp^\{\\mathrm\{tail\}\}=p\_\{t\}^\{\\mathrm\{tail\}\}\. Under the tabular assumption, each logitzvz\_\{v\}is an independent parameter, so the gradient step acts on the logits directly; this assumption is commonly adopted by prior works for analysis\([5](https://arxiv.org/html/2608.14728#bib.bib17);[45](https://arxiv.org/html/2608.14728#bib.bib45);[27](https://arxiv.org/html/2608.14728#bib.bib46)\)\. We writeptail​\(η\)p^\{\\mathrm\{tail\}\}\(\\eta\)for the tail probability after one gradient step with learning rateη\\eta, andΔ​ptail=ptail​\(η\)−ptail​\(0\)\\Delta p^\{\\mathrm\{tail\}\}=p^\{\\mathrm\{tail\}\}\(\\eta\)\-p^\{\\mathrm\{tail\}\}\(0\)for its change\. We proceed in three steps: we first reduceptail​\(η\)p^\{\\mathrm\{tail\}\}\(\\eta\)to a one\-dimensional function of the step size, then identify the derivative of that function as a covariance, and finally expand it to first order inη\\eta\.

#### Step 1: only the top\-kklogits change\.

Split the softmax normalizer into the contributions of the top\-kktokens and of the tail tokens,

A:=∑v∈Sezv,B:=∑v∉Sezv,Z:=A\+B,A:=\\sum\_\{v\\in S\}e^\{z\_\{v\}\},\\qquad B:=\\sum\_\{v\\notin S\}e^\{z\_\{v\}\},\\qquad Z:=A\+B,so thatptail=B/Zp^\{\\mathrm\{tail\}\}=B/Z\. By Proposition[1](https://arxiv.org/html/2608.14728#Thmproposition1), the gradient ofℓnorm\\ell^\{\\mathrm\{norm\}\}is

gv=p¯​\(v\)​\(log⁡p¯​\(v\)q¯​\(v\)−ℓnorm\)for​v∈S,gv=0for​v∉S\.g\_\{v\}=\\bar\{p\}\(v\)\\left\(\\log\\frac\{\\bar\{p\}\(v\)\}\{\\bar\{q\}\(v\)\}\-\\ell^\{\\mathrm\{norm\}\}\\right\)\\ \\ \\text\{for \}v\\in S,\\qquad g\_\{v\}=0\\ \\ \\text\{for \}v\\notin S\.The updatezv←zv−η​gvz\_\{v\}\\leftarrow z\_\{v\}\-\\eta g\_\{v\}therefore leaves tail logit unchanged, and hence leavesBBunchanged\. The change of the tail probability is thus determined byAAalone:

ptail​\(η\)=BA⁡\(η\)\+B,A⁡\(η\)=∑v∈Sezv−η​gv\.p^\{\\mathrm\{tail\}\}\(\\eta\)=\\frac\{B\}\{A\(\\eta\)\+B\},\\qquad A\(\\eta\)=\\sum\_\{v\\in S\}e^\{z\_\{v\}\-\\eta g\_\{v\}\}\.\(13\)

#### Step 2: the derivative ofAAis a covariance\.

DifferentiatingA⁡\(η\)A\(\\eta\)in Eq\. \([13](https://arxiv.org/html/2608.14728#A1.E13)\) with respect toη\\etaand evaluating atη=0\\eta=0gives

A′\(0\)=−∑v∈Sezvgv\.A^\{\\prime\}\(0\)=\-\\sum\_\{v\\in S\}e^\{z\_\{v\}\}g\_\{v\}\.We now rewrite this sum in terms of the normalized distributionp¯\\bar\{p\}\. Usingezv=Z​p​\(v\)e^\{z\_\{v\}\}=Z\\,p\(v\)andp⁡\(v\)=\(1−ptail\)​p¯​\(v\)p\(v\)=\(1\-p^\{\\mathrm\{tail\}\}\)\\,\\bar\{p\}\(v\)forv∈Sv\\in S, and then substituting the expression forgvg\_\{v\}from Step 1,

A′\(0\)=−Z\(1−ptail\)∑v∈Sp¯\(v\)gv=−Z\(1−ptail\)\[∑v∈Sp¯\(v\)2logp¯​\(v\)q¯​\(v\)−ℓnorm∑v∈Sp¯\(v\)2\]\.A^\{\\prime\}\(0\)=\-Z\\bigl\(1\-p^\{\\mathrm\{tail\}\}\\bigr\)\\sum\_\{v\\in S\}\\bar\{p\}\(v\)\\,g\_\{v\}=\-Z\\bigl\(1\-p^\{\\mathrm\{tail\}\}\\bigr\)\\left\[\\sum\_\{v\\in S\}\\bar\{p\}\(v\)^\{2\}\\log\\frac\{\\bar\{p\}\(v\)\}\{\\bar\{q\}\(v\)\}\-\\ell^\{\\mathrm\{norm\}\}\\sum\_\{v\\in S\}\\bar\{p\}\(v\)^\{2\}\\right\]\.\(14\)Each of the two sums in Eq\. \([14](https://arxiv.org/html/2608.14728#A1.E14)\) carries one factorp¯​\(v\)\\bar\{p\}\(v\)that plays the role of a sampling weight, so both are expectations underv∼p¯v\\sim\\bar\{p\}:

∑v∈Sp¯​\(v\)2​log⁡p¯​\(v\)q¯​\(v\)=𝔼v∼p¯​\[p¯​\(v\)​log⁡p¯​\(v\)q¯​\(v\)\],∑v∈Sp¯​\(v\)2=𝔼v∼p¯​\[p¯​\(v\)\]\.\\sum\_\{v\\in S\}\\bar\{p\}\(v\)^\{2\}\\log\\frac\{\\bar\{p\}\(v\)\}\{\\bar\{q\}\(v\)\}=\\mathbb\{E\}\_\{v\\sim\\bar\{p\}\}\\\!\\left\[\\bar\{p\}\(v\)\\log\\frac\{\\bar\{p\}\(v\)\}\{\\bar\{q\}\(v\)\}\\right\],\\qquad\\sum\_\{v\\in S\}\\bar\{p\}\(v\)^\{2\}=\\mathbb\{E\}\_\{v\\sim\\bar\{p\}\}\\bigl\[\\bar\{p\}\(v\)\\bigr\]\.Moreoverℓnorm=𝔼v∼p¯​\[log⁡\(p¯​\(v\)/q¯​\(v\)\)\]\\ell^\{\\mathrm\{norm\}\}=\\mathbb\{E\}\_\{v\\sim\\bar\{p\}\}\\bigl\[\\log\(\\bar\{p\}\(v\)/\\bar\{q\}\(v\)\)\\bigr\]by definition, so the bracket in Eq\. \([14](https://arxiv.org/html/2608.14728#A1.E14)\) is the difference between the expectation of a product and the product of the expectations, which is exactly a covariance:

A′​\(0\)=−Z⁡\(1−ptail\)​Covv∼p¯\(p¯​\(v\),log⁡p¯​\(v\)q¯​\(v\)\)\.A^\{\\prime\}\(0\)=\-Z\\bigl\(1\-p^\{\\mathrm\{tail\}\}\\bigr\)\\operatorname\*\{Cov\}\_\{v\\sim\\bar\{p\}\}\\\!\\left\(\\bar\{p\}\(v\),\\,\\log\\frac\{\\bar\{p\}\(v\)\}\{\\bar\{q\}\(v\)\}\\right\)\.\(15\)The covariance appears because, by Proposition[1](https://arxiv.org/html/2608.14728#Thmproposition1), the gradient on a top\-kklogit is its log\-ratio centered at the mean log\-ratio, while its contribution toAAis weighted by its own probabilityp¯​\(v\)\\bar\{p\}\(v\)\.

#### Step 3: first\-order taylor expansion\.

Differentiatingptail​\(η\)p^\{\\mathrm\{tail\}\}\(\\eta\)in Eq\. \([13](https://arxiv.org/html/2608.14728#A1.E13)\) atη=0\\eta=0and substituting Eq\. \([15](https://arxiv.org/html/2608.14728#A1.E15)\) together withB/Z=ptailB/Z=p^\{\\mathrm\{tail\}\}yields

d​ptaild​η\|η=0=−B\(A\+B\)2​A′​\(0\)=ptail​\(1−ptail\)​Covv∼p¯\(p¯​\(v\),log⁡p¯​\(v\)q¯​\(v\)\)\.\\frac\{\\mathrm\{d\}p^\{\\mathrm\{tail\}\}\}\{\\mathrm\{d\}\\eta\}\\bigg\|\_\{\\eta=0\}=\-\\frac\{B\}\{\(A\+B\)^\{2\}\}\\,A^\{\\prime\}\(0\)=p^\{\\mathrm\{tail\}\}\\bigl\(1\-p^\{\\mathrm\{tail\}\}\\bigr\)\\operatorname\*\{Cov\}\_\{v\\sim\\bar\{p\}\}\\\!\\left\(\\bar\{p\}\(v\),\\,\\log\\frac\{\\bar\{p\}\(v\)\}\{\\bar\{q\}\(v\)\}\\right\)\.A first\-order Taylor expansion inη\\eta, followed by splitting the log\-ratio aslog⁡p¯​\(v\)−log⁡q¯​\(v\)\\log\\bar\{p\}\(v\)\-\\log\\bar\{q\}\(v\)and using the linearity of the covariance in its second argument, gives the stated identity:

Δ​ptail=η​ptail​\(1−ptail\)​\[Covv∼p¯\(p¯​\(v\),log⁡p¯​\(v\)\)−Covv∼p¯\(p¯​\(v\),log⁡q¯​\(v\)\)\]\+O⁡\(η2\)\.\\Delta p^\{\\mathrm\{tail\}\}=\\eta\\,p^\{\\mathrm\{tail\}\}\\bigl\(1\-p^\{\\mathrm\{tail\}\}\\bigr\)\\Bigl\[\\operatorname\*\{Cov\}\_\{v\\sim\\bar\{p\}\}\\bigl\(\\bar\{p\}\(v\),\\log\\bar\{p\}\(v\)\\bigr\)\-\\operatorname\*\{Cov\}\_\{v\\sim\\bar\{p\}\}\\bigl\(\\bar\{p\}\(v\),\\log\\bar\{q\}\(v\)\\bigr\)\\Bigr\]\+O\(\\eta^\{2\}\)\.Sinceptail∈\(0,1\)p^\{\\mathrm\{tail\}\}\\in\(0,1\), the prefactorptail​\(1−ptail\)p^\{\\mathrm\{tail\}\}\(1\-p^\{\\mathrm\{tail\}\}\)is strictly positive, so for sufficiently smallη\>0\\eta\>0the sign ofΔ​ptail\\Delta p^\{\\mathrm\{tail\}\}is determined by the sign of the covariance difference\. The tail probability therefore increases if and only if the student self\-covariance exceeds the student–teacher covariance\. ∎

### A\.3Proof of Proposition[3](https://arxiv.org/html/2608.14728#Thmproposition3)

###### Proof\.

Fix a stepttand a prefixy^<t\\hat\{y\}\_\{<t\}, and omit the subscriptttthroughout, following the convention of Appendix[A\.1](https://arxiv.org/html/2608.14728#A1.SS1): writeS=StkS=S\_\{t\}^\{k\},zv=zt,vz\_\{v\}=z\_\{t,v\},p​\(v\)=pt​\(v\)p\(v\)=p\_\{t\}\(v\),q​\(v\)=qt​\(v\)q\(v\)=q\_\{t\}\(v\), andptail=pttailp^\{\\mathrm\{tail\}\}=p\_\{t\}^\{\\mathrm\{tail\}\},qtail=qttailq^\{\\mathrm\{tail\}\}=q\_\{t\}^\{\\mathrm\{tail\}\}, so that

ℓTA=∑u∈Sp⁡\(u\)​log⁡p⁡\(u\)q⁡\(u\)\+ptail​log⁡ptailqtail,ptail=∑u∉Sp⁡\(u\)\.\\ell^\{\\mathrm\{TA\}\}=\\sum\_\{u\\in S\}p\(u\)\\log\\frac\{p\(u\)\}\{q\(u\)\}\+p^\{\\mathrm\{tail\}\}\\log\\frac\{p^\{\\mathrm\{tail\}\}\}\{q^\{\\mathrm\{tail\}\}\},\\qquad p^\{\\mathrm\{tail\}\}=\\sum\_\{u\\notin S\}p\(u\)\.SinceSSis selected from the teacher distribution,SS,qq, andqtailq^\{\\mathrm\{tail\}\}are constants with respect to the student logits\. Unlike the normalized loss of Proposition[1](https://arxiv.org/html/2608.14728#Thmproposition1),ℓTA\\ell^\{\\mathrm\{TA\}\}depends on the student through the*unnormalized*probabilities, and the tail tokens enter it through their aggregateptailp^\{\\mathrm\{tail\}\}\. We proceed in three steps: we first differentiateℓTA\\ell^\{\\mathrm\{TA\}\}with respect to the probabilities of all tokens, then push the derivative through the full softmax Jacobian, and finally read off the two cases of the claim\.

#### Step 1: the derivative with respect to probability\.

We treatℓTA\\ell^\{\\mathrm\{TA\}\}as a function of the full probability vector\(p⁡\(u\)\)u∈𝒱\\bigl\(p\(u\)\\bigr\)\_\{u\\in\\mathcal\{V\}\}rather than eliminatingptailp^\{\\mathrm\{tail\}\}\. A top\-kktokenu∈Su\\in Sappears only in its own summand, and a tail tokenu∉Su\\notin Sappears only insideptailp^\{\\mathrm\{tail\}\}, on which it acts with∂ptail/∂p⁡\(u\)=1\\partial p^\{\\mathrm\{tail\}\}/\\partial p\(u\)=1\. Hence

∂ℓTA∂p⁡\(u\)=\{log⁡p⁡\(u\)q⁡\(u\)\+1,u∈S,log⁡ptailqtail\+1,u∉S\.\\frac\{\\partial\\ell^\{\\mathrm\{TA\}\}\}\{\\partial p\(u\)\}=\\begin\{cases\}\\log\\dfrac\{p\(u\)\}\{q\(u\)\}\+1,&u\\in S,\\\\\[10\.00002pt\] \\log\\dfrac\{p^\{\\mathrm\{tail\}\}\}\{q^\{\\mathrm\{tail\}\}\}\+1,&u\\notin S\.\\end\{cases\}\(16\)

#### Step 2: pushing through the softmax Jacobian\.

The student distribution is a softmax over the full vocabulary,p⁡\(u\)=ezu/∑w∈𝒱ezwp\(u\)=e^\{z\_\{u\}\}/\\sum\_\{w\\in\\mathcal\{V\}\}e^\{z\_\{w\}\}, whose Jacobian is∂p\(u\)/∂zv=p\(u\)\(𝟏\[u=v\]−p\(v\)\)\\partial p\(u\)/\\partial z\_\{v\}=p\(u\)\\bigl\(\\mathbf\{1\}\[u=v\]\-p\(v\)\\bigr\)for allu,v∈𝒱u,v\\in\\mathcal\{V\}\. Combining it with Eq\. \([16](https://arxiv.org/html/2608.14728#A1.E16)\) through the chain rule gives

∂ℓTA∂zv=∑u∈𝒱∂ℓTA∂p⁡\(u\)​∂p⁡\(u\)∂zv=p⁡\(v\)​∂ℓTA∂p⁡\(u\)\|u=v−p⁡\(v\)​∑u∈𝒱p⁡\(u\)​∂ℓTA∂p⁡\(u\)\.\\frac\{\\partial\\ell^\{\\mathrm\{TA\}\}\}\{\\partial z\_\{v\}\}=\\sum\_\{u\\in\\mathcal\{V\}\}\\frac\{\\partial\\ell^\{\\mathrm\{TA\}\}\}\{\\partial p\(u\)\}\\,\\frac\{\\partial p\(u\)\}\{\\partial z\_\{v\}\}=p\(v\)\\,\\frac\{\\partial\\ell^\{\\mathrm\{TA\}\}\}\{\\partial p\(u\)\}\\bigg\|\_\{u=v\}\-p\(v\)\\sum\_\{u\\in\\mathcal\{V\}\}p\(u\)\\,\\frac\{\\partial\\ell^\{\\mathrm\{TA\}\}\}\{\\partial p\(u\)\}\.\(17\)Note that for the second term, we have

∑u∈𝒱p⁡\(u\)​∂ℓTA∂p⁡\(u\)=∑u∈Sp⁡\(u\)​\(log⁡p⁡\(u\)q⁡\(u\)\+1\)\+\(∑u∉Sp⁡\(u\)\)​\(log⁡ptailqtail\+1\)=ℓTA\+1,\\sum\_\{u\\in\\mathcal\{V\}\}p\(u\)\\,\\frac\{\\partial\\ell^\{\\mathrm\{TA\}\}\}\{\\partial p\(u\)\}=\\sum\_\{u\\in S\}p\(u\)\\left\(\\log\\frac\{p\(u\)\}\{q\(u\)\}\+1\\right\)\+\\left\(\\sum\_\{u\\notin S\}p\(u\)\\right\)\\\!\\left\(\\log\\frac\{p^\{\\mathrm\{tail\}\}\}\{q^\{\\mathrm\{tail\}\}\}\+1\\right\)=\\ell^\{\\mathrm\{TA\}\}\+1,where the second equality uses∑u∉Sp⁡\(u\)=ptail\\sum\_\{u\\notin S\}p\(u\)=p^\{\\mathrm\{tail\}\}to recover the tail term ofℓTA\\ell^\{\\mathrm\{TA\}\}, and∑u∈𝒱p⁡\(u\)=1\\sum\_\{u\\in\\mathcal\{V\}\}p\(u\)=1to collect the constant\.

#### Step 3: the gradient on each logit\.

SubstitutingℓTA\+1\\ell^\{\\mathrm\{TA\}\}\+1for the sum in Eq\. \([17](https://arxiv.org/html/2608.14728#A1.E17)\), we have

∂ℓTA∂zv=\{p⁡\(v\)​\(log⁡p⁡\(v\)q⁡\(v\)−ℓTA\),v∈S,p⁡\(v\)​\(log⁡ptailqtail−ℓTA\),v∉S,\\frac\{\\partial\\ell^\{\\mathrm\{TA\}\}\}\{\\partial z\_\{v\}\}=\\begin\{cases\}p\(v\)\\left\(\\log\\dfrac\{p\(v\)\}\{q\(v\)\}\-\\ell^\{\\mathrm\{TA\}\}\\right\),&v\\in S,\\\\\[11\.99998pt\] p\(v\)\\left\(\\log\\dfrac\{p^\{\\mathrm\{tail\}\}\}\{q^\{\\mathrm\{tail\}\}\}\-\\ell^\{\\mathrm\{TA\}\}\\right\),&v\\notin S,\\end\{cases\}\(18\)which establishes \([5](https://arxiv.org/html/2608.14728#S4.E5)\)\. ∎

### A\.4Proof of Proposition[4](https://arxiv.org/html/2608.14728#Thmproposition4)

###### Proof\.

Fix a stepttand drop the subscriptttas in Appendix[A\.3](https://arxiv.org/html/2608.14728#A1.SS3)\. The two loss functions are

ℓfull=∑v∈Sp⁡\(v\)​log​p⁡\(v\)q⁡\(v\)\+∑v∉Sp⁡\(v\)​log​p⁡\(v\)q⁡\(v\),ℓTA=∑v∈Sp⁡\(v\)​log​p⁡\(v\)q⁡\(v\)\+ptail​log​ptailqtail,\\ell^\{\\mathrm\{full\}\}=\\sum\_\{v\\in S\}p\(v\)\\log\\frac\{p\(v\)\}\{q\(v\)\}\+\\sum\_\{v\\notin S\}p\(v\)\\log\\frac\{p\(v\)\}\{q\(v\)\},\\qquad\\ell^\{\\mathrm\{TA\}\}=\\sum\_\{v\\in S\}p\(v\)\\log\\frac\{p\(v\)\}\{q\(v\)\}\+p\_\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\},where we have splitℓfull=DKL\(p∥q\)\\ell^\{\\mathrm\{full\}\}=D\_\{\\mathrm\{KL\}\}\(p\\\|q\)into the top\-kktokens and the tail tokens so that the two losses can be compared term by term\. The two objectives share the same top\-kkterm and differ only in how they treat the tail, so taking the difference cancels the top\-kkterm and leaves

ℓfull−ℓTA=∑v∉Sp⁡\(v\)​log⁡p⁡\(v\)q⁡\(v\)−ptail​log⁡ptailqtail\.\\ell^\{\\mathrm\{full\}\}\-\\ell^\{\\mathrm\{TA\}\}=\\sum\_\{v\\notin S\}p\(v\)\\log\\frac\{p\(v\)\}\{q\(v\)\}\-p\_\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\.\(19\)
It remains to compute the right\-hand side of \([19](https://arxiv.org/html/2608.14728#A1.E19)\)\. Substitutingp⁡\(v\)=ptail​p~​\(v\)p\(v\)=p\_\{\\mathrm\{tail\}\}\\,\\tilde\{p\}\(v\)andq⁡\(v\)=qtail​q~​\(v\)q\(v\)=q\_\{\\mathrm\{tail\}\}\\,\\tilde\{q\}\(v\)forv∉Sv\\notin Sinto its first term and splitting the logarithm gives

∑v∉Sp⁡\(v\)​log⁡p⁡\(v\)q⁡\(v\)\\displaystyle\\sum\_\{v\\notin S\}p\(v\)\\log\\frac\{p\(v\)\}\{q\(v\)\}=∑v∉Sptail​p~​\(v\)​log⁡ptail​p~​\(v\)qtail​q~​\(v\)\\displaystyle=\\sum\_\{v\\notin S\}p\_\{\\mathrm\{tail\}\}\\,\\tilde\{p\}\(v\)\\log\\frac\{p\_\{\\mathrm\{tail\}\}\\,\\tilde\{p\}\(v\)\}\{q\_\{\\mathrm\{tail\}\}\\,\\tilde\{q\}\(v\)\}=ptail​log⁡ptailqtail​∑v∉Sp~​\(v\)\+ptail​∑v∉Sp~​\(v\)​log⁡p~​\(v\)q~​\(v\)\\displaystyle=p\_\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\\sum\_\{v\\notin S\}\\tilde\{p\}\(v\)\+p\_\{\\mathrm\{tail\}\}\\sum\_\{v\\notin S\}\\tilde\{p\}\(v\)\\log\\frac\{\\tilde\{p\}\(v\)\}\{\\tilde\{q\}\(v\)\}=ptaillogptailqtail\+ptailDKL\(p~∥q~\),\\displaystyle=p\_\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\+p\_\{\\mathrm\{tail\}\}\\,D\_\{\\mathrm\{KL\}\}\(\\tilde\{p\}\\\|\\tilde\{q\}\),where the last equality uses thatp~\\tilde\{p\}is a probability distribution over the tail tokens,∑v∉Sp~​\(v\)=1\\sum\_\{v\\notin S\}\\tilde\{p\}\(v\)=1, and recognizes the remaining sum asDKL\(p~∥q~\)D\_\{\\mathrm\{KL\}\}\(\\tilde\{p\}\\\|\\tilde\{q\}\)\.

Substituting this back into \([19](https://arxiv.org/html/2608.14728#A1.E19)\), the termsptail​log⁡ptailqtailp\_\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}cancel and we obtain

ℓfull−ℓTA=ptailDKL\(p~∥q~\)\.\\ell^\{\\mathrm\{full\}\}\-\\ell^\{\\mathrm\{TA\}\}=p\_\{\\mathrm\{tail\}\}\\,D\_\{\\mathrm\{KL\}\}\(\\tilde\{p\}\\\|\\tilde\{q\}\)\.∎

### A\.5Proof of Proposition[5](https://arxiv.org/html/2608.14728#Thmproposition5)

###### Proof\.

Fix stepttand drop the subscriptttas in Appendix[A\.3](https://arxiv.org/html/2608.14728#A1.SS3): writeS=StkS=S\_\{t\}^\{k\},S\+=St\+S^\{\+\}=S\_\{t\}^\{\+\},p​\(v\)=pt​\(v\)p\(v\)=p\_\{t\}\(v\),q​\(v\)=qt​\(v\)q\(v\)=q\_\{t\}\(v\),y^=y^t\\hat\{y\}=\\hat\{y\}\_\{t\},ptail=pttailp\_\{\\mathrm\{tail\}\}=p\_\{t\}^\{\\mathrm\{tail\}\}, andqtail=qttailq\_\{\\mathrm\{tail\}\}=q\_\{t\}^\{\\mathrm\{tail\}\}\. Recall that the stop\-gradient operatorsg⁡\(⋅\)\\operatorname\{sg\}\(\\cdot\)acts as the identity in the forward pass and has zero derivative, i\.e\.,sg⁡\(x\)=x\\operatorname\{sg\}\(x\)=xin value and∇θ​sg​\(x\)=0\\nabla\_\{\\theta\}\\operatorname\{sg\}\(x\)=0\. Throughout, the expectation is taken with respect toy^∼p\\hat\{y\}\\sim p, where the sampling distribution is fixed at the current parameters\. We will repeatedly use the identity

𝔼y^∼p\[𝟏\[y^∉S\]\]=∑v∉Sp\(v\)=ptail\.\\mathbb\{E\}\_\{\\hat\{y\}\\sim p\}\\bigl\[\\mathbf\{1\}\[\\hat\{y\}\\notin S\]\\bigr\]=\\sum\_\{v\\notin S\}p\(v\)=p\_\{\\mathrm\{tail\}\}\.\(20\)

#### Unbiasedness in value\.

Without considering the gradient,p⁡\(y^\)sg⁡\(p⁡\(y^\)\)=1\\frac\{p\(\\hat\{y\}\)\}\{\\operatorname\{sg\}\(p\(\\hat\{y\}\)\)\}=1, so

ℓSC​\-​TA\\displaystyle\\ell^\{\\mathrm\{SC\\text\{\-\}TA\}\}=∑v∈S\+p\(v\)sg\(logp⁡\(v\)q⁡\(v\)\)\+𝟏\[y^∉S\]p⁡\(y^\)sg⁡\(p⁡\(y^\)\)sg\(logp⁡\(y^\)/ptailq⁡\(y^\)/qtail\)\\displaystyle=\\sum\_\{v\\in S^\{\+\}\}p\(v\)\\operatorname\{sg\}\\\!\\left\(\\log\\frac\{p\(v\)\}\{q\(v\)\}\\right\)\+\\mathbf\{1\}\[\\hat\{y\}\\notin S\]\\,\\frac\{p\(\\hat\{y\}\)\}\{\\operatorname\{sg\}\(p\(\\hat\{y\}\)\)\}\\operatorname\{sg\}\\\!\\left\(\\log\\frac\{p\(\\hat\{y\}\)/p\_\{\\mathrm\{tail\}\}\}\{q\(\\hat\{y\}\)/q\_\{\\mathrm\{tail\}\}\}\\right\)=∑v∈Sp\(v\)logp⁡\(v\)q⁡\(v\)\+ptaillogptailqtail\+𝟏\[y^∉S\]\(logp⁡\(y^\)q⁡\(y^\)−logptailqtail\),\\displaystyle=\\sum\_\{v\\in S\}p\(v\)\\log\\frac\{p\(v\)\}\{q\(v\)\}\+p\_\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\+\\mathbf\{1\}\[\\hat\{y\}\\notin S\]\\left\(\\log\\frac\{p\(\\hat\{y\}\)\}\{q\(\\hat\{y\}\)\}\-\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\\right\),\(21\)where the second equality expandsS\+=S∪\{vtail\}S^\{\+\}=S\\cup\\\{v\_\{\\mathrm\{tail\}\}\\\}usingp⁡\(vtail\)=ptailp\(v\_\{\\mathrm\{tail\}\}\)=p\_\{\\mathrm\{tail\}\}andq⁡\(vtail\)=qtailq\(v\_\{\\mathrm\{tail\}\}\)=q\_\{\\mathrm\{tail\}\}\. The first two terms are deterministic given the prefix\. For the third term, taking the expectation overy^∼p\\hat\{y\}\\sim pand using Eq\. \([20](https://arxiv.org/html/2608.14728#A1.E20)\) gives

𝔼y^∼p\[𝟏\[y^∉S\]\(logp⁡\(y^\)q⁡\(y^\)−logptailqtail\)\]=∑v∉Sp\(v\)logp⁡\(v\)q⁡\(v\)−ptaillogptailqtail\.\\mathbb\{E\}\_\{\\hat\{y\}\\sim p\}\\\!\\left\[\\mathbf\{1\}\[\\hat\{y\}\\notin S\]\\left\(\\log\\frac\{p\(\\hat\{y\}\)\}\{q\(\\hat\{y\}\)\}\-\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\\right\)\\right\]=\\sum\_\{v\\notin S\}p\(v\)\\log\\frac\{p\(v\)\}\{q\(v\)\}\-p\_\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\.Summing the three terms, the two occurrences ofptail​log⁡ptailqtailp\_\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}cancel and we obtain

𝔼y^∼p\[ℓSC​\-​TA\]=∑v∈Sp\(v\)logp⁡\(v\)q⁡\(v\)\+∑v∉Sp\(v\)logp⁡\(v\)q⁡\(v\)=DKL\(p∥q\),\\mathbb\{E\}\_\{\\hat\{y\}\\sim p\}\\bigl\[\\ell^\{\\mathrm\{SC\\text\{\-\}TA\}\}\\bigr\]=\\sum\_\{v\\in S\}p\(v\)\\log\\frac\{p\(v\)\}\{q\(v\)\}\+\\sum\_\{v\\notin S\}p\(v\)\\log\\frac\{p\(v\)\}\{q\(v\)\}=D\_\{\\mathrm\{KL\}\}\(p\\,\\\|\\,q\),which establishes the unbiasedness in value\.

#### Unbiasedness in gradient\.

We first compute the gradient of the exact full\-vocabulary reverse KL\. Sinceqqis fixed,

∇θDKL\(p∥q\)=∑v∈𝒱∇θp\(v\)\(logp⁡\(v\)q⁡\(v\)\+1\)=∑v∈𝒱∇θp\(v\)logp⁡\(v\)q⁡\(v\),\\nabla\_\{\\theta\}\\,D\_\{\\mathrm\{KL\}\}\(p\\,\\\|\\,q\)=\\sum\_\{v\\in\\mathcal\{V\}\}\\nabla\_\{\\theta\}p\(v\)\\left\(\\log\\frac\{p\(v\)\}\{q\(v\)\}\+1\\right\)=\\sum\_\{v\\in\\mathcal\{V\}\}\\nabla\_\{\\theta\}p\(v\)\\,\\log\\frac\{p\(v\)\}\{q\(v\)\},where the last equality uses∑v∈𝒱∇θp\(v\)=∇θ∑v∈𝒱p\(v\)=∇θ1=0\\sum\_\{v\\in\\mathcal\{V\}\}\\nabla\_\{\\theta\}p\(v\)=\\nabla\_\{\\theta\}\\sum\_\{v\\in\\mathcal\{V\}\}p\(v\)=\\nabla\_\{\\theta\}1=0\.

We now differentiateℓSC​\-​TA\\ell^\{\\mathrm\{SC\\text\{\-\}TA\}\}\. Since the log\-ratios are wrapped in stop\-gradients, only the leading probabilities carry gradients in the first term\. Note thatptail=∑v∉Sp⁡\(v\)p\_\{\\mathrm\{tail\}\}=\\sum\_\{v\\notin S\}p\(v\)depends onθ\\theta, so∇θptail=∑v∉S∇θp​\(v\)\\nabla\_\{\\theta\}p\_\{\\mathrm\{tail\}\}=\\sum\_\{v\\notin S\}\\nabla\_\{\\theta\}p\(v\), and therefore

∇θ∑v∈S\+p\(v\)sg\(logp⁡\(v\)q⁡\(v\)\)=∑v∈S∇θp\(v\)logp⁡\(v\)q⁡\(v\)\+\(∑v∉S∇θp\(v\)\)logptailqtail\.\\nabla\_\{\\theta\}\\sum\_\{v\\in S^\{\+\}\}p\(v\)\\operatorname\{sg\}\\\!\\left\(\\log\\frac\{p\(v\)\}\{q\(v\)\}\\right\)=\\sum\_\{v\\in S\}\\nabla\_\{\\theta\}p\(v\)\\,\\log\\frac\{p\(v\)\}\{q\(v\)\}\+\\Bigl\(\\sum\_\{v\\notin S\}\\nabla\_\{\\theta\}p\(v\)\\Bigr\)\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\.For the second term, only the numeratorp⁡\(y^\)p\(\\hat\{y\}\)inp⁡\(y^\)sg⁡\(p⁡\(y^\)\)\\frac\{p\(\\hat\{y\}\)\}\{\\operatorname\{sg\}\(p\(\\hat\{y\}\)\)\}carries gradients, so

∇θ\[𝟏\[y^∉S\]p⁡\(y^\)sg⁡\(p⁡\(y^\)\)sg\(logp⁡\(y^\)/ptailq⁡\(y^\)/qtail\)\]=𝟏\[y^∉S\]∇θp​\(y^\)p⁡\(y^\)\(logp⁡\(y^\)q⁡\(y^\)−logptailqtail\)\.\\nabla\_\{\\theta\}\\\!\\left\[\\mathbf\{1\}\[\\hat\{y\}\\notin S\]\\,\\frac\{p\(\\hat\{y\}\)\}\{\\operatorname\{sg\}\(p\(\\hat\{y\}\)\)\}\\operatorname\{sg\}\\\!\\left\(\\log\\frac\{p\(\\hat\{y\}\)/p\_\{\\mathrm\{tail\}\}\}\{q\(\\hat\{y\}\)/q\_\{\\mathrm\{tail\}\}\}\\right\)\\right\]=\\mathbf\{1\}\[\\hat\{y\}\\notin S\]\\,\\frac\{\\nabla\_\{\\theta\}p\(\\hat\{y\}\)\}\{p\(\\hat\{y\}\)\}\\left\(\\log\\frac\{p\(\\hat\{y\}\)\}\{q\(\\hat\{y\}\)\}\-\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\\right\)\.Taking the expectation overy^∼p\\hat\{y\}\\sim pof the second term gives

𝔼y^∼p\[𝟏\[y^∉S\]∇θp​\(y^\)p⁡\(y^\)\(logp⁡\(y^\)q⁡\(y^\)−logptailqtail\)\]\\displaystyle\\mathbb\{E\}\_\{\\hat\{y\}\\sim p\}\\\!\\left\[\\mathbf\{1\}\[\\hat\{y\}\\notin S\]\\,\\frac\{\\nabla\_\{\\theta\}p\(\\hat\{y\}\)\}\{p\(\\hat\{y\}\)\}\\left\(\\log\\frac\{p\(\\hat\{y\}\)\}\{q\(\\hat\{y\}\)\}\-\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\\right\)\\right\]=∑v∉Sp⁡\(v\)⋅∇θp​\(v\)p⁡\(v\)​\(log⁡p⁡\(v\)q⁡\(v\)−log⁡ptailqtail\)=∑v∉S∇θp​\(v\)​log​p⁡\(v\)q⁡\(v\)−\(∑v∉S∇θp​\(v\)\)​log​ptailqtail\.\\displaystyle=\\sum\_\{v\\notin S\}p\(v\)\\cdot\\frac\{\\nabla\_\{\\theta\}p\(v\)\}\{p\(v\)\}\\left\(\\log\\frac\{p\(v\)\}\{q\(v\)\}\-\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\\right\)=\\sum\_\{v\\notin S\}\\nabla\_\{\\theta\}p\(v\)\\,\\log\\frac\{p\(v\)\}\{q\(v\)\}\-\\Bigl\(\\sum\_\{v\\notin S\}\\nabla\_\{\\theta\}p\(v\)\\Bigr\)\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}\.Combining the two terms, the contributions involvinglog⁡ptailqtail\\log\\frac\{p\_\{\\mathrm\{tail\}\}\}\{q\_\{\\mathrm\{tail\}\}\}cancel exactly and we obtain

𝔼y^∼p\[∇θℓSC​\-​TA\]=∑v∈S∇θp\(v\)logp⁡\(v\)q⁡\(v\)\+∑v∉S∇θp\(v\)logp⁡\(v\)q⁡\(v\)=∇θDKL\(p∥q\),\\mathbb\{E\}\_\{\\hat\{y\}\\sim p\}\\bigl\[\\nabla\_\{\\theta\}\\,\\ell^\{\\mathrm\{SC\\text\{\-\}TA\}\}\\bigr\]=\\sum\_\{v\\in S\}\\nabla\_\{\\theta\}p\(v\)\\,\\log\\frac\{p\(v\)\}\{q\(v\)\}\+\\sum\_\{v\\notin S\}\\nabla\_\{\\theta\}p\(v\)\\,\\log\\frac\{p\(v\)\}\{q\(v\)\}=\\nabla\_\{\\theta\}\\,D\_\{\\mathrm\{KL\}\}\(p\\,\\\|\\,q\),which establishes the unbiasedness in gradient\. This completes the proof\. ∎

## Appendix BDetailed Analysis of the tail probability increase

In this section, we conduct several controlled experiments to identify when the tail probability increase occurs\. We find that it arises under two conditions: a large capability gap between the student and teacher, and an insufficiently largekk\.

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\.#### The tail probability increase occurs when the student–teacher capability gap is large\.

We conduct an ablation study to demonstrate that the tail probability increase occurs when the capability gap between the student and teacher is large\. Specifically, we fix Qwen3\-8B as the teacher and compare two students, Qwen3\-1\.7B\-Base and Qwen3\-8B\-Base, where the latter has a smaller capability gap to the teacher\. For this ablation study, we set the maximum response length to 4096\.

In Figure[7](https://arxiv.org/html/2608.14728#A2.F7), we visualize the student’s and teacher’s tail probabilities on student\-generated prefixes, together with the student’s entropy\. The tail probability increase emerges only under the large capability gap: for Qwen3\-1\.7B\-Base, the student’s tail probability and entropy steadily increase over training steps, whereas for Qwen3\-8B\-Base, both remain low and stable\. Overall, the tail probability increase arises when the student–teacher capability gap is large and is mitigated as the gap narrows\.

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 whenk≤64k\\leq 64\.
#### The tail probability increase occurs whenkkis not sufficiently large\.

We conduct an ablation study on howkkaffects the tail probability increase by fixing the student and teacher, and varying onlykk\. Specifically, we use Qwen2\.5\-Math\-1\.5B as the student and DeepSeek\-R1\-Distill\-Qwen\-7B as the teacher, and sweepk∈\{8,16,32,64\}k\\in\\\{8,16,32,64\\\}\. For this ablation study, we set the training max response length to 4096\.

In Figure[8](https://arxiv.org/html/2608.14728#A2.F8), we visualize the student’s and teacher’s tail probability on student\-generated prefixes, together with the student’s entropy, across different values ofkk\. The tail probability increase emerges whenkkis not sufficiently large \(k≤64k\\leq 64\): over training steps, the student’s tail probability steadily increases together with its entropy\. The effect becomes stronger askkdecreases, with smallerkkyielding higher entropy and a larger increase in tail probability\. In principle, this issue vanishes askkapproaches the vocabulary size, since the top\-kkobjective then reduces to the full\-vocabulary OPD, under which the tail probability is explicitly matched to the teacher’s\. This behavior is expected from Proposition[2](https://arxiv.org/html/2608.14728#Thmproposition2): the magnitude of the single\-step increase is proportional topttail​\(1−pttail\)p\_\{t\}^\{\\mathrm\{tail\}\}\(1\-p\_\{t\}^\{\\mathrm\{tail\}\}\), which shrinks aspttail→0p\_\{t\}^\{\\mathrm\{tail\}\}\\to 0\. Since a largerkkyields a smaller initial tail probability, each update movespttailp\_\{t\}^\{\\mathrm\{tail\}\}less, mitigating the compounding increase over training\. Overall, the tail probability increase arises whenkkis not sufficiently large and is mitigated askkincreases\.

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\.
#### The tail probability increase becomes more pronounced with longer responses\.

We study how the response length affects the tail probability increase by fixing the student and teacher, and varying only the maximum response length\. Specifically, we use Qwen2\.5\-7B\-Instruct as the student and OpenThinker3\-7B as the teacher, and compare maximum response lengths of 4096 and 7168 tokens\.

In Figure[9](https://arxiv.org/html/2608.14728#A2.F9), we visualize the student’s and teacher’s tail probabilities on student\-generated prefixes, together with the student’s entropy, under the two length settings\. The results show that the tail probability increase is amplified under the longer response length: the student’s tail probability and entropy grow faster and reach higher values throughout training\. We attribute this to error compounding along longer trajectories: as the student’s tail probability rises, sampling more tokens per response increases the chance of drifting outside the teacher’s top\-kktokens, driving the student toward uncertain prefixes where the teacher’s supervision is unreliable and further inflating the tail probability\. Overall, the tail probability increase is exacerbated by longer response lengths, suggesting that the issue is particularly concerning for long\-horizon tasks\.

#### The tail probability increase does not occur with the student top\-kkOPD\.

Prior work also adopts the student’s top\-kktokens as the support\([21](https://arxiv.org/html/2608.14728#bib.bib9)\)\. We examine this choice with Qwen3\-1\.7B\-Base as the student and Qwen3\-8B as the teacher, varying only the support betweenTopK⁡\(qt,k\)\\operatorname\{TopK\}\(q\_\{t\},k\)andTopK⁡\(pt,k\)\\operatorname\{TopK\}\(p\_\{t\},k\)\. In each case, the tail probabilities of both models are computed with respect to the support in use, i\.e\. the teacher’s top\-kktokens for the teacher top\-kkvariant and the student’s top\-kktokens for the student top\-kkvariant\.

As shown in Figure[10](https://arxiv.org/html/2608.14728#A2.F10), the tail probability increase is specific to the teacher top\-kksupport\. Two mechanisms explain this\. First, among all token sets of sizekk,TopK⁡\(pt,k\)\\operatorname\{TopK\}\(p\_\{t\},k\)is the one that leaves the least probability mass outside it, so the student’s tail probability starts near zero\. By Proposition[2](https://arxiv.org/html/2608.14728#Thmproposition2), the per\-step change is proportional topttail​\(1−pttail\)p\_\{t\}^\{\\mathrm\{tail\}\}\(1\-p\_\{t\}^\{\\mathrm\{tail\}\}\), and each update therefore movespttailp\_\{t\}^\{\\mathrm\{tail\}\}only slightly\. Second, the student top\-kksupport is changing given a prefix: it is recomputed from the student at every step, so a token whose probability is pushed down simply drops out of the support and is replaced by another token that the student now ranks in its topkk\. The support therefore always consists of the student’skkmost probable tokens, and the probability mass it covers cannot leak away\. The teacher top\-kksupport, in contrast, is fixed by the teacher and does not follow the student, so the mass pushed outside it stays outside and accumulates across steps\.

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\-kksupports\.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\-kksupport\.
#### The teacher top\-kkOPD yields better downstream performance than the student top\-kkOPD\.

Although the student top\-kksupport avoids the tail probability increase, it underperforms the teacher top\-kksupport on downstream benchmarks \(Table[3](https://arxiv.org/html/2608.14728#A2.T3)\)\. We attribute this to two factors\. First, as shown in Figure[11](https://arxiv.org/html/2608.14728#A2.F11), the average rollout length under the student top\-kksupport grows monotonically and saturates at the maximum response length of71687168tokens within the first5050steps, while the teacher top\-kkvariant maintains a healthy response length\. Second, the teacher top\-kksupport directly covers the teacher’s high\-probability tokens, whereas the student’s top\-kktokens may carry little teacher probability, allowing the objective to be reduced without moving the student toward the teacher’s high\-probability modes\. Given the superior downstream performance, we mainly study the teacher top\-kkOPD in this work\.

Figure 11:Average response length across training steps under the student and teacher top\-kksupports\.The student and teacher models are Qwen3\-1\.7B\-Base and Qwen3\-8B, respectively\. Under the student top\-kksupport, the response length saturates at the maximum response length of71687168tokens, while the teacher top\-kkvariant maintains a healthy response length\.Table 3:Performance of normalized top\-kkOPD under the student and teacher top\-kksupports\.TheAvg\.column represents the macro\-average across all six math reasoning benchmarks\. Best results are shown inbold\.

## Appendix CCompared Methods

#### Sampled\-token OPD\.

Sampled\-token OPD\([22](https://arxiv.org/html/2608.14728#bib.bib1);[36](https://arxiv.org/html/2608.14728#bib.bib11)\)supervises only the tokeny^t∼pt\\hat\{y\}\_\{t\}\\sim p\_\{t\}sampled by the student:

ℓtsample=log⁡pt​\(y^t\)−log⁡qt​\(y^t\)\.\\ell\_\{t\}^\{\\mathrm\{sample\}\}=\\log p\_\{t\}\(\\hat\{y\}\_\{t\}\)\-\\log q\_\{t\}\(\\hat\{y\}\_\{t\}\)\.It is an unbiased estimate of the full\-vocabulary reverse KL divergence and requires only the sampled token’s log\-probability from the teacher, but it discards the dense information over the remaining vocabulary and suffers from high variance\.

#### Unnormalized top\-kkOPD\.

Unnormalized top\-kkOPD restricts the divergence computation to the teacher’s top\-kktokensStk=TopK⁡\(qt,k\)S\_\{t\}^\{k\}=\\operatorname\{TopK\}\(q\_\{t\},k\)while keeping the original probabilities:

ℓtunnorm=∑v∈Stkpt​\(v\)​log⁡pt​\(v\)qt​\(v\)\.\\ell\_\{t\}^\{\\mathrm\{unnorm\}\}=\\sum\_\{v\\in S\_\{t\}^\{k\}\}p\_\{t\}\(v\)\\log\\frac\{p\_\{t\}\(v\)\}\{q\_\{t\}\(v\)\}\.It provides dense supervision overStkS\_\{t\}^\{k\}, but is not a well\-defined divergence\.

#### Normalized top\-kkOPD\.

Normalized top\-kkOPD\([21](https://arxiv.org/html/2608.14728#bib.bib9);[7](https://arxiv.org/html/2608.14728#bib.bib3)\)instead renormalizes both distributions onStkS\_\{t\}^\{k\}and minimizes the resulting subset reverse KL divergence:

ℓtnorm=DKL\(p¯t\(Stk\)∥q¯t\(Stk\)\)=∑v∈Stkp¯t\(Stk\)\(v\)logp¯t\(Stk\)​\(v\)q¯t\(Stk\)​\(v\)\.\\ell\_\{t\}^\{\\mathrm\{norm\}\}=D\_\{\\mathrm\{KL\}\}\\\!\\bigl\(\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\\,\\\|\\,\\bar\{q\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\\bigr\)=\\sum\_\{v\\in S\_\{t\}^\{k\}\}\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\(v\)\\log\\frac\{\\bar\{p\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\(v\)\}\{\\bar\{q\}\_\{t\}^\{\(S\_\{t\}^\{k\}\)\}\(v\)\}\.It aligns the student’s relative shape with the teacher’s overStkS\_\{t\}^\{k\}, but the normalization discards the tail probability, which, as we show in Section[3\.1](https://arxiv.org/html/2608.14728#S3.SS1), steadily increases the student’s tail probability and entropy\.

## Appendix DImplementation Details

### D\.1Experimental Details

We implement all methods using the VERL framework\([32](https://arxiv.org/html/2608.14728#bib.bib18)\)and conduct experiments with 6 NVIDIA Pro 6000 GPUs, 8 NVIDIA A100 80 GB GPUs, or 32 NVIDIA A100 40 GB GPUs\. Unless otherwise specified, all experiments use the default settings and hyperparameters listed in Table[4](https://arxiv.org/html/2608.14728#A4.T4)\. For evaluation, we adopt the same prompt as[42](https://arxiv.org/html/2608.14728#bib.bib19), which is shown in Prompt[D\.1](https://arxiv.org/html/2608.14728#A4.SS1)\. For experiments using Llama\-3\.1\-8B as the student, we adopt the tokenizer and chat template of DeepSeek\-R1\-Distill\-Llama\-8B and synchronize the corresponding vocabulary and special\-token configurations\.

Table 4:Default training and evaluation settings\.CategoryItemValueTrainingTraining temperature1\.0Global batch size72Mini batch size36Rollout number4kk\(Number of Top\-kkTokens\)16Top\-pp1\.0Max prompt length1024Max response length7168Learning rate1e\-6Training step300loss aggregationtoken\-meanoptimizerAdamWEvaluationTemperature0\.7Top\-pp0\.95Max new tokens8192Your task is to follow a systematic, thorough reasoning process before providing the final solution\. This involves analyzing, summarizing, exploring, reassessing, and refining your thought process through multiple iterations\. Structure your response into two sections: Thought and Solution\. In the Thought section, present your reasoning using the format: ”`<think\>\\n`thoughts`</think\>\\n`”\. Each thought should include detailed analysis, brainstorming, verification, and refinement of ideas\. After ”`</think\>\\n`” in the Solution section, provide the final, logical, and accurate answer, clearly derived from the exploration in the Thought section\. If applicable, include the answer in`\\boxed\{\}`for closed\-form results like multiple choices or mathematical solutions\. User:`\{QUESTION\}` Assistant:

### D\.2Numerically Stable Computation of TA\-OPD

The per\-token TA\-OPD loss in Eq\. \([4](https://arxiv.org/html/2608.14728#S4.E4)\) depends on the tail log\-probabilitieslog⁡pttail\\log p\_\{t\}^\{\\mathrm\{tail\}\}andlog⁡qttail\\log q\_\{t\}^\{\\mathrm\{tail\}\}\. Below, we describe a failure case of a naive implementation and how we implement the loss\.

#### Numerical issue of a naive implementation\.

The inference engine vLLM only provides access to the teacher’s top\-kklog\-probabilitieslog⁡qt​\(v\)\\log q\_\{t\}\(v\), not the probabilities themselves\. A naive implementation of TA\-OPD loss is to exponentiate these log\-probabilities back to probability space, sum them to obtain the total probability of the top\-kktokens1−qttail=∑v∈Stkqt​\(v\)1\-q\_\{t\}^\{\\mathrm\{tail\}\}=\\sum\_\{v\\in S\_\{t\}^\{k\}\}q\_\{t\}\(v\), and assign the remaining tail probabilityqttailq\_\{t\}^\{\\mathrm\{tail\}\}tovtailv\_\{\\mathrm\{tail\}\}\. However, this is unstable\. When the teacher’s tail probability is close to zero, numerical error can make the sum∑v∈Stkqt​\(v\)\\sum\_\{v\\in S\_\{t\}^\{k\}\}q\_\{t\}\(v\)exceed one\. The tail probabilityqttailq\_\{t\}^\{\\mathrm\{tail\}\}then becomes negative, renderinglog⁡qttail\\log q\_\{t\}^\{\\mathrm\{tail\}\}ill\-defined \(−∞\-\\inftyor NaN\) and producing NaN gradients\.

#### Our log\-space implementation\.

To address the numerical issue, our implementation avoids formingpttailp\_\{t\}^\{\\mathrm\{tail\}\}andqttailq\_\{t\}^\{\\mathrm\{tail\}\}in probability space and carries out the whole computation in log space\. We provide the pseudo code of our implementation in Listing[1](https://arxiv.org/html/2608.14728#LST1)\. In particular, we first obtain the log total probability of the top\-kktokens directly from the top\-kklog\-probabilities via a log\-sum\-exp,

log⁡\(1−pttail\)=logsumexpv∈Stk⁡log⁡pt​\(v\),log⁡\(1−qttail\)=logsumexpv∈Stk⁡log⁡qt​\(v\),\\log\\bigl\(1\-p\_\{t\}^\{\\mathrm\{tail\}\}\\bigr\)=\\operatorname\{logsumexp\}\_\{v\\in S\_\{t\}^\{k\}\}\\log p\_\{t\}\(v\),\\qquad\\log\\bigl\(1\-q\_\{t\}^\{\\mathrm\{tail\}\}\\bigr\)=\\operatorname\{logsumexp\}\_\{v\\in S\_\{t\}^\{k\}\}\\log q\_\{t\}\(v\),which is numerically stable and never overflows\. To further rule out the boundary casepttail→0p\_\{t\}^\{\\mathrm\{tail\}\}\\to 0, we clamp the log total probability such thatlog⁡\(1−pttail\)≤−ϵ\\log\(1\-p\_\{t\}^\{\\mathrm\{tail\}\}\)\\leq\-\\epsilon, ensuring a strictly positive tail probability and a finite log\-domain computation\. The tail log\-probability termlog⁡pttail\\log p\_\{t\}^\{\\mathrm\{tail\}\}is then computed using thelog1mexpprimitive, which directly evaluateslog⁡\(1−exp⁡\(a\)\)\\log\(1\-\\exp\(a\)\)in log space:

log⁡pttail=log1mexp⁡\(log⁡\(1−pttail\)\)\.\\log p\_\{t\}^\{\\mathrm\{tail\}\}=\\operatorname\{log1mexp\}\\\!\\bigl\(\\log\(1\-p\_\{t\}^\{\\mathrm\{tail\}\}\)\\bigr\)\.The teacher termlog⁡qttail\\log q\_\{t\}^\{\\mathrm\{tail\}\}is computed analogously\. In this way, we avoid the numerical issues of the naive implementation\.

Listing 1:Pseudocode for the TA\-OPD loss implementation\.defcompute\_taopd\_loss\(student\_topk\_log\_probs,teacher\_topk\_log\_probs,eps\):

student\_log\_topk\_prob=logsumexp\(student\_topk\_log\_probs,dim=\-1\)

teacher\_log\_topk\_prob=logsumexp\(teacher\_topk\_log\_probs,dim=\-1\)

student\_tail\_log\_prob=log1mexp\(clamp\_max\(student\_log\_topk\_prob,\-eps\)\)

teacher\_tail\_log\_prob=log1mexp\(clamp\_max\(teacher\_log\_topk\_prob,\-eps\)\)

student\_log\_p=cat\(\[student\_topk\_log\_probs,student\_tail\_log\_prob\],dim=\-1\)

teacher\_log\_q=cat\(\[teacher\_topk\_log\_probs,teacher\_tail\_log\_prob\],dim=\-1\)

taopd\_loss=sum\(exp\(student\_log\_p\)\*\(student\_log\_p\-teacher\_log\_q\),dim=\-1\)

returntaopd\_loss

## Appendix EExtensive Study

### E\.1Synthetic experiment

To provide an intuitive understanding of different OPD objectives, we construct a synthetic experiment based on a 30\-armed bandit\. Both the student and teacher policies are parameterized by vectors inℝ1×30\\mathbb\{R\}^\{1\\times 30\}, and the corresponding probability distributions are obtained by applying the softmax function to these vectors\. The teacher policy is kept fixed during training and is defined over the discrete class space𝒱=\{1,…,30\}\\mathcal\{V\}=\\\{1,\\dots,30\\\}as the following bimodal distribution:

q⁡\(v\)∝exp⁡\(−\(v−10\)22⋅22\)\+0\.88⋅exp⁡\(−\(v−20\)22⋅22\),v∈𝒱\.q\(v\)\\propto\\exp\\\!\\left\(\-\\frac\{\(v\-10\)^\{2\}\}\{2\\cdot 2^\{2\}\}\\right\)\\;\+\\;0\.88\\cdot\\exp\\\!\\left\(\-\\frac\{\(v\-20\)^\{2\}\}\{2\\cdot 2^\{2\}\}\\right\),\\quad v\\in\\mathcal\{V\}\.The student policy is parameterized as a categorical distribution over𝒱\\mathcal\{V\}:

pθ​\(v\)=softmax​\(zθ\)v,zθ∈ℝ30,p\_\{\\theta\}\(v\)=\\mathrm\{softmax\}\(z\_\{\\theta\}\)\_\{v\},\\quad z\_\{\\theta\}\\in\\mathbb\{R\}^\{30\},wherezθz\_\{\\theta\}is initialized from𝒩⁡\(0,0\.012\)\\mathcal\{N\}\(0,0\.01^\{2\}\)\. The student is optimized for 20000 steps using AdamW with learning rate1×10−31\\times 10^\{\-3\}\. We compare four objectives: full\-vocabulary OPD, normalized top\-kkOPD, TA\-OPD, and sample\-corrected TA\-OPD\. For top\-kkOPD,kkis set to88\.

Figure 12:Visualization of the student policies under different OPD objectives\. Normalized top\-kkOPD \(Red\) only matches the relative shape of the teacher distribution over the top\-kktokens, but fails to align the tail probability\. In contrast, TA\-OPD \(Blue\) and sample\-corrected TA\-OPD \(Purple\) closely recover the bimodal teacher distribution\.#### TA\-OPD matches both the shape and tail probability of the teacher policy\.

Figure[12](https://arxiv.org/html/2608.14728#A5.F12)compares the final student policies trained with different OPD objectives\. TA\-OPD closely recovers the bimodal teacher distribution\. While normalized top\-kkOPD matches the relative shape over the teacher’s top\-kktokens, it assigns a substantially larger tail probability than the teacher\. Overall, TA\-OPD preserves the missing tail probability information and therefore provides a much closer approximation to the teacher policy than normalized top\-kkOPD\.

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

Figure[13](https://arxiv.org/html/2608.14728#A5.F13)further visualizes the optimization trajectories on the full\-vocabulary reverse\-KL landscape\. In particular, we collect the logit trajectories of all methods together with the teacher logits, center them in logit space, and project them onto the first two principal components\. We then evaluate the full\-vocabulary reverse KL on this two\-dimensional plane and overlay the trajectories of different objectives\. The results show that full\-vocabulary OPD and TA\-OPD move toward the same low\-loss region around the teacher policy, while normalized top\-kkOPD converges to a point that remains far from the teacher under the full\-vocabulary KL\. Overall, TA\-OPD’s objective provides a better top\-kkestimate of the full\-vocabulary reverse KL than the normalized objective\.

### E\.2Top\-kkOverlap Ratio of Different OPD Methods

[21](https://arxiv.org/html/2608.14728#bib.bib9)find that the top\-kkoverlap ratio between the student and the teacher predicts the success of OPD well\. Here we compare the two objectives under this metric\.

#### Setup\.

On each student\-generated prefix, the top\-kkoverlap ratio is the fraction of the teacher’s top\-kktokens that also fall in the student’s top\-kktokens,

Overlaptk=\|TopK⁡\(pt,k\)∩TopK⁡\(qt,k\)\|k,\\mathrm\{Overlap\}\_\{t\}^\{k\}=\\frac\{\\bigl\|\\operatorname\{TopK\}\(p\_\{t\},k\)\\cap\\operatorname\{TopK\}\(q\_\{t\},k\)\\bigr\|\}\{k\},which we average over all tokens in a training batch\. A higher ratio means the student and the teacher agree on which tokens are plausible, so the top\-kkobjective supervises a token set that is meaningful to the student\. We use the three student–teacher pairs of Section[3\.2](https://arxiv.org/html/2608.14728#S3.SS2)and report the ratio across training steps\.

Figure 14:Top\-kkoverlap ratio between the student and the teacher across training steps\.TA\-OPD achieves higher overlap ratios than normalized top\-kkOPD across different model pairs\.
#### TA\-OPD attains a higher top\-kkoverlap ratio\.

Figure[14](https://arxiv.org/html/2608.14728#A5.F14)compares the overlap ratio of the two objectives\. TA\-OPD increases the overlap ratio monotonically and keeps it stable throughout training\. In contrast, normalized top\-kkOPD is consistently lower: for Qwen2\.5\-7B\-Instruct the ratio peaks early and then collapses, and for Llama\-3\.1\-8B it recovers only slowly and remains below TA\-OPD for the entire run\. Overall, TA\-OPD keeps the student and the teacher aligned on the top\-kktokens, which is consistent with its stronger downstream accuracy\.

### E\.3Computational Overhead of TA\-OPD

TA\-OPD introduces no additional computational overhead compared with normalized top\-kkOPD\. Both methods require exactly the same teacher\-side information: the log\-probabilities of the teacher’s top\-kktokens\. Given these quantities, TA\-OPD only additionally computes the tail log\-probabilitieslog⁡pttail\\log p\_\{t\}^\{\\mathrm\{tail\}\}andlog⁡qttail\\log q\_\{t\}^\{\\mathrm\{tail\}\}, which involve a single log\-sum\-exp overkkvalues per token\. This cost is negligible relative to the forward and backward passes of the student model\. Sample\-corrected TA\-OPD additionally requires the teacher’s log\-probability of the sampled tokeny^t\\hat\{y\}\_\{t\}\. Although this appears to be an extra query, it incurs negligible additional cost in practice: when queried for the top\-kklog\-probabilities, inference engines such as vLLM\([18](https://arxiv.org/html/2608.14728#bib.bib35)\)always compute and return the sampled token’s log\-probability as well, even if it falls outside the top\-kktokens\. Therefore, the teacher\-query cost of sample\-corrected TA\-OPD is identical to that of TA\-OPD\.

We empirically verify this by measuring the wall\-clock cost of the sampled\-token OPD, normalized top\-kkOPD, TA\-OPD, and sample\-corrected TA\-OPD under the same setup\. We use Qwen3\-1\.7B as the student and Qwen3\-30B\-A3B\-Instruct\-2507 as the teacher, with all other settings following Table[4](https://arxiv.org/html/2608.14728#A4.T4)\. As shown in Figure[15](https://arxiv.org/html/2608.14728#A5.F15)\(a\), completing the full300300training steps takes2323h22m1818s for TA\-OPD and2222h4747m5656s for normalized top\-kkOPD, i\.e\., a difference of1414m2222s, or1\.05%1\.05\\%of the total runtime\. This 14\-minute gap might stem from the slightly longer average response length under TA\-OPD\. Figure[15](https://arxiv.org/html/2608.14728#A5.F15)\(b\) further compares the per\-step throughput distributions, whose medians are around10001000and10051005tokens/s for TA\-OPD and normalized top\-kkOPD, respectively\. Overall, TA\-OPD incurs negligible computational cost compared to normalized top\-kkOPD\.

\(a\)Total training runtime\.\(b\)Throughput distribution during training\.
Figure 15:Computational overhead comparison between sampled\-token OPD, normalized Top\-kkOPD, 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\.
### E\.4Ablation on the Tail Token

Our analysis attributes the tail probability increase issue to top\-kknormalization\. A natural question is therefore whether simply removing the normalization is already sufficient, making the tail token in TA\-OPD unnecessary\. We show that it is not: without the tail tokenvtailv\_\{\\mathrm\{tail\}\}, training will diverge\.

#### Setup\.

We ablate the tail token by comparing TA\-OPD with unnormalized top\-kkOPD, which removes the tail term in TA\-OPD’s loss function\. Specifically, its loss function is given by:

ℓtunnorm=∑v∈Stkpt​\(v\)​log⁡pt​\(v\)qt​\(v\)=ℓtTA−pttail​log⁡pttailqttail\.\\ell\_\{t\}^\{\\mathrm\{unnorm\}\}=\\sum\_\{v\\in S\_\{t\}^\{k\}\}p\_\{t\}\(v\)\\log\\frac\{p\_\{t\}\(v\)\}\{q\_\{t\}\(v\)\}=\\ell\_\{t\}^\{\\mathrm\{TA\}\}\-p\_\{t\}^\{\\mathrm\{tail\}\}\\log\\frac\{p\_\{t\}^\{\\mathrm\{tail\}\}\}\{q\_\{t\}^\{\\mathrm\{tail\}\}\}\.
The student and teacher are Qwen2\.5\-7B\-Instruct and OpenThinker3\-7B, with all other settings following Table[4](https://arxiv.org/html/2608.14728#A4.T4)\.

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 above0\.90\.9and training collapses\.Table 5:Avg@8 of unnormalized top\-kkOPD and TA\-OPD\.TheAvg\.column represents the macro\-average across all six benchmarks\. Best results are shown inbold\.
#### Removing the tail token collapses training\.

Figure[16](https://arxiv.org/html/2608.14728#A5.F16)compares the training dynamics of the two objectives\. Under unnormalized top\-kkOPD, the student’s tail probability rises above0\.90\.9within 50 steps and its entropy grows monotonically, while the teacher’s tail probability increases in tandem\. The collapse is reflected in downstream accuracy \(Table[5](https://arxiv.org/html/2608.14728#A5.T5)\): its Avg@8 on MATH500 is only 3\.35%, whereas TA\-OPD reaches 77\.88%\.

#### The unnormalized objective is not a divergence\.

The failure is not an optimization artifact but a property of the objective itself: the objective is not minimized atpt=qtp\_\{t\}=q\_\{t\}\. Its minimizer admits the closed form

pt⋆​\(v\)=qt​\(v\)/e∀v∈Stk,hencepttail,⋆=1−1−qttaile≥1−e−1≈0\.63\.p\_\{t\}^\{\\star\}\(v\)=q\_\{t\}\(v\)/e\\quad\\forall v\\in S\_\{t\}^\{k\},\\qquad\\text\{hence\}\\qquad p\_\{t\}^\{\\mathrm\{tail\},\\star\}=1\-\\frac\{1\-q\_\{t\}^\{\\mathrm\{tail\}\}\}\{e\}\\;\\geq\\;1\-e^\{\-1\}\\approx 0\.63\.That is, the objective is minimized by deflating every top\-kkprobability by a factor ofeeand assigning the removed mass to the tail, which drives the loss negative\. The optimum therefore assigns at least63%63\\%of the probability mass outside the teacher’s top\-kktokens, regardless of the teacher distribution\. Adding the tail token contributes exactly the missing termpttail​log⁡\(pttail/qttail\)p\_\{t\}^\{\\mathrm\{tail\}\}\\log\(p\_\{t\}^\{\\mathrm\{tail\}\}/q\_\{t\}^\{\\mathrm\{tail\}\}\), which turns the objective into a reverse KL divergence onSt\+S\_\{t\}^\{\+\}: it is non\-negative and uniquely minimized when the student matches the teacher on both the top\-kktokens and the tail\.

Overall, introducing the tail token is necessary for TA\-OPD\.

### E\.5Discussion of EMA\-PG

Exponential Moving Average\-Policy Gradient \(EMA\-PG\)\([46](https://arxiv.org/html/2608.14728#bib.bib29)\)proposes a top\-kkKL estimator that is related but distinct from our sample\-corrected TA\-OPD\. Its per\-token loss is:

ℓtEMA=∑v∈Stkpt​\(v\)​sg⁡\(log⁡pt​\(v\)qt​\(v\)\)⏟=ℓtunnorm​in value\+𝟏\[y^t∉Stk\]pt​\(y^t\)sg⁡\(pt​\(y^t\)\)sg\(logpt​\(y^t\)qt​\(y^t\)\)\.\\ell\_\{t\}^\{\\mathrm\{EMA\}\}=\\underbrace\{\\sum\_\{v\\in S\_\{t\}^\{k\}\}p\_\{t\}\(v\)\\operatorname\{sg\}\\\!\\left\(\\log\\frac\{p\_\{t\}\(v\)\}\{q\_\{t\}\(v\)\}\\right\)\}\_\{\\textstyle=\\;\\ell\_\{t\}^\{\\mathrm\{unnorm\}\}\\ \\text\{in value\}\}\+\\mathbf\{1\}\[\\hat\{y\}\_\{t\}\\notin S\_\{t\}^\{k\}\]\\,\\frac\{p\_\{t\}\(\\hat\{y\}\_\{t\}\)\}\{\\operatorname\{sg\}\(p\_\{t\}\(\\hat\{y\}\_\{t\}\)\)\}\\operatorname\{sg\}\\\!\\left\(\\log\\frac\{p\_\{t\}\(\\hat\{y\}\_\{t\}\)\}\{q\_\{t\}\(\\hat\{y\}\_\{t\}\)\}\\right\)\.\(22\)In contrast, sample\-corrected TA\-OPD’s per\-token loss is given by:

ℓtSC​\-​TA=∑v∈St\+pt​\(v\)​sg⁡\(log⁡pt​\(v\)qt​\(v\)\)⏟=ℓtTA​in value\+𝟏\[y^t∉Stk\]pt​\(y^t\)sg⁡\(pt​\(y^t\)\)sg\(logpt​\(y^t\)/pttailqt​\(y^t\)/qttail\)\.\\ell\_\{t\}^\{\\mathrm\{SC\\text\{\-\}TA\}\}=\\underbrace\{\\sum\_\{v\\in S\_\{t\}^\{\+\}\}p\_\{t\}\(v\)\\operatorname\{sg\}\\\!\\left\(\\log\\frac\{p\_\{t\}\(v\)\}\{q\_\{t\}\(v\)\}\\right\)\}\_\{\\textstyle=\\;\\ell\_\{t\}^\{\\mathrm\{TA\}\}\\ \\text\{in value\}\}\+\\mathbf\{1\}\[\\hat\{y\}\_\{t\}\\notin S\_\{t\}^\{k\}\]\\,\\frac\{p\_\{t\}\(\\hat\{y\}\_\{t\}\)\}\{\\operatorname\{sg\}\(p\_\{t\}\(\\hat\{y\}\_\{t\}\)\)\}\\operatorname\{sg\}\\\!\\left\(\\log\\frac\{p\_\{t\}\(\\hat\{y\}\_\{t\}\)/p\_\{t\}^\{\\mathrm\{tail\}\}\}\{q\_\{t\}\(\\hat\{y\}\_\{t\}\)/q\_\{t\}^\{\\mathrm\{tail\}\}\}\\right\)\.Both estimators are unbiased in value and in gradient to full\-vocabulary reverse KL divergence, but they differ in where the truncation is placed and in what the sampled token is used to estimate\. EMA\-PG estimates the entire tail contribution∑v∉Stkpt​\(v\)​log⁡pt​\(v\)qt​\(v\)\\sum\_\{v\\notin S\_\{t\}^\{k\}\}p\_\{t\}\(v\)\\log\\frac\{p\_\{t\}\(v\)\}\{q\_\{t\}\(v\)\}with the sampled token\. Instead, our estimator starts fromℓtTA\\ell\_\{t\}^\{\\mathrm\{TA\}\}, which already accounts for the tail probability through the tail token, and uses the sampled token only to estimate the residualpttailDKL\(p~t∥q~t\)p\_\{t\}^\{\\mathrm\{tail\}\}D\_\{\\mathrm\{KL\}\}\(\\tilde\{p\}\_\{t\}\\\|\\tilde\{q\}\_\{t\}\)identified in Proposition[4](https://arxiv.org/html/2608.14728#Thmproposition4)\. In other words, our estimator is obtained by debiasing TA\-OPD, whereas EMA\-PG debiases the unnormalized top\-kkobjective\.

#### Behavior when the sampled token falls inside the top\-kktokens\.

Wheny^t∈Stk\\hat\{y\}\_\{t\}\\in S\_\{t\}^\{k\}, which occurs with probability1−pttail1\-p\_\{t\}^\{\\mathrm\{tail\}\}and thus covers most tokens, the correction term vanishes for both losses\. In this case,ℓtEMA\\ell\_\{t\}^\{\\mathrm\{EMA\}\}reduces to the unnormalized top\-kklossℓtunnorm\\ell\_\{t\}^\{\\mathrm\{unnorm\}\}, whileℓtSC​\-​TA\\ell\_\{t\}^\{\\mathrm\{SC\\text\{\-\}TA\}\}reduces toℓtTA\\ell\_\{t\}^\{\\mathrm\{TA\}\}\. As shown in Appendix[E\.4](https://arxiv.org/html/2608.14728#A5.SS4),ℓtTA\\ell\_\{t\}^\{\\mathrm\{TA\}\}is superior toℓtunnorm\\ell\_\{t\}^\{\\mathrm\{unnorm\}\}\. Consequently, although EMA\-PG is unbiased in expectation, on most tokens its realized objective is the unnormalized top\-kkloss function that we have shown to collapse training, and it relies on infrequent sampled\-token corrections to compensate; our estimator instead falls back to a well\-defined divergence onSt\+S\_\{t\}^\{\+\}that is uniquely minimized atpt=qtp\_\{t\}=q\_\{t\}\.

#### Variance analysis\.

Since both estimators are unbiased, a natural criterion for comparing them is the variance of the estimate\. A direct computation gives

Var\(ℓtEMA\)−Var\(ℓtSC​\-​TA\)=pttail\(1−pttail\)logpttailqttail\(logpttailqttail\+2DKL\(p~t∥q~t\)\)\.\\operatorname\{Var\}\\bigl\(\\ell\_\{t\}^\{\\mathrm\{EMA\}\}\\bigr\)\-\\operatorname\{Var\}\\bigl\(\\ell\_\{t\}^\{\\mathrm\{SC\\text\{\-\}TA\}\}\\bigr\)=p\_\{t\}^\{\\mathrm\{tail\}\}\\bigl\(1\-p\_\{t\}^\{\\mathrm\{tail\}\}\\bigr\)\\,\\log\\frac\{p\_\{t\}^\{\\mathrm\{tail\}\}\}\{q\_\{t\}^\{\\mathrm\{tail\}\}\}\\left\(\\log\\frac\{p\_\{t\}^\{\\mathrm\{tail\}\}\}\{q\_\{t\}^\{\\mathrm\{tail\}\}\}\+2\\,D\_\{\\mathrm\{KL\}\}\(\\tilde\{p\}\_\{t\}\\\|\\tilde\{q\}\_\{t\}\)\\right\)\.\(23\)SinceDKL\(p~t∥q~t\)≥0D\_\{\\mathrm\{KL\}\}\(\\tilde\{p\}\_\{t\}\\\|\\tilde\{q\}\_\{t\}\)\\geq 0, the gap is strictly positive wheneverpttail\>qttailp\_\{t\}^\{\\mathrm\{tail\}\}\>q\_\{t\}^\{\\mathrm\{tail\}\}, i\.e\. whenever the student’s tail probability exceeds the teacher’s, which is likely to happen when using the teacher’s top\-kksupport\. Since the supportStkS\_\{t\}^\{k\}consists of the teacher’s top\-kktokens, the teacher’s tail probability is typically smaller than the student’s\. Figures[4](https://arxiv.org/html/2608.14728#S5.F4)confirms this empirically, whereqttailq\_\{t\}^\{\\mathrm\{tail\}\}is belowpttailp\_\{t\}^\{\\mathrm\{tail\}\}throughout training across all model pairs\. Overall, our loss objective tends to have lower variance than EMA\-PG\.

## Appendix FDetailed Results

### F\.1Training Dynamics of Additional Model Pairs

Figure[17](https://arxiv.org/html/2608.14728#A6.F17)presents additional comparisons between TA\-OPD and normalized top\-kkOPD in terms of tail probability and token\-level entropy during training\. Results are shown on three student–teacher pairs: Qwen3\-1\.7B paired with Qwen3\-30B\-A3B\-Instruct\-2507, Llama\-3\.1\-8B paired with DeepSeek\-R1\-Distill\-Llama\-8B, Qwen3\-1\.7B\-Base paired with Qwen3\-8B\.

Qwen3\-1\.7B \(Student\)→\\rightarrowQwen3\-30B\-A3B\-Instruct\-2507 \(Teacher\)Llama\-3\.1\-8B \(Student\)→\\rightarrowDeepSeek\-R1\-Distill\-Llama\-8B \(Teacher\)Qwen3\-1\.7B\-Base \(Student\)→\\rightarrowQwen3\-8B \(Teacher\)
Figure 17:Tail probability and token\-level entropy over training for TA\-OPD and Normalized top\-kkOPD\.Each row corresponds to a different student\-teacher pair\.Qwen2\.5\-Math\-1\.5B \(Student\)→\\rightarrowJustRL\-DeepSeek\-1\.5B \(Teacher\)DeepSeek\-R1\-Distill\-Qwen\-1\.5B \(Student\)→\\rightarrowJustRL\-DeepSeek\-1\.5B \(Teacher\)
Figure 18:Tail probability and token\-level entropy over training for TA\-OPD and Normalized top\-kkOPD 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\.