TL;DR
- The Switch Transformer replaces the dense feed-forward layer with a sparse mixture of experts, routing each token to exactly one expert (top-1 gating).
- This decouples parameters from compute: adding equal-sized experts grows stored parameters while each accepted token uses one expert FFN. Router and system overhead still grow.
- A capacity factor bounds how many tokens each expert accepts; overflow tokens skip the expert branch and retain their residual, and an auxiliary loss keeps the load balanced.
- The design scaled to trillion-parameter models with the compute budget of a much smaller dense model.
Sparse routing: one expert per token
A dense transformer runs every token through the same feed-forward network. Switch instead keeps N expert FFNs and a small router that picks the single best expert for each token — a top-1 softmax gate.
Parameters without the compute
Because only one expert fires per accepted token, the selected expert FFN computation stays roughly fixed as equal-sized experts are added. Total stored parameters grow. Router scoring grows with the expert count, and padding, memory and communication affect actual cost; the chart below illustrates active parameter counts rather than measured FLOPs.
Capacity and dropped tokens
Sparse routing has a catch: real token distributions are uneven, and experts have fixed-size buffers. The capacity factor sets each expert’s buffer size; tokens routed past it contribute zero from the expert branch while their unchanged representations continue through the residual. Raising the factor reduces drops at the cost of memory and compute.
Implementation
The capacity chart shows which tokens fit. To see what reaches the next layer, keep the selected probability, dispatch accepted tokens into stable slots, weight each expert output and add the residual. The auxiliary loss sees top-1 choices before capacity removes any expert contributions.
Authored offline serial CPU mechanics: eight synthetic tokens, width 4, three experts and FFN width 8. These random reduced weights expose routing, capacity, weighted execution and gradients. This is not a pretrained model, AllToAll implementation, BF16 experiment, trillion-parameter model or throughput benchmark. The equations use authored notation for the cited operations.
The primary reference is the original JMLR §§2.1–2.2 and Appendix F, Figures 14–16. Matching original Mesh TensorFlow operations are pinned to 37553191b0c9. At this revision the single-choice path is named _rand_1_gating; this example uses its argmax convention and exclusive cumsum. It does not use the adjacent no-token-left-behind or top-2 paths, later z-loss recipes or the page’s general T5X link as its source. Capacity uses the pinned integer convention with an explicitly reduced minimum of 1.
1. Compute local router probabilities in float32
Original paper
A small bias-free projection scores every expert for each token.
Math and tensor shapes
Routing is local FP32 matrix multiplication and softmax. This deterministic reference omits training exploration noise. The rest of this CPU example also uses float32; it does not test the paper’s selective BF16 stability or communication savings.
X [T,D], Wᵣ [D,E] → h,p [T,E]
Python implementation
# FP32 local routing; no noise in this deterministic reference.
logits = inputs.float() @ self.router.float() # [T,D] @ [D,E]
if not torch.isfinite(logits).all():
raise ValueError("router logits overflowed float32")
probabilities = logits.softmax(dim=-1) # [T,E]2. Retain both the expert choice and its gate
Original paper
Top-1 selects one index without normalizing the selected probability to one.
Math and tensor shapes
The gate remains differentiable through the selected softmax probability. Argmax is hard; ties choose the first expert in this authored convention. There is no second expert or overflow rerouting.
p [T,E] → index/gate [T], hard choice [T,E]
Python implementation
expert_index = probabilities.argmax(dim=-1) # [T], first on ties
gate = probabilities.gather(1, expert_index[:, None]).squeeze(1)
chosen = F.one_hot(expert_index, num_classes=experts) # [T,E]3. Allocate a buffer and count stable positions
Original paper
Integer capacity must be declared before building a fixed-size dispatch tensor.
Math and tensor shapes
k is the capacity factor. This follows the pinned MTF floor, T clamp, then minimum, with reduced m=1 instead of the upstream default 4. One unpadded token group replaces distributed groups. Exclusive counts begin at zero, so slots 0 through C−1 fit; token order decides who fits.
scalar C; exclusive token positions s [T]
Python implementation
# Original MTF floor, T clamp, then minimum; reduced minimum is 1.
capacity = max(
minimum_capacity,
min(tokens, int(tokens * capacity_factor / experts)),
)
exclusive_counts = chosen.cumsum(dim=0) - chosen
positions = (exclusive_counts * chosen).sum(dim=-1) # [T], zero-based4. Keep fitting tokens and record load before filtering
Original paper
The balancing loss counts original top-1 choices, including overflow.
Math and tensor shapes
Capacity removes an expert contribution, not a token from the sequence. Computing f after filtering would hide overload from the loss. f comes from integer decisions and has no gradient. The pinned overflow predicate uses exclusive positions at lines 1000–1007.
accepted/overflow [T]; hard fractions f [E]
Python implementation
accepted = positions < capacity # [T]; original token order wins slots
overflow = ~accepted
fractions = chosen.float().mean(dim=0) # [E], BEFORE capacity filtering5. Pack tokens into explicit expert slots
Original paper
A binary dispatch tensor separates routing decisions from gate weights.
Math and tensor shapes
Each accepted token has exactly one true slot; each slot holds at most one token. Unused slots contain zeros. Serial CPU einsum makes the same packing algebra visible without any AllToAll or device sharding.
dispatch [T,E,C]; packed inputs U [E,C,D]
Python implementation
dispatch = torch.zeros(tokens, experts, capacity, dtype=torch.bool)
token_ids = torch.arange(tokens)[accepted]
dispatch[token_ids, expert_index[accepted], positions[accepted]] = True
expert_inputs = torch.einsum("tec,td->ecd", dispatch.float(), inputs)6. Apply separate ReLU feed-forward networks
Original paper
Each expert has its own two weight matrices.
Math and tensor shapes
Reduced bias-free ReLU FFNs follow the original MTF operations. Random local initialization and no expert dropout are disclosed reductions. Computing padded slots is visible here; these tensors do not measure accelerator efficiency.
W₁ [E,D,H], W₂ [E,H,D]; Z [E,C,D]
Python implementation
hidden_states = torch.einsum(
"ecd,edh->ech", expert_inputs, self.w1
).relu()
expert_outputs = torch.einsum("ech,ehd->ecd", hidden_states, self.w2)
if not torch.isfinite(expert_outputs).all():
raise ValueError("expert outputs overflowed float32")
# Each expert owns its [D,H] and [H,D] matrices. No biases/dropout.7. Weight expert outputs and preserve the residual
Original paper
The combine tensor carries the selected probability back to the original token position.
Math and tensor shapes
The residual is shown explicitly around the Switch FFN branch. For overflow D is zero, b is zero and Y equals the unchanged X. Accepted outputs retain original token order and the selected probability multiplier. Normalization, attention and full Transformer blocks are outside this reference.
combine [T,E,C] → branch/output [T,D]
Python implementation
combine = dispatch.float() * gate[:, None, None] # [T,E,C]
branch = torch.einsum("tec,ecd->td", combine, expert_outputs)
output = inputs + branch # overflow branch=0, residual token unchanged
if not torch.isfinite(output).all():
raise ValueError("residual output overflowed float32")8. Differentiate the auxiliary balancing objective
Original paper
Hard token fractions multiply mean router probabilities, with expert-count scaling.
Math and tensor shapes
Alpha is 0.01, as in the paper. Uniform f and P give loss alpha; a collapsed hard load with selected mean probability q gives alpha·E·q. Gradients flow through P while f stays fixed, and the expert branch also trains the router through g. Smoke adds an illustrative squared-output objective and one SGD step, not the original language-model training recipe.
f,P [E] → scalar auxiliary loss → router gradients [D,E]
Python implementation
mean_probabilities = probabilities.mean(dim=0) # P [E], differentiable
auxiliary = alpha * experts * (fractions * mean_probabilities).sum()
# Hard argmax/f has no gradient; probabilities and selected gates do.Actual smoke JSON records source SHA-256, seed, reduced shapes, accepted and overflow counts, load fractions, auxiliary loss, token-loop error, unchanged overflow residuals, router/expert gradients and a router update. It reports no FLOP or timing measurement.
switch.pyrequirements-switch.txt
python -m pip install -r requirements-switch.txt python switch.py --smoke --seed 7
Why it mattered
Switch Transformers showed that sparsity is a practical scaling axis, not just a theoretical one. By making routing simple (top-1) and stable (capacity factor + load-balancing loss), it turned mixture-of-experts into a production technique. Modern large models — from open MoE LLMs to frontier systems — inherit its core recipe: scale parameters with experts, keep compute per token bounded.
Related Reading
- Attention Is All You Need — the dense transformer whose feed-forward layer Switch replaces with sparse experts
- LoRA — a complementary efficiency axis: cheap adaptation of a fixed model rather than sparse scaling of a huge one
- FlashAttention — efficiency on the attention side, orthogonal to MoE’s efficiency on the feed-forward side
- BERT — the dense encoder baseline that sparse-expert models scale beyond