Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

The 2017 mixture-of-experts layer sent each token to the top k experts, with k of at least two so the router got useful gradients. Switch Transformers asked whether one is enough. Routing each token to a single expert (the "switch") cuts the router's computation, halves the minimum batch each expert needs, and simplifies the communication of tokens between devices. Inside a Transformer, the Switch layer replaces the feed-forward block; attention stays dense.

One expert per token makes balance and capacity central. Every expert gets a fixed buffer of tokens per batch: (tokens ÷ experts) × a capacity factor. When more tokens choose an expert than it has room for, the extras skip the layer and pass through the residual connection unchanged. An auxiliary loss, α × N × Σ f_i × P_i, pushes the share of tokens each expert receives (f) and the router's average probability for it (P) toward an even split.

Sparse models were also unstable. The authors compute the router alone in float32 while the rest of the model stays in bfloat16, which stopped divergence at bfloat16 speed; scale the default weight initialization down by ten; and fine-tune with dropout of 0.4 inside experts and 0.1 elsewhere.

At the same compute per token as T5-Base, Switch models reached T5-Base quality up to 7 times faster, and gains held across 101 languages. A 1.6-trillion-parameter model trained without instability and reached a fixed perplexity 4 times faster than T5-XXL. Sparse models could also be distilled into small dense ones, keeping about 30% of their gain.

Most deployed mixture-of-experts models descend from choices this paper made, and its vocabulary shows up in serving configurations.

Capacity factor is a quality-versus-memory dial. A low capacity factor saves memory and communication but drops tokens when routing is uneven; a dropped token skips that layer's expert entirely. If an MoE model behaves oddly on unusual inputs, uneven routing with dropped tokens is one cause to check, and some serving stacks expose this setting.

Sparse models are memory-hungry. Switch-C has 1.6 trillion parameters, about 3.1 TB of weights in bfloat16, even though each token touches a small fraction. The same split between parameters stored and parameters used applies to every MoE model you deploy.

Distillation is the exit ramp. If a customer wants an MoE model's quality on modest hardware, the paper shows a dense student can keep a share of the sparse model's gain. Its numbers (about 30%) set an honest expectation: distillation recovers part of the gain, not all of it.