WASD：基于 Wasserstein 的大语言模型知识蒸馏

arXiv:2610.07706v1 公告类型：新 摘要：自回归大语言模型（LLM）的能力取得了快速进展，但其规模不断增大也带来了推理阶段巨大的计算和内存开销。知识蒸馏（KD）通过对齐离散概率分布，将知识从大型教师模型迁移到较小的学生模型，为此提供了一种实用的解决方案。然而，现有面向 LLM 的知识蒸馏方法主要依赖通过词汇表中各索引位置的概率值来评估差异的散度，未显式利用词元级语义信息。我们提出了面向 LLM 的基于 Wasserstein 的知识蒸馏方法（WASD），通过基于 Wasserstein 的距离引入词元级语义信息，该距离使用由词元嵌入构造的代价矩阵。为确保计算可行性，我们采用 Sinkhorn 散度，并推导出一个梯度等价的目标函数，使其能够在不引入额外网络的情况下高效优化。在多个 LLM 系列和不同规模的模型上进行的实验表明，WASD 在包括指令遵循、数学推理和代码生成在内的多种任务上，都能持续提升蒸馏性能。我们的结果凸显了词元空间中编码的语义信息对于 LLM 蒸馏中的有效分布对齐的重要性。实现代码已公开发布于 https://github.com/aailab-kaist/WASD 。