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

Route each token to one expert, and a network can carry a trillion parameters at almost no extra compute per token.

A dense network runs every weight on every token, so more parameters always means more arithmetic. Switch Transformers break that link: each token is sent to a single expert out of hundreds, so the parameter count and the work done per token stop scaling together.

A language model with 1.6 trillion parameters that spends less compute on each token than a dense model over a hundred times smaller.

The dominant story of the last decade of language models is a simple one: make the model bigger and it gets better, in a way that has held across several orders of magnitude (this is the subject of the scaling laws work of Kaplan and collaborators). But in an ordinary Transformer the bill for "bigger" is paid twice. Every parameter you add runs on every token you ever process, so doubling the weights doubles the arithmetic per token, and that arithmetic is what you pay for on every forward pass, forever. Parameters and compute rise together.

Switch Transformers, out of Google, pull the two apart. The idea traces back to a 1991 notion called a mixture of experts: instead of one big network that processes everything, keep many smaller "expert" networks and, for each input, use only the one or two that are relevant. Most of the model sits idle on any given token. The parameter count can balloon while the compute spent on a token stays fixed, because a token only ever touches its own handful of experts. Mixtures of experts had worked before, notably in translation, but they carried a reputation for being fiddly: complex to implement, expensive to communicate across machines, and prone to unstable training. Switch Transformers are the paper that made the idea simple enough to scale, all the way to a model with 1.6 trillion parameters.

The method comes down to a few moving parts: a router that maps each token to an expert, the decision to route to exactly one and why that still trains, a balancing pressure that keeps the router from favoring a few experts, a fixed-capacity scheme that fits dynamic routing onto rigid hardware, and a small bag of tricks that keep it stable. The sections take them in that order.

Every weight fires on every token

Start with what a dense Transformer actually spends its compute on. Each layer has two sublayers: attention, which mixes information across tokens, and a feed-forward network, a plain two-matrix multiply applied to each token on its own. In the models here the feed-forward network holds most of the parameters and does most of the arithmetic. And it is dense in the strict sense: the same weight matrix multiplies every single token that flows through. A token about weather and a token about tax law hit the identical parameters.

A sparse layer breaks that coupling. If every parameter processes every token, the compute per token is tied to the parameter count, so the only way to grow one is to grow the other. Sparsity proposes a different deal: keep many feed-forward networks, call them experts, and route each token to just one of them. Now the parameters can grow with the number of experts while the compute a token pays stays fixed at the cost of one expert. This is conditional computation: which weights run is a function of the input, not a constant.

It helps to hold a concrete picture. Think of a hospital with thousands of specialists. A patient is triaged to the one specialist who fits, so the time any patient spends is the time with one doctor, no matter how large the staff grows. Hiring more specialists deepens the hospital's total expertise without lengthening any single visit. The analogy also exposes the cost: the hospital still has to employ, house, and pay every specialist even though each patient sees one. That cost has a direct echo in the model, and the figure below tracks both sides of it.

Before you drag it, watch one thing: as the expert count climbs, the parameter bar and the memory bar climb with it, while the compute-per-token bar does not move.

The upside comes with fine print. The largest model in the paper, Switch-C, holds 1.6 trillion parameters yet applies about 890 billion floating-point operations per training sequence, fewer than the 11-billion-parameter dense T5-XXL, which applies 6.3 trillion. A model 140 times larger, doing roughly a seventh of the arithmetic per sequence, and so per token. But "constant computational cost" means arithmetic per token and nothing else. The memory to store all those experts grows in lockstep with their number, and, because different tokens want different experts and those experts live on different machines, the tokens have to be physically shipped across the network to reach them. Sparsity buys capacity with memory and bandwidth, not with free compute. (This is the same memory wall that motivates ZeRO, from the same trillion-parameter era.)

Mixture of experts: pick a specialist

Routing a token starts with the router, a learned component about as simple as they get: a single weight matrix that scores every expert. Write $x$ for a token's representation (a vector of length $d_{\text{model}}$) and $W_r$ for the router's weights. It produces one score per expert,

and a softmax turns those $N$ scores into a probability distribution over the $N$ experts. The number $p_i(x)$, the gate value for expert $i$, is the probability the router assigns to expert $i$ for this token. The classic mixture-of-experts layer, from Shazeer and collaborators in 2017, then keeps the top $k$ experts and combines their outputs, each weighted by its gate value:

