WASD: 基于 Wasserstein 距离的大语言模型知识蒸馏¶

会议: NeurIPS 2026
arXiv: 2610.07706
OpenReview: https://openreview.net/forum?id=FGlwR3s9v3
代码: https://github.com/aailab-kaist/WASD
领域: 模型压缩
关键词: 知识蒸馏、大语言模型、最优传输、Sinkhorn散度、语义感知

一句话总结¶

提出基于 Sinkhorn 散度的 WASD 蒸馏框架，利用词嵌入几何代价矩阵捕获词元级语义关系，并通过对偶势函数无偏闭式梯度实现高效无额外网络的自回归 LLM 知识对齐。

研究背景与动机¶

自回归大语言模型（LLM）的性能突破严重依赖于参数规模与海量训练数据的扩展，但随之而来的高昂推理计算与显存开销极大限制了其在延迟敏感或资源受限场景下的实际部署。知识蒸馏（Knowledge Distillation, KD）作为主流压缩手段，旨在通过对齐自回归生成过程中教师与学生模型在词表空间上的离散概率分布，将大模型的推理与表征能力迁移至轻量学生模型。然而，现有的 LLM 蒸馏方法（如正向/逆向 KL 散度、全变差 TV、广义 JS 散度及 $\alpha$-$\beta$ 散度等）均属于基于概率索引取值直接计算的 $f$-散度变体，仅机械比对同一词表索引处的概率标量差异，完全忽略了词元所承载的几何语义结构。

这种基于索引的概率匹配在语言建模中暴露出严重缺陷：在长尾词表预测场景下，不同词元间的语义关联存在本质差异。例如面对问句“你在客厅坐在什么上面？”，词表中的候选词“sofa”与“couch”高度语义同义，而“apple”则完全无关。现有 $f$-散度在优化时仅根据索引位置惩罚概率差异，若学生模型预测“couch”而教师倾向“sofa”，学生将遭受与预测“apple”完全相同量级的严厉惩罚。这不仅惩罚了合理的语义等价表达、抑制了生成的多样性，更因 LLM 词表分布高度稀疏且存在大量近零概率值，极易引发密度比导致的数值不稳定与优化退化。

解决该瓶颈的核心挑战在于如何将词元间的连续语义距离引入离散分布对齐，同时维持反向传播的计算可行性与目标收敛一致性。为此，本文采用最优传输理论中的 Wasserstein 距离，以教师模型的词嵌入距离构建传输代价矩阵，并引入去偏 Sinkhorn 散度消除熵正则化偏差。核心 idea：利用词嵌入诱导的传输代价矩阵构建去偏 Sinkhorn 散度目标，借助对偶势函数闭式解析梯度在停止梯度（stop-gradient）下实现高效优化，使学生模型在对齐概率分布的同时充分保留词元级语义几何关系。

方法详解¶

整体框架¶

WASD（Wasserstein-based Knowledge Distillation）的整体管线构建在自回归逐步离散概率对齐的范式之上。对于给定输入提示词与前序上下文，教师模型与学生模型分别输出整个词表上的下一步预测概率分布。为了赋予概率转移以语义感知能力，WASD 首先从预训练教师词嵌入中离线提取几何代价矩阵并做局部稀疏化。在训练过程中，算法在每一步通过轻量矩阵缩放迭代求解跨传输（教师-学生）与自传输（学生-学生）的最优对偶势函数，最终将其势差作为梯度权重直接赋给学生模型的对数似然梯度。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列与前序词元上下文"] --> B["词嵌入几何代价矩阵构建<br/>提取教师嵌入计算余弦距离并做 Top-k 稀疏截断"]
    B --> C["去偏 Sinkhorn 散度对齐<br/>交替求解跨传输与自传输矩阵缩放迭代"]
    C --> D["对偶势差停止梯度优化<br/>将对偶差值赋权至学生模型对数似然梯度"]
    D --> E["学生模型参数更新"]

关键设计¶

1. 词嵌入几何代价矩阵构建：引入词元语义距离并保持计算稀疏度

