🤖 machine learning

DOT-MoE: Differentiable Optimal Transport for MoEfication

DOT-MoE 是一个新颖的框架，它通过将神经元分解公式化为一个可微的最优传输问题，将预训练的稠密大语言模型转换为高效的混合专家模型，从而实现了分配过程的端到端学习，在保留原始性能 90% 的同时减少了 50% 的活跃参数，其表现显著优于现有的启发式方法。

原作者： Udbhav Bamba, Arnav Chavan, Aryamaan Thakur, Steve Teig, Deepak Gupta

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Udbhav Bamba, Arnav Chavan, Aryamaan Thakur, Steve Teig, Deepak Gupta

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你拥有一个庞大且极其聪明的图书馆（大语言模型），每当你提出一个问题时，里面每一本书都在同时被阅读。虽然这让图书馆变得非常聪明，但运行起来也非常缓慢且昂贵，因为你在为那些实际上并不需要的书籍消耗能量。

为了解决这个问题，科学家们创建了一个“专家混合”（Mixture of Experts, MoE）系统。你可以把它想象成聘请了一支专家团队。与其让整个图书馆去阅读你的问题，你只需询问少数几位了解答案的特定专家。这样可以节省能量并提高速度。

问题所在：如何挑选专家？
论文解释说，将一个标准的“全读全看”式图书馆转变为“专家”团队是非常棘手的。

旧方法就像是随机地将书籍分装到不同的房间里，并寄希望于合适的人最终能出现在合适的地方。或者，他们会观察书的封面（权重）并猜测哪些书应该归为一类。
问题在于： 这些方法通常将“对书籍进行分类”的过程与“训练图书管理员”的过程分离开来。他们先对书籍进行分类，然后再尝试教图书管理员如何寻找它们。这往往会导致一个混乱的团队，专家们并没有实现真正的专业化，或者图书管理员会感到困惑。

解决方案：DOT-MoE（智能分类器）
作者提出了一种名为 DOT-MoE 的新方法。他们将把书籍分拣到专家室的问题视为一个被称为“最优传输”（Optimal Transport）的物流谜题。

以下是类比：
想象你有一个装满数千名工人（神经元）的仓库，以及一组建筑工地（专家）。每个工地需要恰好相同数量的工人，且每位工人必须分配到一个工地。

目标： 你希望将正确的工人送到正确的工地，从而使建筑物的建造效果与原始仓库完全一致。
创新点： DOT-MoE 并没有使用猜测或随机列表，而是使用了一个数学上的“交通控制器”（可微最优传输）。它计算出移动每一位工人的最有效路径，确保没有工地过于拥挤，也没有任何一名工人被遗漏。
“秘密武器”： 该系统在分类的同时进行学习。它不仅仅是先对书籍进行分类再去雇佣图书管理员，而是在教图书管理员如何为每个新问题挑选正确房间的同时，也在计算哪些书该去哪个房间。它们协同学习，不断调整彼此的位置，直到一切都完美契合。

为什么这种方法更好？
论文声称，通过使用这种“智能物流”方法：

保留了智慧： 这个新的专家团队保留了原图书馆约 90% 的智能。
降低了成本： 在回答问题时，它仅使用了 50% 的活跃“脑力”（参数）。
更稳定： 不同于其他在尝试拆分团队时可能会崩溃或表现不佳的方法，这种方法确保了从一开始就是一个平衡且功能完备的团队。

结果
在测试中，这种新方法比以往拆分图书馆或随机挑选专家的做法表现得更好。它证明了，如果你将神经网络的组织形式视为一个精确、可解的物流问题，你就可以让巨大的 AI 模型运行得更快、更便宜，且不会丧失其理解和生成语言的能力。

简而言之： 他们找到了一种方法，通过使用一个在分类的同时还能教导系统如何使用的数学“交通控制器”，将一个庞大、缓慢的 AI 大脑重新组织成一支快速、高效的专家团队，并且整个过程一气呵成。

技术摘要：DOT-MoE：用于 MoEfication 的可微最优传输

1. 问题陈述

大语言模型（LLMs）的快速规模化在推理效率方面造成了显著的瓶颈。虽然混合专家（Mixture-of-Experts, MoE）架构通过每个 token 仅激活一部分参数，实现了模型规模与推理成本的解耦，但从头开始训练 MoE 是极其不稳定、耗费数据且计算成本高昂的。

因此，MoEfication（将预训练稠密模型转换为稀疏 MoE）已成为一种可行的替代方案。然而，现有的将预训练稠密前馈网络（FFNs）转换为稀疏专家的方法过度依赖启发式策略（例如，随机拆分、基于权重的聚类或基于激活值的聚类）。这些方法存在三个主要局限性：

关注点分离： 它们通常将神经元分配（划分 FFN）和路由训练视为两个独立的、顺序进行的阶段，导致两者无法进行协同演化。
代理优化： 它们基于中间统计量（权重或激活值）而非神经元的实际输出贡献进行优化，这往往导致对稠密模型行为的重建效果不佳。
缺乏可微性： 神经元分配的离散性质阻碍了端到端的梯度流，使得难以联合优化结构分解和路由策略。

2. 方法论：DOT-MoE

作者提出了 DOT-MoE，该框架将稠密层向专家的分解形式化为一个**可微最优传输（Differentiable Optimal Transport, DOT）**问题。这种方法实现了专家分解和路由机制的端到端学习。