Here $E_i(x)$ is expert $i$'s feed-forward network applied to the token, and $\mathcal{T}$ is the small set of chosen experts. The gate value out front is doing two jobs at once, and the two are separate: it determines which experts pass the top-$k$ selection, and it scales how much each winner contributes to the output. That second job, the gate multiplying the expert's output, matters enormously in a moment.

The prior work took $k \geq 2$ as a given. Shazeer's 2017 paper had argued you needed at least two experts per token: only by comparing two options, the reasoning went, could the router get a useful gradient and learn to route at all. Later work pushed $k$ higher still in the lower layers. Routing to a single expert was assumed to be a dead end, a router that could never learn. Switch Transformers exist to show that assumption is wrong.

The Switch layer: route to one

The Switch layer takes the mixture-of-experts idea and sets $k = 1$. Each token goes to exactly one expert, the argmax of the router distribution, and no others run. This is the maximally sparse choice, and it buys three concrete things: the router does less work, each expert needs a smaller share of the batch (since no token is sent to two places), and the cross-machine communication that dominates a distributed MoE roughly halves versus the usual two-expert routing, since each token now travels to one expert instead of two. All of that is upside. The open question is whether the router can still learn when it only ever picks one.

It can, and walking through exactly why is worth the care, because the method rests on it. Look again at the output for a single routed token, now with $k = 1$:

The selection $i^\ast$ is a hard argmax, a step function with a flat derivative everywhere. Nothing useful comes back through it, and nothing needs to. The gradient the router needs rides on the other factor, the gate value $p_{i^\ast}(x)$ sitting in front. Because that scalar multiplies the expert's output, the loss depends on it smoothly: nudging the router weights changes $p_{i^\ast}(x)$, which changes the output, which changes the loss. Concretely, when the chosen expert produced a good output the gradient pushes the router to raise that expert's gate value on tokens like this one, and when the output was bad it pushes the value down. The router learns which tokens belong to which expert without ever needing a second expert to compare against. The decision of which expert is not differentiable; the magnitude of how much is, and that is all the signal the router requires.

So why did the earlier work believe $k = 1$ could not work? The difference is a subtle ordering, and it is the one place a reader can get the mechanism exactly backwards. Shazeer's gating masked all but the top $k$ logits to negative infinity before the softmax. With $k = 1$, the lone survivor of that masking is the only thing left to normalize, so its softmax value is a constant 1, and a constant has no gradient. Switch takes the softmax over all $N$ experts first and only then picks the argmax, so the gate it keeps is a genuine fraction below 1 with a live derivative. Same arithmetic, opposite order, and that order separates a router that learns from one that cannot.

Drag the token below along its representation axis. The router's softmax over six experts updates live; the single highest-probability expert lights up and runs while the other five stay idle. Watch two things: the assignment jumps from one expert to the next as you cross a border, and the gate value on the winner is never 1, which is exactly what leaves the router a gradient to learn from.

Make it concrete with real shapes. In Switch-Base a token is a 768-dimensional vector and there are 128 experts, so the router $W_r$ is a $768 \times 128$ matrix and produces 128 logits per token. Suppose the softmax picks expert 57 with gate value 0.31. Then expert 57's feed-forward network runs on this token, its output is scaled by 0.31 and added back to the token through the residual connection, and the other 127 experts never touch it. One token, one expert, 127 sets of weights sitting out the play. The complete layer, dropped tokens and all, looks like this:

# one Switch feed-forward layer, top-1 routing
# x: token representations, shape [tokens, d_model]
logits = x @ W_r                  # [tokens, N experts]   W_r: d_model x N
logits = logits.astype(f32)       # selective precision: softmax in float32
p      = softmax(logits, axis=-1) # gate values p_i(x), eq (1)
expert = argmax(p, axis=-1)       # top-1: WHICH expert (not differentiable)
gate   = p.max(axis=-1)           # the winning probability p_i(x)

cap = int(tokens / N * capacity_factor)      # fixed seats per expert, eq (3)
y = x.copy()                      # default is identity (a dropped token)
for i in range(N):                # each expert has its own FFN weights
    seat = where(expert == i)[:cap]           # first cap tokens routed to i
    y[seat] = x[seat] + gate[seat] * E[i](x[seat])   # eq (2), k=1, residual

