SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs

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

Modern language models are becoming increasingly large because their performance strongly correlates with the number of parameters. One way to increase network size while maintaining near-constant inference cost is to use Mixture of Experts (MoE) architecture, which activates only a small fraction of the network, called experts. This architecture is becoming an ubiquitous solution in modern Large Language Models. That said, expert selection in MoE networks is typically based on non-differentiable top-$k$ routing, which introduces significant limitations and hinders optimization of the routing process. We propose SoftMoE an expert selection mechanism that replaces the hard top-k function with its soft, differentiable version based on the LapSum operator. Instead of fixing the number of activated experts a priori, SoftMoE allows the model to adaptively adjust the activation size depending on the input token. Furthermore, we propose a mechanism that enables the model to learn how to allocate active experts between network layers, while maintaining a global computational budget constraint. Our experiments show that SoftMoE achieves comparable or better results than classical expert routing methods, while activating significantly fewer experts on average. Moreover, the analysis of the learned budget allocation reveals an interesting phenomenon: the model naturally shifts a larger part of the experts' budget to deeper layers of the transformer architecture. This may provide new insights into how information is processed in Large Language Models.