Switch Transformer,
explained.

A mixture of experts replaces a transformer's single feed-forward layer with many parallel ones — the experts — and a small router that decides which expert each token should use. The payoff is sparse activation: the model can carry an enormous number of parameters, but each token only touches a few of them, so the compute per token barely rises. Capacity grows without the usual proportional cost in FLOPs.

The catch, before this paper, was that MoE was hard to train. Routing to two or more experts per token added complexity and communication overhead, gradients were noisy, and — worst of all — routers collapsed, funneling most tokens to a handful of favored experts while the rest sat idle. Switch Transformer's contribution was to make MoE simple and stable enough to scale reliably, and its central move was counterintuitive: route less, not more.

Earlier MoE work insisted you needed at least two experts per token — the belief was that routing needed a comparison to produce useful gradients. Switch Transformer challenged that and used top-1 routing: the router scores the experts for a token, and the token goes to the single highest-scoring one. Nothing else.

The simplification pays off three ways. Routing computation halves or better. The all-to-all communication that shuffles tokens to their experts across devices shrinks. And the implementation gets dramatically cleaner — one expert per token means fixed, predictable tensor shapes. The router still produces a probability for the chosen expert, which is used to weight the expert's output so gradients flow back to the router. Less machinery, same learning signal.

Left to itself, a router develops favorites — a rich-get-richer loop where a few experts get most tokens, train faster, and get chosen even more, while the rest waste their parameters. Switch counters this with a load-balancing auxiliary loss, added to the training objective, that is minimized only when tokens spread evenly across experts. With N experts, f_i the fraction of tokens routed to expert i, and P_i the router's mean probability for expert i:

The product f_i·P_i couples how many tokens an expert actually gets to how confident the router is about it, and the sum is smallest when both are spread uniformly (each 1/N), giving N·N·(1/N)² = 1. Pile everything onto one expert and it climbs toward N. Because it is differentiable through P_i, minimizing it during training actively pushes the router toward even utilization. The runnable version below computes it for balanced and skewed loads.

The balancing loss is a soft nudge; hardware needs a hard guarantee. On real accelerators every expert must be given a fixed-size buffer, so Switch sets an expert capacity — the maximum tokens one expert will process in a batch — as tokens-per-expert times a tunable capacity factor:

If routing still sends an expert more tokens than its capacity, the overflow tokens are dropped — they skip the expert entirely and pass through unchanged via the residual connection. A capacity factor above 1 adds slack to reduce dropping, at the cost of wasted compute on padding; below 1 saves compute but drops more. It is the explicit lever that trades a little accuracy for guaranteed, hardware-friendly balance.

With routing simplified and stabilized, the numbers got large. The Switch Transformer scaled to 1.6 trillion parameters — a headline figure at the time — while keeping the compute per token comparable to a far smaller dense model, because only one expert fires per token. That is the whole point of sparsity made concrete: parameters and FLOPs, usually chained together, come apart.

The trade is honest: you pay in memory and cross-device communication for parameters you activate only sparsely. But for large-scale training, spending memory to buy quality-per-FLOP was exactly the right exchange — and it reframed how people thought about scaling.

Switch Transformer took mixture-of-experts from a finicky research idea to a scalable, practical technique, and the modern MoE models that followed — from open trillion-scale releases to the sparse layers rumored inside frontier systems — inherit its recipe: simple top-1-ish routing, an auxiliary balancing loss, and capacity limits. Its specific choices are refined in newer work, but the skeleton is the same.

Its lasting idea is the decoupling itself. For years, "bigger model" meant "more compute," full stop. Sparse expert routing broke that identity, letting parameter count and per-token compute scale on separate dials. That is why mixture-of-experts is now a central tool for building very large models economically — and Switch Transformer is the paper that made the tool dependable.