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

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

• The paper presents a differentiable, input-dependent routing mechanism using LapSum relaxation to optimize expert selection via gradient-based learning.
• It introduces a learnable, per-layer expert budget that adaptively allocates compute resources, concentrating capacity in deeper layers and reducing expert activations by up to 24%.
• Empirical results demonstrate that SoftMoE achieves lower language modeling loss and competitive downstream performance while maintaining efficient and stable training.

SoftMoE: Differentiable Routing and Adaptive Expert Allocation for Sparse Mixture-of-Experts LLMs

Introduction and Motivation

Scaling LLMs via Mixture-of-Experts (MoE) architectures enables an increase in parameter count without commensurate inference-time compute growth by selectively activating a subset of experts per token. Standard sparse MoE methods employ hard, discrete top-$k$ routing to select active experts, maintaining autoregressive compatibility and preserving model causality. However, the discrete nature of top-$k$ routing introduces fundamental optimization challenges due to non-differentiability. This necessitates the use of surrogate gradients, imposes a rigid, fixed count of active experts per token and layer, and precludes adaptive capacity reallocation across the model. These constraints result in suboptimal computational resource utilization and limit architectural flexibility, especially as scale and expert counts increase.

The SoftMoE architecture addresses these core limitations by introducing a differentiable, input-dependent routing mechanism based on the LapSum relaxation of the top-$k$ operator. This enables end-to-end gradient-based optimization of expert selection, integration of a learnable global capacity budget, and input- and depth-dependent expert allocation. The result is an MoE architecture that retains autoregressive compatibility and computational efficiency while adaptively focusing compute where it is most beneficial.

Differentiable Soft Routing via LapSum Relaxation

SoftMoE replaces the hard top-$k$ selection with a truncated, differentiable soft top-$k$ mechanism constructed using the LapSum function, which is a sum of Laplace cumulative distribution functions. Given routing scores $r \in \mathbb{R}^n$ produced by the gating network for $n$ experts, the LapSum approach computes a threshold such that the sum of soft selection weights (as probabilities) equals a global, possibly non-integer, parameter $k$. These soft weights are interpreted as probabilistic activations of each expert per token and are subsequently thresholded to enforce efficient, sparse computation.

Distinctive properties of LapSum-based routing include:

• Continuous and efficient computation with closed-form inversion, avoiding sorting or iterative optimization.
• Linear time and space complexity with respect to the number of experts.
• Full differentiability not only in the gating scores but also in the expert count parameter $k$, which is pivotal for learning compute allocations across layers.

Overall, this mechanism generalizes the softmax operator to enforce exact cardinality constraints, providing sparse, gradient-compatible expert selection.

Adaptive, Learnable Per-Layer Expert Budgets

A key innovation of SoftMoE is the integration of a learnable, globally-allocated expert budget. Instead of setting a fixed number of active experts per layer, SoftMoE explicitly parameterizes the mean number of active experts for each layer and optimizes these values end-to-end, subject to a global constraint on total expert activations. The per-layer allocations are obtained via a softmax-reparametrized set of allocation parameters that guarantee at least one active expert per layer and a fixed total allocation across the network.

This mechanism introduces competitive expert allocation, wherein layers with higher marginal utility for extra compute can claim more of the global budget, while layers with lower utility allocate less to themselves. Empirically, this results in highly non-uniform capacity distributions—later layers typically garner more experts, aligning with findings that deeper model layers encode more semantically complex or abstract information.

Experimental Evaluation

SoftMoE and baseline models were compared using 1.6B parameter GPT-2-style architectures, with decoders containing 10 transformer blocks and 32 experts per block. Experiments were conducted on the C4 and OpenWebText corpora for pretraining and evaluated on PIQA, HellaSwag, and ARC-E zero-shot downstream benchmarks.

Main results are as follows:

• Efficiency: Across all active expert budgets, SoftMoE consistently achieves lower language modeling loss than sparse MoE baselines, while activating fewer experts on average. Notably, with $k \leq 2$, SoftMoE exhibits a 17–24% reduction in average expert activations during training while improving loss.
• Downstream Tasks: SoftMoE matches or surpasses sparse MoE in downstream accuracies on HellaSwag and PIQA, and is competitive on ARC-E.
• Emergent Allocation Patterns: When learning per-layer allocations, SoftMoE rapidly (within 25k steps) concentrates capacity towards deeper layers, with the final three layers absorbing about 50% of the expert budget.

Crucially, these empirical advantages are achieved without introducing instability or appreciable additional training cost. Training remained stable, and the differentiable routing provided robust gradient flow and convergence properties.

Theoretical and Practical Implications

The differentiable, LapSum-based soft top-$k$0 operator resolves the long-standing issue of mismatched objectives between training (requiring differentiability) and inference (requiring sparsity and efficiency) in MoE LLMs. The framework provides a formal apparatus for end-to-end optimization of both routing and budget allocation, facilitating:

• Efficient compute utilization—flexibly modulating compute on a per-token and per-layer basis.
• Empirical and theoretical investigation into depth-aware capacity allocation, potentially uncovering new strategies for model design at scale.
• Seamless scaling to large expert pools and model sizes due to the linear complexity of the LapSum operator.

From a practical engineering standpoint, SoftMoE allows practitioners to decouple architectural design from rigid capacity settings. Dynamic compute budgeting naturally yields more resource-efficient models, enabling training and deployment under stricter inference constraints without loss of accuracy.

Limitations and Future Work

Evaluation was limited to English-language textual tasks and 1.63B parameter models. It remains an open question how the method performs at trillion-scale model sizes, across diverse modalities, or in multilingual/extensive downstream evaluations. Future research directions include exploring hierarchical or token-level adaptive budget allocation, integration with retrieval-augmented and multimodal models, and systematic study of allocation dynamics as a function of data distribution and task complexity.

Conclusion

SoftMoE demonstrates that differentiable, input- and depth-adaptive expert routing is both feasible and beneficial for scaling LLMs within fixed computational budgets. The introduction of globally-constrained, learnable expert allocation produces non-uniform, semantically-aligned compute allocation across model depth, yielding improved compute efficiency with equal or superior performance relative to conventional sparse MoE models. This paradigm lays a foundation for future work on scalable, conditional computation in large-scale neural architectures, with implications for both model architecture research and real-world deployment efficiency.

Reference: "SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs" (2606.17952).