DOT-MoE:用于MoE化的可微最优传输
DOT-MoE: Differentiable Optimal Transport for MoEfication
AI总结:
提出DOT-MoE框架,通过可微最优传输将密集层分解为专家,联合学习神经元分配和路由策略,在减少50%活跃参数的同时保留90%原始性能。
摘要:
大型语言模型(LLMs)的扩展带来了显著的性能提升,但也造成了推理效率方面的重大挑战。虽然混合专家(MoEs)架构通过将模型大小与推理成本解耦来解决这一问题,但从头训练MoEs通常不稳定且计算密集。将预训练的密集模型转换为稀疏MoEs已成为一种替代方案;然而,现有方法通常依赖启发式神经元聚类或随机分割来将前馈网络(FFN)划分为专家。在这项工作中,我们提出了DOT-MoE,一种新颖的框架,将密集层的分解建模为可微最优传输(DOT)问题。与静态启发式方法不同,我们将神经元分配建模为平衡传输问题,利用可微的Sinkhorn-Knopp迭代来强制执行严格的专家容量约束。此外,我们利用直通估计器(STE)来联合学习离散的神经元到专家的分配和令牌到专家的路由策略。跨多个架构和基准的大量实验表明,DOT-MoE显著优于结构化剪枝、启发式聚类和随机分割基线,在减少50%活跃参数的同时保留了原始密集模型90%的性能。
英文摘要:
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%.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。