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