Mixture of experts: routing, load balancing and capacity

Before this lesson: lesson 04 (the SwiGLU feed-forward network), lesson 05 (the residual path, which carries dropped tokens), lesson 09 (the softmax and its Jacobian), lesson 11 (the feed-forward sublayer and parameter counting) and lesson 12 (bf16 and the z-loss). The all-to-all collective is covered in the ML-systems track. Time: about 75 minutes, plus the notebook.

---

Why this matters #

A dense transformer uses every parameter for every token: a forward pass costs about $2N$ FLOPs per token for $N$ parameters (lesson 11). An MoE layer breaks this link. Parameters grow with the number of experts $E$; compute per token grows only with the number $k$ of experts each token uses.

The gains are large. With 64 experts, Switch-Base reached the quality of the FLOP-matched dense T5-Base in one seventh of the training time (Fedus et al. 2021). Mixtral 8×7B has about 47B parameters but uses about 13B per token, and matches or beats Llama 2 70B on most benchmarks. DeepSeek-V3 has 671B parameters and uses about 37B per token. Interviews ask about each part: "count the active parameters", "why does the router collapse?", "what is the capacity factor?".

1. Predict first #

2. The picture #

One token, two experts. Take a token vector $x \in \mathbb{R}^D$ and two experts $E_1, E_2$, each a small feed-forward network. The router is a linear map: logits $z = xW_r \in \mathbb{R}^2$ and probabilities $p = \operatorname{softmax}(z)$. Top-1 routing picks $i^* = \arg\max_i p_i$ and outputs $y = p_{i^*}\, E_{i^*}(x)$. Only one expert runs. The layer has twice the parameters of one network but the same compute, plus a tiny router.

From soft to sparse. The classic mixture of experts (Jacobs et al. 1991) runs every expert and mixes the outputs: $y = \sum_i p_i(x)\, E_i(x)$. A sparse MoE keeps only the $k$ largest terms of this sum, so the other $E - k$ experts never run for this token.

Why multiply by $p_{i^*}$? The choice $i^*$ is a step function of $z$: it has zero gradient almost everywhere. The factor $p_{i^*}$ is smooth in $z$. Through it, the loss can tell the router "this expert helped, give it more weight". Without it, the router would get no gradient from the main loss.

Where it sits. In a transformer the feed-forward sublayer acts on each token separately (lesson 11). So an MoE layer can replace it with no other change: each token in the batch picks its own experts. Attention stays dense. Mixtral replaces every feed-forward sublayer; GShard and most Switch models replace every other one; DeepSeek-V3 replaces all but the first three.

Interactive figure: loading…

Four new problems, one per part of section 3:

• balance: the router can send most tokens to a few experts (3.3);
• fixed buffers: an expert can receive more tokens than it has room for (3.4);
• stability: the router's exponentials amplify bf16 rounding (3.5);
• communication: experts live on different GPUs, so tokens must travel (3.6).

3. Definitions and main results #

3.1 The top-k MoE layer #

Notation, used in the rest of the lesson:

• $T$ tokens in the batch (batch × sequence, flattened), model width $D$, $E$ experts, $k$ experts per token;
• expert $i$ is a function $E_i: \mathbb{R}^D \to \mathbb{R}^D$ (with an index; without one, $E$ is the number of experts);
• $x_t = \operatorname{LN}(h_t)$ is the normalized residual stream $h_t$ of token $t$ (pre-LN, lesson 05);
• the router has a weight $W_r \in \mathbb{R}^{D \times E}$ and no bias; logits $z_t = x_t W_r \in \mathbb{R}^E$; probabilities $p_t = \operatorname{softmax}(z_t)$ with entries $p_{t,i}$;
• $S_t$ is the set of the $k$ largest $p_{t,i}$;
• gate weights $g_{t,i}$ for $i \in S_t$; the output is $y_t = \sum_{i \in S_t} g_{t,i}\, E_i(x_t)$, and the block outputs $h_t + y_t$.

Two choices of gate weights are common:

• Switch ($k = 1$): $g_{t,i} = p_{t,i}$, the raw probability.
• Renormalized (GShard top-2, Mixtral): $g_{t,i} = p_{t,i} \big/ \sum_{j \in S_t} p_{t,j}$, so the $k$ weights sum to 1.

