尾感知的 Top-$k$ 在策略蒸馏

# 尾部感知的Top-k在线策略蒸馏  
来源：https://arxiv.org/html/2608.14728  
黄惠鹏  
隶属机构：南方科技大学统计与数据科学系  
魏宏欣  
致谢：通讯作者 \(weihx@sustech\.edu\.cn\)  
隶属机构：南方科技大学统计与数据科学系  

###### 摘要  
在线策略蒸馏已成为语言模型间知识迁移的有效范式，学生模型通过自身轨迹训练以使其下一词元分布与教师模型对齐。为在可行成本下提供稠密监督，许多研究通过最小化学生与教师在教师Top-k词元上的归一化分布间的反向Kullback-Leibler散度。然而，这一归一化目标丢弃了尾部概率信息：教师Top-k词元之外的总概率。因此，优化过程会持续增加学生的尾部概率与熵，经验上导致下游准确率下降。为解决此问题，我们提出尾部感知的Top-k在线策略蒸馏，这是一种恢复缺失尾部概率信号的新蒸馏方法。TA-OPD通过最小化Top-k词元及携带尾部概率的尾部词元上的反向KL散度，从而更好地对齐学生与教师的下一词元分布，避免Top-k归一化导致的尾部概率与熵增加。大量实验证明TA-OPD的优越性，在常用基准上最高提升Avg@8指标8.05分。代码已开源：https://github.com/HuipengHuang/TA-OPD。

## 1 引言  
知识蒸馏是将大语言模型能力迁移至小模型的有效途径。标准蒸馏采用离线策略：学生模型在教师生成的序列上训练，但推理时需基于自身生成内容进行条件推断。训练与推理间的这种差异会导致误差沿生成序列快速累积，并阻碍有效学习。这催生了在线策略蒸馏，其中学生模型通过自身轨迹训练以对齐下一词元分布。流行的归一化Top-k OPD方法通过最小化学生与教师在教师Top-k词元上的归一化分布间的反向KL散度，但丢弃了尾部概率信息。理论分析表明，当学生分布未对齐时，优化会增加学生尾部概率；经验上，尾部概率持续增长会导致学生更频繁采样教师Top-k之外的词元，偏离教师监督可靠的前缀区域，最终降低下游性能。为解决此问题，我们提出TA-OPD——通过恢复尾部概率信号的新蒸馏方法。TA-OPD最小化Top-k词元及携带尾部概率的尾部词元上的反向KL散度，有效对齐学生尾部概率，防止归一化导致的尾部概率与熵增长。理论证明TA-OPD目标是全词表反向KL散度的紧下界。此外，当可获取采样词元概率时，我们推导了TA-OPD的采样变体，该变体通过采样词元去偏并获得全词表反向KL散度的无偏估计。在数学基准上的大量实验表明TA-OPD优于基线方法。特别是在将OpenThinker3-7B蒸馏至Qwen2.5-7B-Instruct时，归一化Top-k OPD训练出现坍塌：训练熵增至约6，MATH500的Avg@8准确率降至68.78%。而TA-OPD稳定蒸馏过程，保持训练熵低于1.5，并将MATH500的Avg@8准确率提升至77.88%。进一步分析表明TA-OPD在师生能力差距较大时效果最佳，且可采用较小k值应用。本文贡献总结如下：  
- • 证明归一化Top-k OPD丢弃尾部概率，最小化该目标会持续增加学生尾部概率与熵，导致下游准确率下降。  
- • 提出TA-OPD这一恢复尾部概率信号的新蒸馏方法，证明其解决了Top-k归一化导致的尾部概率与熵增长。推导了可获取采样词元概率时的无偏估计变体。  
- • 通过大量实验证明TA-OPD性能优于归一化Top-k OPD，在常用基准上最高提升Avg@8准确率8.05分。

## 2 前提知识  
在线策略蒸馏旨在将教师LLM的知识蒸馏至学生LLM。给定提示，学生采样响应序列。在每步，在线策略蒸馏的目标是使学生下一词元分布与教师分布对齐。

