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

2026-08-12

AI Summary

This study addresses the issue in on-policy distillation where Top-k normalization discards tail probabilities, leading to increased student model entropy and degraded accuracy. To mitigate this, we propose TA-OPD, a novel method that introduces virtual tokens into the Top-k distribution to preserve tail information. By minimizing reverse KL divergence augmented with these tail signals, TA-OPD explicitly recovers missing supervision and effectively suppresses entropy inflation in the student model. Extensive experiments demonstrate that our approach significantly enhances alignment and performance in language model distillation. Notably, TA-OPD achieves a maximum improvement of 8.05 points on the Avg@8 metric across multiple benchmarks, validating its effectiveness in addressing probability truncation issues during knowledge transfer.

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

Problem

• On-policy distillation
• Top-k normalization
• Tail probability
• Reverse KL divergence
• Language models

Innovation

• On-Policy Distillation
• Tail-Aware
• Reverse KL Divergence
• Top-k Normalization
• Knowledge Distillation