Load balancing in MoE: Exercises — AI-ML deep dive

Quick symbol refresher (so nothing is used before it is named):

Before the exercises, one picture to keep in your head. The figure below plots both quantities side by side — cyan bars are the hard dispatch fractions $f_{i}$ , amber bars are the soft router probabilities $P_{i}$ — first when the layer is perfectly balanced, then when it has collapsed onto expert 1. Notice in the collapsed panel how $f_{1}$ towers over the rest (the amber arrow marks the overloaded expert whose $P_{1}$ the gradient will shove down) while expert 3 is starved to $f_{3} = 0$ (cyan arrow). Refer back to this whenever a problem talks about "overloaded" or "starved" experts.

Level 1 — Recognition

L1.1

A Top-2 MoE layer has $N = 8$ experts. In a perfectly balanced batch, what fraction $f_{i}$ of tokens does each expert receive, and what value should each router probability $P_{i}$ settle to?

Recall Solution

Perfect dispatch balance means $f_{i} = K / N$ . With $K = 2, N = 8$ : $f_{i} = 2/8 = 0.25$ . Perfect probability balance means $P_{i} = 1/ N = 1/8 = 0.125$ . Note these are different numbers — $f$ sums to $K = 2$ , $P$ sums to $1$ . Confusing them is the classic slip.

L1.2

State, without deriving, which of $f_{i}$ and $P_{i}$ carries a gradient during backprop, and why the other does not.

Recall Solution

$P_{i}$ carries the gradient — it is a smooth softmax output of $W_{g} x$ . $f_{i}$ does not — it is built from an indicator function $1 [\cdot]$ (a hard yes/no count). Its derivative is zero almost everywhere, so backprop treats $f_{i}$ as a constant.

L1.3

Write the Switch Transformer auxiliary loss and identify each factor.

Recall Solution

$L_{aux} = α N \sum_{i = 1}^{N} f_{i} P_{i}$

$α$ : small weight (typically $0.01$ ), already inside the loss.
$N$ : normalizer so the balanced value is scale-free.
$f_{i}$ : constant load factor (measured dispatch).
$P_{i}$ : differentiable probability the gradient acts on.

Level 2 — Application

L2.1

$N = 4$ , $B = 8$ , Top-1. Dispatch: tokens 1–4 → E1, tokens 5,6 → E2, token 7 → E3, token 8 → E4. Router averages $P = (0.6, 0.2, 0.1, 0.1)$ . With $α = 0.01$ , compute $L_{aux}$ .

Recall Solution

Fractions: $f = (0.5, 0.25, 0.125, 0.125)$ , and indeed $\sum f_{i} = 1 = K$ . ✓ Dot product: $\sum_{i} f_{i} P_{i} = (0.5) (0.6) + (0.25) (0.2) + (0.125) (0.1) + (0.125) (0.1)$ $= 0.30 + 0.05 + 0.0125 + 0.0125 = 0.375$ Loss: $L_{aux} = α N \sum_{i} f_{i} P_{i} = 0.01 \cdot 4 \cdot 0.375 = 0.015$ .

L2.2

Same layer, now perfectly balanced: $f_{i} = 0.25$ , $P_{i} = 0.25$ for all $i$ . Compute $L_{aux}$ and confirm it equals the fixed-point prediction $α K$ .

Recall Solution

$\sum_{i} f_{i} P_{i} = 4 \times (0.25) (0.25) = 4 \times 0.0625 = 0.25$ . $L_{aux} = 0.01 \cdot 4 \cdot 0.25 = 0.01$ . Fixed-point formula: $α K = 0.01 \cdot 1 = 0.01$ . ✓ Matches. The imbalanced case (L2.1) gave $0.015 > 0.01$ — imbalance always sits above the floor.

L2.3

Capacity factor. $N = 4$ , $B = 8$ , $C = 1.0$ . What is each expert's capacity? In L2.1's dispatch, how many tokens does E1 drop?

Recall Solution

Capacity $= C \cdot \frac{B}{N} = 1.0 \cdot \frac{8}{4} = 2$ tokens per expert. E1 was assigned 4 tokens but can hold only 2 → it drops $4 - 2 = 2$ tokens. Those two tokens skip the expert and pass through as a residual.

Level 3 — Analysis

L3.1

Show that the gradient of $L_{aux}$ with respect to $P_{i}$ is $α N f_{i}$ , and explain in one sentence why this makes the loss a feedback controller.

Recall Solution

