WASD：基于 Wasserstein 距离的语言模型蒸馏

自回归大语言模型 (LLM) 的能力迅速提高，但其规模的不断扩大伴随着推理时的大量计算和记忆成本。知识蒸馏 (KD) 提供了一种实用的解决方案，通过离散概率分布的对齐将知识从大型 教师模型转移到较小的 学生模型。然而，LLM 的现有 KD 方法主要依赖于通过每个词汇索引的概率值来评估差异的散度，而没有明确利用 token 级别的语义信息。我们为 LLM 提出了基于 Wasserstein 的知识蒸馏（WASD），它通过基于 Wasserstein 的距离与从 token 嵌入导出的成本矩阵合并了 token 级语义信息。为了确保计算的易处理性，我们采用 Sinkhorn 散度并派生出一个与梯度等效的目标，该目标可以在不引入额外网络的情况下进行有效优化。跨多个 LLM 系列和规模的实验表明，WASD 持续提高了各种任务的蒸馏性能，包括指令跟踪、数学推理和代码生成。我们的结果强调了 token 空间中编码的语义信息对于 LLM 蒸馏中有效分布对齐的重要性。该实现可在 https://github.com/aailab-kaist/WASD 上公开获得。