Mixture of Experts (MoE) from Scratch — Routing and Load Balancing in the Switch Transformer

A ground-up explanation of Mixture of Experts, the sparse architecture behind today's largest LLMs, built strictly from the Switch Transformer paper: the router math, expert capacity, the load-balancing loss, and the three tricks that make sparse training stable.

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

More Parameters, Same Compute

An ordinary neural network uses all of its parameters for every input. To make it smarter you add parameters, and every parameter you add makes each forward pass proportionally more expensive. That is the bargain dense models are stuck with.

Mixture of Experts (MoE) breaks the bargain: it selects different parameters for each incoming example. The result is a sparsely-activated model — an outrageous number of parameters, but a constant computational cost per token.

The 2021 paper Switch Transformer (Fedus, Zoph, Shazeer) simplified MoE to its bare minimum, trained a model with 1.6 trillion parameters stably, and reported up to 7x faster pre-training than T5 with the same computational resources (§Abstract). This article builds MoE from the ground up using that paper as the sole primary source.

The Metaphor: A Hospital Front Desk

• Patients = tokens (the words in your sentence)
• Specialists = experts (small networks, each with its own separate weights)
• The front desk = the router (decides which patient sees which doctor)

A dense model is a hospital where one enormous generalist sees every patient — and the smarter you make the doctor, the slower every visit gets. In an MoE hospital, the front desk sends each patient to one specialist. You can grow the staff to 100 doctors while each patient still gets exactly one consultation. That is the whole trick.

The Intuition: Decouple Knowledge from Compute

The most parameter-heavy part of a Transformer block is the FFN (feed-forward network). The Switch Transformer replaces it with N copies — the experts — and passes each token through exactly one of them (Figure 2).

Scaling laws had studied three axes: model size, data, and compute. This paper proposes a fourth axis: increase the parameter count while keeping FLOPs per token constant (§2). The only added compute is the router itself, a lightweight $O(d_{model} \times \text{num experts})$ operation (§3).

Mechanism, Part 1: The Router Picks Experts with a Softmax

Start with the standard MoE that preceded Switch (Shazeer et al. 2017, reviewed in §2.1). Given a token representation $x$, the router computes logits $h(x) = W_r \cdot x$ with a weight matrix $W_r$, then turns them into probabilities:

Here $p_i(x)$ is "how much the router wants to send token $x$ to expert $i$," and the N values sum to 1 — the same softmax you met in attention. Put in words: score every expert, make each score positive by exponentiating it, then give each expert its share of the total. Because the scores pass through an exponential first, an expert that scores only slightly higher ends up with a much larger share — the front desk forms strong preferences from weak evidence. Classic MoE keeps the top-$k$ experts and blends their outputs, weighted by these probabilities:

$\mathcal{T}$ is the set of chosen experts and $E_i(x)$ is what expert $i$ computes on the token. Read it as "a weighted average of the chosen specialists' opinions, weighted by the front desk's confidence." Multiplying by $p_i(x)$ is what lets gradients flow into the router so it can learn (§2.1). Now look at which experts appear in that sum: only the chosen ones. The other N−k never enter the equation at all — a detail which says, in a single line, why sparsity is free. An expert you never evaluate costs you nothing, exactly like a doctor who was never paged.

Mechanism, Part 2: Switch's Answer — Route to Just One

The 2017 conventional wisdom held that the router could only learn by comparing at least two experts, so $k > 1$ was considered necessary. The Switch paper contradicts this head-on and routes each token to a single expert, $k = 1$ (§2.1). Three benefits: (1) less routing computation, (2) each expert's batch — its capacity, defined below — can be at least halved, and (3) the implementation gets simpler and communication costs drop. The gate value from Equation (2) survives, so differentiability is preserved.

In Table 1 the paper runs a head-to-head against top-2 MoE under identical conditions (128 experts, 32 TPUv3 cores): Switch is faster, and at small capacity factors it wins on quality too (§2.3). The intuition that "you can't learn without a comparison" simply did not survive the experiment.

