WASD: Wasserstein-based Knowledge Distillation for Large Language Models

2026-10-05

AI Summary

This study addresses the limitation of existing knowledge distillation methods for large language models (LLMs), which often overlook token-level semantic information. To this end, we propose a novel distillation framework based on the Wasserstein distance. Specifically, our approach constructs a cost matrix using token embeddings to incorporate semantic priors and introduces the Sinkhorn divergence into LLM distillation for the first time. This formulation enables efficient optimization of gradient-equivalent objectives without requiring any additional network components. Experimental results demonstrate that the proposed method significantly improves distillation performance across instruction following, mathematical reasoning, and code generation tasks.

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 https://github.com/aailab-kaist/WASD .

Problem

• Knowledge Distillation
• Large Language Models
• Token-level Semantics
• Distribution Alignment

Innovation

• Knowledge Distillation
• Wasserstein Distance
• Sinkhorn Divergence
• Large Language Models
• Token-level Semantics