2.1. 最优传输公式化

其核心洞察是将神经元分配视为一个平衡传输问题，其中“质量”（神经元）必须在严格的容量约束下，从源分布（稠密 FFN 神经元）传输到目标分布（专家）。

源（Source）： $dffn$ 个神经元，每个具有单位质量。
目标（Target）： $E$ 个专家，每个专家需要恰好 $s$ 个神经元（满足 $s \cdot E = dffn$ ）。
代价函数： 该方法没有使用固定的代价矩阵，而是使用了一个可学习的亲和力矩阵 $A$ ，其中 $A_{i,e}$ 表示将神经元 $i$ 分配给专家 $e$ 的亲和力。目标是在满足边缘约束的同时最大化总亲和力。

2.2. 通过 Sinkhorn-Knopp 实现的可微松弛

为了使离散分配问题变得可微，DOT-MoE 采用了熵正则化（Entropic Regularization）。

最优传输问题通过 Sinkhorn-Knopp 算法进行松弛，该算法迭代地对传输计划矩阵的行和列进行归一化。
这产生了一个软分配矩阵 $M_{soft} \in [0, 1]^{dffn \times E}$ ，它在期望上满足容量约束。
当温度参数 $\tau \to 0$ 时，解趋于离散最优解。

2.3. 使用直通估计器（STE）进行联合优化

一个关键挑战是，部署时需要离散分配，而路由涉及离散的 top-k 选择，两者都是不可微的。DOT-MoE 使用**直通估计器（Straight-Through Estimators, STE）**来解决这一问题：

前向传播： 使用硬决策（对 $M_{soft}$ 进行贪婪舍入以进行分配；使用 top-k 选择进行路由）。
反向传播： 通过软对应项（ $M_{soft}$ 和路由对数概率 $P$ ）传递梯度。
这使得路由权重（ $W_r$ ）和分配亲和力矩阵（ $A$ ）可以在训练期间进行协同演化。

2.4. 训练目标

该框架通过**对齐阶段（Alignment Phase）**进行训练，在此阶段稠密模型的权重是冻结的。模型最小化以下损失组合：

输出重建： 通过稠密教师模型与稀疏 MoE 学生模型输出之间的 KL 散度来衡量。
语言建模： 基于语言建模目标的交叉熵损失。
稳定性辅助项： 路由 $z$ -loss（防止对数概率爆炸）和负载均衡损失（防止专家崩溃）。

3. 核心贡献

最优传输框架： 本文引入了一种新颖的将稠密到 MoE 的转换形式化为平衡传输问题的方法，并通过可微 Sinkhorn 迭代来强制执行严格的专家容量约束。
双层协同演化： 不同于先固定分配再训练路由的现有方法，DOT-MoE 使用互补的 STE 联合优化神经元到专家的分解以及 token 到专家的路由。
输出感知构建： 通过直接建模神经元对最终输出的贡献（而非优化权重或中间激活值等代理指标），该方法实现了更高的重建保真度。

4. 实验结果

作者在三个模型家族（LLaMA-2, LLaMA-3, Qwen2.5）和六个基准测试（ARC-Challenge, Winogrande, HellaSwag, PIQA, SciQ, BoolQ）上评估了 DOT-MoE。

性能 vs. 剪枝： 在 LLaMA-2 7B 拥有 50% 参数预算的情况下，DOT-MoE 在 WikiText-2 上的困惑度为 7.99，显著优于最先进的结构化剪枝方法 DISP-LLM（9.84）。
性能 vs. 现有 MoEfication： 在零样本设置下，DOT-MoE 显著超越了 CMoE 和 LLaMA-MoE-v2 等方法。例如，在 LLaMA-2 7B 上，DOT-MoE 实现了 61.5% 的平均准确率，而 CMoE 为 44.5%。
效率： DOT-MoE 保留了原始稠密模型约 90% 的性能，同时减少了 50% 的激活参数。
微调： 仅通过 1.2B token 的连续微调，DOT-MoE 模型即可恢复接近稠密模型的性能（例如，LLaMA-2 7B 上为 66.6% vs 76.6%），其表现优于需要相似或更大数据预算的基准模型。
消融研究：
- 粒度： 性能随专家粒度的增加而提高，直到达到饱和点，之后收益递减。
- 推理速度： 尽管专家数量增加，但使用融合 MoE 内核（如 vLLM）时，推理吞吐量保持稳定，因为总的融合中间维度保持不变。
- 稀疏性： 在更高稀疏度水平（75%）下训练的模型，在推理时的更极端稀疏场景下仍具有鲁棒的泛化能力。

5. 重要性与主张

本文声称，DOT-MoE 通过提供一种原则性的、可微的方法将稠密模型转换为稀疏 MoE，代表了 LLM 推理效率的一次重大进步。

范式转变： 它从启发式的、静态的划分转向了动态的、可学习的分解，这种分解尊重专家的功能连贯性。
实用性： 该方法实现了对预训练模型进行高效的后验稀疏化，无需从头开始进行完整的重新训练，在大幅减少激活参数的同时保留了高容量。
鲁棒性： 该框架证明，保持 FFN 保真度需要显式建模神经元对输出的贡献，而现有的基于代理的方法未能捕捉到这一原则。

作者总结道，DOT-MoE 为稀疏模型提供了卓越的初始化，以极少的额外训练数据缩小了稀疏架构与稠密架构之间的性能差距。