Renormalized top-$k$ = softmax over the top-$k$ logits. This is Mixtral's $G(x) = \operatorname{Softmax}(\operatorname{TopK}(x W_g))$, with the logits outside the top $k$ set to $-\infty$. Proof: $p_{t,i} / \sum_{j \in S_t} p_{t,j} = e^{z_{t,i}} / \sum_{j \in S_t} e^{z_{t,j}}$, because the full softmax denominator $\sum_{\text{all } j} e^{z_{t,j}}$ cancels. The softmax is monotone, so ranking the logits or the probabilities selects the same $S_t$.

Parameter count. Let each expert be a SwiGLU network with hidden width $d_e$: three matrices, $3Dd_e$ weights (lesson 04). Per MoE layer:

• total expert parameters: $E \cdot 3Dd_e$;
• active expert parameters per token: $k \cdot 3Dd_e$;
• router: $DE$;
• forward FLOPs per token: about $2k \cdot 3Dd_e$ for the experts and $2DE$ for the router.

Everything else (attention, norms, embeddings) is shared and always active. So total parameters = shared + all experts, and active parameters = shared + $k$ experts per MoE layer.

Interactive figure: loading…

3.2 How the router learns from the main loss #

Let $\mathcal{L}$ be the main loss and $a_{t,i} = \partial \mathcal{L} / \partial g_{t,i} = (\partial \mathcal{L} / \partial y_t)\, E_i(x_t)^\top$, a number: how fast the loss changes when the weight of expert $i$ grows. The softmax Jacobian is $\partial p_{t,i} / \partial z_{t,j} = p_{t,i}(\mathbb{1}[i = j] - p_{t,j})$, where $\mathbb{1}[\cdot]$ is 1 if the condition holds and 0 otherwise (lesson 09).

• Switch ($k = 1$, raw gate), with $i^*$ the chosen expert: $\partial \mathcal{L} / \partial z_{t,j} = a_{t,i^*}\, p_{t,i^*}(\mathbb{1}[j = i^*] - p_{t,j})$. If the chosen expert helped ($a_{t,i^*} < 0$), gradient descent raises $z_{t,i^*}$ and lowers the other logits. Shazeer et al. (2017) thought that $k \ge 2$ was needed for useful router gradients; Switch showed that $k = 1$ works.
• Top-1 with renormalization. Then $g_{t,i^*} = p_{t,i^*}/p_{t,i^*} = 1$ for every $z$. The main loss sends no gradient to the router at all.
• Top-2, renormalized. $g_{t,1} = s(z_{t,1} - z_{t,2})$ (with $s$ the logistic sigmoid) and $g_{t,2} = 1 - g_{t,1}$. The gradient with respect to the logit difference is $(a_{t,1} - a_{t,2})\, g_{t,1} g_{t,2}$: the router learns which of the two chosen experts is better. It learns nothing about the experts it did not choose.

3.3 Collapse and the Switch load-balancing loss #

Collapse. Shazeer et al. (2017) observed that the router "tends to converge to a state where it always produces large weights for the same few experts. This imbalance is self-reinforcing": the favored experts get more tokens, so more gradient updates, so they improve and are chosen even more. Unused experts are wasted parameters. On many GPUs the busiest expert also sets the step time.

Definition (Switch load-balancing loss). For a batch of $T$ tokens and top-1 routing,

$f_i$ is the share of tokens dispatched to expert $i$: a count. $P_i$ is the mean router probability of expert $i$. Switch uses $\alpha = 10^{-2}$, chosen from a sweep over $10^{-1}$ to $10^{-5}$: large enough to balance quickly, small enough not to disturb the main loss.

Top-$k$ version. $f_i$ is the share of the $Tk$ assignments that go to expert $i$, so $\sum_i f_i = 1$ again. Conventions differ by a factor $k$: the Hugging Face transformers Mixtral code counts all $k$ choices and divides by $T$, so its uniform value is $k$.

Facts.

1. Uniform routing gives $\alpha$. With $f_i = P_i = 1/E$: $\alpha E \cdot E \cdot (1/E^2) = \alpha$. The factor $E$ makes this value independent of the number of experts.
2. A collapsed router gives $\alpha E$. If every token has $p_t = (1, 0, \dots, 0)$, then $f_1 = P_1 = 1$ and the loss is $\alpha E$.
3. Why both $f$ and $P$. $\sum_i P_i = 1$ always, so a loss linear in $P$ alone would be constant. And $f$ alone has no gradient: it counts argmax decisions. The product weights each expert's mean probability by its real load, and the gradient flows through $P$ only.

Theorem (gradient of the load-balancing loss). Treat $f$ as a constant (it has zero gradient almost everywhere). Then

Proof.

