PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Huipeng Huang, Hongxin Wei

arXiv:2608.14728v1cs.LGcs.AI

TL;DR

Normalized top-k on-policy distillation discards tail probability, which can increase student tail probability and entropy. TA-OPD restores this signal with a tail token and improves Avg@8 accuracy by up to 8.05 points on common benchmarks.

  • Problem

    Normalized top-k OPD discards the teacher’s tail probability, limiting alignment between student and teacher distributions during on-policy distillation.

  • Method

    TA-OPD minimizes reverse KL over the teacher’s top-k tokens plus a tail token carrying the remaining probability mass.

  • Results

    TA-OPD achieves superior performance across benchmarks, improving Avg@8 accuracy by up to 8.05 points over normalized top-k OPD.

  • Takeaways & Limitations

    Restoring tail probability provides a practical way to address the tail probability and entropy increases caused by top-k normalization.

  • Takeaways & Limitations

    TA-OPD’s performance gains diminish as k increases or the student–teacher capability gap decreases, and experiments are limited to models up to 8B parameters.

Abstract

from arXiv · show

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 https://github.com/HuipengHuang/TA-OPD.

1 INTRODUCTION

Standard off-policy distillation creates a training–inference mismatch that causes errors to accumulate, while normalized top-k OPD discards tail probability and can increase student tail probability and entropy. TA-OPD restores this signal by adding a tail token, improving alignment and stabilizing distillation.

  • Off-policy distillation trains students on teacher-generated sequences, although inference conditions on the students’ own generations, causing errors to accumulate.
  • Normalized top-k OPD discards probability mass outside the teacher’s top-k tokens, allowing student tail probability and entropy to increase.
  • TA-OPD minimizes reverse KL over the teacher’s top-k tokens plus a tail token carrying tail probability.
  • In the OpenThinker3-7B to Qwen2.5-7B-Instruct setting, normalized top-k OPD reaches training entropy around 6 and MATH500 Avg@8 accuracy of 68.78%.
  • TA-OPD keeps training entropy below 1.5 and raises MATH500 Avg@8 accuracy to 77.

2 PRELIMINARIES

OPD distills a teacher LLM into a student using student-sampled trajectories and reverse KL alignment of next-token distributions. Existing approximations trade off computational cost and supervision density, while normalized teacher-top-k OPD omits tail probability and can allow substantial distributional divergence.

  • On-policy distillation: OPD distills teacher πte into student πθ on trajectories sampled from the student.For prompt x ∼ Dx, the student samples a response and supplies the prefixes used for distillation.
  • On-policy distillation: At each student-generated prefix, OPD defines student and teacher next-token distributions pt and qt over vocabulary V and minimizes reverse KL divergence.The distributions are pt = πθ(· | x, ŷ<t) and qt = πte(· | x, ŷ<t).
  • OPD approximations: Sampled-token OPD is cheap but discards dense vocabulary information, whereas full-vocabulary OPD provides exact dense supervision at prohibitive LLM cost.Sampled-token OPD uses ℓsample_t = log pt(ŷt) − log qt(ŷt), while full-vocabulary OPD computes DKL(pt∥qt) over V.
  • Normalized top-k OPD: Normalized top-k OPD restricts divergence computation to a subset, with prior work using either student top-k or teacher top-k tokens.The paper primarily focuses on the teacher top-k variant, which reduces teacher-query cost while retaining dense multi-token supervision.
  • Limitation: Normalized teacher-top-k OPD discards tail probability outside the selected tokens, so its objective can be minimized even when the student’s tail probability is much higher.This mismatch can make the student’s next-token distribution diverge substantially from the teacher’s.

3 MOTIVATION

The normalized top-k OPD objective lacks a mechanism to reduce the student’s tail probability and can systematically increase tail probability and entropy. This increase is empirically associated with substantial accuracy degradation, especially under large capability gaps, small k, and longer response lengths.

  • Mechanism: The normalized objective can steadily increase the student’s tail probability because tail logits receive zero gradient and top-k logit gradients sum to zero.Reducing tail probability requires raising top-k logits relative to tail logits, which this objective cannot explicitly do.
  • Mechanism: The student’s tail probability increases after one optimization step when student self-covariance exceeds student–teacher covariance.The student self-covariance is nonnegative, while student–teacher covariance can be small or negative when distributions are poorly aligned.
  • Empirical evidence: Across three student–teacher pairs, normalized top-k OPD steadily increases the student’s tail probability and entropy during on-policy training.These experiments use DAPO-MATH-17K with k = 16 and validate the theoretical analysis.
  • Empirical evidence: Accuracy falls from above 0.7 to below 0.2 as tail probability grows from near 0 to around 0.7 on MATH500.The relationship is measured across checkpoints from normalized top-k OPD with Qwen2.5-7B-Instruct as student and OpenThinker3-7B as teacher.
  • Conditions: The tail-probability increase is more likely with a large capability gap, k ≤64, or longer maximum training response lengths.Appendix analysis finds the phenomenon is more pronounced when the maximum training response length becomes longer.

