Switch Transformer (Fedus, Zoph & Shazeer, 2021) took the mixture-of-experts layer and made one aggressive simplification: send each token to exactly one expert. Every earlier sparse MoE routed to at least two, on the belief that a single choice would starve the router of gradient. Switch showed that belief was wrong, and cashed the simplification in for a cheaper router, half the dispatch traffic, and a model scaled to 1.6 trillion parameters. But k = 1 also sharpens every failure mode of sparse routing: a token that misses its expert’s buffer gets nothing at all. This piece works through the top-1 layer end to end — the gradient argument, the capacity formula, the balancing loss, the communication, the instabilities, and the accounting that makes a sparse model cheap per token and expensive in RAM.
The Switch layer: one expert per token
A Switch layer replaces one feed-forward block with N copies of it plus a router. The router is the smallest thing in the layer: a single weight matrix W_r: [d, N] followed by a softmax.
h(x) = W_r^T x logits, shape [N]
p_i(x) = exp(h_i) / Σ_j exp(h_j)
i* = argmax_i p_i(x)
y = p_i*(x) · E_i*(x)Two details carry the whole design. First, only E_i* is evaluated — the other N-1 experts do no work for this token, which is where the FLOP saving lives. Second, the expert output is scaled by its own gate probability p_i*(x), a number in (0, 1]. That scalar looks decorative next to a full FFN, but it is the only path by which the router receives any learning signal at all, and it is what makes top-1 viable.
Does k = 1 break the router gradient?
The 2017 sparsely-gated MoE argued that k ≥ 2 was necessary: with one expert, the reasoning went, there is nothing to compare against and so no meaningful gradient to the gate. Switch’s counter is a one-line derivative. Because the output is y = p_i*(x) · E_i*(x), the loss reaches the router through the multiplicative scalar:
∂L/∂W_r = (∂L/∂y · E_i*(x)) · ∂p_i*/∂W_rAnd p_i* is a softmax output, so it depends on every logit through the shared denominator — ∂p_i*/∂h_j ≠ 0 for all j. One selected expert still updates all N router rows: if the chosen expert helped, its logit rises and the rest fall. Fedus et al. list three payoffs for the simplification — less router computation, at least halved per-expert buffers, and reduced communication.
Expert capacity: a fixed buffer per expert
Routing is data-dependent, but the hardware is not. Experts live on different devices, and an all-to-all needs tensors of a shape known before the routing decisions exist. So each expert gets a fixed buffer, sized by a capacity factor:
C = CF · (tokens_per_batch / N) [Switch, k = 1]
CF = 1.0 → exactly the perfectly balanced share
CF > 1.0 → slack for skew, at the cost of idle slotsThis is the generic MoE capacity CF · T·k/N with k collapsed to 1 — the reason Switch can afford buffers half the size of a top-2 model at the same slack. The capacity factor is a pure memory-for-robustness dial: N · C slots are allocated and padded whether or not tokens fill them, so CF = 2.0 means half the expert compute is spent on padding even when routing is perfect.
Token dropping and what it costs
When more than C tokens choose the same expert, the overflow is dropped: those tokens are simply not processed by any FFN. They are not re-routed and not queued. Their representation passes to the next layer through the residual connection unchanged, as though the layer had been skipped for them.
At k = 2 this is a partial loss — a dropped token usually still has its other expert, and the combine just reweights. At k = 1 there is no second expert, so a dropped token receives zero feed-forward contribution from that layer. This is the sharpest edge of the top-1 design and the reason capacity and balancing matter more here than in any other MoE variant. A small drop rate is survivable and even acts a little like structured dropout; a large one silently removes a whole layer’s worth of computation from a biased subset of the vocabulary.
Worked example: capacity, skew, drops
Take one Switch layer with N = 64 experts serving a batch of T = 8192 tokens across the expert-parallel group (global, not per-device). At k = 1 there are 8192 dispatches.
perfect share μ = 8192 / 64 = 128
CF = 1.25 → C = 1.25 · 128 = 160
skewed routing: 4 experts ← 300 tokens each
8 experts ← 200 tokens each
52 experts ← ~104 tokens each
dropped = 4·(300-160) + 8·(200-160) = 560 + 320 = 880
drop_rate = 880 / 8192 = 10.7%Nearly eleven percent of tokens skip the FFN entirely. Raise CF to 2.0 and C = 256: only the four hottest experts overflow, dropped = 4·44 = 176, a 2.1% drop rate — but you now allocate 64 × 256 = 16384 slots for 8192 tokens, wasting half the expert compute on padding. That trade is the whole tuning problem.
The load-balancing auxiliary loss
Capacity handles skew; the auxiliary loss tries to prevent it. Switch adds a product of two per-expert fractions over each batch B:
f_i = fraction of tokens dispatched to expert i (hard counts)
P_i = mean_{x in B} p_i(x) (soft mass)
L_aux = α · N · Σ_i f_i · P_i, α = 1e-2
uniform: N · Σ (1/N)(1/N) = 1 → the minimumThe product form is what makes it usable. f_i comes from an argmax and a count — piecewise constant, gradient identically zero. P_i is a mean of softmax outputs and is perfectly smooth. Treating f_i as a constant weight leaves ∇L_aux = αN Σ_i f_i · ∇P_i: the observed overload becomes the coefficient, and the gradient pushes probability mass off exactly the experts that were overloaded. The scale α must be large enough to balance and small enough not to fight the task loss.
All-to-all: what sparsity costs on the wire
Experts are sharded across devices, so a Switch layer is not a local operation. Each layer performs two all-to-all collectives: a dispatch that ships every token’s hidden vector to whichever device owns its expert, and a combine that ships the outputs back to the token’s original position. Both move roughly T · k · d activation elements.
Two consequences follow. First, top-1 halves that volume relative to top-2 — the concrete form of Fedus et al.’s communication benefit, and a large part of why Switch is faster per step than a top-2 model of the same quality. Second, all-to-all is a synchronizing collective: it finishes when the slowest shard finishes, so imbalance costs wall-clock even when nothing is dropped. On a single machine there is no wire, and the entire cost reverts to a gather over memory.
Training instability and the three fixes
Sparse models at scale diverge in ways dense ones do not, and the router is the culprit: it contains an exponential, and it is the one place where a tiny numerical wobble changes a discrete decision. Switch reports three targeted fixes.
Selective precision. Cast the router input to fp32, compute the logits and softmax in fp32, then recast to bfloat16 before anything leaves the router. Because the fp32 tensor is never communicated, this costs almost nothing while removing the round-off that made routing decisions flip between steps. Smaller initialization. With weights drawn from a truncated normal of std = sqrt(s / n_in), reduce the scale s from 1.0 to 0.1 — smaller initial logits mean a flatter, more forgiving early router. Expert dropout. When fine-tuning on small datasets, apply much higher dropout inside the expert layers than elsewhere, since the sparse parameters overfit fastest.
Cheap per token, expensive in memory
The accounting is the point of the whole architecture. Because exactly one expert fires, the FLOPs of a Switch FFN equal the FLOPs of a single dense FFN — roughly 2 · d · d_ff per token per matmul pair — independent of N. Adding experts adds parameters at constant compute per token, which is exactly how Switch-C reaches 1.6T parameters with 2048 experts while costing what its dense counterpart costs per step.
Memory does not cooperate. Every expert must be resident, so weight storage grows linearly in N, and the capacity buffers add N · C · d activations on top. For a CPU SLM this is the decisive asymmetry: sparsity buys you nothing on a machine that is bound by RAM and bandwidth rather than arithmetic, and the honest use of a Switch model there is to distill it back into a dense one, accepting that only part of the gain survives.
N logits — the 2017 claim that k ≥ 2 was required is simply false. Top-1 halves buffers and halves all-to-all volume, but it removes the safety net: a token that overflows C = CF · (T / N) gets no FFN at all, so capacity factor and the α·N·Σ f_i P_i balancing loss carry more weight here than anywhere else — the product form works precisely because the non-differentiable count multiplies the differentiable probability mass. Stabilize with an fp32 router, a 10× smaller init scale, and heavy expert dropout when fine-tuning. And remember the shape of the bargain: constant FLOPs per token, linear memory in the expert count — a great trade on an accelerator cluster, a bad one on a CPU.