DOT-MoE: Differentiable Optimal Transport for MoEfication

📄 arXiv: 2606.01666v1 📥 PDF

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

分类: cs.LG, cs.AI

发布日期: 2026-06-01

备注: Accepted at ICML 2026

💡 一句话要点

提出DOT-MoE，通过可微最优传输实现高效MoE化，提升大模型推理效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 模型压缩 最优传输 可微学习 大语言模型

📋 核心要点

  1. 现有MoE方法依赖启发式聚类或随机分割，导致专家划分不合理，影响模型性能。
  2. DOT-MoE将稠密层分解建模为可微最优传输问题，实现神经元到专家的最优分配。
  3. 实验表明，DOT-MoE在减少50%活跃参数的同时，保留了原始模型90%的性能。

📝 摘要（中文）

大规模语言模型(LLMs)的扩展虽然带来了显著的性能提升，但也给推理效率带来了巨大挑战。混合专家模型(MoEs)通过将模型大小与推理成本解耦来解决这个问题，但从头开始训练MoEs通常不稳定且计算密集。将预训练的稠密模型转换为稀疏MoEs已成为一种替代解决方案；然而，现有方法通常依赖于启发式神经元聚类或随机分割来将前馈网络(FFN)划分为专家。在这项工作中，我们提出了DOT-MoE，这是一种新颖的框架，它将稠密层的分解公式化为可微最优传输(DOT)问题。我们没有采用静态启发式方法，而是将神经元分配建模为一个平衡的传输问题，利用可微Sinkhorn-Knopp迭代来强制执行严格的专家容量约束。此外，我们利用Straight-Through Estimators (STE)来端到端地联合学习离散的神经元到专家的分配和token到专家的路由策略。在多种架构和基准测试中进行的大量实验表明，DOT-MoE显著优于结构化剪枝、启发式聚类和随机分割基线，在减少50%的活跃参数的同时，保留了原始稠密模型90%的性能。

🔬 方法详解

问题定义：现有将稠密模型转化为MoE的方法，如启发式聚类或随机分割，无法保证神经元到专家的最优分配，导致模型性能下降。这些方法缺乏可学习性，无法根据数据自适应地调整专家结构。

核心思路：DOT-MoE的核心思想是将神经元到专家的分配问题建模为一个可微的最优传输问题。通过最小化神经元特征和专家特征之间的传输成本，实现神经元的最优分配，从而提高MoE模型的性能。这种方法允许模型学习专家结构，而不是依赖于预定义的启发式规则。

技术框架：DOT-MoE框架包含以下主要模块：1) 稠密层分解模块：将稠密层分解为多个专家层。2) 最优传输模块：使用Sinkhorn-Knopp迭代解决神经元到专家的最优传输问题，得到软分配矩阵。3) Straight-Through Estimator (STE)：用于联合学习离散的神经元到专家的分配和token到专家的路由策略。4) MoE层：根据学习到的分配策略，将token路由到相应的专家进行处理。

关键创新：DOT-MoE的关键创新在于将神经元分配问题建模为可微最优传输问题，并使用Sinkhorn-Knopp迭代进行求解。与传统的启发式方法相比，DOT-MoE能够学习到更优的专家结构，从而提高MoE模型的性能。此外，使用STE进行端到端训练，可以同时优化神经元分配和token路由策略。

关键设计：DOT-MoE的关键设计包括：1) 使用Sinkhorn-Knopp迭代来强制执行严格的专家容量约束，保证每个专家处理的神经元数量大致相同。2) 使用Straight-Through Estimators (STE)来解决离散分配带来的梯度消失问题，实现端到端训练。3) 定义合适的传输成本函数，例如神经元特征之间的距离，以指导神经元的最优分配。4) 平衡专家容量约束和传输成本之间的权重，以获得最佳的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，DOT-MoE在多种架构和基准测试中显著优于结构化剪枝、启发式聚类和随机分割等基线方法。DOT-MoE在减少50%活跃参数的同时，保留了原始稠密模型90%的性能。这表明DOT-MoE能够有效地将稠密模型转换为稀疏MoE，并在保持性能的同时显著降低计算成本。

🎯 应用场景

DOT-MoE可应用于各种大规模语言模型的MoE化，尤其适用于资源受限的场景。通过将预训练的稠密模型转换为稀疏MoE，可以在保持性能的同时显著降低推理成本，从而加速LLM在边缘设备和移动端的部署。该方法还可用于模型压缩和加速，提高计算效率。

📄 摘要（原文）

The scaling of Large Language Models (LLMs) has driven significant performance gains but created substantial challenges in inference efficiency. While Mixture of Experts (MoEs) architectures address this by decoupling model size from inference cost, training MoEs from scratch is often unstable and compute intensive. Conversion of pre-trained dense models into sparse MoEs has emerged as an alternative solution; however, existing methods typically rely on heuristic neuron clustering or random splitting to partition the Feed-Forward Network (FFN) into experts. In this work, we propose DOT-MoE, a novel framework that formulates the decomposition of dense layers as a Differentiable Optimal Transport (DOT) problem. Instead of static heuristics, we model neuron assignment as a balanced transport problem, utilizing differentiable Sinkhorn-Knopp iterations to enforce strict expert capacity constraints. Furthermore, we utilize Straight-Through Estimators (STE) to jointly learn the discrete neuron-to-expert assignment and the token-to-expert routing policy end-to-end. Extensive experiments across multiple architectures and benchmarks demonstrate that DOT-MoE significantly outperforms structured pruning, heuristic clustering, and random-split baselines, retaining 90% of the original dense model's performance while reducing active parameters by 50%.