## 附录A 理论分析  
### A.3 命题3的证明  
###### 证明  
固定步长与前缀... 尾部概率增加当且仅当学生自协方差超过师生协方差。∎  

## 附录D 数值稳定计算  
### D.2 TA-OPD的数值稳定计算  
每词元TA-OPD损失依赖于尾部对数概率。朴素实现存在数值问题：当教师尾部概率接近零时，Top-k概率求和可能超过1，导致负尾部概率和未定义的对数运算。我们的对数空间实现通过避免概率空间构造，在对数域进行完整计算。我们首先通过log-sum-exp直接获取Top-k总对数概率，并钳位确保严格正尾部概率。尾部对数概率通过log1mexp原语计算。教师项同理处理。Listing 1展示了伪代码实现。

Listing 1: TA-OPD损失实现伪代码  
def compute_taopd_loss(student_topk_log_probs, teacher_topk_log_probs, eps):  
    student_log_topk_prob = logsumexp(student_topk_log_probs, dim=-1)  
    teacher_log_topk_prob = logsumexp(teacher_topk_log_probs, dim=-1)  
    student_tail_log_prob = log1mexp(clamp_max(student_log_topk_prob, -eps))  
    teacher_tail_log_prob = log1mexp(clamp_max(teacher_log_topk_prob, -eps))  
    student_log_p = cat([student_topk_log_probs, student_tail_log_prob], dim=-1)  
    teacher_log_q = cat([teacher_topk_log_probs, teacher_tail_log_prob], dim=-1)  
    taopd_loss = sum(exp(student_log_p) * (student_log_p - teacher_log_q), dim=-1)  
    return taopd_loss  

## 附录E 深入研究  
### E.1 合成实验  
基于30臂老虎机构建合成实验。教师策略在训练中固定，定义为双模态分布。学生策略初始化为类别分布，使用AdamW优化。比较全词表OPD、归一化Top-k OPD、TA-OPD及采样校正TA-OPD四种目标。Top-k OPD中k设为8。  

图12: 不同OPD目标下学生策略可视化。归一化Top-k OPD仅匹配教师分布Top-k词元的相对形状，未能对齐尾部概率；TA-OPD则紧密复现双模态分布。  

#### TA-OPD同时匹配教师策略的形状与尾部概率  
图12比较最终学生策略，TA-OPD紧密复现双模态教师分布。归一化Top-k OPD虽匹配Top-k词元相对形状，但分配了显著更大的尾部概率。总体而言，TA-OPD保留了缺失的尾部概率信息，提供更接近教师策略的近似。  

图13: 不同OPD目标的学习动态。TA-OPD更好地近似全词表OPD。  

#### TA-OPD更好地近似全词表OPD  
图13可视化全词表反向KL景观上的优化轨迹。收集所有方法与教师的logit轨迹，在logit空间中心化后投影至前两个主成分，在该二维平面上评估全词表反向KL。结果显示全词表OPD与TA-OPD均向教师策略周围的低损失区域移动，而归一化Top-k OPD收敛于在全词表KL下仍远离教师的点。TA-OPD目标提供了全词表反向KL的更优Top-k估计。  

### E.2 不同OPD方法的Top-k重叠比例  
文献发现师生Top-k重叠比例能良好预测OPD成功率。  

#### 设置  
在每个学生生成前缀上，重叠比例定义为教师Top-k词元中同时出现在学生Top-k词元中的比例。我们对训练批次中所有词元平均该比例。更高比例意味着师生对合理词元的判断一致，Top-k目标监督对学生有意义的词元集。使用第3.2节的三组师生对报告训练过程中的比例。  

图14: 训练过程中师生Top-k重叠比例。TA-OPD在不同模型对上实现更高重叠比例。  

#### TA-OPD实现更高Top-k重叠比例  
图14比较两种目标的重叠比例。TA-OPD单调增加并保持重叠比例稳定；归一化Top-k OPD...