arXiv：cs.LG· Byeonghu Na, Donghyeok Shin, Yeongmin Kim, Mina Kang, Il-Chul Moon·· 4 小时前AI 评分40
WASD：面向大语言模型的 Wasserstein 知识蒸馏
WASD: Wasserstein-based Knowledge Distillation for Large Language Models
AI 导读
研究提出 WASD，一种基于 Wasserstein 距离的大语言模型知识蒸馏方法，通过由 token 嵌入构建的代价矩阵引入 token 级语义信息，并采用 Sinkhorn 散度推导出梯度等价目标以高效优化。在多个 LLM 系列与规模上的实验显示，WASD 在指令遵循、数学推理和代码生成等任务上持续提升蒸馏性能。实现已公开，论文被 NeurIPS 2026 接收。
正文
Abstract: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 this https URL .
| Comments: | Accepted at NeurIPS 2026 |
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2610.07706 [cs.LG] |
| (or arXiv:2610.07706v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.07706 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Byeonghu Na [view email]
[v1]
Tue, 6 Oct 2026 03:58:54 UTC (5,240 KB)
来源：arXiv：cs.LG · arxiv.org