Treat each $f_{i}$ as constant (L1.2). Then $\frac{\partial L _{aux}}{\partial P _{i}} = \frac{\partial}{\partial P _{i}} [α N \sum_{j} f_{j} P_{j}] = α N f_{i} .$ The push to lower $P_{i}$ is proportional to the current measured load $f_{i}$ . Overloaded experts (big $f_{i}$ ) get their probability shoved down hardest; because $\sum_{i} P_{i} = 1$ , that freed mass flows to under-used experts. That load-proportional correction is the feedback loop.

L3.2

Fact 1 from the parent note claims the bare dot product $\sum_{i} f_{i} P_{i}$ , under only $\sum f_{i} = K$ and $\sum P_{i} = 1$ , is minimized at a corner, not at uniformity. Demonstrate this: with $N = 4$ , $K = 1$ , fix $P = (0.6, 0.2, 0.1, 0.1)$ and find the $f$ (subject to $\sum f_{i} = 1$ , $f_{i} \geq 0$ ) that minimizes $\sum_{i} f_{i} P_{i}$ .

Recall Solution

With $P$ fixed, $\sum_{i} f_{i} P_{i}$ is a weighted average of the $P_{i}$ using weights $f_{i}$ . To minimize a weighted average, dump all weight on the smallest $P_{i}$ . The smallest are $P_{3} = P_{4} = 0.1$ . So put $f = (0, 0, 1, 0)$ (or $(0, 0, 0, 1)$ ): all tokens to expert 3. Value $= 1 \cdot 0.1 = 0.1$ . Compare uniform $f = (0.25, \dots)$ giving $\sum f_{i} P_{i} = 0.25 (0.6 + 0.2 + 0.1 + 0.1) = 0.25$ . The corner ( $0.1$ ) beats uniform ( $0.25$ ). So the bare dot product does not prefer uniformity — this confirms Fact 1. Balance emerges only because $f$ is a measured constant and gradients act on $P$ (L3.1), not because the dot product itself is uniform-loving.

L3.3

Suppose the router collapses: $f = (0.9, 0.033, 0.033, 0.034)$ and $P = (0.9, 0.033, 0.033, 0.034)$ (router mirrors the load), $N = 4$ , $α = 0.01$ , $K = 1$ . Compute $L_{aux}$ and compare to the balanced floor $α K$ .

Recall Solution

$\sum_{i} f_{i} P_{i} = 0. 9^{2} + 0.03 3^{2} + 0.03 3^{2} + 0.03 4^{2} = 0.81 + 0.001089 + 0.001089 + 0.001156 \approx 0.813334$ . $L_{aux} = 0.01 \cdot 4 \cdot 0.813334 \approx 0.03253$ . Floor: $α K = 0.01$ . The collapsed loss ( $\approx 0.0325$ ) is over triple the balanced floor — exactly the strong corrective signal we want.

Level 4 — Synthesis

L4.1

Design a mini training step. $N = 4$ , $B = 8$ , Top-1. Current router logits give $P = (0.7, 0.1, 0.1, 0.1)$ , and the current hard dispatch gives $f = (0.75, 0.125, 0.0, 0.125)$ (E3 got starved this batch). Using $\frac{\partial L _{aux}}{\partial P _{i}} = α N f_{i}$ with $α = 0.01$ : (a) compute the gradient on each $P_{i}$ ; (b) state which expert's probability gets pushed down hardest and which experts get no auxiliary push, and explain the consequence.

Recall Solution

(a) $\frac{\partial L _{aux}}{\partial P _{i}} = 0.01 \cdot 4 \cdot f_{i} = 0.04 f_{i}$ : $g = (0.04 \cdot 0.75, 0.04 \cdot 0.125, 0.04 \cdot 0, 0.04 \cdot 0.125) = (0.03, 0.005, 0, 0.005) .$ (b) E1 ( $g = 0.03$ ) gets pushed down hardest — it's the overloaded chef. E3 gets gradient $0$ : because it was starved ( $f_{3} = 0$ ), the aux loss does not lower its probability at all. Since gradient descent reduces $P_{1}$ and softmax renormalizes ( $\sum P_{i} = 1$ ), the freed probability mass flows toward E3 (and the others). Next batch E3 receives more tokens — the controller self-corrects.

L4.2

Combine load balancing with capacity. $N = 4$ , $B = 16$ , Top-1, $C = 1.25$ . Router dispatches $f_{1} = 0.5$ (8 tokens), $f_{2} = 0.25$ (4), $f_{3} = 0.125$ (2), $f_{4} = 0.125$ (2). (a) Compute per-expert capacity. (b) How many tokens does each expert drop? (c) What is the total number of dropped tokens, and what fraction of the batch is that?

Recall Solution

(a) Capacity $= C \cdot \frac{B}{N} = 1.25 \cdot \frac{16}{4} = 1.25 \cdot 4 = 5$ tokens per expert. (b) Drops $= max (assigned - 5, 0)$ :

