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

Technical Summary: Tail-Aware Top-k On-Policy Distillation (TA-OPD)

Problem Statement
On-policy distillation (OPD) trains a student language model to align its next-token distribution with a teacher's along the student's own trajectories, addressing the training-inference discrepancy inherent in off-policy distillation. A prevalent approach, Normalized Top-k OPD, minimizes the reverse Kullback-Leibler (KL) divergence between the student and teacher distributions restricted to the teacher's top-$k$ tokens.

The paper identifies a critical flaw in this normalized objective: it discards information regarding tail probability (the total probability mass outside the top-$k$ tokens). Theoretical analysis and empirical evidence demonstrate that minimizing the normalized objective creates a systematic bias that steadily increases the student's tail probability and entropy. As the student's tail probability grows, it increasingly samples tokens outside the teacher's top-$k$ set, drifting toward prefixes where the teacher's supervision is unreliable. This phenomenon leads to a degradation in downstream accuracy, particularly when there is a large capability gap between the student and teacher.

Methodology: Tail-Aware Top-k OPD (TA-OPD)
To resolve the loss of tail probability signals, the authors propose TA-OPD, a method that restores this missing information by introducing an explicit tail token ($v_{tail}$).

1. Augmented Token Set: Instead of normalizing probabilities over only the top-$k$ tokens, TA-OPD defines an augmented set $S^+_t = S^k_t \cup \{v_{tail}\}$, where $S^k_t$ represents the teacher's top-$k$ tokens.
2. Probability Assignment:
   • Tokens in $S^k_t$ retain their original probabilities $p_t(v)$.
   • The tail token $v_{tail}$ is assigned the total tail probability: $p_t(v_{tail}) = p^{tail}_t = \sum_{v \notin S^k_t} p_t(v)$.
   • Similarly, the teacher's distribution is augmented with $q^{tail}_t$.
3. Loss Function: The method minimizes the reverse KL divergence over this augmented set:
   $$\ell^{TA}_t = \sum_{v \in S^k_t} p_t(v) \log \frac{p_t(v)}{q_t(v)} + p^{tail}_t \log \frac{p^{tail}_t}{q^{tail}_t}$$
   This formulation explicitly compares the student's and teacher's tail probabilities, driving the student's tail probability toward the teacher's.

Theoretical Properties

• Gradient Dynamics: Unlike the normalized objective, which provides zero gradient to logits outside the top-$k$ set, TA-OPD updates tail logits via the tail log-ratio. Furthermore, the gradients on top-$k$ logits no longer sum to zero, allowing the method to explicitly raise top-$k$ logits relative to tail logits when the student over-weights the tail.
• Lower Bound: The TA-OPD loss is proven to be a tight lower bound of the full-vocabulary reverse KL divergence. The difference is a non-negative residual term proportional to the tail probability. In practice, as TA-OPD drives tail probabilities toward zero, this residual vanishes, making the objective a tight approximation.
• Unbiased Variant (SC-TA-OPD): When the probability of the sampled token is available, the authors derive a Sample-Corrected TA-OPD. This variant adds a bias correction term based on the sampled token, yielding an unbiased estimate of the full-vocabulary reverse KL divergence in both value and gradient.

Key Contributions

1. Identification of Tail Probability Increase: The paper theoretically and empirically demonstrates that normalized top-k OPD causes a steady increase in student tail probability and entropy, leading to performance degradation.
2. Proposal of TA-OPD: A novel distillation method that restores the tail probability signal via a tail token, preventing the entropy increase associated with top-k normalization.
3. Theoretical Guarantees: Proof that TA-OPD is a tight lower bound of the full-vocabulary objective and the derivation of an unbiased sample-corrected variant.
4. Empirical Superiority: Extensive experiments showing TA-OPD outperforms baselines, particularly in scenarios with large student-teacher capability gaps.

Experimental Results
The authors evaluated TA-OPD on mathematical reasoning benchmarks (MATH500, Minerva, OlympiadBench, AMC, AIME24/25) and out-of-distribution (OOD) tasks (ARC-c, MMLU-Pro) using three student-teacher pairs (e.g., Qwen2.5-7B-Instruct distilled from OpenThinker3-7B).

• Performance Gains: TA-OPD consistently achieved the best average accuracy. Notable improvements include:
  • +8.05 points on the average of math benchmarks for Llama-3.1-8B (compared to normalized top-k OPD).
  • +19.53 points on MATH500 for Llama-3.1-8B (improving from 31.30% to 50.83%).
  • Significant gains in OOD performance (e.g., ARC-c accuracy improved from 23.07% to 76.30% for Qwen2.5-7B-Instruct).
• Training Dynamics: TA-OPD stabilized training, keeping student entropy below 1.5 and tail probabilities near 0, whereas normalized top-k OPD caused entropy to rise to ~6 and tail probabilities to ~0.6.
• Capability Gap Sensitivity: The method is most effective when the capability gap between student and teacher is large. When the gap is small, performance gains diminish.
• Parameter Sensitivity: TA-OPD remains effective even with small $k$ (e.g., $k=2$), showing insensitivity to $k$ once $k \ge 2$.

Significance and Claims
The paper claims that TA-OPD addresses a fundamental flaw in existing top-k distillation methods by preserving the tail probability signal, which is crucial for maintaining the student's alignment with the teacher. The method is presented as a practical improvement that requires no additional teacher queries beyond the top-$k$ probabilities and is straightforward to implement. The authors note that while TA-OPD offers substantial gains in large-gap scenarios, its advantage diminishes as $k$ increases or the capability gap narrows. The work is limited to models up to 8B parameters due to computational constraints.