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.