The Switch Transformer: How Google Scaled to Trillion Parameters by Routing Each Token to One Expert

Every major language model you can name today runs the same computation for every token. GPT-2, LLaMA, Mistral, the T5 family: when a token passes through the feed-forward layer, all model weights participate every time. That design is simple and stable, but expensive. Doubling the parameters doubles the compute per token. The Switch Transformer, introduced in January 2021 by William Fedus, Barret Zoph, and Noam Shazeer at Google Brain, broke that constraint. Published as 'Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity' (arXiv:2101.03961) and later accepted to JMLR, the paper showed that routing each token to exactly one specialized sub-network called an expert made it possible to scale total parameter count without scaling compute proportionally. The result was a 7x speedup in pre-training speed over T5 on the same hardware budget, and the first language models trained at trillion-parameter scale.

Why dense models hit a compute wall

In a standard dense Transformer, every parameter fires on every token in every forward pass. Adding parameters costs proportionally more at inference and training time. GPT-3 with 175 billion parameters costs roughly 117 times as much compute per token as GPT-2 with 1.5 billion. Mixture of Experts (MoE) was a known approach to break this constraint. The core idea dates back to Jacobs et al. in 1991, and Shazeer et al. revisited it in their 2017 paper 'Outrageously Large Neural Networks,' which replaced dense feed-forward layers with a gated network that dispatched each token to the top-2 highest-scoring experts from a pool. That 2017 system showed strong improvements in machine translation quality but carried real costs: the routing logic was complex, all-to-all communication across distributed GPUs was expensive, and training instability was common. Widespread adoption stalled. The Switch Transformer's primary contribution was to ask whether MoE's complexity was load-bearing, or whether a simpler design could match or beat it.

The switch: one expert per token

The answer was simpler than expected. Instead of routing each token to the top-2 experts as Shazeer et al. 2017 did, Switch sends each token to exactly ONE expert. A small learned linear layer (the router) computes a probability distribution over the available experts for every token in the sequence. The token is dispatched to whichever expert scores highest. That is the switch in Switch Transformer: a hard, discrete dispatch to a single expert rather than a weighted sum across multiple. This simplification delivered compounding benefits. Communication per token dropped by half, since each token now crosses at most one GPU boundary. Gradient flow simplified, since backpropagation traces one expert per token rather than two. The router could still learn meaningful routing behavior despite the hard selection, because the router weights themselves receive gradients through the router probability scores. In practice, experts quickly specialize during pre-training, with different experts gravitating toward different syntactic patterns, domains, or token types.

Capacity factors and load balancing

Sparse routing introduces a problem: nothing prevents all tokens in a batch from being routed to the same expert. The paper addresses this with two mechanisms. First, a capacity factor limits how many tokens each expert will accept per batch, set to capacity_factor times (total_tokens divided by num_experts). A capacity factor of 1.0 means experts are perfectly balanced. In practice the paper uses values of 1.25 to 2.0 to absorb natural imbalances without excessively dropping tokens. Tokens routed to a full expert skip that layer and pass through unchanged (their residual stream is preserved, but the expert computation is skipped for them). Second, an auxiliary load-balancing loss added to the training objective penalizes routing distributions that concentrate too many tokens on too few experts. This loss uses a small coefficient (the paper experiments with values around 0.01 to 0.1) and is computed from the fraction of tokens dispatched to each expert versus the fraction of router probability mass assigned to each expert. Multiplying these two fractions and summing across experts gives a differentiable measure of imbalance that the optimizer can minimize without disrupting the language model loss.

Training stability and bfloat16

Sparse models have a reputation for training instability, and for good reason: early in training, if routing is highly unbalanced, some experts receive no tokens and thus no gradient signal, while overloaded experts may diverge. The Switch Transformer paper introduced two targeted fixes. The first is a small initialization for router weights, keeping initial routing probabilities close to uniform and preventing early specialization that locks in a bad routing pattern before the experts have learned anything useful. The second is selective precision training: the router computation runs in float32 for numerical stability, while the rest of the model uses bfloat16. This combination produced, according to the paper, the first successful training of large sparse language models in bfloat16 precision, reducing memory footprint and memory bandwidth requirements versus float32. The paper also introduced a fallback routing strategy called No-Token-Left-Behind: when a token is dropped by its top-1 expert due to capacity limits, it is rerouted to its second-highest expert rather than discarded entirely, preventing complete information loss for dropped tokens.

Results: speed, scale, and multilingual gains

The benchmark results span pre-training speed, fine-tuning quality, scale, and multilingual transfer.

• Pre-training speed: Switch Transformer variants based on T5-Base and T5-Large architecture achieved up to 7x faster pre-training on the same compute budget, measured in steps to reach a fixed perplexity on the C4 dataset (Colossal Clean Crawled Corpus).
• Trillion-parameter scale: the largest Switch models pre-trained reached up to a trillion parameters, achieving a 4x speedup over the T5-XXL model (which has 11 billion parameters) on a per-time-basis comparison.
• Multilingual: Switch Transformer showed improvements over mT5-Base across all 101 languages tested in the multilingual evaluation, including low-resource languages where data scarcity makes every parameter update count.
• Distillation: the paper showed that a large Switch model could be compressed into a small dense model using knowledge distillation, achieving a reduction of up to 99% in model size while retaining approximately 30% of the quality gains the sparse model had over a dense baseline.
• Low-compute regimes: even with as few as two experts (the minimum meaningful configuration), Switch models outperformed equally sized dense baselines, suggesting the routing mechanism adds value even at very small scale.

The legacy: MoE at every frontier lab

Switch Transformer established that top-1 sparse routing was both practical and stable at scale, and the wave of MoE models that followed shows how thoroughly that lesson was absorbed. Google's GLaM model (2021) used MoE routing and demonstrated strong few-shot performance at lower training cost than GPT-3. Mixtral-8x7B from Mistral AI, released in 2023, uses top-2 routing with 8 experts and 46.7 billion total parameters while matching GPT-3.5-class performance at the active compute cost of roughly a 12 billion parameter model. OpenAI has not officially confirmed GPT-4's architecture, but the model is widely believed by researchers to use a mixture of experts design. Google's Gemini 1.5 Pro uses MoE. Each of these systems traces its practical feasibility back to the problems Switch Transformer solved: a simpler routing algorithm, bfloat16 training stability, and proven expert-dispatch communication patterns at scale. The paper's core claim proved correct. The fastest path to more capable models is not always building denser networks. Sometimes it is building smarter routing. PromptingIndex covers Switch Transformer alongside QLoRA, ZeRO, Mamba, and multi-head latent attention as part of its series on the infrastructure and architectural choices that define how frontier models are built.