WASD: Wasserstein-based Knowledge Distillation for Large Language Models

📄 arXiv: 2610.07706v1 📥 PDF

作者: Byeonghu Na, Donghyeok Shin, Yeongmin Kim, Mina Kang, Il-Chul Moon

分类: cs.LG, cs.AI

发布日期: 2026-10-06

备注: Accepted at NeurIPS 2026

🔗 代码/项目: GITHUB

💡 一句话要点

提出WASD以解决大语言模型蒸馏效率低下问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 知识蒸馏 大语言模型 Wasserstein距离 Sinkhorn散度 语义信息 模型压缩 自然语言处理

📋 核心要点

  1. 现有的知识蒸馏方法未能充分利用令牌级的语义信息，导致蒸馏效果不理想。
  2. WASD通过引入Wasserstein距离和令牌嵌入的成本矩阵，增强了知识蒸馏过程中的语义信息利用。
  3. 实验结果显示，WASD在指令跟随、数学推理和代码生成等多项任务上均显著提升了蒸馏性能。

📝 摘要（中文）

自回归的大语言模型（LLMs）在能力上迅速提升，但其规模的增加在推理时带来了巨大的计算和内存成本。知识蒸馏（KD）通过将知识从大型教师模型转移到较小的学生模型，提供了一种实用的解决方案。然而，现有的KD方法主要依赖于通过每个词汇索引的概率值评估差异，而未能显式利用令牌级的语义信息。我们提出了基于Wasserstein距离的知识蒸馏（WASD），通过从令牌嵌入派生的成本矩阵引入令牌级的语义信息。为了确保计算的可行性，我们采用了Sinkhorn散度，并推导出一个梯度等效的目标，可以高效优化而无需引入额外的网络。实验结果表明，WASD在多种任务上持续提升了蒸馏性能，强调了令牌空间中编码的语义信息在LLM蒸馏中的重要性。

🔬 方法详解

问题定义：论文旨在解决现有知识蒸馏方法在大语言模型蒸馏中未能有效利用令牌级语义信息的问题，导致蒸馏效果不佳。

核心思路：提出WASD方法，通过Wasserstein距离引入令牌级语义信息，利用Sinkhorn散度确保计算的可行性，从而提升蒸馏效果。

技术框架：整体架构包括教师模型和学生模型，采用Wasserstein距离计算令牌嵌入之间的差异，优化过程中使用Sinkhorn散度来处理概率分布的对齐。

关键创新：WASD的核心创新在于引入了令牌级的语义信息和Wasserstein距离的结合，显著改善了现有方法中对概率值的依赖，提升了蒸馏效果。

关键设计：在损失函数设计上，采用Sinkhorn散度作为优化目标，确保在不增加额外网络的情况下实现高效优化，同时保持了模型的计算效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，WASD在多个大语言模型家族和规模上均显著提升了蒸馏性能，尤其在指令跟随、数学推理和代码生成任务中，蒸馏效果提升幅度达到XX%（具体数据待补充），显示出其在多样化任务中的有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、代码生成和智能助手等，能够有效降低大语言模型的计算和内存成本，提升其在实际应用中的可用性和效率。未来，WASD方法可能推动更小型、智能化的语言模型的发展，使其在资源受限的环境中也能发挥强大的功能。

📄 摘要（原文）

Autoregressive large language models (LLMs) have rapidly advanced in capability, but their increasing scale comes with substantial computational and memory costs at inference time. Knowledge distillation (KD) offers a practical solution by transferring knowledge from a large teacher model to a smaller student model via alignment of discrete probability distributions. However, existing KD methods for LLMs primarily rely on divergences that evaluate discrepancies through probability values at each vocabulary index, without explicitly leveraging token-level semantic information. We propose Wasserstein-based knowledge distillation (WASD) for LLMs, which incorporates token-level semantic information via the Wasserstein-based distance with a cost matrix derived from token embeddings. To ensure computational tractability, we adopt the Sinkhorn divergence and derive a gradient-equivalent objective that can be efficiently optimized without introducing additional networks. Experiments across multiple LLM families and scales show that WASD consistently improves distillation performance on diverse tasks, including instruction following, mathematical reasoning, and code generation. Our results highlight the importance of semantic information encoded in the token space for effective distribution alignment in LLM distillation. The implementation is publicly available at https://github.com/aailab-kaist/WASD .