Every Expert Has a Fixed Seat Count

Now the unglamorous part that makes or breaks real MoE systems. TPUs require statically-shaped tensors, so the number of tokens each expert processes — its expert capacity — must be fixed in advance (§2.2):

Divide the batch evenly across experts, then multiply by a slack ratio called the capacity factor. A factor of 1.0 means "exactly enough seats if routing is perfectly balanced"; 1.25 reserves 25% spare seats. Put in words, the right-hand side is just "patients per batch ÷ number of doctors × slack." Notice what is absent from it: no router probabilities, no token content. The seat count is decided before the model has looked at a single input, which is precisely why overflow is possible at all.

When routing is skewed and an expert overflows, the excess tokens skip the expert entirely and pass through the residual connection unprocessed — the paper calls these dropped tokens (§2.2). Raising the factor reduces drops but wastes compute and memory on empty seats. The paper stresses that keeping drop rates low matters for scaling, yet finds Switch performs best at small factors like 1.0–1.25 (§2.3).

The Load-Balancing Loss: Taxing the Popular Doctors

Left alone, the router converges on whichever few experts happened to be good early in training — every patient queuing for the same doctor. The fix is an auxiliary loss added to the model's main loss (§2.2). For a batch $\mathcal{B}$ of $T$ tokens and $N$ experts:

$f_i$ is the fraction of tokens actually dispatched to expert $i$; $P_i$ is the average router probability assigned to expert $i$. The dot product is minimized when both are uniform ($1/N$ everywhere) — in other words, the more skewed the routing, the bigger the fine. Put in words, the sum multiplies "how many patients this doctor really got" by "how badly the front desk wanted to send them there," and adds that up over all doctors. A doctor who is high on only one of the two barely moves the total; it is the doctor who is high on both that makes the product explode — which is the concentration the loss is there to punish. The $f$ vector (an argmax) is not differentiable, but $P$ is, so gradients still reach the router. The factor $N$ keeps the loss scale constant as you vary the expert count. The coefficient $\alpha = 10^{-2}$ was chosen from a sweep of $10^{-1}$ to $10^{-5}$ as large enough to balance load quickly without interfering with the main language-modeling loss (§2.2).

The whole mechanism fits in ten lines of numpy pseudocode:

import numpy as np

def switch_route(x, W_r):            # x: (T, d_model)
    logits = x @ W_r                 # (T, N) score per expert
    probs  = softmax(logits)         # Eq (1): dispatch probabilities
    expert = probs.argmax(-1)        # Eq (2) with k=1: one destination
    gate   = probs.max(-1)           # gate value applied to the output
    f = np.bincount(expert, minlength=N) / len(x)   # actual load
    P = probs.mean(0)                                # mean probability
    aux_loss = alpha * N * (f * P).sum()             # Eq (4)
    return expert, gate, aux_loss

Three Tricks That Make Training Stable (§2.4)

1. Selective precision. Run the router in bfloat16 and its softmax blows up: the all-bfloat16 model diverged at a negative log perplexity of -3.780, while all-float32 was stable at -1.718 but slow (1,160 examples/sec, Table 2). Casting only the router's internals to float32 delivered both: -1.716 quality at 1,390 examples/sec. For background on these formats, see number formats.

2. Initialize at one tenth the scale. Dropping the initialization scale from $s = 1.0$ to $0.1$ improved a 32-expert model's early quality from -3.60 to -2.72 and collapsed the run-to-run standard deviation from 0.68 to 0.01 (Table 3). The same scheme worked from a 223M baseline up past one trillion parameters.

3. Dropout the experts, hard. With far more parameters than a FLOP-matched dense model, Switch overfits small fine-tuning sets faster. Using dropout 0.1 on non-expert layers but 0.4 inside the experts — expert dropout — gave the best results on four downstream tasks, e.g. GLUE 85.2 (Table 4).

So How Much Faster Is It?