1. Only $P_i$ depends on $z_{t,j}$, and only through token $t$: $\partial P_i / \partial z_{t,j} = \frac{1}{T}\, \partial p_{t,i} / \partial z_{t,j}$.
2. Softmax Jacobian: $\partial p_{t,i} / \partial z_{t,j} = p_{t,i}(\mathbb{1}[i = j] - p_{t,j})$.
3. Sum over $i$: $\frac{\partial \mathcal{L}_{\text{aux}}}{\partial z_{t,j}} = \frac{\alpha E}{T}\sum_i f_i\, p_{t,i}(\mathbb{1}[i = j] - p_{t,j}) = \frac{\alpha E}{T}\big(f_j p_{t,j} - p_{t,j}\sum_i f_i p_{t,i}\big)$.
4. Factor out $p_{t,j}$. $\blacksquare$

Reading the formula.

• $\bar f_t$ is the load that token $t$ "sees", averaged with its own probabilities. A descent step lowers $z_{t,j}$ when expert $j$ is busier than that ($f_j > \bar f_t$) and raises it when expert $j$ is less busy. Probability moves from overloaded to underloaded experts.
• The factor $p_{t,j}$ means that a token mainly changes the logits of experts it already likes.
• $\sum_j \partial \mathcal{L}_{\text{aux}} / \partial z_{t,j} = \frac{\alpha E}{T}(\bar f_t - \bar f_t) = 0$, as for any loss of a softmax. For the router weights (row convention, $Z = XW_r$): $\partial \mathcal{L}_{\text{aux}} / \partial W_r = X^\top G$, where row $t$ of $G$ is $\partial \mathcal{L}_{\text{aux}} / \partial z_t$.

Interactive figure: loading…

It is a surrogate. Switch says the loss "is minimized under a uniform distribution", but this is not exactly true: a few very confident tokens on an underused expert can push it below $\alpha$. In practice the gradient above still drives the routing toward balance. (Shazeer et al. 2017 used noisy top-$k$ gating and two separate balancing losses; Switch merged them into this one.)

3.4 Expert capacity and token dropping #

Why a fixed capacity. Switch was built for TPUs, which need static tensor shapes. Each expert gets a buffer of $C$ token slots, fixed before the batch is routed. The expert capacity is

where cf is the capacity factor: $\mathrm{cf} = 1$ gives exactly enough slots for perfectly balanced routing.

Dropping. When an expert's buffer is full, further assignments to it are dropped: that expert's computation is skipped for the token. If all its assignments are dropped, the token reaches the next layer unchanged, through the residual connection ($h_t + 0$). For top-$k$, a common priority order (GShard-style) is: all first choices in token order, then all second choices.

The trade-off. The layer computes $E \cdot C \approx kT\,\mathrm{cf}$ slots, filled or not, so compute, activation memory and communication all grow with cf. A larger cf drops fewer tokens but pays for padding. Switch found the best speed–quality trade-off at cf 1.0 and 1.25, and with its auxiliary loss it typically dropped less than 1% of the tokens.

Balanced on average is not enough. The load of one expert is random. For $T = 4096$ and $E = 8$ it has mean 512 and standard deviation $\sqrt{4096 \cdot \tfrac18 \cdot \tfrac78} \approx 21$, and the busiest of 8 experts is about 1.5 standard deviations above the mean. So cf $= 1$ drops tokens even then.

Dropless MoE. MegaBlocks (Gale et al. 2022) writes the expert computation as block-sparse matrix products that accept any number of tokens per expert, so no token is dropped. DeepSeek-V3 reports no dropped tokens in training or inference, thanks to its balancing (section 3.7).

3.5 The router z-loss: stability in low precision #

The problem. bf16 has 7 stored fraction bits, so neighboring bf16 numbers in $[2^m, 2^{m+1})$ are $2^{m-7}$ apart: 0.0078 in $[1, 2)$, but 0.5 in $[64, 128)$ (lesson 12). A rounding error $\varepsilon$ in a logit multiplies the probability ratio of two experts by $e^{\varepsilon}$. In exact arithmetic the softmax ignores a common shift of the logits, but not after rounding: large logits carry large errors. ST-MoE (Zoph et al. 2022) links the instabilities of large sparse models to these extra exponentials in the router.

Definition (router z-loss, ST-MoE). For $T$ tokens with router logits $z_t$,

The total loss is $\mathcal{L}_{\text{CE}} + c_B \mathcal{L}_{\text{aux}} + c_z \mathcal{L}_z$, where $\mathcal{L}_{\text{CE}}$ is the cross-entropy and $c_B$ plays the role of $\alpha$; ST-MoE uses $c_z = 0.001$.

