Contribution
It shows that the simplest sparse routing, one expert per token, works as well as or better than top-2 mixture-of-experts at lower cost, and supplies the engineering (capacity factor, auxiliary balancing loss, selective float32, small initialization, expert dropout, expert parallelism) that makes it train at scale. The empirical case is that parameter count at fixed computation per token is a separate axis of scaling, from 2 experts on a few cores to 2048 experts and 1.6 trillion parameters.
Key insight
Parameters and computation can be decoupled: a token sees one expert’s weights, so adding experts adds parameters but not FLOPs per token, and the language-modelling loss keeps improving with the number of experts. The price is a routing decision that must be kept balanced, both for hardware (fixed capacity per expert, overflow tokens skip the layer) and for learning (experts that receive no tokens do not train).
Assumptions
- Encoder–decoder T5 models on C4 (masked span prediction), TPU v3, Mesh TensorFlow; static tensor shapes, hence a fixed expert capacity (tokens per batch / experts × capacity factor).
- Experts replace every other feed-forward layer (Table 1); homogeneous experts throughout.
- Comparisons are FLOP-matched to T5 baselines; speed is wall-clock on identical hardware, which the authors note depends on implementation.
Key results
- Switch vs. top-2 MoE (Table 1, 128 experts). At capacity factor 1.0–1.25 Switch-Base is faster and reaches the −1.50 threshold sooner; a wider Switch-Base+ matched to MoE’s speed is best per step.
- Stability (Tables 2–3). bfloat16 throughout diverged (−3.780); casting only the router to float32 matched float32 quality (−1.716 vs −1.718) at bfloat16 speed. A 0.1× initialization scale gave −2.72 ± 0.01 against −3.60 ± 0.68 at 3.5k steps over 3 seeds.
- Scaling (§3). Loss improves with experts at fixed FLOPs; Switch-Base 64e matches T5-Base’s 450k-step quality at 60k steps (7.5× by step) and in one-seventh the time; Switch-Base beats T5-Large, which uses 3.5× more FLOPs, by 2.5× in time.
- Fine-tuning (Table 5). SuperGLUE +4.4 (Base) and +2.0 (Large); large gains on Winogrande, closed-book TriviaQA and XSum; losses on ARC Challenge (Base, 32.8 vs 35.5) and ARC Easy (Large, 66.0 vs 68.8).
- Distillation (Tables 6–8). Initializing from non-expert weights and mixing 0.25 teacher with 0.75 hard labels keeps about 30% of the quality gain at 95–99% compression.
- Multilingual (§4.3). Gains on all 101 languages over mT5-Base; mean step speed-up 5×, 91% of languages at least 4×.
- Trillion scale (Table 9). Switch-XXL (395B, FLOP-matched to T5-XXL) and Switch-C (1.6T, 2048 experts) beat T5-XXL in negative log perplexity at 250k steps by 0.061 and 0.051; Switch-XXL by 0.087 at 500k. Switch-XXL was sometimes unstable; Switch-C was not. Downstream, Switch-XXL sets closed-book QA records but trails on SuperGLUE.
- Appendices. Expert attention layers help in float32 but diverge in bfloat16 (A); rerouting overflow tokens (“No-Token-Left-Behind”) brought no benefit (B); of four exploration schemes, input jitter was best and sampling from the softmax worst (C); 2–8 experts already help (D); for fixed upstream perplexity Switch does better than dense on TriviaQA and worse on SuperGLUE at the largest scale (E).
Claims
| id | claim | strength | support |
|---|---|---|---|
| C1 | Top-1 routing preserves quality relative to top-2 at lower cost | moderate | Table 1, one setting, 128 experts |
| C2 | More experts at fixed FLOPs per token improve pre-training loss and speed | strong | Figs. 4–6, repeated across sizes |
| C3 | Router-local float32 suffices for bfloat16 stability | moderate | Table 2, one model, early in training |
| C4 | The auxiliary loss with α = 10⁻² balances load without hurting training loss | weak | stated from a sweep 10⁻¹…10⁻⁵ whose results are not shown |
| C5 | Upstream gains translate better to knowledge tasks than to reasoning tasks | weak | Appendix E, few points; the authors call it preliminary |
Method
Router logits h(x) = W_r x, softmax p(x); each token goes to argmax p and the expert’s output is scaled by its gate value p_i(x), which keeps the router differentiable. Expert capacity fixes how many tokens each expert takes; overflow tokens pass through the residual. The auxiliary loss is α·N·Σ_i f_i P_i, f_i the fraction of tokens dispatched to expert i (not differentiable) and P_i the mean router probability for i (differentiable), computed per core in the pseudocode.
Concepts
- Switch layer: a top-1 routed mixture-of-experts feed-forward layer.
- Expert capacity / capacity factor: the static batch size per expert and its slack over an even split.
- Dropped tokens: overflow tokens that skip the layer (typically < 1%).
- Expert dropout: a higher dropout rate (0.4) inside experts at fine-tuning.
Connections
It simplifies Shazeer et al.’s sparsely-gated MoE (LIT-1507), which had conjectured that k > 1 was needed for router gradients, and follows GShard’s simplified auxiliary loss and overflow handling. Wang et al. (LIT-1414) replace its auxiliary loss with a bias controller and use its loss as their baseline; Li et al. (LIT-1411) borrow the loss’s form for a gating network in continual learning.
Bearing on the record
- CASE-117. The LIT says the exchange cited it as “the standard case of routing among experts with an explicit load-balancing term, the imposed alternative to diversity that emerges”. The reading confirms that the term is in the paper (eq. 4), with a sharper description: the balancing is imposed on the router’s dispatch fractions and mean probabilities, not on what experts learn, and the paper never measures specialization or diversity among experts. So it is the standard case of imposed balance of load, and evidence about emergent diversity only in the negative sense that it was not studied. Appendix C’s finding that sampling the router is worse than argmax, and jitter slightly better, bears on CASE-117’s exploration regimes. Appendix B’s result, that rerouting overflow tokens to a second-choice expert did not help and that “if this association is changed … performance could be degraded”, is a small datum for the case’s question of practical exclusion.
The paper is machine-learning practice throughout and belongs to the anthology; it carries the anthology-candidate flag.
Limitations
- Single runs for nearly every comparison; seeds only in Table 3.
- Wall-clock comparisons depend on the authors’ TPU implementation, as they say.
- Errors and slips. (i) §2.2 says the auxiliary loss “is minimized under a uniform distribution”. It is not only so minimized: with two experts, if 90% of tokens have p = (0.5 + δ, 0.5 − δ) and 10% have p = (0, 1), then f = (0.9, 0.1), P → (0.45, 0.55) and N·Σf_iP_i → 0.92 as δ → 0 (0.934 at δ = 0.01, computed), below 1, the uniform value. The loss pushes towards balance in practice but its uniform point is not its minimum in general. (ii) §5.6 says “both Switch Transformer variants improve over the T5-XXL … by over 0.061” at 250k steps; by Table 9 Switch-C’s margin is 0.051. (iii) Table 11’s caption says “lower is better” for negative log perplexity, where the text and the numbers mean higher is better. (iv) The router pseudocode adds Uniform(1 − ε, 1 + ε) noise to the logits, while Appendix C describes multiplicative jitter on the incoming representation. (v) The abstract’s “up to 7x” is by time; §3.1’s 7.5× is by step.
Open questions
- What do the experts specialize in, and does balancing load suppress or produce specialization? Not addressed.
- Why does Switch-C, with more parameters and fewer FLOPs, fine-tune worse than Switch-XXL (SQuAD 87.7 vs 89.6)? The authors call the dependence “poorly understood”.