让专家模型路由可优化,自动分配计算资源,提升效率。
SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs

- 用软化路由替代离散选择,实现梯度反传优化。
- 相同算力下激活专家更少,语言任务性能相当或更好。
- 模型自学习各层专家使用量,后期层更依赖专家。
稀疏专家混合(MoE)架构通过 top-$k$ 路由仅激活少量专家,在固定推理预算下扩展大模型参数。但离散的 top-$k$ 操作不可微,导致每输入激活专家数固定,计算效率低下。本文提出 SoftMoE,以截断的软 top-$k$ LapSum 逼近替代离散路由,支持梯度优化路由策略。进一步对每层平均激活专家数参数化,并施加全局预算约束,使模型可学习跨层专家容量分配。该方法保持与自回归建模完全兼容,在语言建模及下游任务上表现相当于或优于传统稀疏 MoE,同时显著减少激活专家数量。值得注意的是,学习到的分配高度非均匀,后期层激活更多专家。源代码已公开。
原文摘要 · Abstract (English)
Sparse Mixture-of-Experts (MoE) architectures enable scaling LLM parameters under a fixed inference budget by activating only a small subset of experts via top-$k$ routing. While this preserves causality and suits autoregressive language models, the discrete top-$k$ operator is not differentiable, forcing a fixed number of active experts per input and resulting in inefficient use of computation. We propose SoftMoE, which replaces discrete routing with a truncated soft top-$k$ LapSum relaxation, allowing gradient-based optimization of expert routing. We further parameterize the mean number of active experts per layer and impose a global budget constraint, enabling the model to learn how to allocate expert capacity across layers. SoftMoE remains fully compatible with autoregressive modeling and achieves performance comparable to or better than sparse MoE on language modeling and downstream tasks, while activating significantly fewer experts. Notably, the learned allocation is highly non-uniform, with later layers activating more experts. The source code is publicly available$^\dagger$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。