What it does.

• Gradient: $\partial \mathcal{L}_z / \partial z_{t,j} = \frac{2}{T}\operatorname{lse}_t\, p_{t,j}$, because $\partial \operatorname{lse}_t / \partial z_{t,j} = p_{t,j}$.
• It fixes the free shift of the softmax. Since $\max_j z_{t,j} \le \operatorname{lse}_t \le \max_j z_{t,j} + \log E$, keeping $\operatorname{lse}_t$ near 0 keeps the largest logit in $[-\log E, 0]$, where bf16 is precise.
• When $\operatorname{lse}_t > 0$ a descent step lowers all logits, most for the most likely expert. It also makes the router a little less confident.

In ST-MoE the z-loss made all runs stable and slightly improved quality; tighter update clipping also stabilized training, but hurt quality. Earlier, Switch used selective precision (the router in float32, cast back to bf16 before the dispatch) and a 10 times smaller initialization scale. ST-MoE found selective precision alone not enough at its largest scale.

Interactive figure: loading…

3.6 Expert parallelism #

Layout. $G$ GPUs form an expert-parallel group. Each GPU holds $E/G$ experts and its own $T_g$ tokens. The other layers (attention, norms, router) are copies on every GPU (data parallelism).

One MoE layer, forward pass:

1. Route. Each GPU runs the router and the top-$k$ on its own tokens.
2. Dispatch (all-to-all 1). Each GPU sends every token copy to the GPU that holds its expert. An all-to-all is a collective in which every GPU sends a different piece of data to every other GPU. With a capacity, the buffers have the static shape $(E, C, D)$.
3. Expert compute. Each GPU runs its experts, each on the batch of tokens it received.
4. Combine (all-to-all 2). The outputs go back to the tokens' home GPUs, where they are multiplied by the gates, summed over the $k$ copies and added to the residual.

The backward pass runs the same two all-to-alls in reverse (output gradients to the experts, input gradients back). So a training step has four all-to-alls per MoE layer.

Volume. In one all-to-all a GPU sends about $k\,T_g\,D\,b$ bytes ($b$ bytes per value). With uniform routing, the fraction $(G-1)/G$ leaves the GPU. Example: $G = 8$, $T_g = 4096$, $k = 2$, $D = 4096$, bf16 ($b = 2$): $2 \cdot 4096 \cdot 4096 \cdot 2 = 67.1$ MB per GPU, of which $\tfrac78 \cdot 67.1 = 58.7$ MB go to other GPUs, per all-to-all. This traffic scales with tokens and $k$ (activations), not with parameters, unlike the gradient all-reduce of data parallelism.

What makes it slow. Stragglers: all GPUs wait for the busiest one, so imbalance costs time even without dropping. Slow links: the all-to-all crosses nodes (InfiniBand instead of NVLink). DeepSeek-V3 limits each token to at most 4 nodes ("node-limited routing") and overlaps communication with computation.

3.7 Balancing without an auxiliary loss, and other routers #

DeepSeek-V3. Each MoE layer has $N_s = 1$ shared expert, used by every token, and $N_r = 256$ routed experts, of which $K_r = 8$ are active. The router computes affinities $s_{i,t} = \operatorname{sigmoid}(u_t^\top e_i)$, with $u_t$ the input of token $t$ and $e_i$ a learned vector per expert. Then:

• selection uses biased scores: the top $K_r$ of $s_{i,t} + b_i$;
• gate values use the unbiased scores: $g_{i,t} = s_{i,t} / \sum_{j \text{ selected}} s_{j,t}$;
• after each step, $b_i \leftarrow b_i - \gamma$ if expert $i$ was overloaded in the batch, and $b_i \leftarrow b_i + \gamma$ if it was underloaded ($\gamma = 0.001$ for most of training, then 0).

Why: an auxiliary loss sends gradients into the router that compete with the language-model loss (a large $\alpha$ hurts quality, a small one leaves imbalance). The bias has no gradient. It acts like a feedback controller on the load: it changes which experts are picked, and the mixture weights stay undistorted. DeepSeek-V3 keeps only a sequence-wise balance loss with a tiny $\alpha = 0.0001$, against extreme imbalance inside one sequence.

Fine-grained and shared experts (DeepSeekMoE). Split each expert into $m$ smaller ones and activate $mk$ of them: the same parameters and FLOPs, but many more combinations. Top-2 of 16 experts gives $\binom{16}{2} = 120$ combinations; after a 4-way split, top-8 of 64 gives $\binom{64}{8} \approx 4.4 \times 10^9$. Shared experts hold common knowledge, so the routed experts can specialize.