Two details in that listing point to later sections. The line that casts the logits to float32 stabilizes training, a selective-precision fix covered below. The $\text{cap}$ variable, a fixed number of seats per expert, is how dynamic routing fits on hardware that demands static shapes. Both get their own section. First, the pressure that keeps routing balanced.

Keeping the experts busy

Hard top-1 routing has a failure mode you can guess. Early in training the router is close to random, but any small preference feeds on itself: an expert that gets slightly more tokens gets slightly more gradient, gets slightly better, attracts more tokens, and the routing collapses onto a favored few while the rest starve. Collapse defeats the model's purpose, since idle experts are parameters you paid for and never use. Something has to push the router toward spreading tokens evenly.

An auxiliary loss, added alongside the usual language-modeling objective, supplies that push. For each Switch layer, over a batch of $T$ tokens and $N$ experts, it is a single dot product:

It takes two vectors, one a hard count and one a soft probability. The first is the actual traffic:

the fraction of tokens whose argmax landed on expert $i$, a plain head-count. The second is the router's intent:

the average probability mass the router assigned to expert $i$ across the batch. The balanced target for both is $1/N$, an even share. The loss uses both a hard count and a soft mass, rather than one of them twice, because only one of them can be trained through: the head-count $f_i$, like the $i^\ast$ selection earlier, comes from an argmax and carries no gradient, so it enters the loss as a fixed constant. All of the loss's gradient flows through the soft mass $P_i$. Working it out, the gradient on each expert's probability is $\alpha N f_i$, a pressure proportional to how loaded that expert already is. An overloaded expert gets its router probability pushed down hardest, so the next batch diverts toward the quiet ones. It behaves like a congestion toll: the busier a route, the more it costs to send more traffic there, and the toll is priced off current load while only ever charging the soft probability.

The constant $\alpha$ sets how hard the toll bites; the paper uses $\alpha = 10^{-2}$ after sweeping from $10^{-1}$ down to $10^{-5}$, large enough to balance quickly and small enough not to drown the language objective. The factor of $N$ out front looks cosmetic but is not. When routing is perfectly uniform, $f_i = P_i = 1/N$, so the sum is $N \cdot (1/N)^2 = 1/N$, and the whole loss becomes $\alpha N \cdot 1/N = \alpha$. The $N$ cancels the $1/N$, which means the loss bottoms out at the same value $\alpha$ no matter how many experts there are. Its strength stays fixed as you scale from 8 experts to 2048. Notice the floor is $\alpha$, not zero. If the router tried to cheat by jamming most tokens onto one expert, both that expert's head-count $f_j$ and its soft mass $P_j$ would spike together, since they measure the same argmax, so the product $f_j P_j$ rises rather than vanishes. Because $f$ and $P$ are two views of the same routing decision, they cannot be pushed apart, and uniform routing is genuinely the best they can jointly do.

Drag the routing below from collapsed toward uniform and watch the dot product fall to its floor. The amber bars are the hard counts $f$, the teal bars the soft mass $P$; the loss meter reads its value against the minimum $\alpha$ and the maximum $\alpha N$.

In code the loss is four lines, and the one comment that matters is the last one:

# auxiliary load-balancing loss for one Switch layer, eq (4)
f   = one_hot(expert, N).mean(axis=0)   # dispatch fraction f_i, eq (5)  HARD
P   = p.mean(axis=0)                     # router prob mass P_i, eq (6)   SOFT
aux = alpha * N * (f * P).sum()          # alpha = 1e-2; floor at uniform = alpha
# f is a head-count with no gradient. All of aux's gradient flows through P.

With this pressure in place the router keeps the experts within a hair of an even load. The paper reports the fraction of tokens that overflow their expert is typically under one percent. That number only makes sense once you know what "overflow" is, which is the next piece.

Fixed seats and dropped tokens

Routing is dynamic: which expert a token wants depends on the token, and it is different every batch. But the accelerators these models train on, TPUs, demand tensors of a size fixed at compile time. You cannot allocate "however many tokens happen to want expert 57 this step." So every expert gets a fixed number of seats, the expert capacity, computed once:

The first factor is the even share each expert would get under perfectly balanced routing. The capacity factor is a slack multiplier on top. At exactly 1.0 each expert has room for its uniform share and not one token more, which sounds sufficient but is not, because real routing is never perfectly even: some expert always draws a few extra. Push the factor above 1.0 and you add buffer seats to absorb the imbalance.

