Soft MoE: Differentiable Routing

Routing mechanisms, for example top-k and switch gating, perform a 'hard' assignment. A token is routed to a discrete, small set of experts, while all other experts are ignored for that specific token's computation. This hard selection is the source of the computational savings in sparse models, but it also introduces challenges like non-differentiability and the need for auxiliary load-balancing losses.

Soft MoE offers a different approach by replacing this discrete selection with a "soft," weighted combination of all experts. Instead of choosing which experts to use, the gating network determines a weight for every expert, and the final output is a weighted sum of the outputs from all experts. This makes the entire MoE layer fully differentiable and elegantly sidesteps the training instabilities associated with hard gating.

The Mechanism of Soft Routing

In a Soft MoE layer, the gating network operates similarly to a standard router by producing a logit for each expert. However, instead of using these logits to select the top-k experts, we apply a softmax function across them. This converts the logits into a set of positive weights that sum to one, effectively forming a probability distribution over the experts.

The final output for an input token $x$ is not the output of a few selected experts, but a linear combination of the outputs from all $N$ experts. The contribution of each expert $E_i(x)$ is scaled by its corresponding softmax weight $w_i$ .

The mathematical formulation is direct. Given an input $x$ , the gating network $G$ computes a vector of logits $h(x)$ . The weights $w$ are then calculated as:

w = \text{softmax}(h(x))

The final output $y$ of the Soft MoE layer is the weighted sum:

y = \sum_{i=1}^{N} w_i \cdot E_i(x)

This formulation might look familiar. It closely resembles the attention mechanism, where a query attends to a set of keys to produce weights, which are then used to compute a weighted sum of values. In Soft MoE, you can think of the token's representation as the query and the experts as the keys and values.

The diagrams below illustrate the difference between the data flow in a hard-gating MoE and a Soft MoE.

In hard routing, the gating network selects a discrete expert (Expert 1), and all computation flows through it. Other experts remain inactive for this token.

In Soft MoE, the gating network computes a weight for every expert. The final output is a weighted combination of all expert outputs.

The Trade-off: Differentiability vs. Computation

The primary advantage of Soft MoE is its elegant solution to the training challenges of sparse models.

Full Differentiability: The entire layer, including the routing decision, is a continuous function. This allows gradients to flow smoothly through the entire network during backpropagation.
Implicit Load Balancing: Because every expert receives a non-zero weight (however small), every expert participates in the forward and backward passes for every token. This naturally prevents the "expert collapse" phenomenon where some experts receive no training signals. There is no need for an explicit auxiliary loss to encourage load balancing.

However, this elegance comes at a significant and often prohibitive cost.

Computational Expense: The fundamental benefit of a sparse MoE is activating only a fraction of the model's parameters for any given input, thereby decoupling the parameter count from the computational cost (FLOPs). Soft MoE completely eliminates this advantage. Since every token is processed by every expert, the computational cost is equivalent to running an ensemble of $N$ models and then averaging their outputs. This makes it as expensive as a dense model with a parameter count equal to the sum of all experts.
Memory Footprint: In a sparse model, only the active experts' weights need to be readily available in high-speed memory (e.g., GPU HBM). In Soft MoE, all expert weights must be loaded simultaneously, leading to a massive memory footprint that is impractical for models with a large number of experts.

When to Consider Soft MoE

Given its computational demands, a "pure" Soft MoE is rarely used in large-scale language models where computational efficiency is a primary design goal. Its formulation serves more as a theoretical benchmark and an analytical tool.

However, the core idea of soft, differentiable assignments has influenced the design of more practical, hybrid systems. For example, some approaches might use a top-k router to select a small subset of experts and then compute a soft, weighted combination within that subset. This can provide some of the training stability of soft routing while retaining most of the computational benefits of sparsity.

Understanding Soft MoE is important because it clearly delineates the trade-off between mathematical simplicity in training and the computational sparsity required for scaling. It represents one end of the spectrum in MoE design, where training stability is maximized at the expense of inference efficiency. This provides a valuable contrast to mechanisms like Switch Transformers, which occupy the other end of the spectrum by prioritizing computational efficiency above all else.

Was this section helpful?

References

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean, 2017 arXiv preprint arXiv:1701.06538 DOI: 10.48550/arXiv.1701.06538 - This foundational paper introduces the sparsely-gated Mixture-of-Experts (MoE) layer, providing context by contrasting it with dense (soft) gating and discussing the computational advantages of sparsity.
Learning to Route: A Differentiable Approach to Mixture of Experts, Clemens Rosenbaum, Chetan Sanan, Charith Gunasekara, Josh Trani, Andrew Gordon Wilson, Kyunghyun Cho, 2018 Proceedings of the 35th International Conference on Machine Learning (ICML), Vol. 80 (PMLR) DOI: 10.5555/3326938.3326955 - This paper presents a methodology for achieving differentiable routing in Mixture of Experts models, aligning directly with the section's discussion of soft routing.
Attention Is All You Need, Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin, 2017 Advances in Neural Information Processing Systems (NeurIPS) 30 DOI: 10.48550/arXiv.1706.03762 - Introduces the Transformer architecture and its self-attention mechanism, which serves as an excellent analogy for the weighted sum computation in soft routing.
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, William Fedus, Barret Zoph, Noam Shazeer, 2021 arXiv DOI: 10.48550/arXiv.2101.03961 - Presents Switch Transformers, a prominent example of hard-gating MoE, demonstrating the computational efficiency of sparsity that Soft MoE diverges from.