Expert-choice routing (Zhou et al. 2022). Each expert picks its top-$C$ tokens from the batch: perfect balance, and a token can get zero or several experts. But the choice uses later tokens of the batch, so it does not directly apply to autoregressive generation.

3.8 MoE against dense #

	Dense FFN	MoE FFN ($E$ experts of the same size, top-$k$)
Parameters	$P$	$E \cdot P$ (+ router $DE$)
Forward FLOPs per token	$2P$	$\approx 2kP$
Weights to store	$P$	$E \cdot P$: all experts, always
Extra communication	none	4 all-to-alls per layer per training step
Weights read per decoding step, batch $B$	$P$	about $E\big(1 - (1 - k/E)^B\big) \cdot P$
New failure modes	none	imbalance, drops, router instability, overfitting when fine-tuning

• Quality: better per FLOP (the Switch and Mixtral numbers above), worse per parameter: a dense model with the same total parameters is an upper bound for an MoE (DeepSeekMoE).
• Fine-tuning: sparse models overfit small tasks more; Switch used dropout 0.4 inside the experts and 0.1 elsewhere ("expert dropout").
• Decoding: $E(1 - (1 - k/E)^B)$ is the expected number of distinct experts that $B$ tokens touch under uniform routing: $k$ at batch 1, nearly all $E$ at a large batch. Decoding is limited by reading weights from memory, so a small batch reads far fewer bytes per step than a dense model of the same total size.
• Rule of thumb (Hugging Face blog): MoE for high throughput on many accelerators; dense when memory is the limit.

Answers to section 1. (1) The router collapses onto a few experts, and the effect feeds itself (3.3); every MoE needs a balancing mechanism (3.3, 3.7). (2) About 2%: a simulation (section 5) gives 1.65%, because the busiest expert gets about 544 tokens on average. With capacity factor 1.25 (640 slots) the drops almost vanish; with $E = 64$ at cf 1.0 they are about 5%.

4. Worked example #

Setup. $E = 3$ experts, numbered 0, 1, 2, and $T = 4$ tokens. To make the arithmetic simple, the logits are $z_t = \log p_t$, so the router probabilities are:

token	$p_{t,0}$	$p_{t,1}$	$p_{t,2}$	top-1
1	0.7	0.2	0.1	0
2	0.6	0.3	0.1	0
3	0.5	0.1	0.4	0
4	0.2	0.5	0.3	1

(a) Gates. Switch (top-1, raw): token 1 uses $0.7\,E_0(x_1)$, and token 4 uses $0.5\,E_1(x_4)$. Top-2 renormalized, token 3: experts 0 and 2, with weights $0.5/0.9 = 0.556$ and $0.4/0.9 = 0.444$. These equal $e^{z_0}/(e^{z_0} + e^{z_2})$ and $e^{z_2}/(e^{z_0} + e^{z_2})$: the softmax over the two chosen logits.

(b) Load-balancing loss ($\alpha = 1$).

1. Top-1 counts: expert 0 gets 3 tokens, expert 1 gets 1, expert 2 gets 0. So $f = (0.75, 0.25, 0)$.
2. Column means: $P = \big(\tfrac{0.7+0.6+0.5+0.2}{4}, \tfrac{0.2+0.3+0.1+0.5}{4}, \tfrac{0.1+0.1+0.4+0.3}{4}\big) = (0.5, 0.275, 0.225)$.
3. $\sum_i f_i P_i = 0.75 \cdot 0.5 + 0.25 \cdot 0.275 + 0 = 0.375 + 0.06875 = 0.44375$.
4. $\mathcal{L}_{\text{aux}} = 3 \times 0.44375 = 1.331$. Compare: 1 for uniform routing, 3 for a fully collapsed router.

(c) Gradient for token 1.

1. $\bar f_1 = 0.75 \cdot 0.7 + 0.25 \cdot 0.2 + 0 \cdot 0.1 = 0.525 + 0.05 = 0.575$.
2. $f - \bar f_1 = (0.175, -0.325, -0.575)$, and $\alpha E/T = 3/4$.
3. $$\frac{\partial \mathcal{L}_{\text{aux}}}{\partial z_1} = \tfrac34\big(0.7 \cdot 0.175,\; 0.2 \cdot (-0.325),\; 0.1 \cdot (-0.575)\big) = (0.0919,\; -0.0488,\; -0.0431).$$

