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

2026-06-16

AI Summary

This work addresses the inefficiency of traditional sparse Mixture-of-Experts (MoE) models, which rely on non-differentiable top-k routing that rigidly activates a fixed number of experts per input, leading to suboptimal computational resource utilization. The authors propose SoftMoE, the first end-to-end differentiable soft routing mechanism for MoE. It replaces discrete routing with a truncated soft top-k LapSum relaxation and incorporates a global computation budget constraint alongside hierarchical expert capacity parameterization, enabling adaptive, dynamic expert allocation across layers. Compatible with autoregressive language modeling, SoftMoE matches or exceeds the performance of conventional MoE approaches on language modeling and downstream tasks while significantly reducing the number of activated experts—particularly allocating more capacity in deeper layers—to enhance computational efficiency.

Abstract

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

Problem

• Mixture-of-Experts
• differentiable routing
• sparse activation
• large language models
• computation efficiency

Innovation

• SoftMoE
• differentiable routing
• Mixture-of-Experts
• soft top-k
• expert allocation