arXiv:2606.01666cs.LGcs.AI2026-06中稿 · ICML

用可微最优传输方法,高效将大模型转为稀疏专家网络。

DOT-MoE: Differentiable Optimal Transport for MoEfication

论文配图:DOT-MoE: Differentiable Optimal Transport for MoEfication
图 1 · 摘自论文原文
  • 将神经元分配建模为可微最优传输问题,实现动态均衡分组。
  • 在多个模型上保持原性能90%,激活参数减少50%。
  • 适合需要高效推理的大模型部署场景。

大型语言模型的规模扩展带来了显著性能提升,但也引发推理效率难题。虽然混合专家(MoE)架构通过解耦模型大小与推理开销缓解了该问题,但从头训练MoE通常不稳定且计算成本高。将预训练密集模型转换为稀疏MoE成为替代方案;然而,现有方法多依赖启发式神经元聚类或随机分割来划分前馈网络(FFN)。本文提出DOT-MoE,将密集层分解建模为可微最优传输(DOT)问题。不同于静态启发式方法,我们以平衡运输问题建模神经元分配,利用可微Sinkhorn-Knopp迭代强制执行严格的专家容量约束。同时,采用直通估计器(STE)联合学习离散神经元-专家分配与令牌-专家路由策略,实现端到端训练。在多种架构和基准上的大量实验表明,DOT-MoE显著优于结构化剪枝、启发式聚类和随机分割基线,在保留原模型90%性能的同时,将活跃参数减少50%。

原文摘要 · Abstract (English)

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%.

MoE模型压缩可微优化大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。