传统散度仅将词表视为无序离散索引集合。为了让概率质量的转移代价反映真实语言语义，本文利用教师模型的预训练词嵌入向量 $e_i \in \mathbb{R}^d$，构建静态成对距离矩阵 $C_{ij} = 1 - \cos(e_i, e_j)$（通用评测使用余弦距离，部分任务特定场景使用 $L_2$ 距离）。由于现代大语言模型词表规模庞大（$|\mathcal{V}| \ge 32000$ 甚至更大），计算全词表间的全连接传输代价矩阵会导致显存爆炸与迭代过慢。为此，WASD 在词嵌入几何空间中采用 $k$ 近邻截断策略：对每个词元仅保留几何距离最近的 $k$ 个非自身近邻，对传输图进行对称化处理并增加自身零代价自环。这使得代价矩阵在捕获局部语义等价项（如同义词、近义词）的同时保持极高稀疏度，大幅降低了后续最优传输的计算复杂度。

2. 去偏 Sinkhorn 散度对齐：消除熵正则化带来的固有分布漂移

直接求解标准离散 Wasserstein 距离 $D_{\text{W}}(p, q_{\bm{\theta}}) := \min_{\bm{P} \in U(p, q_{\bm{\theta}})} \langle \bm{P}, \bm{C} \rangle$ 需面临大规模线性规划计算瓶颈。虽然引入熵正则化 $-\epsilon \mathcal{H}(\bm{P})$ 后能借助经典的 Sinkhorn-Knopp 矩阵缩放算法高效求解，但熵项的引入破坏了距离度量的自相似性，即自传输代价 $\tilde{D}_{\text{W}}^{\epsilon}(p, p) \ne 0$。这导致正则化传输损失在 $q_{\bm{\theta}} = p$ 时无法取得最小值，产生不可忽略的优化偏差（Entropic Bias），无法保证学生分布真正收敛至教师目标。针对此痛点，WASD 采用去偏 Sinkhorn 散度作为真正的蒸馏目标： $$D_{\text{S}}^{\epsilon}(p, q_{\bm{\theta}}) := \tilde{D}_{\text{W}}^{\epsilon}(p, q_{\bm{\theta}}) - \frac{1}{2}\tilde{D}_{\text{W}}^{\epsilon}(q_{\bm{\theta}}, q_{\bm{\theta}}) - \frac{1}{2}\tilde{D}_{\text{W}}^{\epsilon}(p, p)$$ 该散度在对称代价矩阵下严格保证非负性，且当且仅当 $q_{\bm{\theta}} = p$ 时取零，从数学上确保了蒸馏目标的无偏自洽，并在 $\epsilon \to 0$ 时退化为标准 Wasserstein 距离、在 $\epsilon \to \infty$ 时渐近逼近最大均值差异（MMD）。

3. 对偶势差停止梯度优化：无需展开求解器的高效闭式梯度反传

在传统最优传输神经网络训练中，若直接对 Sinkhorn 算法多步迭代过程使用自动微分求导，将占用极其庞大的计算图显存并大幅拖慢训练速度。WASD 基于凸对偶理论与包络定理（Envelope Theorem）证明：在最优点处，对偶变量对学生分布参数的隐式导数无需显式反传。记跨传输问题 $\tilde{D}_{\text{W}}^{\epsilon}(p, q_{\bm{\theta}})$ 的最优对偶变量为 $\bm{\phi}^{*(p, q_{\bm{\theta}})}$，学生自传输问题 $\tilde{D}_{\text{W}}^{\epsilon}(q_{\bm{\theta}}, q_{\bm{\theta}})$ 的对称最优对偶变量为 $\bm{\phi}^{*(q_{\bm{\theta}}, q_{\bm{\theta}})}$，二者分别通过标准 Sinkhorn-Knopp 迭代与自对偶不动点迭代快速收敛获得。本文提出梯度等价目标： $$\mathcal{L}_{\text{WASD}}^{\epsilon}(\bm{\theta}) := \mathbb{E}_{(\mathbf{x}, \mathbf{y}) \sim \mathcal{D}} \left[ \sum_{l=1}^{L} \sum_{i=1}^{d} \operatorname{sg}\left(\phi_i^{*(p, q_{\bm{\theta}})}(\mathbf{x}, \mathbf{y}_{<l}) - \phi_i^{*(q_{\bm{\theta}}, q_{\bm{\theta}})}(\mathbf{x}, \mathbf{y}_{<l})\right) q_{\bm{\theta}}(y_l = i | \mathbf{x}, \mathbf{y}_{<l}) \right]$$ 其中 $\operatorname{sg}(\cdot)$ 表示停止梯度算子。该损失对参数 $\bm{\theta}$ 的梯度形式为： $$\nabla_{\bm{\theta}} D_{\text{S}}^{\epsilon} = \sum_{i=1}^d q_{\bm{\theta}}(y_l = i | \mathbf{x}, \mathbf{y}_{<l}) \left(\phi_i^{*(p, q_{\bm{\theta}})} - \phi_i^{*(q_{\bm{\theta}}, q_{\bm{\theta}})}\right) \nabla_{\bm{\theta}} \log q_{\bm{\theta}}(y_l = i | \mathbf{x}, \mathbf{y}_{<l})$$ 与传统前向 KL 散度直接使用教师概率标量 $-p(i)$ 赋权不同，WASD 将包含几何上下文传输信息的对偶势差作为权重，无需通过迭代求解器进行反向传播，完全省去了辅助对齐网络或复杂可微求解模块。