The entries sum to 0. A descent step lowers token 1's logit for the overloaded expert 0 and raises the other two. Token 4 ($\bar f_4 = 0.275$) gets $(0.0713, -0.0094, -0.0619)$: it is pushed away from expert 0 as well, and mostly toward the unused expert 2.

(d) Capacity (top-1).

• cf $= 1.0$: $C = \lceil 4/3 \rceil = 2$. Expert 0 receives tokens 1, 2, 3 in order, so token 3 is dropped: its MoE output is 0 and only its residual $h_3$ goes on. Slots: $3 \times 2 = 6$, of which 3 are used and 3 are padding (50%).
• cf $= 2.0$: $C = \lceil 8/3 \rceil = 3$. Nothing is dropped, but 5 of the 9 slots (56%) are padding. A larger capacity factor is an expensive fix for imbalance.

(e) Router z-loss. Here $\operatorname{lse}_t = \log\sum_j p_{t,j} = \log 1 = 0$ for every token, so $\mathcal{L}_z = 0$. Now add 10 to every logit. The probabilities, the main loss and $\mathcal{L}_{\text{aux}}$ do not change, but $\operatorname{lse}_t = 10$ and $\mathcal{L}_z = 100$. Its gradient for token 1 is $\frac{2}{4} \cdot 10 \cdot (0.7, 0.2, 0.1) = (3.5, 1.0, 0.5)$: it pulls the logits back down, most for expert 0.

(f) Counting DeepSeek-V3. From its configuration: $D = 7168$, expert width 2048 (SwiGLU), 256 routed experts and 1 shared expert, top-8, MoE in 58 of 61 layers.

1. One expert: $3 \cdot 7168 \cdot 2048 = 44{,}040{,}192 \approx 44.0$M.
2. All routed experts: $58 \cdot 256 \cdot 44.04\text{M} \approx 653.9$B.
3. Active expert parameters per token: $58 \cdot (8 + 1) \cdot 44.04\text{M} \approx 23.0$B.
4. The rest (attention, the 3 dense layers, embeddings, routers) is about 14.6B, used by every token. Total $\approx 671$B, and active about 37B, as reported. Only $9/257 = 3.5\%$ of the expert parameters work on any one token.

5. In code #

PyTorch's autograd checks the two gradient formulas on the worked example, and NumPy reproduces the drop rate and the counts:

import numpy as np
import torch

p = torch.tensor([[0.7, 0.2, 0.1], [0.6, 0.3, 0.1],
                  [0.5, 0.1, 0.4], [0.2, 0.5, 0.3]], dtype=torch.float64)
T, E = p.shape

# (b), (c) the Switch loss and its gradient; f is a count, so it carries no gradient
z = p.log().requires_grad_()
probs = z.softmax(-1)
f = torch.bincount(probs.argmax(-1), minlength=E).double() / T
aux = E * (f * probs.mean(0)).sum()
aux.backward()
fbar = (p * f).sum(1, keepdim=True)
print(round(aux.item(), 4), z.grad[0].numpy().round(4))
print(torch.allclose(z.grad, (E / T) * p * (f - fbar)))
# 1.3312 [ 0.0919 -0.0488 -0.0431]
# True

# (e) the router z-loss after a shift of +10
z = (p.log() + 10).requires_grad_()
z_loss = (torch.logsumexp(z, dim=1) ** 2).mean()
z_loss.backward()
print(round(z_loss.item(), 6), z.grad[0].numpy().round(6))
# 100.0 [3.5 1.  0.5]

# large router logits in bf16: same softmax in exact arithmetic, different after rounding
big = torch.tensor([100.3, 100.0, 99.0])
for zz in (big, big - 99.0):
    zb = zz.to(torch.bfloat16).float()
    print(zb.tolist(), torch.softmax(zb, 0).numpy().round(3))
# [100.5, 100.0, 99.0] [0.547 0.331 0.122]
# [1.296875, 1.0, 0.0] [0.496 0.369 0.136]      (exact: [0.497 0.368 0.135])

# token drops of a router that is balanced on average (T = 4096, E = 8, top-1)
rng = np.random.default_rng(0)
for C in [512, 640]:                                 # capacity factor 1.0 and 1.25
    loads = np.stack([np.bincount(rng.integers(0, 8, 4096), minlength=8) for _ in range(2000)])
    dropped = np.maximum(loads - C, 0).sum(axis=1) / 4096
    print(f"C={C}: mean busiest load {loads.max(axis=1).mean():.0f}, dropped {100 * dropped.mean():.2f}%")
# C=512: mean busiest load 544, dropped 1.65%
# C=640: mean busiest load 544, dropped 0.00%