Slack costs you in both directions. Seats an expert does not fill are padding: allocated memory and wasted arithmetic on nothing. Tokens that arrive after an expert's seats are full overflow, which sounds worse than it is. An overflowed token is not an error and is not lost. Its expert is skipped: the combine weight is set to zero, so the expert contributes nothing, and the token's own representation passes through unchanged to the next layer along the residual connection that already wraps every sublayer. It misses one feed-forward update and nothing more. This is graceful degradation, borrowed from the earlier GShard work, not a crash.

Drag the two sliders below. Capacity factor sets how many seats each expert has; routing skew piles tokens onto one expert. At a low factor with skewed routing, the overflow appears as dropped tokens above the full column; crank the factor up and the drops vanish but empty padding seats appear.

The two dials interact in a way the paper leans on. A higher capacity factor makes drops rarer but costs more memory and compute per step, which at trillion scale is exactly the resource that is scarce. The load-balancing loss from the previous section resolves this: keep routing near-uniform and you barely overflow even at a tight capacity factor. That is why the paper finds Switch works best at low factors of 1.0 and 1.25, the opposite of what you would reach for if you were trying to be safe. Good balancing lets you run lean.

Making hard routing trainable

A hard argmax at every layer is a jagged thing to optimize, and left alone it makes large sparse models unstable. Three fixes, none of them heavy, turn it into something that trains reliably from 200 million parameters to a trillion. Each one is a small idea with a clean reason, and two of them refute a tradeoff you would expect to pay.

Selective precision. Large models train in bfloat16 to save memory and bandwidth, but the router's softmax specifically misbehaves in bfloat16. It is a precision problem, not a range one: bfloat16 uses the same exponent as float32 and covers the same magnitudes, so nothing overflows, and the softmax subtracts its max before exponentiating anyway. What bfloat16 gives up is mantissa bits, keeping about 7 where float32 keeps 23, sixteen fewer, so its rounding steps are coarser by a factor of $2^{16}$, tens of thousands. That coarse rounding, run through the exponentials of a softmax and then a hard argmax that can flip on a rounding error, is enough to destabilize routing. So the model casts only the router's local computation to float32 and leaves everything else in bfloat16. Crucially, the float32 values never leave the router: the dispatch and combine tensors are cast back to bfloat16 before any token is shipped across the network, so no expensive high-precision tensor crosses machines. This is not a speed-for-stability trade. Selective precision reaches float32's quality at essentially bfloat16's speed, while training everything in bfloat16 diverges.

Smaller initialization. The paper draws weights from a truncated normal with standard deviation $\sigma = \sqrt{s/n}$, where $n$ is the number of inputs to the weight and $s$ is a scale it sets to 0.1, a tenth of the usual default of 1.0. Starting the network smaller, closer to the identity, keeps the early gradients tame while the router is still flailing. Like selective precision, this is not a compromise: the smaller scale both improves the average model quality and sharply cuts the run-to-run variance across random seeds, so training is both better and more repeatable.

Expert dropout. Fine-tuning is where a sparse model's huge parameter count turns into a liability, because the downstream datasets are small and the model overfits. The wrinkle is that each expert only ever sees the tokens routed to it, a fraction of the data, so the experts overfit faster than the shared trunk. The remedy matches the diagnosis: use ordinary dropout at a low rate of 0.1 in the shared layers, but a much higher rate of 0.4 inside the expert feed-forward networks. Raising dropout uniformly everywhere hurts; the asymmetric setting, heavy only where overfitting is worst, is what helps.

A fourth axis to scale

The scaling-laws work established three levers on model quality: parameters, data, and compute. Switch adds a fourth that the others do not have, the number of experts, and it is nearly free in compute. Add experts and you add parameters at a fixed cost per token, and the paper finds the reliable result: at a matched compute budget, more experts means faster learning. A 64-expert Switch-Base reaches in roughly 60,000 steps a quality that dense T5-Base needs about 450,000 to hit, about 7.5 times fewer steps for the same computation per token.

One qualification applies to that number. Fewer steps is not the same as less wall-clock time, because a Switch step does extra work a dense step does not: it computes routing and it ships tokens across the network to their experts. Those costs shave the advantage from 7.5 times in steps down to about 7 times in wall-clock, still a large speedup, though not a free one. And compared against the stronger dense baseline T5-Large, which spends 3.5 times the compute per token, Switch-Base is still 2.5 times faster to a given quality despite its far smaller per-token cost.