损失函数 / 训练策略¶

训练时，总损失即为自回归序列上 WASD 目标期望，模型基于 AdamW 优化器进行端到端更新。超参数配置方面，熵正则化系数固定为 $\epsilon = 0.001$，跨传输与自传输矩阵缩放迭代步数均仅设为 10 步；稀疏邻居数在大多数模型（GPT-2、OpenLLaMA2、Gemma）上设为 $k=8$，在词表规模较大的 Qwen2.5 上采用 $k=2$ 以平衡吞吐量。此外，训练期配合学生端温度缩放策略（通常取温度为 2.0）调节概率分布熵，避免学生模型在自由解码时因同义替代过于宽泛而损害基准 ROUGE-L 指标。

实验关键数据¶

主实验¶

在指令遵循任务评测中，针对不同模型族（GPT-2、OpenLLaMA2）与参数规模设置进行了广泛对比，在 Dolly Eval、Self-Instruct、Vicuna、Super-Natural Instructions（Super NI）与 Unnatural Instructions（UnNI）五个基准上报告生成质量指标 ROUGE-L：

蒸馏设置	对齐方法	Dolly Eval	Self Inst	Vicuna	Super NI	UnNI	平均 ROUGE-L (↑)
GPT-2 XL (1.5B)	Teacher 上界	26.93	14.77	16.43	27.19	31.43	23.35
GPT-2 (1.5B → 0.1B)	GKD (ICLR'24)	23.86	11.21	14.63	19.43	22.28	18.28
	TAID (ICLR'25)	25.44	12.55	16.65	24.45	26.70	21.16
	DistiLLM-SKL (ICML'24)	25.39	12.42	16.21	25.13	26.94	21.22
	ABKD (ICML'25)	26.15	12.56	17.24	27.07	29.03	22.41
	CSD (ICLR'26)	25.91	12.97	16.80	26.90	28.51	22.22
	AMiD (ICLR'26)	26.39	14.18	16.81	29.11	30.78	23.46
	WASD (本文)	27.07	14.34	17.25	29.60	31.86	24.02
GPT-2 (1.5B → 0.3B)	AMiD (ICLR'26)	27.82	14.84	16.55	30.30	34.17	24.74
	WASD (本文)	28.37	15.45	18.41	30.24	35.06	25.51
OpenLLaMA2 (7B → 3B)	AMiD (ICLR'26)	30.21	20.65	20.51	37.49	37.61	29.30
	WASD (本文)	29.83	20.33	21.34	38.83	37.74	29.61

在下游特定任务评测中，针对 Gemma-7B-IT 到 2B-IT（翻译、摘要、算术推理）以及 Qwen2.5-Coder-7B 到 1.5B（代码生成 pass@1）同样验证了显著优势： - 机器翻译（Flores-200, COMET）：AMiD 73.70 vs. WASD 74.53（教师 79.33）； - 算术推理（GSM8k, Accuracy）：AMiD 24.49% vs. WASD 24.87%（学生零样本 0.00%）； - 代码生成（HumanEval / MBPP 均值）：AMiD 74.1% vs. WASD 74.9%。

消融实验¶

针对传输代价矩阵类型与去偏机制在 GPT-2（1.5B → 0.1B）设置下的消融对比（五任务平均 ROUGE-L）：

代价矩阵配置 / 目标变体	平均 ROUGE-L (↑)	说明
基线 AMiD（无最优传输）	23.36	采用前向 KL 散度与辅助分布
Uniform Cost（均匀代价矩阵）	22.08	移除了词间分级的几何信息，退化为常数惩罚，掉点 1.28
Permuted Cost（置换打乱代价矩阵）	23.47	保留传输框架但随机重排词间对应关系，表现仅持平 AMiD
有偏熵正则 Wasserstein 距离	23.65	未扣除自传输修正项，受 Entropic Bias 影响次于完整模型
Semantic Cost (WASD 完整模型)	23.94	完整保留词嵌入几何距离与去偏 Sinkhorn 修正

超参数与度量消融发现： - 距离度量替换：将余弦距离换为欧氏距离（$L_2$），平均 ROUGE-L 保持在 22.45（余弦为 22.47），证明算法对底层连续距离定义具备强鲁棒性； - 邻域大小 $k$：当 $k$ 从 2 增至 8 时性能稳步攀升并在 $k=8$ 处达到峰值（22.47），继续增大到 16 或 64 时由于引入长距离弱相关词元扰动，性能维持平稳（22.45）而无额外增益。

关键发现¶

语义几何优于机械平滑：置换代价矩阵实验表明，最优传输的增益绝非单纯来自平滑正则化效应，真实反映语言语义的词嵌入拓扑结构是知识成功迁移的关键决定因素。
质量与多样性前沿显著改善：在不同解码温度下绘制 ROUGE-L 与 Self-BLEU 的权衡曲线（Trade-off Frontier），WASD 始终显著外扩于所有对比基线。其在鼓励发散采样时能够稳定抑制离题输出，保证生成的多样性与高语义连贯性。
收敛加速抵消微小时间损耗：尽管每次迭代求解对偶变量使单步耗时约为标准基线的 2 倍并增加约 20% 显存，但在相同物理时间墙钟（Wall-clock time）预算下，WASD 凭借更优的梯度流更快达到收敛，实际训练效率极具竞争力。

亮点与洞察¶

将连续语义几何自然嫁接至离散词表蒸馏：通过最优传输对偶理论巧妙跨越了离散 token 空间与连续语义表征的鸿沟，从根本上破除了 $f$-散度无法感知词表同义结构的固有缺陷。
包络定理支撑的零额外显存求导：利用包络定理证明了最优对偶变量在梯度计算时可作为常数处理，配合停止梯度算子构造等价目标，避免了通过 Sinkhorn 迭代展开反向传播或额外训练辅助判别网络。
模块化即插即用：该散度形式在工程实现上仅表现为对输出 logits 梯度的对偶加权，能直接嵌入现有的各类高级蒸馏流水线（如辅助分布蒸馏 AMiD、对比蒸馏 DistiLLM-2、跨分词器蒸馏等）。

局限与展望¶

训练端额外矩阵迭代开销：虽然无需反传求解器，但前向传播阶段的矩阵缩放迭代仍然增加了显存占用与前向延迟，在大规模上下文或更大词表模型（如 Gemma 256k 词表）上需要更激进的并行或稀疏工程优化。
静态词嵌入代价矩阵未建模上下文动态：当前传输代价由静态预训练词嵌入离线计算得出，未能动态反映一词多义及特定输入上下文下的细微语境变动。
未来方向：探索由中间层隐藏状态或动态注意力图诱导的动态上下文传输代价矩阵，以及将该框架拓展至跨模态 Token 几何对齐。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将去偏 Sinkhorn 散度与词元语义代价应用于自回归大模型蒸馏，理论优雅且闭式推导严谨。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 GPT-2、OpenLLaMA2、Gemma、Qwen2.5 多模型族及通用/特定任务，包含代价矩阵置换消融与 GPT-4/GPT-4o-mini 双重评测。
写作质量: ⭐⭐⭐⭐⭐ 逻辑脉络清晰，数学证明完备，从局限发现到理论推导层层推进。
价值: ⭐⭐⭐⭐⭐ 为大语言模型压缩与高效分发提供了通用、原则性强的全新损失设计范式。