Switch Transformer › Top-1 Routing: The Simplification Insight

Prior MoE models, including the influential GShard architecture, used top-2 routing: each token was sent to two experts, with the outputs combined as a weighted average. The reasoning was intuitive: combining perspectives from multiple specialists should produce richer representations. After all, if one expert knows about syntax and another knows about semantics, shouldn't we benefit from consulting both? The original Shazeer MoE paper from 2017 used top-2 routing as a default, and subsequent work largely inherited this choice without questioning it deeply. Switch Transformer was specific partly for its willingness to challenge this received wisdom empirically rather than theoretically.

Switch Transformer challenged the top-2 assumption with empirical evidence that contradicted conventional wisdom. The authors found that routing to a single expert reaches comparable or better performance while giving several practical advantages that compound at scale. The finding held across multiple model sizes, pre-training objectives, and evaluation benchmarks. Its consistency indicated that top-1 routing reflects MoE optimization rather than a dataset-specific artifact.

The effectiveness of top-1 routing can be understood through several lenses, each revealing a different facet of why simpler is better in this context. These arguments reinforce each other, and together they explain why the empirical results so consistently favor single-expert assignment over the seemingly more powerful two-expert approach.

Reduced communication cost. In distributed training, each expert typically lives on a different device. This means that routing a token to an expert requires sending that token's representation across the network to wherever the expert resides, processing it there, and returning the result. This all-to-all communication is a basic bottleneck in distributed MoE training: it requires every device to communicate with every other device, and the total data transferred scales with the number of tokens per batch times the dimensionality of the token representations. Top-2 routing means every token must participate in two such communication rounds, effectively doubling the network traffic. Top-1 cuts this communication cost in half, which becomes increasingly important as models scale across hundreds or thousands of accelerators. At the scale of Switch-C, which used 2048 experts across a correspondingly large device mesh, this factor-of-two communication reduction translates to a major practical advantage.

Simpler gradient flow. With top-2 routing, the gradient for a token flows through two expert networks weighted by their respective probabilities. This creates interdependencies that complicate optimization: the gradient with respect to one expert's parameters depends on what the other expert produced, because the router probability determines how much each expert's output contributes to the final result. When both experts influence the output, neither expert receives a clean signal telling it how to improve. Top-1 gives cleaner gradients, where each expert receives full responsibility for the tokens it processes. When expert 7 handles a token and the model produces a poor output, the gradient flows cleanly back through expert 7 and through the router's probability for expert 7. This clarity in the optimization objective helps the model converge more reliably and prevents situations where two experts compete for the same tokens without either fully committing to specialization.

Higher effective capacity per expert. If each token visits two experts, and each expert has a fixed capacity (more on this shortly), then top-2 routing effectively halves the number of tokens any expert can accommodate. Top-1 routing means twice as many tokens can be processed per expert for the same capacity setting. This translates to larger effective batch sizes per expert, which improves the quality of gradient estimates and allows each expert to see more diverse examples during training. Better gradient quality per expert means faster and more reliable specialization.

Empirical validation at multiple scales. The Switch Transformer paper demonstrated that top-1 routing with the same total compute reaches better perplexity than top-2 routing across multiple model scales. This wasn't a marginal improvement; the gains were consistent and significant, validating the theoretical arguments with concrete evidence. The perplexity gap was largest at the largest scales, consistent with the hypothesis that the communication and gradient benefits compound as scale increases.

The mathematical formulation for top-1 routing is straightforward, which is part of its elegance. Given a token representation $x$, the router computes a probability distribution over all available experts by first projecting the input into a logit vector and then applying softmax:

$$p(x) = \text{softmax}(W_r \cdot x)$$

where:

• $p(x)$: the probability distribution over experts for the input token, with $p_i(x)$ representing the probability of routing to expert $i$
• $W_r$: the learnable router weight matrix of shape $[d_{\text{model}} \times E]$, projecting the input to a vector of expert logits
• $x$: the input token representation, a vector of dimension $d_{\text{model}}$

This probability distribution is the router's assessment of how well-suited each expert is for processing this particular token. The selected expert is simply the one with the highest probability: $e^* = \arg\max_i p_i(x)$.

Why does this formula make sense? Notice that the linear projection $W_r \cdot x$ produces a scalar score for each expert, and the softmax converts these scores into a proper probability distribution that sums to 1. The expert with the highest score gets the highest probability, and the argmax selects that expert as the routing destination. The simplicity of this computation is what allows the router to be extremely fast, adding minimal overhead to the forward pass.

Once the expert is selected, the output of the Switch layer combines the expert's processing with a scaling factor based on the router's probability:

$$y = p_{e^*}(x) \cdot E_{e^*}(x)$$

where:

• $y$: the output of the Switch layer for the token, which replaces the dense FFN output in the transformer block
• $e^*$: the index of the selected expert (the one with the highest probability)
• $p_{e^*}(x)$: the probability assigned to the selected expert, acting as a multiplicative gating factor
• $E_{e^*}(x)$: the output of the selected expert network after processing the input token through its feed-forward layers

The multiplicative weighting by $p_{e^*}(x)$ serves an important purpose beyond just combining the expert output. It gives a differentiable signal that allows gradients to flow back through the router, letting the router to learn which experts work best for which tokens. Without this multiplicative term, the argmax operation would be non-differentiable: the gradient would have no information about how to adjust the router's parameters. By preserving the router probability in the output, we create a path through which gradients can flow during backpropagation, letting end-to-end training of both the router and the experts.

A potential concern with top-1 routing is that hard assignment creates discontinuities. Small changes in a token's representation could flip it to a different expert, potentially causing unstable training dynamics. If a token is hovering on the decision boundary between two experts, tiny perturbations might cause it to oscillate between them during training. This could create a noisy gradient signal where the expert receiving a token changes unpredictably from step to step, preventing stable specialization.

However, the multiplicative weighting smooths this effect in an elegant way. If a token is marginally assigned to expert 3 with probability 0.51 (vs. 0.49 for expert 2), the output is scaled by 0.51. This scaling gives a soft transition that prevents gradient spikes. When the token is near a decision boundary, the output magnitude is reduced, naturally dampening the impact of near-ties. As the router becomes more confident about a token's assignment, the probability approaches 1.0 and the full expert output is used.

This design ensures that tokens near decision boundaries contribute less strongly to the gradient signal, letting the model to focus its learning on tokens where the routing decision is clear and real. Think of it as the model saying: when I am not sure which expert should handle a token, I will use a softer, more conservative output rather than committing fully to an uncertain choice. This prevents spurious gradient updates driven by borderline routing decisions and encourages the router to become increasingly confident over the course of training, pushing probabilities toward 0 or 1 for most tokens as the experts develop clear specializations.