改进语言模型蒸馏方法,防止学生模型过度发散。
Tail-Aware Top-$k$ On-Policy Distillation
- 引入尾部概率信号,优化时兼顾top-k与尾部词
- 实验显示平均准确率提升最高达8.05点
- 适合追求高精度的模型压缩与知识迁移场景
基于策略的蒸馏(OPD)已成为在语言模型间转移知识的有效范式,其中学生模型在其自身轨迹上学习对齐教师模型的下一个词分布。为在可接受成本下提供密集监督,许多工作最小化学生与教师在教师top-k词上的归一化分布之间的反向KL散度。然而,这种归一化目标会丢失尾部概率信息——即教师top-k词之外的总概率。结果导致优化过程持续增加学生模型的尾部概率和熵,实证中降低下游任务准确率。为此,我们提出尾部感知的top-k OPD(TA-OPD),一种新蒸馏方法,以恢复缺失的尾部概率信号。具体而言,TA-OPD在top-k词及携带尾部概率的尾部词上最小化反向KL散度。有效提升了学生模型与教师模型下一个词分布的对齐程度,防止由top-k归一化引起的尾部概率和熵的增加。大量实验表明TA-OPD性能更优,在常见基准上平均准确率(Avg@8)提升最高达8.05点。代码已开源:https://github.com/HuipengHuang/TA-OPD。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。