让大模型蒸馏更关注低概率但有用的信息，提升小模型性能。
Don't Ignore the Tail: Decoupling top-K Probabilities for Efficient Language Model Distillation
- 将教师模型的前K个高概率预测与低概率部分解耦，避免被主导项淹没
- 在多个数据集上达到与传统方法相当甚至更好的小模型性能
- 计算开销不变，适合学术级资源完成大规模模型蒸馏
语言模型蒸馏的核心学习信号通常是学生与教师分布之间的标准KL散度。传统KL散度容易被教师输出中概率最高的几个词（即模式）主导，从而弱化了低概率但可能有信息量的部分。本文提出一种新的尾部感知散度，将教师模型前K个最高概率预测与其余低概率预测的贡献解耦，同时保持与标准KL散度相同的计算复杂度。该方法降低了教师模式的影响，增强了分布尾部的贡献。实验表明，该改进蒸馏方法在多种数据集上的预训练和有监督蒸馏任务中均表现优异，且整个蒸馏过程高效，仅需学术级算力即可完成大规模数据处理，无需工业级计算资源。
原文摘要 · Abstract (English)
The core learning signal used in language model distillation is the standard Kullback-Leibler (KL) divergence between the student and teacher distributions. Traditional KL divergence tends to be dominated by the next tokens with the highest probabilities, i.e., the teacher's modes, thereby diminishing the influence of less probable yet potentially informative components of the output distribution. We propose a new tail-aware divergence that decouples the contribution of the teacher model's top-K predicted probabilities from that of lower-probability predictions, while maintaining the same computational profile as the KL Divergence. Our decoupled approach reduces the impact of the teacher modes and, consequently, increases the contribution of the tail of the distribution. Experimental results demonstrate that our modified distillation method yields competitive performance in both pre-training and supervised distillation of decoder models across various datasets. Furthermore, the distillation process is efficient and can be performed with a modest academic budget for large datasets, eliminating the need for industry-scale computing.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。