Paper deep dive
Tail-Aware Top-$k$ On-Policy Distillation
Huipeng Huang, Hongxin Wei
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%
Last extracted: 8/18/2026, 4:27:18 AM
Summary
The paper introduces Tail-Aware Top-k On-Policy Distillation (TA-OPD), a method to improve knowledge transfer from large language models to smaller ones. It addresses a flaw in normalized top-k OPD, where discarding tail probability causes the student's entropy and tail probability to increase, degrading performance. TA-OPD restores this signal by including a tail token in the divergence calculation, resulting in superior accuracy on math reasoning benchmarks.
Entities (9)
Relation Signals (6)
Normalized Top-kk OPD ā discards ā Tail Probability
confidence 95% Ā· However, this normalized objective discards the information about tail probability: the total probability outside the teacherās top-kk tokens.
TA-OPD ā improves ā Normalized Top-kk OPD
confidence 95% Ā· Extensive experiments demonstrate the superiority of TA-OPD, improving Avg@8 by up to 8.05 points on common benchmarks.
TA-OPD ā restores ā Tail Probability
confidence 95% Ā· TA-OPD, a novel distillation method that restores the missing tail probability signal.
Normalized Top-kk OPD ā causes ā Increased Entropy
confidence 90% Ā· As a result, the optimization can steadily increase the studentās tail probability and entropy, empirically degrading downstream accuracy.
Qwen2.5-7B-Instruct ā isstudentof ā OpenThinker3-7B
confidence 90% Ā· when distilling OpenThinker3-7B (10) into Qwen2.5-7B-Instruct (26)
TA-OPD ā minimizes ā Reverse KL Divergence
confidence 90% Ā· TA-OPD minimizes the reverse KL divergence over the top-kk tokens plus a tail token that carries the tail probability.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:On-policy distillation (OPD) has emerged as an effective paradigm for transferring knowledge between language models, where a student is trained to align its next-token distribution with the teacher's along its own trajectories. To provide dense supervision at tractable cost, many works minimize the reverse Kullback-Leibler (KL) divergence between the student and teacher's normalized distributions over the teacher's top-$k$ tokens. However, this normalized objective discards the information about tail probability: the total probability outside the teacher's top-$k$ tokens. As a result, the optimization can steadily increase the student's tail probability and entropy, empirically degrading downstream accuracy. To address this issue, we propose Tail-Aware Top-$k$ OPD (\textbf{TA-OPD}), a novel distillation method that restores the missing tail probability signal. In particular, TA-OPD minimizes the reverse KL divergence over the top-$k$ tokens plus a tail token that carries the tail probability. In effect, TA-OPD better aligns the student's next-token distribution with the teacher's, preventing the increase in tail probability and entropy caused by top-$k$ normalization. Extensive experiments demonstrate the superiority of TA-OPD, improving Avg@8 by up to 8.05 points on common benchmarks. Our code is available at this https URL.
Tags
Links
- Source: https://arxiv.org/abs/2608.14728v1
- Canonical: https://arxiv.org/abs/2608.14728v1
Trouble viewing inline? Open PDF directly ā
Full Text
101,427 characters extracted from source content.
Expand or collapse full text
Tail-Aware Top-k On-Policy Distillation Huipeng Huang Affiliation: Department of Statistics and Data Science, Southern University of Science and Technology Hongxin Wei Thanks: 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-k tokens. However, this normalized objective discards the information about tail probability: the total probability outside the teacherās top-k tokens. As a result, the optimization can steadily increase the studentās tail probability and entropy, empirically degrading downstream accuracy. To address this issue, we propose Tail-Aware Top-k OPD (TA-OPD), a novel distillation method that restores the missing tail probability signal. In particular, TA-OPD minimizes the reverse KL divergence over the top-k tokens plus a tail token that carries the tail probability. In effect, TA-OPD better aligns the studentās next-token distribution with the teacherās, preventing the increase in tail probability and entropy caused by top-k normalization. Extensive experiments demonstrate the superiority of TA-OPD, improving Avg@8 by up to 8.05 points on common benchmarks. Our code is available at https://github.com/HuipengHuang/TA-OPD. 1 Introduction Knowledge distillation (15) 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; 9), 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; 2) and hinders effective learning (41). This motivates on-policy distillation (OPD) (1; 22; 40), where the student is trained to align its next-token distribution with the teacherās along its own trajectories. A popular idea, normalized top-k OPD (7; 38), minimizes the reverse Kullback-Leibler (KL) divergence between the student and teacherās normalized distributions over the teacherās top-k tokens. However, it discards the information about tail probability: the total probability mass outside the teacherās top-k tokens. 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-k tokens, 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-k OPD (TA-OPD), a novel distillation method that restores the missing tail probability signal. In particular, TA-OPD minimizes the reverse KL divergence over the top-k tokens plus a tail token that carries the tail probability. In effect, TA-OPD aligns the studentās tail probability with the teacherās, preventing the tail probability and entropy increase caused by top-k normalization. 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) into Qwen2.5-7B-Instruct (26), normalized top-k OPDās training collapses: its training entropy increases to around 6, and its Avg@8 accuracy on MATH500 (14) 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 small k. We summarize our contributions as follows: ā¢ We show that normalized top-k OPD 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-k normalization. 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-k OPD, improving Avg@8 accuracy by up to 8.05 points on common benchmarks. 2 Preliminaries On-policy distillation (OPD) aims to distill knowledge from a teacher LLM Ļte _te to a student LLM ĻĪø _Īø on trajectories sampled from the student. Given a prompt xā¼x _x, the student samples a response y^=(y^1,ā¦,y^T)ā¼ĻĪø(ā ā£x) y=( y_1,ā¦, y_T) _Īø(Ā· x). At each step t, on the prefix y^<t=(y^1,ā¦,y^tā1) y_<t=( y_1,ā¦, y_t-1), we define the studentās and teacherās next-token distributions over a vocabulary V as pt=ĻĪø(ā ā£x,y^<t)p_t= _Īø(Ā· x, y_<t) and qt=Ļte(ā ā£x,y^<t)q_t= _te(Ā· x, y_<t). OPD minimizes reverse KL divergence over student-sampled trajectories: āOPD(Īø)=xā¼x,y^ā¼ĻĪø(ā ā£x)[āt=1TDKL(ptā„qt)],L_OPD(Īø)=E_x _x,\; y _Īø(Ā· x) [ _t=1^TD_KL(p_t\|q_t) ], (1) where DKL(ptā„qt)=āvāpt(v)logptā(v)qtā(v)D_KL(p_t\|q_t)= _v p_t(v) p_t(v)q_t(v). Different OPD methods vary in how they approximate DKL(ptā„qt)D_KL(p_t\|q_t) with their per-token losses āt _t. Given a student-sampled token y^tā¼pt y_t p_t, sampled-token OPD (22; 36) computes ātsample=logā”ptā(y^t)ālogā”qtā(y^t) _t^sample= p_t( y_t)- q_t( y_t), which is cheap to compute but discards the dense information over the remaining vocabulary. Full-vocabulary OPD (39) computes ātfull=DKL(ptā„qt) _t^full=D_KL(p_t\|q_t) exactly over V. It provides dense supervision but is prohibitively expensive for LLMs. Normalized top-k OPD 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-k tokens (21) or the teacherās top-k tokens (16; 7). In this work, we primarily focus on the teacher top-k variant. Formally, we define the teacherās top-k tokens as Stk=TopKā”(qt,k)S_t^k=TopK(q_t,k). The normalized distributions on StkS_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), p_t^(S_t^k)(v)= p_t(v)\,1[vā S_t^k] _uā S_t^kp_t(u), q_t^(S_t^k)(v)= q_t(v)\,1[vā S_t^k] _uā S_t^kq_t(u), where ā[ā ]1[Ā·] denotes the indicator function. Distillation is then performed by minimizing the subset reverse KL divergence ātnorm=DKL(pĀÆt(Stk)ā„qĀÆt(Stk)) _t^norm=D_KL\! ( p_t^(S_t^k)\,\|\, q_t^(S_t^k) ), yielding the trajectory-level objective: āOPDnorm(Īø)=xā¼x,y^ā¼ĻĪø(ā ā£x)[āt=1TDKL(pĀÆt(Stk)ā„qĀÆt(Stk))].L_OPD^norm(Īø)=E_x _x,\; y _Īø(Ā· x) [ _t=1^TD_KL\! ( p_t^(S_t^k)\,\|\, q_t^(S_t^k) ) ]. (2) Normalized top-k OPD reduces the teacher-query cost while retaining dense, multi-token supervision over StkS_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-k tokens: pttail=āvāStkptā(v),qttail=āvāStkqtā(v).p_t^tail= _vā S_t^kp_t(v), q_t^tail= _vā S_t^kq_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. 3 Motivation In this section, we investigate how the normalized objective in Eq. (2) changes the studentās tail probability. We find that minimizing the objective can steadily increase the studentās tail probability. 3.1 Theoretical Analysis on the Tail Probability Fix a step t and let zt,vz_t,v denote the student logit for token v under the prefix, so that ptā(v)=ezt,vāuāezt,up_t(v)= e^z_t,v _u e^z_t,u. We have the following propositions for the optimization of the normalized objective. Proposition 1. For every step t, the gradient of ātnorm _t^norm with respect to the student logit zt,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. ā _tā z_t,v= cases p_t^(S_t^k)(v)\! [ p_t^(S_t^k)(v) q_t^(S_t^k)(v)- _t ],&vā S_t^k,\\[10.00002pt] 0,&vā S_t^k. cases Moreover, the gradients on the logits of the top-k tokens sum to zero: āvāStkāātāzt,v=0 _vā S_t^k ā _tā z_t,v=0. The proof is provided in Appendix A.1. 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-k logits relative to the tail logits. However, the normalized objective can do neither: the tail logits receive zero gradient, and the gradients on the top-k logits sum to zero. In the following, we further explore how the zero-sum gradient within the top-k tokens affects the studentās tail probability. Proposition 2. Assume that the student ĻĪø _Īø is a tabular softmax policy, where each token v at each step t is associated with an independent logit parameter zt,v=Īøt,vz_t,v= _t,v. Write pĀÆ:=pĀÆt(Stk) p:= p_t^(S_t^k) and qĀÆ:=qĀÆt(Stk) q:= q_t^(S_t^k) for the normalized distributions on StkS_t^k, and let Covvā¼r(f,g) *Cov_v r(f,g) denote the covariance of fā”(v)f(v) and gā”(v)g(v) under vā¼rv r. After one gradient descent step ĪøāĪøāĪ·āāātnormāĪøāĪø-Ī·\, ā _t^normāĪø with learning rate Ī·, the studentās tail probability change Īāpttail p_t^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). p_t^tail=Ī·\,p_t^tail (1-p_t^tail ) [ *Cov_v p ( p(v), p(v) )_ student self-covariance- *Cov_v p ( p(v), q(v) )_student-teacher covariance ]+O(Ī·^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. The proposition11 1 We 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, since logā”pĀÆ p is monotone with pĀÆ p. In contrast, the student-teacher covariance is positive only when the student concentrates its probability on the teacherās high-probability tokens within StkS_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.2 Empirical Study on the Tail Probability Increase Setup. We conduct experiments with three student-teacher model pairs: Qwen3-1.7B (43) with Qwen3-30B-A3B-Instruct-2507, Qwen2.5-7B-Instruct with OpenThinker3-7B, and Llama-3.1-8B (8) with DeepSeek-R1-Distill-Llama-8B (11). All models are trained on DAPO-MATH-17K (44) with k=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-k OPD consistently increases the studentās tail probability and entropy across all three model pairs. Normalized top-k OPD steadily increases the studentās tail probability. Figure 1 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-k tokens. 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; 37), as the teacher itself exhibits significantly reduced accuracy given these prefixes. Consequently, normalized top-k OPD progressively shifts probability mass away from the teacherās top-k tokens, 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 plots validation accuracy against the studentās tail probability. Each point is a checkpoint from a normalized top-k OPD 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 is 11. The validation accuracy degrades significantly from above 0.70.7 to below 0.20.2 as the tail probability grows from near 00 to around 0.70.7, with a Pearson correlation coefficient of ā0.91-0.91. In Appendix B, 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 when k is not sufficiently large (e.g., kā¤64kā¤ 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-k OPD): The normalized objective discards the tail probability, leading to increased tail probability. (c) Decreased Tail Probability (Tail-Aware Top-k OPD, Ours): TA-OPD restores the missing tail probability signal, matching the studentās tail probability to the teacherās. 4 Method 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 (b)). To address this problem, our key idea is to restore the missing tail probability signal. Tail-Aware Top-k OPD. We propose Tail-Aware Top-k OPD (TA-OPD), which explicitly aligns the studentās tail probability with the teacherās (Figure 3 (c)). In particular, we introduce a tail token vtailv_tail that carries the tail probability. We then minimize the reverse KL divergence over the top-k tokens plus the tail token. Formally, we define the augmented token set St+=StkāŖvtailS_t^+=S_t^kāŖ\v_tail\. On St+S_t^+, each top-k token retains its original probability, ptā(v)=ĻĪøā(vā£x,y^<t)p_t(v)= _Īø(v x, y_<t), while the tail token carries the tail probability, ptā(vtail)=pttailp_t(v_tail)=p_t^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. _t^TA= _vā S_t^+p_t(v) p_t(v)q_t(v)= _vā S_t^kp_t(v) p_t(v)q_t(v)+p_t^tail p_t^tailq_t^tail. (4) By introducing the tail token vtailv_tail, the term pttailālogā”pttailqttailp_t^tail p_t^tailq_t^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 _t^TA. Proposition 3. For every step t, the gradient of ātTA _t^TA with respect to the student logit zt,vz_t,v is āātTAāzt,v=ptā(v)ā(logā”ptā(v)qtā(v)āātTA),vāStk,ptā(v)ā(logā”pttailqttailāātTA),vāStk. ā _t^TAā z_t,v= casesp_t(v) ( p_t(v)q_t(v)- _t^TA ),&vā S_t^k,\\[11.99998pt] p_t(v) ( p_t^tailq_t^tail- _t^TA ),&vā S_t^k. cases (5) The proposition reveals two key properties of TA-OPD. First, while the normalized objective gives no gradient to logits outside the top-k tokens, TA-OPD updates them through the tail log-ratio logā”pttailqttail p_t^tailq_t^tail, which compares the studentās tail probability with the teacherās. Second, the gradients on the top-k logits no longer sum to zero. The top-k gradient sum āvāStkāātTAāzt,v=pttailā(ātTAālogā”pttailqttail) _vā S_t^k ā _t^TAā z_t,v=p_t^tail ( _t^TA- p_t^tailq_t^tail ) is negative whenever the tail log-ratio is larger than the average log-ratio over StkS_t^k. It therefore can explicitly raise the top-k logits relative to the tail whenever the student over-weights the tail relative to the teacher. In the following, we show that ātTA _t^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). Let p~tā(v)=ptā(v)/pttail p_t(v)=p_t(v)/p_t^tail and q~tā(v)=qtā(v)/qttail q_t(v)=q_t(v)/q_t^tail for vāStkvā S_t^k be the normalized tail distributions. Then ātTA=ātfullāpttailDKL(p~tā„q~t)ā¤ātfull. _t^TA= _t^full-p_t^tail\,D_KL( p_t\| q_t)ā¤ _t^full. (6) The proofs of the above propositions are presented in Appendix A.3 and A.4. This proposition shows that the loss functions of TA-OPD and full-vocabulary OPD differ only by a non-negative residual term pttailDKL(p~tā„q~t)p_t^tail\,D_KL( p_t\| q_t). In practice, TA-OPD drives both the studentās and teacherās tail probabilities close to zero (see Figure 4 (a) and (b)). As a result, the residual term becomes small, making ātTA _t^TA a tight lower bound to the ideal full-vocabulary reverse KL objective. 5 Experiments 5.1 Experimental Setup Models and Training Dataset. For main experiments, we use the three student-teacher model pairs as in section 3.2: 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), OlympiadBench (13), AMC (20) and AIME24/25 (20)āand two out-of-distribution benchmarks, ARC-c (4) and MMLU-Pro (35). 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-k OPD. In Appendix E.4, we conduct an additional comparison with unnormalized top-k OPD. We provide a detailed introduction to these methods in Appendix C. Implementation. All the OPD experiments use the same training setup: 300 training steps with a learning rate of 1Ć10ā61Ć 10^-6. Unless otherwise specified, we set k=16k=16 for top-k OPD. 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. 5.2 Results Table 1: Results on math reasoning and out-of-distribution (OOD) benchmarks. We compare TA-OPD with sampled-token OPD (Sampled-token) and normalized top-k OPD (Norm. top-k). Each Avg. column represents the macro-average. Best results are shown in bold. Methods In-Distribution Performance OOD Performance MATH500 Minerva Olympiad AMC AIME 24 AIME 25 Avg. ARC-c MMLU-Pro Avg. Student: Qwen2.5-7B-Instruct Teacher: OpenThinker3-7B Sampled-token 74.60 31.25 41.22 45.48 14.58 16.25 37.23 74.99 47.66 61.33 Norm. top-k 68.78 25.55 36.02 41.42 13.75 15.00 33.42 23.07 37.38 30.23 TA-OPD (Ours) 77.88 32.58 42.41 45.03 16.67 17.50 38.68 76.30 50.35 63.33 Student: Llama-3.1-8B Teacher: DeepSeek-R1-Distill-Llama-8B Sampled-token 43.45 12.78 16.56 17.62 2.08 1.25 15.62 41.88 28.99 35.44 Norm. top-k 31.30 9.24 9.96 11.45 1.25 0.42 10.60 10.01 18.33 14.17 TA-OPD (Ours) 50.83 14.94 20.65 20.48 2.50 2.50 18.65 65.70 34.57 50.14 Student: Qwen3-1.7B Teacher: Qwen3-30B-A3B-Instruct-2507 Sampled-token 82.98 49.49 34.93 52.41 26.25 19.58 44.27 83.91 50.08 67.00 Norm. top-k 81.80 47.20 33.28 50.60 23.33 17.08 42.22 82.48 49.66 66.07 TA-OPD (Ours) 83.48 49.82 34.50 52.86 26.25 20.83 44.62 83.59 50.42 67.01 TA-OPD achieves the best performance across studentāteacher model pairs. Table 1 compares TA-OPD with sampled-token OPD and normalized top-k OPD. Evaluated on six math reasoning benchmarks, TA-OPD attains the best average accuracy on every studentāteacher pair, outperforming normalized top-k OPD by +5.26 points on Qwen2.5-7B-Instruct and +8.05 points on Llama-3.1-8B. Notably, TA-OPD demonstrates a significantly greater advantage on MATH500 with Llama-3.1-8B over normalized top-k OPD, improving Avg@8 from 31.30 to 50.83 (+19.53 points). 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-k OPD 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-k OPD 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, we compare the training dynamics of TA-OPD and normalized top-k OPD on the Qwen2.5-7B-Instruct and OpenThinker3-7B model pair. With normalized top-k OPD, the studentās tail probability, the teacherās tail probability, and the studentās entropy increase to approximately 0.60.6, 0.50.5, and 66 over the course of training. In contrast, TA-OPD keeps the two tail probabilities close to 00 and the studentās entropy below 1.51.5. The same trends are consistently observed for three alternative model pairs in Figure 17. In short, these results show that TA-OPD addresses the tail probability and entropy increase caused by top-k normalization. Figure 5: Ablation study on k. TA-OPD is effective with a small k. In Figure 5, we ablate how the number of top-k tokens 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 vary kā1,2,4,8,16,32kā\1,2,4,8,16,32\. For each k, we report the average accuracy macro-averaged across the six math reasoning benchmarks. The results show that the accuracy is insensitive to k once kā„2kā„ 2: k=2k=2 attains an average accuracy of 30.4330.43, only 0.430.43 points below the best result of 30.8630.86 at k=32k=32. Overall, TA-OPD can be applied with a small k. Additional results. Due to space constraints, we defer additional analyses to the appendix. Appendix E.1 conducts a synthetic experiment to provide an intuitive understanding of TA-OPD and normalized top-k OPDās optimization. Appendix E.2 shows that TA-OPD consistently achieves higher student-teacher top-k overlap ratios compared to normalized top-k OPD. Appendix E.3 shows that TA-OPD incurs negligible computational overhead over normalized top-k OPD. 6 Discussion Table 2: Comparison between TA-OPD and sample-corrected TA-OPD on math reasoning and out-of-distribution (OOD) benchmarks. Best results are shown in bold. Methods In-Distribution Performance OOD Performance MATH500 Minerva Olympiad AMC AIME 24 AIME 25 Avg. ARC-c MMLU-Pro Avg. Student: DeepSeek-R1-Distill-Qwen-1.5B Teacher: JustRL-DeepSeek-1.5B TA-OPD 86.13 30.97 57.04 76.20 41.67 29.73 53.62 35.40 25.98 30.69 SC-TA-OPD 87.00 32.08 57.65 76.51 42.92 30.17 54.39 35.55 25.76 30.66 Student: Qwen3-1.7B Teacher: Qwen3-30B-A3B-Instruct-2507 TA-OPD 83.48 49.82 34.50 52.86 26.25 20.83 44.62 83.59 50.42 67.01 SC-TA-OPD 84.00 49.54 34.93 53.16 26.67 19.58 44.65 83.91 50.28 67.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, the bias is pttailDKL(p~tā„q~t)p_t^tail\,D_KL( p_t\,\|\, q_t). When the sampled tokenās probability is additionally available, we can estimate the bias with the sampled token y^t y_t, and add this estimate back to ātTA _t^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= _t^SC -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), _vā S_t^+p_t(v)sg\! ( p_t(v)q_t(v) )+1[ y_tā S_t^k]\, p_t( y_t)sg(p_t( y_t))sg\! ( p_t( y_t)/p_t^tailq_t( y_t)/q_t^tail ), (7) where sgā”(ā )sg(Ā·) denotes the stop-gradient operator. We next characterize its theoretical property. Proposition 5 (Unbiasedness). For every step t, taking the randomness over y^tā¼pt y_t p_t, we have y^tā¼pt[ātSCā-āTA]=DKL(ptā„qt),y^tā¼pt[āĪøātSCā-āTA]=āĪøDKL(ptā„qt).E_ y_t p_t [ _t^SC -TA ]=D_KL(p_t\,\|\,q_t), _ y_t p_t [ _Īø\, _t^SC -TA ]= _Īø\,D_KL(p_t\,\|\,q_t). The proof is provided in Appendix A.5. Proposition 5 shows that ātSCā-āTA _t^SC -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-k OPDās dense supervision. Table 2 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 from 53.6253.62 to 54.3954.39 for DeepSeek-R1-Distill-Qwen-1.5B and from 44.6244.62 to 44.6544.65 for Qwen3-1.7B, while the OOD averages remain essentially unchanged in both settings (30.6930.69 vs. 30.6630.66 and 67.0167.01 vs. 67.1067.10). We attribute this to the tightness of TA-OPD: its bias to full-vocabulary OPD pttailDKL(p~tā„q~t)p_t^tail\,D_KL( p_t\,\|\, 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) 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 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-k OPD by 13.50 points. 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, 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. 7 Related Work On-policy Distillation. On-policy distillation (OPD) is an effective post-training paradigm that has attracted a surge of interest (33). To provide dense supervision at tractable cost, top-k OPD has become a popular research direction (47; 38). Normalized top-k OPD (21; 7) directly uses normalized top-k reverse KL as the loss function. Entropy-aware OPD (16) improves sampled-token OPD by applying normalized top-k forward KL at high-entropy positions. vOPD (23) uses normalized top-k reverse KL as a reward baseline to reduce variance of sampled-token OPD. While these methods use normalized top-k KL for different purposes, they share a common limitation: the normalized objective does not faithfully approximate the full-vocabulary KL. We note that 6 study a related tail probability issue, but it considers off-policy full-vocabulary distillation and is not applicable to top-k OPD. 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-k objectives 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; 24). In OPD, the student is optimized by minimizing the reverse KL divergence to the teacher (22), while in reinforcement learning (RL), a KL constraint against the base model is commonly imposed to prevent the policy from drifting away (48; 31). For efficiency, most existing methods estimate the divergence using only the sampled tokens (28; 30). 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). Existing top-k OPD methods predominantly adopt a normalized top-k formulation (21; 7), 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) 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) with a principled constraint. In contrast, we directly use the top-k KL approximation as the optimization objective and establish its theoretical properties. Besides, EMA-PG (46) 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. 8 Conclusion 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-k tokens 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-k normalization. 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-k probabilities. 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-k OPD diminishes as k increases 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 Agarwal et al. (2024) R. Agarwal, N. Vieillard, Y. Zhou, P. Stanczyk, S. Ramos Garea, M. Geist, and O. Bachem On-policy distillation of language models: learning from self-generated mistakes. In International Conference on Learning Representations, Vol. 2024, p. 21246ā21263. Cited by: Ā§1. Arora et al. (2022) K. Arora, L. El Asri, H. Bahuleyan, and J. C. K. Cheung Why exposure bias matters: an imitation learning perspective of error accumulation in language generation. In Findings of the Association for Computational Linguistics: ACL 2022, p. 700ā710. Cited by: Ā§1. Bengio et al. (2015) S. Bengio, O. Vinyals, N. Jaitly, and N. Shazeer Scheduled sampling for sequence prediction with recurrent neural networks. Advances in Neural Information Processing Systems 28. Cited by: Ā§1. Clark et al. (2018) P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457. Cited by: Ā§5.1. Cui et 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. Dasgupta et al. (2026) S. Dasgupta, T. Cohn, and T. Baldwin Donāt ignore the tail: decoupling top-k probabilities for efficient language model distillation. arXiv preprint arXiv:2602.20816. Cited by: Ā§7. Fu et al. (2026) Y. Fu, H. Huang, K. Jiang, J. Liu, Z. Jiang, Y. Zhu, and D. Zhao Revisiting on-policy distillation: empirical failure modes and simple fixes. arXiv preprint arXiv:2603.25562. Cited by: Appendix C, Ā§1, Ā§2, Ā§3.2, Ā§7, Ā§7. Grattafiori et 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. Gu et al. (2024) Y. Gu, L. Dong, F. Wei, and M. Huang Minillm: knowledge distillation of large language models. In International Conference on Learning Representations, Vol. 2024, p. 32694ā32717. Cited by: Ā§1. Guha et 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. Schmidt OpenThoughts: data recipes for reasoning models. External Links: 2506.04178, Link Cited by: Ā§1. Guo et 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. He et 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. He et 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. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), p. 3828ā3850. Cited by: Ā§5.1. Hendrycks et al. (2021) D. Hendrycks, C. Burns, S. Kadavath, A. Arora, S. Basart, E. Tang, D. Song, and J. Steinhardt Measuring mathematical problem solving with the math dataset. Advances in Neural Information Processing Systems. Cited by: Ā§1. Hinton et al. (2015) G. Hinton, O. Vinyals, and J. Dean Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531. Cited by: Ā§1. Jin et al. (2026) W. Jin, T. Min, Y. Yang, S. R. Kadhe, Y. Zhou, D. Wei, N. Baracaldo, and K. Lee Entropy-aware on-policy distillation of language models. arXiv preprint arXiv:2603.07079. Cited by: Ā§2, Ā§7. Kim and Rush (2016) Y. Kim and A. M. Rush Sequence-level knowledge distillation. In Proceedings of the 2016 conference on empirical methods in natural language processing, p. 1317ā1327. Cited by: Ā§1. Kwon et al. (2023) W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. Gonzalez, H. Zhang, and I. Stoica Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th symposium on operating systems principles, p. 611ā626. Cited by: Ā§E.3. Lewkowycz et 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 Systems 35, p. 3843ā3857. Cited by: Ā§5.1. Li et 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 repository 13 (9), p. 9. Cited by: Ā§5.1. Li et 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, Appendix C, Ā§E.2, Ā§2, Ā§7, Ā§7. Lu and Lab (2025) K. Lu and T. M. Lab On-policy distillation. Thinking Machines Lab: Connectionism. Note: https://thinkingmachines.ai/blog/on-policy-distillation External Links: Document Cited by: Appendix C, Ā§1, Ā§2, Ā§7. Oh et al. (2026) M. Oh, S. Song, G. Choi, Y. Choi, and Y. Jo KL for a kl: on-policy distillation with control variate baseline. arXiv preprint arXiv:2605.07865. Cited by: Ā§7. Ouyang et 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 Systems 35, p. 27730ā27744. Cited by: Ā§7. Qi et al. (2026) P. Qi, X. Zhou, Z. Liu, T. Pang, C. Du, M. Lin, and W. S. Lee Rethinking the trust region in llm reinforcement learning. arXiv preprint arXiv:2602.04879. Cited by: Ā§7. Qwen et 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. Qiu Qwen2.5 technical report. External Links: 2412.15115, Link Cited by: Ā§1. Ren and Sutherland (2025) Y. Ren and D. Sutherland Learning dynamics of llm finetuning. In International Conference on Learning Representations, Vol. 2025, p. 70523ā70563. Cited by: Ā§A.2. Schulman et al. (2015) J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz Trust region policy optimization. In International Conference on Machine Learning, p. 1889ā1897. Cited by: Ā§7. Schulman et al. (2017) J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. Cited by: Ā§7. Schulman (2020) J. Schulman Approximating kl divergence. John Schulmanās Homepage 5. Cited by: Ā§7. Shah et 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. Sheng et al. (2025) G. Sheng, C. Zhang, Z. Ye, X. Wu, W. Zhang, R. Zhang, Y. Peng, H. Lin, and C. Wu Hybridflow: a flexible and efficient rlhf framework. In Proceedings of the Twentieth European Conference on Computer Systems, p. 1279ā1297. Cited by: Ā§D.1. Song and Zheng (2026) M. Song and M. Zheng A survey of on-policy distillation for large language models. arXiv preprint arXiv:2604.00626. Cited by: Ā§7. Vieillard et al. (2020) N. Vieillard, T. Kozuno, B. Scherrer, O. Pietquin, R. Munos, and M. Geist Leverage the average: an analysis of kl regularization in reinforcement learning. Advances in Neural Information Processing Systems 33, p. 12163ā12174. Cited by: Ā§7. Wang et 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 Systems 37, p. 95266ā95290. Cited by: Ā§5.1. Xiao et 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, Ā§2. Xie et al. (2026) Y. Xie, S. Zhu, T. Wen, B. Chen, and Y. Wang On the position bias of on-policy distillation. arXiv preprint arXiv:2606.22600. Cited by: Ā§3.2. Xing et al. (2026) X. Xing, H. Wang, B. Gao, Z. Li, and Y. Tang Trust region on-policy distillation. arXiv preprint arXiv:2606.01249. Cited by: Ā§1, Ā§7. Xu et 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. Xu et al. (2025a) H. Xu, Q. Zhu, H. Deng, J. Li, L. Hou, Y. Wang, L. Shang, R. Xu, and F. Mi Kdrl: post-training reasoning llms via unified knowledge distillation and reinforcement learning. arXiv preprint arXiv:2506.02208. Cited by: Ā§1. Xu et al. (2025b) W. Xu, R. Han, Z. Wang, L. Le, D. Madeka, L. Li, W. Wang, R. Agarwal, C. Lee, and T. Pfister Speculative knowledge distillation: bridging the teacher-student gap through interleaved sampling. In International Conference on Learning Representations, Vol. 2025, p. 64616ā64646. Cited by: Ā§1. Yan et al. (2026) J. Yan, Y. Li, Z. Hu, Z. Wang, G. Cui, X. Qu, Y. Cheng, and Y. Zhang Learning to reason under off-policy guidance. Advances in Neural Information Processing Systems 38, p. 117157ā117186. Cited by: Ā§D.1. Yang et 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. Yu et 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 Systems 38, p. 113222ā113244. Cited by: Ā§3.2. Zhang et al. (2026) H. Zhang, Y. Li, Z. Wang, Z. Wang, S. Zhang, X. Qu, and Y. Cheng Characterizing, evaluating, and optimizing complex reasoning. In Forty-third International Conference on Machine Learning, External Links: Link Cited by: Ā§A.2. Zhang and Ba (2026) L. Zhang and J. Ba EMA policy gradient: taming reinforcement learning for llms with ema anchor and top-k kl. arXiv preprint arXiv:2602.04417. Cited by: Ā§E.5, Ā§7. Zhu et al. (2026) S. Zhu, X. Ye, H. Lu, W. Shi, and G. Liu The many faces of on-policy distillation: pitfalls, mechanisms, and fixes. arXiv preprint arXiv:2605.11182. Cited by: Ā§7. Ziegler et al. (2019) D. M. Ziegler, N. Stiennon, J. Wu, T. B. Brown, A. Radford, D. Amodei, P. Christiano, and G. Irving Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.08593. Cited by: Ā§7. Appendix A Proof A.1 Proof of Proposition 1 Proof. Fix a step t and a prefix y^<t y_<t, and omit the subscript t throughout: write S=Stk=TopKā”(qt,k)S=S_t^k=TopK(q_t,k) for the teacherās top-k tokens, zv=zt,vz_v=z_t,v for the student logit of token v, and pĀÆ=pĀÆt(Stk) p= p_t^(S_t^k), qĀÆ=qĀÆt(Stk) q= q_t^(S_t^k) for the normalized distributions on S, so that ānorm=āuāSpĀÆā(u)ālogā”pĀÆā(u)qĀÆā(u). ^norm= _uā S p(u) p(u) q(u). Since the token set S is selected from the teacher distribution, both S and qĀÆ q are constants with respect to the student logits. We proceed in three steps: we first show that pĀÆ p is a softmax over the logits restricted to top-k tokens S, then differentiate the loss through it, and finally sum the resulting gradients over S. Step 1: pĀÆ p can be expressed as a softmax restricted to top-k tokens. Writing the student distribution as pā”(v)=ezv/āuāezup(v)=e^z_v/ _u e^z_u, the vocabulary-level normalizer appears in both the numerator and the denominator of pĀÆ p and therefore cancels, pĀÆā(v)=pā”(v)āwāSpā”(w)=ezvāwāSezw,vāS. p(v)= p(v) _wā Sp(w)= e^z_v _wā Se^z_w, vā S. (8) Two consequences of Eq. (8) drive the whole proof: pĀÆ p is a softmax over the sub-vector (zw)wāS(z_w)_wā S, and it does not depend on any logit outside S. Its Jacobian is therefore the standard softmax Jacobian on S and vanishes elsewhere, āpĀÆā(u)āzv=pĀÆ(u)([u=v]āpĀÆ(v)),u,vāS,0,vāS. ā p(u)ā z_v= cases p(u) (1[u=v]- p(v) ),&u,vā S,\\[3.99994pt] 0,&vā S. cases (9) Step 2: the gradient with respect to logit. Because qĀÆ q is fixed with respect to the student logits, the loss depends on the student only through pĀÆā(u)uāS\ p(u)\_uā S, with āānormāpĀÆā(u)=logā”pĀÆā(u)qĀÆā(u)+1,uāS. ā ^normā p(u)= p(u) q(u)+1, uā 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, ā ^normā z_v= _uā S ā ^normā p(u)Ć ā p(u)ā z_v= _uā S ( p(u) q(u)+1 ) ā p(u)ā z_v, (10) where the terms with uāSuā S are already dropped because ānorm ^norm does not involve them. For vāSvā S, every Jacobian entry in Eq. (9) is zero, hence āānorm/āzv=0ā ^norm/ā z_v=0, which is the second case of the claim: the loss is blind to the logits outside the teacherās top-k tokens. For vāSvā S, substituting Eq. (9) into Eq. (10) and separating the diagonal term u=vu=v gives āānormāzv=pĀÆā(v)ā(logā”pĀÆā(v)qĀÆā(v)+1)āpĀÆā(v)āāuāSpĀÆā(u)ā(logā”pĀÆā(u)qĀÆā(u)+1). ā ^normā z_v= p(v) ( p(v) q(v)+1 )- p(v) _uā S p(u) ( p(u) q(u)+1 ). (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. _uā S p(u) ( p(u) q(u)+1 )= _uā S p(u) p(u) q(u)+ _uā S p(u)= ^norm+1. By substituding it into Eq 11, we obtain āānormāzv=pĀÆā(v)ā[logā”pĀÆā(v)qĀÆā(v)āānorm],vāS, ā ^normā z_v= p(v) [ p(v) q(v)- ^norm ], vā S, (12) which is the first case of the claim. Since ānorm=vā¼pĀÆā[logā”(pĀÆā(v)/qĀÆā(v))] ^norm=E_v p [ ( p(v)/ q(v)) ] by definition, Eq. (12) says that the gradient on a top-k logit is its log-ratio centered at the pĀÆ p-weighted mean log-ratio. Step 3: the top-k gradients sum to zero. Summing Eq. (12) over vāSvā S and using āvāSpĀÆā(v)ālogā”(pĀÆā(v)/qĀÆā(v))=ānorm _vā S p(v) ( p(v)/ q(v) )= ^norm together with āvāSpĀÆā(v)=1 _vā S p(v)=1 yields āvāSāānormāzv=ānormāānorm=0. _vā S ā ^normā z_v= ^norm- ^norm=0. This is precisely the mean-centering exhibited in Eq. (12): the update redistributes logit mass among the top-k tokens 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 within S, and never on how much probability the student assigns to S as a whole. ā A.2 Proof of Proposition 2 Proof. Fix a step t and a prefix y^<t y_<t, and omit the subscript t throughout, following the convention of Appendix A.1: write S=StkS=S_t^k, zv=zt,vz_v=z_t,v, pĀÆ=pĀÆt(Stk) p= p_t^(S_t^k), qĀÆ=qĀÆt(Stk) q= q_t^(S_t^k), and ptail=pttailp^tail=p_t^tail. Under the tabular assumption, each logit zvz_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; 45; 27). We write ptailā(Ī·)p^tail(Ī·) for the tail probability after one gradient step with learning rate Ī·, and Īāptail=ptailā(Ī·)āptailā(0) p^tail=p^tail(Ī·)-p^tail(0) for its change. We proceed in three steps: we first reduce ptailā(Ī·)p^tail(Ī·) 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 Ī·. Step 1: only the top-k logits change. Split the softmax normalizer into the contributions of the top-k tokens and of the tail tokens, A:=āvāSezv,B:=āvāSezv,Z:=A+B,A:= _vā Se^z_v, B:= _vā Se^z_v, Z:=A+B, so that ptail=B/Zp^tail=B/Z. By Proposition 1, the gradient of ānorm ^norm is gv=pĀÆā(v)ā(logā”pĀÆā(v)qĀÆā(v)āānorm)for āvāS,gv=0for āvāS.g_v= p(v) ( p(v) q(v)- ^norm )\ \ for vā S, g_v=0\ \ for vā S. The update zvāzvāĪ·āgvz_vā z_v-Ī· g_v therefore leaves tail logit unchanged, and hence leaves B unchanged. The change of the tail probability is thus determined by A alone: ptailā(Ī·)=BAā”(Ī·)+B,Aā”(Ī·)=āvāSezvāĪ·āgv.p^tail(Ī·)= BA(Ī·)+B, A(Ī·)= _vā Se^z_v-Ī· g_v. (13) Step 2: the derivative of A is a covariance. Differentiating Aā”(Ī·)A(Ī·) in Eq. (13) with respect to Ī· and evaluating at Ī·=0Ī·=0 gives Aā²(0)=āāvāSezvgv.A (0)=- _vā Se^z_vg_v. We now rewrite this sum in terms of the normalized distribution pĀÆ p. Using ezv=Zāpā(v)e^z_v=Z\,p(v) and pā”(v)=(1āptail)āpĀÆā(v)p(v)=(1-p^tail)\, p(v) for vāSvā S, and then substituting the expression for gvg_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 (0)=-Z (1-p^tail ) _vā S p(v)\,g_v=-Z (1-p^tail ) [ _vā S p(v)^2 p(v) q(v)- ^norm _vā S p(v)^2 ]. (14) Each of the two sums in Eq. (14) carries one factor pĀÆā(v) p(v) that plays the role of a sampling weight, so both are expectations under vā¼pĀÆv 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)]. _vā S p(v)^2 p(v) q(v)=E_v p\! [ p(v) p(v) q(v) ], _vā S p(v)^2=E_v p [ p(v) ]. Moreover ānorm=vā¼pĀÆā[logā”(pĀÆā(v)/qĀÆā(v))] ^norm=E_v p [ ( p(v)/ q(v)) ] by definition, so the bracket in Eq. (14) 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 (0)=-Z (1-p^tail ) *Cov_v p\! ( p(v),\, p(v) q(v) ). (15) The covariance appears because, by Proposition 1, the gradient on a top-k logit is its log-ratio centered at the mean log-ratio, while its contribution to A is weighted by its own probability pĀÆā(v) p(v). Step 3: first-order taylor expansion. Differentiating ptailā(Ī·)p^tail(Ī·) in Eq. (13) at Ī·=0Ī·=0 and substituting Eq. (15) together with B/Z=ptailB/Z=p^tail yields dāptaildāĪ·|Ī·=0=āB(A+B)2āAā²ā(0)=ptailā(1āptail)āCovvā¼pĀÆ(pĀÆā(v),logā”pĀÆā(v)qĀÆā(v)). dp^taildĪ· |_Ī·=0=- B(A+B)^2\,A (0)=p^tail (1-p^tail ) *Cov_v p\! ( p(v),\, p(v) q(v) ). A first-order Taylor expansion in Ī·, followed by splitting the log-ratio as logā”pĀÆā(v)ālogā”qĀÆā(v) p(v)- 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). p^tail=Ī·\,p^tail (1-p^tail ) [ *Cov_v p ( p(v), p(v) )- *Cov_v p ( p(v), q(v) ) ]+O(Ī·^2). Since ptailā(0,1)p^tailā(0,1), the prefactor ptailā(1āptail)p^tail(1-p^tail) is strictly positive, so for sufficiently small Ī·>0Ī·>0 the sign of Īāptail p^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.3 Proof of Proposition 3 Proof. Fix a step t and a prefix y^<t y_<t, and omit the subscript t throughout, following the convention of Appendix A.1: write S=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), and ptail=pttailp^tail=p_t^tail, qtail=qttailq^tail=q_t^tail, so that āTA=āuāSpā”(u)ālogā”pā”(u)qā”(u)+ptailālogā”ptailqtail,ptail=āuāSpā”(u). ^TA= _uā Sp(u) p(u)q(u)+p^tail p^tailq^tail, p^tail= _uā Sp(u). Since S is selected from the teacher distribution, S, q, and qtailq^tail are constants with respect to the student logits. Unlike the normalized loss of Proposition 1, āTA ^TA depends on the student through the unnormalized probabilities, and the tail tokens enter it through their aggregate ptailp^tail. We proceed in three steps: we first differentiate āTA ^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 ^TA as a function of the full probability vector (pā”(u))uā (p(u) )_u rather than eliminating ptailp^tail. A top-k token uāSuā S appears only in its own summand, and a tail token uāSuā S appears only inside ptailp^tail, on which it acts with āptail/āpā”(u)=1ā p^tail/ā p(u)=1. Hence āāTAāpā”(u)=logā”pā”(u)qā”(u)+1,uāS,logā”ptailqtail+1,uāS. ā ^TAā p(u)= cases p(u)q(u)+1,&uā S,\\[10.00002pt] p^tailq^tail+1,&uā S. 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/ _w e^z_w, whose Jacobian is āp(u)/āzv=p(u)([u=v]āp(v))ā p(u)/ā z_v=p(u) (1[u=v]-p(v) ) for all u,vāu,v . Combining it with Eq. (16) 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). ā ^TAā z_v= _u ā ^TAā p(u)\, ā p(u)ā z_v=p(v)\, ā ^TAā p(u) |_u=v-p(v) _u p(u)\, ā ^TAā 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, _u p(u)\, ā ^TAā p(u)= _uā Sp(u) ( p(u)q(u)+1 )+ ( _uā Sp(u) )\! ( p^tailq^tail+1 )= ^TA+1, where the second equality uses āuāSpā”(u)=ptail _uā Sp(u)=p^tail to recover the tail term of āTA ^TA, and āuāpā”(u)=1 _u p(u)=1 to collect the constant. Step 3: the gradient on each logit. Substituting āTA+1 ^TA+1 for the sum in Eq. (17), we have āāTAāzv=pā”(v)ā(logā”pā”(v)qā”(v)āāTA),vāS,pā”(v)ā(logā”ptailqtailāāTA),vāS, ā ^TAā z_v= casesp(v) ( p(v)q(v)- ^TA ),&vā S,\\[11.99998pt] p(v) ( p^tailq^tail- ^TA ),&vā S, cases (18) which establishes (5). ā A.4 Proof of Proposition 4 Proof. Fix a step t and drop the subscript t as in Appendix A.3. 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, ^full= _vā Sp(v) p(v)q(v)+ _vā Sp(v) p(v)q(v), ^TA= _vā Sp(v) p(v)q(v)+p_tail p_tailq_tail, where we have split āfull=DKL(pā„q) ^full=D_KL(p\|q) into the top-k tokens and the tail tokens so that the two losses can be compared term by term. The two objectives share the same top-k term and differ only in how they treat the tail, so taking the difference cancels the top-k term and leaves āfullāāTA=āvāSpā”(v)ālogā”pā”(v)qā”(v)āptailālogā”ptailqtail. ^full- ^TA= _vā Sp(v) p(v)q(v)-p_tail p_tailq_tail. (19) It remains to compute the right-hand side of (19). Substituting pā”(v)=ptailāp~ā(v)p(v)=p_tail\, p(v) and qā”(v)=qtailāq~ā(v)q(v)=q_tail\, q(v) for vāSvā S into its first term and splitting the logarithm gives āvāSpā”(v)ālogā”pā”(v)qā”(v) _vā Sp(v) p(v)q(v) =āvāSptailāp~ā(v)ālogā”ptailāp~ā(v)qtailāq~ā(v) = _vā Sp_tail\, p(v) p_tail\, p(v)q_tail\, q(v) =ptailālogā”ptailqtailāāvāSp~ā(v)+ptailāāvāSp~ā(v)ālogā”p~ā(v)q~ā(v) =p_tail p_tailq_tail _vā S p(v)+p_tail _vā S p(v) p(v) q(v) =ptaillogptailqtail+ptailDKL(p~ā„q~), =p_tail p_tailq_tail+p_tail\,D_KL( p\| q), where the last equality uses that p~ p is a probability distribution over the tail tokens, āvāSp~ā(v)=1 _vā S p(v)=1, and recognizes the remaining sum as DKL(p~ā„q~)D_KL( p\| q). Substituting this back into (19), the terms ptailālogā”ptailqtailp_tail p_tailq_tail cancel and we obtain āfullāāTA=ptailDKL(p~ā„q~). ^full- ^TA=p_tail\,D_KL( p\| q). ā A.5 Proof of Proposition 5 Proof. Fix step t and drop the subscript t as in Appendix A.3: write S=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 y= y_t, ptail=pttailp_tail=p_t^tail, and qtail=qttailq_tail=q_t^tail. Recall that the stop-gradient operator sgā”(ā )sg(Ā·) acts as the identity in the forward pass and has zero derivative, i.e., sgā”(x)=xsg(x)=x in value and āĪøāsgā(x)=0 _Īøsg(x)=0. Throughout, the expectation is taken with respect to y^ā¼p y 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.E_ y p [1[ yā S] ]= _vā Sp(v)=p_tail. (20) Unbiasedness in value. Without considering the gradient, pā”(y^)sgā”(pā”(y^))=1 p( y)sg(p( y))=1, so āSCā-āTA ^SC -TA =āvāS+p(v)sg(logpā”(v)qā”(v))+[y^āS]pā”(y^)sgā”(pā”(y^))sg(logpā”(y^)/ptailqā”(y^)/qtail) = _vā S^+p(v)sg\! ( p(v)q(v) )+1[ yā S]\, p( y)sg(p( y))sg\! ( p( y)/p_tailq( y)/q_tail ) =āvāSp(v)logpā”(v)qā”(v)+ptaillogptailqtail+[y^āS](logpā”(y^)qā”(y^)ālogptailqtail), = _vā Sp(v) p(v)q(v)+p_tail p_tailq_tail+1[ yā S] ( p( y)q( y)- p_tailq_tail ), (21) where the second equality expands S+=SāŖvtailS^+=SāŖ\v_tail\ using pā”(vtail)=ptailp(v_tail)=p_tail and qā”(vtail)=qtailq(v_tail)=q_tail. The first two terms are deterministic given the prefix. For the third term, taking the expectation over y^ā¼p y p and using Eq. (20) gives y^ā¼p[[y^āS](logpā”(y^)qā”(y^)ālogptailqtail)]=āvāSp(v)logpā”(v)qā”(v)āptaillogptailqtail.E_ y p\! [1[ yā S] ( p( y)q( y)- p_tailq_tail ) ]= _vā Sp(v) p(v)q(v)-p_tail p_tailq_tail. Summing the three terms, the two occurrences of ptailālogā”ptailqtailp_tail p_tailq_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),E_ y p [ ^SC -TA ]= _vā Sp(v) p(v)q(v)+ _vā Sp(v) p(v)q(v)=D_KL(p\,\|\,q), which establishes the unbiasedness in value. Unbiasedness in gradient. We first compute the gradient of the exact full-vocabulary reverse KL. Since q is fixed, āĪøDKL(pā„q)=āvāāĪøp(v)(logpā”(v)qā”(v)+1)=āvāāĪøp(v)logpā”(v)qā”(v), _Īø\,D_KL(p\,\|\,q)= _v _Īøp(v) ( p(v)q(v)+1 )= _v _Īøp(v)\, p(v)q(v), where the last equality uses āvāāĪøp(v)=āĪøāvāp(v)=āĪø1=0 _v _Īøp(v)= _Īø _v p(v)= _Īø1=0. We now differentiate āSCā-āTA ^SC -TA. Since the log-ratios are wrapped in stop-gradients, only the leading probabilities carry gradients in the first term. Note that ptail=āvāSpā”(v)p_tail= _vā Sp(v) depends on Īø, so āĪøptail=āvāSāĪøpā(v) _Īøp_tail= _vā S _Īø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. _Īø _vā S^+p(v)sg\! ( p(v)q(v) )= _vā S _Īøp(v)\, p(v)q(v)+ ( _vā S _Īøp(v) ) p_tailq_tail. For the second term, only the numerator pā”(y^)p( y) in pā”(y^)sgā”(pā”(y^)) p( y)sg(p( 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). _Īø\! [1[ yā S]\, p( y)sg(p( y))sg\! ( p( y)/p_tailq( y)/q_tail ) ]=1[ yā S]\, _Īøp( y)p( y) ( p( y)q( y)- p_tailq_tail ). Taking the expectation over y^ā¼p y p of the second term gives y^ā¼p[[y^āS]āĪøpā(y^)pā”(y^)(logpā”(y^)qā”(y^)ālogptailqtail)] _ y p\! [1[ yā S]\, _Īøp( y)p( y) ( p( y)q( y)- p_tailq_tail ) ] =ā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. = _vā Sp(v)Ā· _Īøp(v)p(v) ( p(v)q(v)- p_tailq_tail )= _vā S _Īøp(v)\, p(v)q(v)- ( _vā S _Īøp(v) ) p_tailq_tail. Combining the two terms, the contributions involving logā”ptailqtail p_tailq_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),E_ y p [ _Īø\, ^SC -TA ]= _vā S _Īøp(v)\, p(v)q(v)+ _vā S _Īøp(v)\, p(v)q(v)= _Īø\,D_KL(p\,\|\,q), which establishes the unbiasedness in gradient. This completes the proof. ā Appendix B Detailed 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 large k. 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, 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 when kā¤64kā¤ 64. The tail probability increase occurs when k is not sufficiently large. We conduct an ablation study on how k affects the tail probability increase by fixing the student and teacher, and varying only k. Specifically, we use Qwen2.5-Math-1.5B as the student and DeepSeek-R1-Distill-Qwen-7B as the teacher, and sweep kā8,16,32,64kā\8,16,32,64\. For this ablation study, we set the training max response length to 4096. In Figure 8, we visualize the studentās and teacherās tail probability on student-generated prefixes, together with the studentās entropy, across different values of k. The tail probability increase emerges when k is not sufficiently large (kā¤64kā¤ 64): over training steps, the studentās tail probability steadily increases together with its entropy. The effect becomes stronger as k decreases, with smaller k yielding higher entropy and a larger increase in tail probability. In principle, this issue vanishes as k approaches the vocabulary size, since the top-k objective 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: the magnitude of the single-step increase is proportional to pttailā(1āpttail)p_t^tail(1-p_t^tail), which shrinks as pttailā0p_t^tailā 0. Since a larger k yields a smaller initial tail probability, each update moves pttailp_t^tail less, mitigating the compounding increase over training. Overall, the tail probability increase arises when k is not sufficiently large and is mitigated as k increases. 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, 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-k tokens, 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-k OPD. Prior work also adopts the studentās top-k tokens as the support (21). We examine this choice with Qwen3-1.7B-Base as the student and Qwen3-8B as the teacher, varying only the support between TopKā”(qt,k)TopK(q_t,k) and TopKā”(pt,k)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-k tokens for the teacher top-k variant and the studentās top-k tokens for the student top-k variant. As shown in Figure 10, the tail probability increase is specific to the teacher top-k support. Two mechanisms explain this. First, among all token sets of size k, TopKā”(pt,k)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, the per-step change is proportional to pttailā(1āpttail)p_t^tail(1-p_t^tail), and each update therefore moves pttailp_t^tail only slightly. Second, the student top-k support 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 top k. The support therefore always consists of the studentās k most probable tokens, and the probability mass it covers cannot leak away. The teacher top-k support, 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-k supports. The student and teacher models are Qwen3-1.7B-Base and Qwen3-8B. The tail probabilities are computed with respect to the support in use. The tail probability increase occurs only under the teacher top-k support. The teacher top-k OPD yields better downstream performance than the student top-k OPD. Although the student top-k support avoids the tail probability increase, it underperforms the teacher top-k support on downstream benchmarks (Table 3). We attribute this to two factors. First, as shown in Figure 11, the average rollout length under the student top-k support grows monotonically and saturates at the maximum response length of 71687168 tokens within the first 5050 steps, while the teacher top-k variant maintains a healthy response length. Second, the teacher top-k support directly covers the teacherās high-probability tokens, whereas the studentās top-k tokens 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-k OPD in this work. Figure 11: Average response length across training steps under the student and teacher top-k supports. The student and teacher models are Qwen3-1.7B-Base and Qwen3-8B, respectively. Under the student top-k support, the response length saturates at the maximum response length of 71687168 tokens, while the teacher top-k variant maintains a healthy response length. Table 3: Performance of normalized top-k OPD under the student and teacher top-k supports. The Avg. column represents the macro-average across all six math reasoning benchmarks. Best results are shown in bold. Method MATH500 Minerva Olympiad AMC AIME 24 AIME 25 Avg. Student: Qwen3-1.7B-Base Teacher: Qwen3-8B Student top-k 68.58 25.32 33.35 36.00 9.17 6.25 29.78 Teacher top-k 69.93 25.46 33.41 36.00 9.17 7.08 30.18 Appendix C Compared Methods Sampled-token OPD. Sampled-token OPD (22; 36) supervises only the token y^tā¼pt y_t p_t sampled by the student: ātsample=logā”ptā(y^t)ālogā”qtā(y^t). _t^sample= p_t( y_t)- q_t( 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-k OPD. Unnormalized top-k OPD restricts the divergence computation to the teacherās top-k tokens Stk=TopKā”(qt,k)S_t^k=TopK(q_t,k) while keeping the original probabilities: ātunnorm=āvāStkptā(v)ālogā”ptā(v)qtā(v). _t^unnorm= _vā S_t^kp_t(v) p_t(v)q_t(v). It provides dense supervision over StkS_t^k, but is not a well-defined divergence. Normalized top-k OPD. Normalized top-k OPD (21; 7) instead renormalizes both distributions on StkS_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). _t^norm=D_KL\! ( p_t^(S_t^k)\,\|\, q_t^(S_t^k) )= _vā S_t^k p_t^(S_t^k)(v) p_t^(S_t^k)(v) q_t^(S_t^k)(v). It aligns the studentās relative shape with the teacherās over StkS_t^k, but the normalization discards the tail probability, which, as we show in Section 3.1, steadily increases the studentās tail probability and entropy. Appendix D Implementation Details D.1 Experimental Details We implement all methods using the VERL framework (32) 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. For evaluation, we adopt the same prompt as 42, which is shown in Prompt D.1. 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. Category Item Value Training Training temperature 1.0 Global batch size 72 Mini batch size 36 Rollout number 4 k (Number of Top-k Tokens) 16 Top-p 1.0 Max prompt length 1024 Max response length 7168 Learning rate 1e-6 Training step 300 loss aggregation token-mean optimizer AdamW Evaluation Temperature 0.7 Top-p 0.95 Max new tokens 8192 Your 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> thoughts </think> ā. Each thought should include detailed analysis, brainstorming, verification, and refinement of ideas. After ā</think> ā 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 for closed-form results like multiple choices or mathematical solutions. User: QUESTION Assistant: D.2 Numerically Stable Computation of TA-OPD The per-token TA-OPD loss in Eq. (4) depends on the tail log-probabilities logā”pttail p_t^tail and logā”qttail q_t^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-k log-probabilities logā”qtā(v) 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-k tokens 1āqttail=āvāStkqtā(v)1-q_t^tail= _vā S_t^kq_t(v), and assign the remaining tail probability qttailq_t^tail to vtailv_tail. However, this is unstable. When the teacherās tail probability is close to zero, numerical error can make the sum āvāStkqtā(v) _vā S_t^kq_t(v) exceed one. The tail probability qttailq_t^tail then becomes negative, rendering logā”qttail q_t^tail ill-defined (āā-ā or NaN) and producing NaN gradients. Our log-space implementation. To address the numerical issue, our implementation avoids forming pttailp_t^tail and qttailq_t^tail in probability space and carries out the whole computation in log space. We provide the pseudo code of our implementation in Listing 1. In particular, we first obtain the log total probability of the top-k tokens directly from the top-k log-probabilities via a log-sum-exp, logā”(1āpttail)=logsumexpvāStkā”logā”ptā(v),logā”(1āqttail)=logsumexpvāStkā”logā”qtā(v), (1-p_t^tail )=logsumexp_vā S_t^k p_t(v), (1-q_t^tail )=logsumexp_vā S_t^k q_t(v), which is numerically stable and never overflows. To further rule out the boundary case pttailā0p_t^tailā 0, we clamp the log total probability such that logā”(1āpttail)ā¤āĻµ (1-p_t^tail)ā¤-Īµ, ensuring a strictly positive tail probability and a finite log-domain computation. The tail log-probability term logā”pttail p_t^tail is then computed using the log1mexp primitive, which directly evaluates logā”(1āexpā”(a)) (1- (a)) in log space: logā”pttail=log1mexpā”(logā”(1āpttail)). p_t^tail=log1mexp\! ( (1-p_t^tail) ). The teacher term logā”qttail q_t^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. ā¬ def compute_taopd_loss(student_topk_log_probs, teacher_topk_log_probs, eps): # Input: # student_topk_log_probs: student log-probs on the teacherās top-k tokens # teacher_topk_log_probs: teacher top-k log-probs, log q(v) # eps: small constant to keep the tail prob positive # Return: # loss: per-token TA-OPD reverse KL divergence on the augmented token set # 1. log total probability of the top-k tokens via logsumexp student_log_topk_prob = logsumexp(student_topk_log_probs, dim=-1) # log(1 - p_tail) teacher_log_topk_prob = logsumexp(teacher_topk_log_probs, dim=-1) # log(1 - q_tail) # 2. tail-token log-prob log(p_tail) via log1mexp student_tail_log_prob = log1mexp(clamp_max(student_log_topk_prob, -eps)) teacher_tail_log_prob = log1mexp(clamp_max(teacher_log_topk_prob, -eps)) # 3. append the tail token and compute reverse KL on the augmented token set 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) return taopd_loss Appendix E Extensive Study E.1 Synthetic 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Ć30R^1Ć 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,ā¦,30V=\1,ā¦,30\ as the following bimodal distribution: qā”(v)āexpā”(ā(vā10)22ā 22)+ 0.88ā expā”(ā(vā20)22ā 22),vā.q(v) \! (- (v-10)^22Ā· 2^2 )\;+\;0.88Ā· \! (- (v-20)^22Ā· 2^2 ), v . The student policy is parameterized as a categorical distribution over V: pĪøā(v)=softmaxā(zĪø)v,zĪøāā30,p_Īø(v)=softmax(z_Īø)_v, z_Īø ^30, where zĪøz_Īø is initialized from ā”(0,0.012)N(0,0.01^2). The student is optimized for 20000 steps using AdamW with learning rate 1Ć10ā31Ć 10^-3. We compare four objectives: full-vocabulary OPD, normalized top-k OPD, TA-OPD, and sample-corrected TA-OPD. For top-k OPD, k is set to 88. Figure 12: Visualization of the student policies under different OPD objectives. Normalized top-k OPD (Red) only matches the relative shape of the teacher distribution over the top-k tokens, but fails to align the tail probability. In contrast, TA-OPD (Blue) and sample-corrected TA-OPD (Purple) closely recover the bimodal teacher distribution. TA-OPD matches both the shape and tail probability of the teacher policy. Figure 12 compares the final student policies trained with different OPD objectives. TA-OPD closely recovers the bimodal teacher distribution. While normalized top-k OPD matches the relative shape over the teacherās top-k tokens, 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-k OPD. 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 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-k OPD converges to a point that remains far from the teacher under the full-vocabulary KL. Overall, TA-OPDās objective provides a better top-k estimate of the full-vocabulary reverse KL than the normalized objective. E.2 Top-k Overlap Ratio of Different OPD Methods 21 find that the top-k overlap 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-k overlap ratio is the fraction of the teacherās top-k tokens that also fall in the studentās top-k tokens, Overlaptk=|TopKā”(pt,k)ā©TopKā”(qt,k)|k,Overlap_t^k= |TopK(p_t,k) (q_t,k) |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-k objective supervises a token set that is meaningful to the student. We use the three studentāteacher pairs of Section 3.2 and report the ratio across training steps. Figure 14: Top-k overlap ratio between the student and the teacher across training steps. TA-OPD achieves higher overlap ratios than normalized top-k OPD across different model pairs. TA-OPD attains a higher top-k overlap ratio. Figure 14 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-k OPD 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-k tokens, which is consistent with its stronger downstream accuracy. E.3 Computational Overhead of TA-OPD TA-OPD introduces no additional computational overhead compared with normalized top-k OPD. Both methods require exactly the same teacher-side information: the log-probabilities of the teacherās top-k tokens. Given these quantities, TA-OPD only additionally computes the tail log-probabilities logā”pttail p_t^tail and logā”qttail q_t^tail, which involve a single log-sum-exp over k values 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 token y^t y_t. Although this appears to be an extra query, it incurs negligible additional cost in practice: when queried for the top-k log-probabilities, inference engines such as vLLM (18) always compute and return the sampled tokenās log-probability as well, even if it falls outside the top-k tokens. 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-k OPD, 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. As shown in Figure 15 (a), completing the full 300300 training steps takes 2323h22m1818s for TA-OPD and 2222h4747m5656s for normalized top-k OPD, i.e., a difference of 1414m2222s, or 1.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 (b) further compares the per-step throughput distributions, whose medians are around 10001000 and 10051005 tokens/s for TA-OPD and normalized top-k OPD, respectively. Overall, TA-OPD incurs negligible computational cost compared to normalized top-k OPD. (a) Total training runtime. (b) Throughput distribution during training. Figure 15: Computational overhead comparison between sampled-token OPD, normalized Top-k OPD, TA-OPD, and sample-corrected TA-OPD. Left: total training runtime. Right: throughput distribution across training steps. The student and teacher models are Qwen3-1.7B and Qwen3-30B-A3B-Instruct-2507, respectively. E.4 Ablation on the Tail Token Our analysis attributes the tail probability increase issue to top-k normalization. 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 token vtailv_tail, training will diverge. Setup. We ablate the tail token by comparing TA-OPD with unnormalized top-k OPD, 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. _t^unnorm= _vā S_t^kp_t(v) p_t(v)q_t(v)= _t^TA-p_t^tail p_t^tailq_t^tail. The student and teacher are Qwen2.5-7B-Instruct and OpenThinker3-7B, with all other settings following Table 4. Figure 16: The studentās tail probability, the teacherās tail probability, and the studentās token-level entropy across training steps. Without the tail token, the studentās tail probability rises above 0.90.9 and training collapses. Table 5: Avg@8 of unnormalized top-k OPD and TA-OPD. The Avg. column represents the macro-average across all six benchmarks. Best results are shown in bold. Methods MATH500 Minerva Olympiad AMC AIME 24 AIME 25 Avg. Student: Qwen2.5-7B-Instruct Teacher: OpenThinker3-7B UnNorm. top-k 3.35 1.70 1.42 3.00 0.42 0.42 1.72 TA-OPD (Ours) 77.88 32.58 42.41 45.03 16.67 17.50 38.68 Removing the tail token collapses training. Figure 16 compares the training dynamics of the two objectives. Under unnormalized top-k OPD, the studentās tail probability rises above 0.90.9 within 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): 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 at pt=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 (v)=q_t(v)/e ā vā S_t^k, p_t^tail, =1- 1-q_t^taile\;ā„\;1-e^-1ā 0.63. That is, the objective is minimized by deflating every top-k probability by a factor of e and assigning the removed mass to the tail, which drives the loss negative. The optimum therefore assigns at least 63%63\% of the probability mass outside the teacherās top-k tokens, regardless of the teacher distribution. Adding the tail token contributes exactly the missing term pttailālogā”(pttail/qttail)p_t^tail (p_t^tail/q_t^tail), which turns the objective into a reverse KL divergence on St+S_t^+: it is non-negative and uniquely minimized when the student matches the teacher on both the top-k tokens and the tail. Overall, introducing the tail token is necessary for TA-OPD. E.5 Discussion of EMA-PG Exponential Moving Average-Policy Gradient (EMA-PG) (46) proposes a top-k KL 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)). _t^EMA= _vā S_t^kp_t(v)sg\! ( p_t(v)q_t(v) )_ =\; _t^unnorm\ in value+1[ y_tā S_t^k]\, p_t( y_t)sg(p_t( y_t))sg\! ( p_t( y_t)q_t( y_t) ). (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). _t^SC -TA= _vā S_t^+p_t(v)sg\! ( p_t(v)q_t(v) )_ =\; _t^TA\ in value+1[ y_tā S_t^k]\, p_t( y_t)sg(p_t( y_t))sg\! ( p_t( y_t)/p_t^tailq_t( y_t)/q_t^tail ). 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) _vā S_t^kp_t(v) p_t(v)q_t(v) with the sampled token. Instead, our estimator starts from ātTA _t^TA, which already accounts for the tail probability through the tail token, and uses the sampled token only to estimate the residual pttailDKL(p~tā„q~t)p_t^tailD_KL( p_t\| q_t) identified in Proposition 4. In other words, our estimator is obtained by debiasing TA-OPD, whereas EMA-PG debiases the unnormalized top-k objective. Behavior when the sampled token falls inside the top-k tokens. When y^tāStk y_tā S_t^k, which occurs with probability 1āpttail1-p_t^tail and thus covers most tokens, the correction term vanishes for both losses. In this case, ātEMA _t^EMA reduces to the unnormalized top-k loss ātunnorm _t^unnorm, while ātSCā-āTA _t^SC -TA reduces to ātTA _t^TA. As shown in Appendix E.4, ātTA _t^TA is superior to ātunnorm _t^unnorm. Consequently, although EMA-PG is unbiased in expectation, on most tokens its realized objective is the unnormalized top-k loss 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 on St+S_t^+ that is uniquely minimized at pt=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)).Var ( _t^EMA )-Var ( _t^SC -TA )=p_t^tail (1-p_t^tail )\, p_t^tailq_t^tail ( p_t^tailq_t^tail+2\,D_KL( p_t\| q_t) ). (23) Since DKL(p~tā„q~t)ā„0D_KL( p_t\| q_t)ā„ 0, the gap is strictly positive whenever pttail>qttailp_t^tail>q_t^tail, i.e. whenever the studentās tail probability exceeds the teacherās, which is likely to happen when using the teacherās top-k support. Since the support StkS_t^k consists of the teacherās top-k tokens, the teacherās tail probability is typically smaller than the studentās. Figures 4 confirms this empirically, where qttailq_t^tail is below pttailp_t^tail throughout training across all model pairs. Overall, our loss objective tends to have lower variance than EMA-PG. Appendix F Detailed Results F.1 Training Dynamics of Additional Model Pairs Figure 17 presents additional comparisons between TA-OPD and normalized top-k OPD 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) ā Qwen3-30B-A3B-Instruct-2507 (Teacher) Llama-3.1-8B (Student) ā DeepSeek-R1-Distill-Llama-8B (Teacher) Qwen3-1.7B-Base (Student) ā Qwen3-8B (Teacher) Figure 17: Tail probability and token-level entropy over training for TA-OPD and Normalized top-k OPD. Each row corresponds to a different student-teacher pair. Qwen2.5-Math-1.5B (Student) ā JustRL-DeepSeek-1.5B (Teacher) DeepSeek-R1-Distill-Qwen-1.5B (Student) ā JustRL-DeepSeek-1.5B (Teacher) Figure 18: Tail probability and token-level entropy over training for TA-OPD and Normalized top-k OPD under different student-teacher capability gaps. The tail probability increase only occurs for Qwen2.5-Math-1.5B but not for DeepSeek-R1-Distill-Qwen-1.5B.