# (f) DeepSeek-V3 expert parameters
D, d_e, n_moe = 7168, 2048, 58
expert = 3 * D * d_e                                 # SwiGLU expert: three matrices
print(expert, round(n_moe * 256 * expert / 1e9, 1), round(n_moe * (8 + 1) * expert / 1e9, 1))
# 44040192 653.9 23.0

The notebook builds the full layer and checks these formulas with exercises.

6. Where it shows up #

• Language models. Sparse MoE layers in Switch (up to 1.6T parameters and 2,048 experts), ST-MoE, Mixtral, DeepSeek-V2 and V3, and many other recent models. Vision MoE (V-MoE, Riquelme et al. 2021) applies the same layer to the tokens of a vision transformer. Hybrid models such as Jamba combine MoE layers with state-space layers (lesson 16).
• ML for science. Physical fields mix regimes (ocean and land, tropics and poles). Spatial Mixture-of-Experts (Dryden and Hoefler, NeurIPS 2022) routes each grid location to experts and was applied to medium-range weather prediction. In an Earth-system foundation model, MoE adds capacity for many variables and regimes at a fixed FLOP cost. But a per-batch balancing loss pushes all experts to the same load, even when the regimes are not equally frequent.
• Practice. Log per layer: the busiest expert's load divided by the mean load, the fraction of dropped tokens, the auxiliary loss and the z-loss (or the mean $\operatorname{lse}_t$). In a bf16 MoE run with loss spikes, check the router logits, the router precision and the z-loss first.

7. Common mistakes #

• "Mixtral 8×7B has 56B parameters." Attention, norms and embeddings are shared. It has about 46.7B parameters, and 12.9B are active per token.
• Renormalizing a top-1 gate. The gate becomes exactly 1, and the main loss no longer trains the router.
• Mixing conventions for top-$k$. $f$ normalized by $Tk$ gives 1 for uniform routing; normalized by $T$ it gives $k$. Check the uniform value of your implementation.
• "cf = 1 is enough when routing is balanced on average." Random fluctuations alone drop about 2% of the tokens at $T = 4096$, $E = 8$, and more with more experts. Raising cf to fix imbalance pays for padding; fix the balance instead.
• "MoE saves memory." It saves FLOPs. All experts must be stored.
• A bf16 router with no protection. Use a float32 router, a z-loss, or both.
• A balancing coefficient that is too large. It competes with the main loss and hurts quality.

8. Check yourself #

1. A Switch-style layer with top-1 routing renormalizes its gate. What gradient does the main loss send to the router, and what does the router learn?
   Answer

   None. The gate is $p_{i^*}/p_{i^*} = 1$ for all logits, and the argmax is piecewise constant. Only the auxiliary loss trains the router, so it learns balance but not which expert is good for which token. Switch uses the raw $p_{i^*}$ for this reason.
2. With $E = 64$ experts, what is $E\sum_i f_i P_i$ for uniform routing, and for a router that gives every token the probability vector $(1, 0, \dots, 0)$?
   Answer

   Uniform: $64 \cdot 64 \cdot (1/64)^2 = 1$. Collapsed: $f_1 = P_1 = 1$, so $64 \cdot 1 \cdot 1 = 64$.
3. $T = 2048$ tokens, $E = 64$ experts, top-2, capacity factor 1.25. Give the capacity, the total number of slots, and the smallest possible fraction of empty slots.
   Answer

   $C = \lceil 2 \cdot 2048 \cdot 1.25 / 64 \rceil = 80$. Slots: $64 \cdot 80 = 5120$ for $2 \cdot 2048 = 4096$ assignments. Even with perfect balance, 1024 slots (20%) are empty.
4. Derive $\partial \mathcal{L}_z / \partial z_{t,j}$ for $\mathcal{L}_z = \frac{1}{T}\sum_t \operatorname{lse}_t^2$. Why does the z-loss fix the free shift of the softmax?
   Answer

   $\partial \operatorname{lse}_t/\partial z_{t,j} = p_{t,j}$, so the gradient is $\frac{2}{T}\operatorname{lse}_t\, p_{t,j}$. Adding $c$ to all logits of token $t$ leaves $p_t$ and the other losses unchanged but turns $\operatorname{lse}_t^2$ into $(\operatorname{lse}_t + c)^2$. The z-loss prefers $\operatorname{lse}_t \approx 0$, which keeps the largest logit in $[-\log E, 0]$, where bf16 rounding is small.