E1: $max (8 - 5, 0) = 3$ dropped.
E2: $max (4 - 5, 0) = 0$ .
E3: $max (2 - 5, 0) = 0$ .
E4: $max (2 - 5, 0) = 0$ . (c) Total dropped $= 3$ . Fraction $= 3/16 = 0.1875 = 18.75%$ . Interpretation: raising $C$ from $1.0$ (capacity $4$ , would drop $4$ ) to $1.25$ (capacity $5$ , drops $3$ ) reduces token loss but costs more memory per expert — the core capacity trade-off.

Level 5 — Mastery

L5.1

Prove the fixed-point value. Show that when both dispatch and probability are perfectly balanced ( $f_{i} = K / N$ , $P_{i} = 1/ N$ for all $i$ ), the auxiliary loss equals exactly $α K$ , independent of $N$ .

Recall Solution

$L_{aux} = α N \sum_{i = 1}^{N} f_{i} P_{i} = α N \sum_{i = 1}^{N} \frac{K}{N} \cdot \frac{1}{N} = α N \cdot N \cdot \frac{K}{N} \cdot \frac{1}{N} .$ The $N \cdot N$ from the sum-of- $N$ -equal-terms and the outer $N$ combine: $α N \cdot (N \cdot \frac{K}{N ^{2}}) = α N \cdot \frac{K}{N} = α K$ . The $N$ 's cancel completely → floor is $α K$ , a constant chosen so different-sized MoE layers report comparable balanced losses.

L5.2

Sensitivity / worst case. For fixed $\sum_{i} f_{i} = K = 1$ and $\sum_{i} P_{i} = 1$ with $N$ experts, and the worst collapse where a single expert takes everything ( $f = P = (1, 0, \dots, 0)$ ), express $L_{aux}$ and compute the ratio of worst-case to balanced-floor loss for $N = 8$ . Then say what this ratio tells you about the loss's dynamic range.

Recall Solution

Worst case: $\sum_{i} f_{i} P_{i} = 1 \cdot 1 = 1$ , so $L_{aux}^{m a x} = α N \cdot 1 = α N$ . Floor: $α K = α$ (since $K = 1$ ). Ratio $= \frac{α N}{α K} = \frac{N}{K} = \frac{8}{1} = 8$ . So the loss ranges over a factor of $N / K$ between full collapse and perfect balance. With more experts, the dynamic range grows — the aux loss shouts louder (relative to its floor) precisely when there are more experts to potentially waste, giving proportionally stronger correction where it matters most.

L5.3

Full pipeline. $N = 4$ , $B = 8$ , Top-1, $α = 0.01$ , $C = 1.0$ . A batch dispatches $f = (0.5, 0.25, 0.125, 0.125)$ with router probabilities $P = (0.55, 0.25, 0.10, 0.10)$ , and the task loss for this batch is $L_{task} = 2.000$ . Compute (a) $L_{aux}$ , (b) $L_{total}$ , (c) capacity and tokens dropped by the busiest expert.

Recall Solution

(a) $\sum_{i} f_{i} P_{i} = (0.5) (0.55) + (0.25) (0.25) + (0.125) (0.10) + (0.125) (0.10)$ $= 0.275 + 0.0625 + 0.0125 + 0.0125 = 0.3625$ . $L_{aux} = 0.01 \cdot 4 \cdot 0.3625 = 0.0145$ . (b) $L_{total} = L_{task} + L_{aux} = 2.000 + 0.0145 = 2.0145$ (no double $α$ ). (c) Capacity $= 1.0 \cdot \frac{8}{4} = 2$ . Busiest expert E1 has $f_{1} \cdot B = 0.5 \cdot 8 = 4$ tokens → drops $max (4 - 2, 0) = 2$ tokens.

Recall Self-test recap

Balanced dispatch fraction with Top-2, N=8 ::: $f_{i} = K / N = 0.25$ Balanced router probability with N=8 ::: $P_{i} = 1/ N = 0.125$ Which quantity carries the gradient ::: $P_{i}$ (differentiable softmax) Gradient of aux loss w.r.t. $P_{i}$ ::: $α N f_{i}$ Balanced fixed-point loss value ::: $α K$ Worst-case (full collapse) loss value ::: $α N$ Expert capacity formula ::: $C \cdot B / N$ Tokens dropped by an expert ::: $max (assigned - capacity, 0)$

Related: Switch Transformer · GShard · Top-K Routing in MoE · Expert Capacity and Token Dropping · MoE Training Dynamics · Gradient Routing in MoE · Mixture of Experts (MoE) Architecture