SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs
This work addresses the inefficiency of discrete top-k routing in MoE for LLMs, offering a differentiable alternative that improves computational efficiency without sacrificing performance.
SoftMoE introduces a differentiable soft top-k routing mechanism for Mixture-of-Experts in LLMs, enabling gradient-based optimization of expert allocation. It achieves comparable or better performance than sparse MoE while activating significantly fewer experts, with learned non-uniform allocation across layers.
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$.