5. DeepSeek-V3-style balancing with 4 experts. In one step the loads are 40%, 30%, 20% and 10% of the assignments, and $\gamma = 0.001$. How do the biases change, and why is the bias not used in the gate value?
   Answer

   The mean load is 25%. Experts 1 and 2 are overloaded: $b_i \leftarrow b_i - 0.001$. Experts 3 and 4 are underloaded: $b_i \leftarrow b_i + 0.001$. The bias only changes which experts are selected; the gate uses the original affinity, so the mixture weights are not distorted and no balancing gradient competes with the main loss.
6. A Mixtral-style layer ($E = 8$, top-2) decodes with batch 1 and with batch 16. How many experts per layer are read on average (uniform routing)? What does that mean for memory bandwidth?
   Answer

   Batch 1: 2 experts, a quarter of the expert weights. Batch 16: $8(1 - 0.75^{16}) \approx 7.92$, nearly all of them. Decoding is limited by reading weights, so at batch 1 the MoE reads far fewer bytes per step than a dense model of the same total size; at batch 16 it reads almost all of its weights, but this cost is shared by 16 tokens.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write the output of a top-$k$ MoE layer for token $t$.

   Show answer

2. Why is an expert's output multiplied by its gate value?

   Show answer

3. Mixtral 8×7B: total and active parameters?

   Show answer

4. Write the Switch load-balancing loss.

   Show answer

5. What is the difference between $f_i$ and $P_i$ in the balance loss?

   Show answer

6. Gradient of the Switch loss with respect to the router logit $z_{t,j}$?

   Show answer

7. Formula for the expert capacity?

   Show answer

8. What happens to a token whose expert is full?

   Show answer

9. What does the router z-loss penalize, and what is its gradient?

   Show answer

10. How many all-to-alls does one MoE layer need per training step, and why?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Mixtral is 8×7B: how many parameters, how many active?" — About 47B and 13B; only the feed-forward sublayers are replicated.
• "Why does a router collapse, and how do you prevent it?" — Rich-get-richer feedback; the Switch loss, noisy top-$k$, DeepSeek-V3's biases, expert-choice routing.
• "Derive the gradient of the Switch loss." — $\frac{\alpha E}{T} p_{t,j}(f_j - \bar f_t)$; $f$ is a constant count.
• "What limits MoE training on many GPUs?" — Four all-to-alls per layer per step, stragglers, bandwidth between nodes.
• "Your bf16 MoE run has loss spikes. What do you check?" — Router logit size ($\operatorname{lse}$), the z-loss, a float32 router, the initialization scale, per-expert load and drop rate.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution, and two demos:

1. Write top_k_gating (raw and renormalized gates) and check it on the worked example.
2. Write a sparse MoE forward pass with SwiGLU experts and check it against a dense reference.
3. Write the Switch load-balancing loss and check its gradient against the formula of section 3.3.
4. Write capacity-limited dispatch with the GShard priority order (top-1 and top-2).
5. Write the router z-loss and check its value and gradient.

Demos: a tiny MoE trained with and without the balancing loss (watch one expert die without it), and the parameter counts of Mixtral 8×7B and DeepSeek-V3.

11. Watch and read #

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient SparsityYannic Kilcher · 34 min

Top-1 routing, the capacity factor and the load-balancing loss of sections 3.1–3.4.

Stanford CS25: V1 I Mixture of Experts (MoE) paradigm and the Switch TransformerStanford Online · CS25 Transformers United · 1 h 6 min

The Switch authors on routing, capacity, selective precision and fine-tuning sparse models.

Mixtral of Experts (Paper Explained)Yannic Kilcher · 35 min

The top-2 renormalized router and the parameter count of section 3.1.

• Fedus, Zoph and Shazeer, Switch Transformers (2021): top-1 routing, capacity, the load-balancing loss, selective precision and expert parallelism.
• Hugging Face, Mixture of Experts Explained: a readable overview of routing, load balancing, the z-loss and when to use MoE.
• Zoph et al., ST-MoE: Designing Stable and Transferable Sparse Expert Models (2022): the router z-loss and the effect of rounding in the router.
• Jiang et al., Mixtral of Experts (2024): the top-2 router and the results against dense models.
• DeepSeek-AI, DeepSeek-V3 Technical Report (2024): auxiliary-loss-free balancing, shared and fine-grained experts, node-limited routing.
• Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (2017): noisy top-$k$ gating and the first observation of collapse.
• Dai et al., DeepSeekMoE (2024): fine-grained and shared experts; Gale et al., MegaBlocks (2022): dropless MoE; Zhou et al., Mixture-of-Experts with Expert Choice Routing (2022).