Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
By William Fedus*, Barret Zoph*, Noam Shazeer
TL;DR
"We simplify the Mixture of Experts (MoE) routing algorithm to route each token to only a single expert. This simple change yields a 7x speedup in pre-training, mitigates training instabilities, and allows us to scale language models up to 1.6 trillion parameters on standard hardware budgets."
Why This Matters
Modern deep learning has been dominated by a simple rule: bigger models perform better. However, scaling up traditional dense models (like standard Transformers) means that every single parameter is computed for every single token. This creates a massive computational bottleneck, requiring astronomical budgets and energy to train state-of-the-art models.
Mixture of Experts (MoE) models offered a promising alternative. By routing tokens to different sub-networks (experts), they allow the model's parameter count to grow independently of the FLOPs (floating-point operations) per token. Yet, MoE models historically suffered from intense training instability, high communication overhead between devices, and complex routing algorithms that hindered widespread adoption.
The Big Idea
The Switch Transformer simplifies MoE by making a radical design choice: route each token to exactly one expert (Top-1 Routing), rather than two or more.
This simple shift has massive downstream engineering benefits:
- Reduced Router Computation: Selecting the single best expert is mathematically simpler than sorting and balancing multiple.
- Halved Expert Capacity: The buffer size needed on each device to process tokens is dramatically reduced.
- Simplified Communication: All-to-all cross-device communication is minimized, matching hardware limits perfectly.
Interactive Demo 1: Top-1 vs. Top-2 Routing
Toggle between standard Mixture of Experts (Top-2) and Switch Transformer (Top-1) to see how routing paths, compute cost, and cross-device communication scales.
How It Works
1. Simplifying Sparse Routing
Let $x$ be the input representation of a token. The router produces a probability distribution over $N$ experts using a softmax layer:
Where $h(x) = W_r \cdot x$ represents the router's logits. In standard MoE, the top-$k$ (where $k \ge 2$) indices are selected. In Switch Routing, we set $k=1$. The output of the layer is simply:
where $i$ is the index of the single expert with the highest probability.
2. Expert Capacity & Token Dropping
To ensure static tensor shapes required by hardware accelerators (like TPUs), each expert has a fixed capacity:
If more tokens are routed to an expert than its capacity allows, the excess tokens are dropped. Dropped tokens bypass the expert computation entirely and are sent directly to the next layer via a residual connection.
Interactive Demo 2: Capacity Factor & Token Dropping Simulator
Simulate how uneven token routing can cause experts to overflow, and see how increasing the Capacity Factor or applying a Load Balancing Loss prevents dropped tokens.
Higher capacity factors reduce dropped tokens but increase memory padding and communication overhead.
3. Differentiable Load Balancing Loss
To prevent all tokens from routing to just a few "popular" experts, Switch Transformers use an auxiliary loss during training. For each Switch layer, this loss is added to the total model loss:
Where $f_i$ is the fraction of tokens dispatched to expert $i$, and $P_i$ is the fraction of the router probability allocated to expert $i$. Minimizing this dot product encourages a uniform distribution of tokens across all experts.
Interactive Demo 3: Scaling & Parameter Calculator
See how increasing the number of experts scales the total model parameters while keeping the FLOPs per token completely constant.
Headline Results
The Switch Transformer sets new benchmarks for training efficiency, showing consistent improvements across pre-training speed, downstream task fine-tuning, and distillation.
| Model | Parameters | FLOPs/seq | Neg. Log Perplexity @500k steps |
|---|---|---|---|
| T5-Base (Dense) | 223M | 124B | -1.556 |
| Switch-Base (Sparse) | 7.4B | 124B | -1.306 |
| T5-Large (Dense) | 739M | 425B | -1.350 |
| Switch-Large (Sparse) | 26.3B | 425B | -1.177 |
Distillation: Shrinking Trillion-Parameter Models
Deploying models with billions or trillions of parameters can be highly impractical. To solve this, the authors successfully distilled these massive sparse models back into compact dense models. By combining weight initialization from non-expert layers and a mixture of hard/soft losses, they preserved 30% of the sparse model's quality gains while reducing parameter count by up to 99%.
Limitations & Open Questions
- Fine-Tuning Instabilities: While Switch Transformers pre-train exceptionally well, scaling to extreme parameter sizes (like Switch-XXL) sometimes introduces sporadic training instability during downstream fine-tuning.
- Task-Type Discrepancy: The model translates upstream gains significantly better to knowledge-heavy tasks (e.g., TriviaQA) than to complex reasoning tasks (e.g., SuperGLUE).
- Hardware Co-adaptation: The static tensor requirements of modern accelerators (TPUs/GPUs) force the use of capacity factors, occasionally leading to dropped tokens or wasted memory padding.
Glossary
Cite This Work
@article{fedus2022switch,
title={Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity},
author={Fedus, William and Zoph, Barret and Shazeer, Noam},
journal={Journal of Machine Learning Research},
volume={23},
number={120},
pages={1--40},
year={2022}
}