Only numbers stated in the paper:

• Against a FLOP-matched T5-Base, the 64-expert Switch-Base reaches the same quality 7.5x faster per step and in one-seventh the wall-clock time (§3.1–3.2)
• Even against T5-Large, which spends 3.5x more FLOPs per token, Switch-Base keeps a 2.5x speedup (§3.3)
• Multilingual pre-training on 101 languages improved every single language, with a mean 5x speedup and at least 4x for 91% of languages (§4.3)
• Switch-C (1.6 trillion parameters, 2,048 experts) reaches fixed quality 4x faster than T5-XXL (§5.6)
• Distillation into a dense student keeps about 30% of the quality gain while cutting parameters by 95%+ (§4.2)
• No supercomputer required: gains appear with as few as 2 experts (§7, Appendix D)

What the Paper Admits It Hasn't Solved

• Instability at the largest scale is open. Switch-XXL (395B, high FLOPs per token) was sporadically unstable; the stabilization tricks weren't sufficient and it never completed the planned 1M steps (§5.6, §8).
• Upstream gains don't fully transfer downstream. Despite 1.6T parameters, Switch-C scores 87.7 on SQuAD versus 89.6 for the smaller Switch-XXL. The paper explicitly calls the relation between FLOPs, parameter count, and fine-tuning quality "poorly understood" (§8).
• Expert attention layers were shelved. Replacing the QKV projections with Switch layers improved quality but diverged under bfloat16 (Appendix A).
• Rescuing dropped tokens didn't help. No-Token-Left-Behind, which reroutes overflow to the second-choice expert, showed no empirical benefit (Appendix B).
• The engineering targets TPUs and Mesh-TensorFlow, and all baselines are T5-era models — read the numbers with that context.

How This Is Used on the Job

From here on we leave the paper; this is the author's practical reading.

Who touches this, and when. ML engineers choosing an architecture for LLM pre-training, and inference-platform engineers designing serving for MoE models. Whenever a model card lists "total parameters" and "active parameters" separately, you are looking at exactly this article's $N$ experts and $k$-way selection.

The parameter names you'll actually type. The paper's Appendix F pseudocode uses the vocabulary that stuck: num_experts, capacity_factor, expert_capacity, the auxiliary loss coefficient $\alpha$, and the expert dropout rate. In training logs, always watch the dropped-token rate and the load skew across experts. The paper's defaults — $\alpha = 10^{-2}$, capacity factor 1.0–1.25, float32 router — are still the standard starting point.

Pitfalls that cause incidents.

1. Running the router in half precision. Table 2's divergence is not hypothetical. Even in a mixed-precision setup, keep the router softmax in float32.
2. Misreading "constant compute" as "constant memory." FLOPs stay flat, but every expert's weights must sit in memory, and serving involves all-to-all communication across devices (§5.4). Distributed training covers the prerequisites.
3. Fine-tuning small datasets with dense-model instincts. Overfitting arrives early; reach for strong regularization like expert dropout first.

If it comes up in an interview. "Why doesn't MoE cost more compute?" — each token passes through one expert (in Switch), and the router is only $O(d_{model} \times N)$. "Then why wasn't it everywhere already?" — complexity, communication cost, and training instability; Switch is the paper that answered all three at once with k=1 routing, reduced communication, and the stability recipe. If you can also sketch the load-balancing loss of Equation (4), you're done.

Takeaways

• MoE swaps parameters per input, scaling a fourth axis — parameter count — at constant compute
• Switch's core claim: k=1 routing works, beating top-2 MoE on speed, simplicity, and communication
• The machinery that makes it real is unglamorous: expert capacity and the load-balancing loss
• The stability kit: float32 router, 0.1x initialization, expert dropout
• The honest limits: instability at the very largest scale, and an unresolved gap between pre-training quality and downstream transfer

If you haven't read the Transformer itself yet, start with attention from scratch; for the serving-side story of giant models, continue to the KV cache.