4 METHOD

TA-OPD restores the tail-probability signal lost by normalized top-k OPD by adding a tail token and minimizing reverse KL over the top-k tokens plus that token. This mechanism updates tail logits, adjusts top-k logits when the student overweights the tail, and provides a tight lower-bound approximation to full-vocabulary reverse KL.

  • Tail-Aware Top-k OPD: TA-OPD adds a tail token carrying the remaining probability and minimizes reverse KL over the top-k tokens plus this token.Each top-k token retains its original probability, while the tail token represents the student or teacher tail probability.
  • Gradient signal: TA-OPD updates logits outside the top-k through a tail log-ratio comparing student and teacher tail probabilities.Unlike normalized top-k OPD, which gives no gradient to outside logits, this signal drives the student tail probability toward the teacher’s.
  • Gradient signal: TA-OPD can raise top-k logits relative to the tail when the student overweights the tail relative to the teacher.Its top-k gradient sum need not be zero, enabling this corrective behavior.
  • Relation to full-vocabulary OPD: TA-OPD differs from full-vocabulary reverse KL by a non-negative residual term involving the normalized tail distributions.The residual is ptail DKL(˜pt∥˜qt).
  • Relation to full-vocabulary OPD: TA-OPD becomes a tight lower bound to full-vocabulary reverse KL as both student and teacher tail probabilities approach zero.The residual term becomes small in practice under this condition.

5 EXPERIMENTS

Section 5 evaluates TA-OPD against sampled-token and normalized top-k OPD across math and out-of-distribution benchmarks. TA-OPD achieves the strongest average accuracy, stabilizes tail probability and entropy, and remains effective with small k.

  • Main results: TA-OPD attains the best average accuracy across all three student–teacher pairs, exceeding normalized top-k OPD by +5.26 points on Qwen2.5-7B-Instruct and +8.05 points on Llama-3.1-8B.The comparison covers six math reasoning benchmarks and includes sampled-token OPD and normalized top-k OPD.
  • Training dynamics: Normalized top-k OPD raises student tail probability, teacher tail probability, and student entropy to approximately 0.6, 0.5, and 6, respectively, during training.These dynamics are reported for Qwen2.5-7B-Instruct and OpenThinker3-7B.
  • Training dynamics: TA-OPD keeps both tail probabilities close to 0 and student entropy below 1.5 throughout training.This contrasts with normalized top-k OPD on the Qwen2.5-7B-Instruct and OpenThinker3-7B pair.
  • Sensitivity to k: With Qwen3-1.7B-Base as student and Qwen3-8B as teacher, k = 2 achieves an average accuracy of 30.43, only 0.43 points below the best result.Accuracy is insensitive to k once k ≥2 across six macro-averaged math reasoning benchmarks.
  • Additional results: Additional analyses find higher student–teacher top-k overlap for TA-OPD and negligible computational overhead compared with normalized top-k OPD.A synthetic experiment further analyzes the methods’ optimization behavior.

6 DISCUSSION

The discussion introduces sample-corrected TA-OPD, which is unbiased for the full-vocabulary reverse KL in both value and gradient while retaining dense supervision. Empirically, bias correction offers little improvement, whereas TA-OPD is especially effective when the student–teacher capability gap is large.

  • Sample-corrected TA-OPD: Sample-corrected TA-OPD adds a sampled-token bias estimate to TA-OPD’s loss to obtain an unbiased full-vocabulary reverse-KL estimate.The correction uses the sampled token’s probability and a stop-gradient operator.
  • Sample-corrected TA-OPD: Proposition 5 establishes that sample-corrected TA-OPD is unbiased in both loss value and gradient, combining sampled-token OPD’s unbiasedness with top-k OPD’s dense supervision.The proof is provided in Appendix A.5.
  • Empirical comparison: 53.62 to 54.39 and 44.62 to 44.65 are the mathematical-reasoning macro-average changes after bias correction, while OOD averages remain essentially unchanged.The corresponding OOD averages are 30.69 vs. 30.66 and 67.01 vs. 67.10; the authors attribute the limited gains to TA-OPD’s negligible practical bias.
  • Capability gap: 75.20 is TA-OPD’s Avg@8 on MATH500 for Qwen2.5-Math-1.5B, outperforming normalized top-k OPD by 13.50 points; on DeepSeek-R1-Distill-Qwen-1.5B, sample-corrected TA-OPD reaches 87.00.The larger capability gap corresponds to Qwen2.5-Math-1.5B, and the discussion links its stronger benefit to a more pronounced tail-probability increase.