Drag the expert count below. The dense T5-Base curve is fixed; the Switch curve pulls left as experts are added, reaching the same quality in fewer and fewer steps, and the bracket reports the speedup, which grows toward 7.5 times at 64 experts.

The gains carry past pretraining. Fine-tuned on a spread of language tasks, the FLOP-matched Switch models beat their dense counterparts on SuperGLUE by 4.4 points over T5-Base and 2 points over T5-Large, with large jumps on knowledge-heavy tasks like closed-book TriviaQA. The exceptions belong in the result too: on the reasoning-flavored ARC benchmarks the dense models sometimes hold their ground or win. Across 101 languages, a multilingual Switch improves over mT5 on every single one, and reaches mT5's quality about 5 times faster on average.

Reaching a trillion parameters

Scaling to a trillion parameters means slicing the model across thousands of accelerator cores, and there are three ways to slice, which compose. Data parallelism puts a full copy of the model on each core and feeds each a different slice of the batch, syncing gradients once per step. Model parallelism, the approach of Megatron-LM, splits a single layer's matrices across cores that must pool their partial results within every layer. Expert parallelism is the one sparsity adds: place different experts on different cores, and route each token to the core holding its expert via an all-to-all shuffle. Experts are a natural unit to spread out, since they are independent by construction.

The paper builds two trillion-era models that spend a fixed budget in opposite ways, and comparing them corrects a tempting oversimplification. Switch-C is 1.6 trillion parameters and 2048 experts, built with expert parallelism alone, so its individual layers stay narrow and cheap. Switch-XXL is smaller at 395 billion parameters but is FLOP-matched to the dense T5-XXL, using all three kinds of parallelism to afford much wider layers and about seven times the compute per token that Switch-C spends. Both beat dense T5-XXL on pretraining, with Switch-C reaching a fixed quality about four times faster.

Raw parameter count, that comparison shows, is not the figure of merit it looks like. Per training step the smaller Switch-XXL is the stronger model and beats the four-times-larger Switch-C on downstream tasks, because what it lacks in parameters it more than makes up in compute spent per token. And instability does not track size the way you might assume: the 1.6-trillion Switch-C trains with no instability at all, while the smaller Switch-XXL is sometimes unstable, because instability follows compute-per-token, layer width, and precision, not the parameter count. The trillion is a real milestone, but the 395-billion model is arguably the better one, and it is worth keeping the two claims separate.

Shipping it: distill back to dense

A trillion-parameter model is a research artifact, not something you serve to users: all those experts have to be stored and shuffled, which is exactly the memory and bandwidth cost sparsity trades for. So the paper closes the loop by distilling the sparse teacher back into a small dense student, in the classic sense of training a small model to match a large one's outputs.

Two moves make it work, and the larger of the two exploits a coincidence of the setup. Because the sparse and dense models are FLOP-matched, their non-expert layers, all the attention and the shared machinery, have identical shapes. So you can initialize the dense student by copying the teacher's non-expert weights directly, so it starts with the shared skeleton already in place and has only the collapsed expert knowledge left to learn. On top of that, the training target blends the teacher's soft probabilities with the ground-truth labels, weighted 0.75 toward the hard labels and 0.25 toward the teacher, the reverse of the soft-heavy mix Hinton's original distillation favored.

The distilled dense model preserves about 30 percent of the quality gain of the sparse teacher while shedding 95 percent or more of the parameters. That 30 percent is the fraction of the gap between the dense baseline and the sparse teacher that the student recovers, not "30 percent as good as the teacher." A student that fully matched the teacher would score 100 percent by this measure; matching the original dense baseline scores 0. Compress the 1.1-billion-parameter teacher by 82 percent and you keep 37 percent of its gain; even the far larger 14.7-billion-parameter teacher, compressed by 99 percent, still preserves 28. You train the enormous sparse model to learn efficiently, then pour a third of what it learned into a model small enough to deploy.

A mixture of experts had been a decade-old idea with a reputation for being more trouble than it was worth. Switch Transformers pared it down to its simplest form, one expert per token, showed that the router still learns because the gate value it keeps carries a gradient, and wrapped it in the balancing, capacity, and precision machinery that let it train stably at unprecedented size. What comes out the other side is a parameter count decoupled from compute, and the first demonstration that the decoupling holds all the way to a trillion.