7 RELATED WORK

Related work centers on on-policy distillation, particularly top-k variants that provide dense supervision at tractable cost. KL divergence is also central to LLM post-training, serving as a reverse-KL objective in OPD and a base-model constraint in reinforcement learning.

  • On-policy Distillation: On-policy distillation has attracted growing interest as an effective post-training paradigm, with top-k OPD emerging to provide dense supervision at tractable cost.The passage identifies top-k OPD as a popular research direction.
  • On-policy Distillation: Normalized top-k OPD uses normalized top-k reverse KL directly as its loss function, while entropy-aware OPD applies normalized topThe supplied passage truncates the description of entropy-aware OPD after “normalized top.”
  • KL computation for LLM post-training: KL divergence is a crucial component of LLM post-training, with OPD minimizing reverse KL to the teacher and reinforcement learning commonly constraining KL against the base model.The reinforcement-learning constraint is used to prevent policy drift away from the base model.

8 CONCLUSION · A PROOF

TA-OPD restores the tail-probability signal omitted by normalized top-k OPD by adding a tail token to reverse-KL distillation. Its gains diminish with larger k or smaller student–teacher capability gaps, and experiments covered models up to 8B parameters without assuming transfer to 30B+ models.

  • 8 CONCLUSION: TA-OPD restores the missing tail-probability signal in on-policy distillation.It is introduced as a novel distillation method addressing the information discarded by top-k normalization.
  • 8 CONCLUSION: TA-OPD minimizes reverse KL over the teacher’s top-k tokens plus a tail token carrying the tail probability.The added tail token represents the probability mass outside the teacher’s top-k tokens.
  • 8 CONCLUSION: TA-OPD explicitly aligns the student’s tail probability with the teacher’s.This mechanism targets the tail-probability mismatch created by top-k normalization.
  • 8 CONCLUSION: TA-OPD addresses the tail-probability and entropy increases caused by top-k normalization.The method’s tail-aware objective is designed to correct both effects.
  • 8 CONCLUSION: TA-OPD’s performance gain over normalized top-k OPD diminishes as k increases.The limitation concerns comparisons against normalized top-k OPD.
  • 8 CONCLUSION: TA-OPD’s performance gain over normalized top-k OPD diminishes as the student–teacher capability gap decreases.The reported limitation applies when the student and teacher become more similar in capability.
  • 8 CONCLUSION: Experiments were restricted to models with up to 8B parameters because of limited computational resources.The authors describe these sizes as standard for research-stage OPD studies.
  • 8 CONCLUSION: The authors do not presume automatic transfer of the reported results to 30B+ models.This caution follows the experiments’ parameter-size restriction.

A.1 PROOF OF PROPOSITION 1 … A.5 PROOF OF PROPOSITION 5

Appendix A proves that top-k normalization updates only the student’s distribution shape within the teacher’s top-k set, leaving tail allocation unconstrained, while TA-OPD restores tail-sensitive gradients. It also establishes the resulting tail-probability condition, the relation to full-vocabulary reverse KL, and unbiased value and gradient estimators.

  • A.1 PROOF OF PROPOSITION 1: The normalized top-k loss is a restricted softmax over teacher-selected tokens, so its Jacobian vanishes for logits outside that set.The teacher top-k set and normalized teacher distribution are fixed with respect to student logits.
  • A.1 PROOF OF PROPOSITION 1: Its top-k logit gradients sum to zero, redistributing mass within the set without changing the student’s total top-k probability or tail allocation.Thus the normalized objective acts only on the within-set distribution shape.
  • A.2 PROOF OF PROPOSITION 2: Under the tabular assumption, only top-k logits change after a normalized-loss update, so the tail-probability change is governed by the top-k normalizer.The tail contribution to the softmax normalizer remains unchanged.
  • A.2 PROOF OF PROPOSITION 2: For sufficiently small η > 0, tail probability increases if and only if student self-covariance exceeds student–teacher covariance.The sign is determined by the covariance difference because ptail(1−ptail) is strictly positive.
  • A.3 PROOF OF PROPOSITION 3: TA-OPD differentiates through the full softmax and includes the aggregate student tail probability, giving gradients for both top-k and tail logits.Unlike the normalized loss, TA-OPD depends on unnormalized probabilities and treats tail tokens through ptail.
  • A.4 PROOF OF PROPOSITION 4: The TA-OPD and full-vocabulary reverse-KL objectives share the same top-k term and differ only in their treatment of the tail.Splitting the full reverse KL into top-k and tail contributions cancels the common top-k component.
  • A.5 PROOF OF PROPOSITION 5: The stop-gradient construction is unbiased in value: its expected sampled objective equals the exact tail-aware objective.The proof expands the sampled tail token using p(vtail) = ptail and q(vtail) = qtail, after which matching ptail log ptail terms cancel.
  • A.5 PROOF OF PROPOSITION 5: The same stop-gradient construction is unbiased in gradient because the log-ratio contributions involving log ptail and log qtail cancel exactly.Only the leading probabilities carry gradients, yielding the exact full-vocabulary reverse-KL gradient.

B DETAILED ANALYSIS OF THE TAIL PROBABILITY INCREASE

The tail probability increase arises with a large student–teacher capability gap, insufficiently large k, and longer responses. It is specific to teacher top-k support, while student top-k support avoids the increase but underperforms downstream.

  • Capability gap: The increase emerges only with a large student–teacher capability gap, while narrowing the gap keeps tail probability and entropy low and stable.Qwen3-1.7B-Base shows steadily increasing tail probability and entropy, whereas Qwen3-8B-Base does not.
  • Top-k size: The increase occurs when k ≤64, becoming stronger as k decreases because smaller k yields higher entropy and greater tail probability growth.Larger k reduces the initial tail probability and mitigates compounding increases over training.
  • Response length: Longer maximum response lengths amplify the increase, causing tail probability and entropy to grow faster and reach higher values.The experiments compare maximum response lengths of 4096 and 7168 tokens and attribute the effect to error compounding along longer trajectories.
  • Support choice: The increase is specific to teacher top-k support; student top-k support avoids it because its dynamically recomputed support cannot lose covered probability mass.Teacher support is fixed, whereas student support follows the student’s k most probable tokens at every step.
  • Support choice: Student top-k OPD avoids the tail probability increase but underperforms teacher top-k OPD on downstream benchmarks, so the study focuses mainly on teacher top-k OPD.Under student top-k support, average rollout length saturates at the maximum response length of 7168 tokens within the first 50 steps.

C COMPARED METHODS … E.3 COMPUTATIONAL OVERHEAD OF TA-OPD

The paper compares OPD objectives, details a numerically stable TA-OPD implementation, and evaluates its distributional alignment, overlap, and computational cost. TA-OPD restores tail supervision while adding negligible overhead relative to normalized top-k OPD.

  • C COMPARED METHODS: Sampled-token OPD is an unbiased full-vocabulary reverse-KL estimate but discards dense vocabulary information and suffers from high variance.It supervises only the student-sampled token and requires that token’s teacher log-probability.
  • C COMPARED METHODS: Normalized top-k OPD aligns relative probabilities over the teacher’s top-k tokens but discards tail probability, increasing the student’s tail probability and entropy.Unnormalized top-k OPD retains original probabilities and provides dense top-k supervision, but it is not a well-defined divergence.
  • D.1 EXPERIMENTAL DETAILS: Experiments use VERL with 6 NVIDIA Pro 6000 GPUs, 8 NVIDIA A100 80 GB GPUs, or 32 NVIDIA A100 40 GB GPUs, alongside default Table 4 settings.Evaluation follows Yan et al. (2026)’s prompt, and Llama-3.1-8B student experiments use the DeepSeek-R1-Distill-Llama tokenizer and chat template.
  • D.2 NUMERICALLY STABLE COMPUTATION OF TA-OPD: TA-OPD computes tail probabilities in log space using log-sum-exp, clamping, and log1mexp to avoid numerical instability.The naive approach can make the summed top-k probability exceed one, producing negative tail probabilities and invalid logarithms.
  • E.1 SYNTHETIC EXPERIMENT: In a 30-armed bandit, TA-OPD matches both the teacher’s top-k shape and tail probability, whereas normalized top-k OPD assigns substantially larger tail probability.TA-OPD closely recovers the teacher’s bimodal distribution and provides a closer approximation to the teacher policy.
  • E.1 SYNTHETIC EXPERIMENT: TA-OPD better approximates full-vocabulary OPD during optimization trajectories.The comparison evaluates full-vocabulary reverse KL on a two-dimensional principal-component projection of centered logit trajectories.
  • E.2 TOP-k OVERLAP RATIO OF DIFFERENT OPD METHODS: TA-OPD increases top-k overlap monotonically and keeps it stable, while normalized top-k OPD remains lower across the reported model pairs.For Qwen2.5-7B-Instruct, normalized top-k overlap peaks early then collapses; for Llama-3.1-8B, it recovers slowly and remains below TA-OPD.
  • E.3 COMPUTATIONAL OVERHEAD OF TA-OPD: 14m22s, or 1.05% of total runtime, separates TA-OPD from normalized top-k OPD over 300 training steps.TA-OPD and normalized top-k OPD require the same teacher top-k log-probabilities, while TA-OPD adds only one log-sum-exp over k values per token.

E.4 ABLATION ON THE TAIL TOKEN … F.1 TRAINING DYNAMICS OF ADDITIONAL MODEL PAIRS

The tail token is necessary for stable TA-OPD training because removing it causes collapse, while sample-corrected TA-OPD remains a well-defined divergence with lower variance than EMA-PG. Additional model-pair comparisons further examine tail probability and entropy dynamics across capability gaps.

  • E.4 ABLATION ON THE TAIL TOKEN: Removing the tail token causes training collapse: the student’s tail probability exceeds 0.9 within 50 steps, entropy grows monotonically, and MATH500 Avg@8 falls to 3.35% versus 77.88% for TA-OPD.The comparison uses unnormalized top-k OPD with Qwen2.5-7B-Instruct as student and OpenThinker3-7B as teacher.
  • E.4 ABLATION ON THE TAIL TOKEN: The unnormalized top-k objective is not minimized at pt = qt; its minimizer deflates every top-k probability by e and assigns the removed mass to the tail.This drives the loss negative and assigns at least 63% of probability mass outside the teacher’s top-k tokens.
  • E.4 ABLATION ON THE TAIL TOKEN: Adding the tail term turns the objective into a non-negative reverse KL divergence uniquely minimized when student and teacher match on top-k tokens and the tail.The analysis concludes that introducing the tail token is necessary for TA-OPD.
  • E.5 DISCUSSION OF EMA-PG: EMA-PG and sample-corrected TA-OPD are both unbiased in value and gradient, but EMA-PG debiases the unnormalized top-k objective whereas TA-OPD debiases its tail-aware divergence.TA-OPD uses the sampled token to estimate residual tail probability, while EMA-PG estimates the entire tail contribution.
  • E.5 DISCUSSION OF EMA-PG: When the sampled token lies in the top-k, EMA-PG’s correction vanishes, leaving the collapsing unnormalized top-k loss on most tokens; TA-OPD instead retains a well-defined divergence.TA-OPD’s divergence is uniquely minimized at pt = qt.
  • E.5 DISCUSSION OF EMA-PG: Because the student’s tail probability typically exceeds the teacher’s, sample-corrected TA-OPD generally has lower estimate variance than EMA-PG throughout training.The cited analysis attributes this comparison to the teacher’s top-k support and reports the pattern across model pairs.
  • F.1 TRAINING DYNAMICS OF ADDITIONAL MODEL PAIRS: Figure 17 compares TA-OPD with normalized top-k OPD using tail probability and token-level entropy across three additional student–teacher pairs.The pairs are Qwen3-1.7B/Qwen3-30B-A3B-Instruct-2507, Llama-3.1-8B/DeepSeek-R1-Distill-Llama-8B, and Qwen3-1.7B-Base/Qwen3-8B.
  • F.1 TRAINING DYNAMICS OF ADDITIONAL MODEL PAIRS: Across different capability gaps, Figure 18 reports that tail-probability increase occurs for Qwen2.5-Math-1.5B but not for DeepSeek-R1-Distill-Qwen-1.5B.The figure compares TA-OPD and normalized top-k OPD using tail probability and token-level entropy during training.
Loading 2608.14728v1…