Load balancing in MoE: Worked examples — AI-ML deep dive

This is the worked-example companion to the parent note on load balancing. There we derived the Switch Transformer auxiliary loss

L_{aux} = α N i = 1 \sum N f_{i} P_{i}

and explained why only the routing probability $P_{i}$ carries gradients while the dispatch fraction $f_{i}$ is a frozen constant. Here we run the machinery on every kind of input it will ever meet — balanced, collapsed, degenerate, over-capacity, and a couple of exam traps — so no scenario surprises you.

Recall Symbols we will reuse (all built in the parent, restated here)

$N$ ::: number of experts. $B$ ::: number of tokens in the batch. $K$ ::: how many experts each token is routed to (Top- $K$ ). $f_{i}$ ::: fraction of tokens actually dispatched to expert $i$ ; a hard count, so $\sum_{i} f_{i} = K$ . Non-differentiable (frozen constant in backprop). $P_{i}$ ::: average softmax routing probability the router gives expert $i$ ; differentiable, and $\sum_{i} P_{i} = 1$ . $α$ ::: strength knob on the auxiliary loss (typically $0.01$ ). $C$ ::: capacity factor; expert capacity $= C \cdot B / N$ tokens.

Before any numbers, one picture of what the loss actually measures so every example lands on the same mental image.

The scenario matrix

Every situation this topic can throw at you falls into one of these cells. The examples below are labelled with the cell(s) they cover — together they fill the whole table.

Cell	Scenario class	What is stressed	Example
A	Perfectly balanced, $f_{i} = K / N$ , $P_{i} = 1/ N$	The fixed-point value $α K$	Ex 1
B	Fully collapsed, one expert takes almost all	Upper end of the loss	Ex 2
C	Top- $K$ with $K > 1$	$\sum f_{i} = K \neq = 1$ bookkeeping	Ex 3
D	Degenerate: $f_{i} = 0$ for a dead expert	A zero factor kills its term	Ex 4
E	Gradient sign — which $P_{i}$ goes up vs down	The feedback controller in action	Ex 5
F	Capacity + token dropping	Hard constraint, dropped-token count	Ex 6
G	Real-world word problem (the chef restaurant)	Translating words → $f_{i}, P_{i}$	Ex 7
H	Exam twist: the "double- $α$ " trap	Total-loss composition	Ex 8

Worked example Example 1 — Cell A: the perfectly balanced fixed point

Setup: $N = 4$ experts, Top- $1$ ( $K = 1$ ), and training has converged to perfect balance: every expert receives an equal share of tokens and the router assigns equal probability. So $f_{i} = K / N = 1/4$ and $P_{i} = 1/ N = 1/4$ for all $i$ . Take $α = 0.01$ . What is $L_{aux}$ ?

Forecast: guess before computing — will the loss be zero, or some small positive constant? (Hint: the terms $f_{i} P_{i}$ are all positive, so it can't be zero.)

Write each term: $f_{i} P_{i} = \frac{1}{4} \cdot \frac{1}{4} = \frac{1}{16}$ . Why this step? The loss is a sum over experts of $f_{i} P_{i}$ ; at balance every term is identical, so I only need one.
Sum over $N = 4$ experts: $\sum_{i} f_{i} P_{i} = 4 \cdot \frac{1}{16} = \frac{1}{4}$ . Why this step? The formula sums all experts; four equal terms.
Multiply by $α N$ : $L_{aux} = 0.01 \cdot 4 \cdot \frac{1}{4} = 0.01$ . Why this step? That is the full formula $α N \sum_{i} f_{i} P_{i}$ .

Verify: The parent note predicts the fixed point equals $α K = 0.01 \cdot 1 = 0.01$ . ✓ It matches — and note it is not zero, so a converged MoE always carries this small residual balancing loss.

Worked example Example 2 — Cell B: full collapse (the loss blows up)

Setup: $N = 4$ , $K = 1$ , but the router has collapsed. Expert 1 grabs almost everything: $f_{1} = 0.9, f_{2} = 0.1, f_{3} = f_{4} = 0$ , and $P_{1} = 0.85, P_{2} = 0.10, P_{3} = 0.03, P_{4} = 0.02$ . $α = 0.01$ .

Forecast: relative to Ex 1's $0.01$ , will the collapsed loss be larger or smaller? By how many times, roughly?

Big term first: $f_{1} P_{1} = 0.9 \cdot 0.85 = 0.765$ . Why this step? The dominant expert contributes the overwhelming share; the picture's giant rectangle.
Remaining terms: $f_{2} P_{2} = 0.1 \cdot 0.10 = 0.01$ ; $f_{3} P_{3} = 0$ ; $f_{4} P_{4} = 0$ . Why this step? Dead experts ( $f = 0$ ) contribute nothing — a preview of Cell D.
Sum: $0.765 + 0.01 = 0.775$ .
$L_{aux} = 0.01 \cdot 4 \cdot 0.775 = 0.031$ . Why this step? Full formula.

Verify: $0.031 > 0.01$ — collapse costs ≈3.1× the balanced fixed point, exactly the "imbalance raises it above $α K$ " claim. The gradient on $P_{1}$ is $α N f_{1} = 0.01 \cdot 4 \cdot 0.9 = 0.036$ , the largest of all four → the router is pushed hardest to reduce $P_{1}$ . ✓

Worked example Example 3 — Cell C: Top-

$K$ with $K = 2$ (bookkeeping trap) Setup: $N = 4$ , Top-2 routing, $B = 10$ tokens. Each token now picks 2 experts, so $\sum_{i} f_{i} = K = 2$ . Dispatch counts: expert 1 chosen by 8 tokens, expert 2 by 6, expert 3 by 4, expert 4 by 2. Probabilities $P = (0.4, 0.3, 0.2, 0.1)$ , $α = 0.01$ .

Forecast: What must $\sum_{i} f_{i}$ equal — 1 or 2? Get this wrong and every later number is off.

Convert counts to fractions: $f_{1} = 8/10 = 0.8, f_{2} = 6/10 = 0.6, f_{3} = 4/10 = 0.4, f_{4} = 2/10 = 0.2$ . Why this step? $f_{i}$ is a fraction of tokens, so divide by $B = 10$ .
Check bookkeeping: $\sum_{i} f_{i} = 0.8 + 0.6 + 0.4 + 0.2 = 2.0 = K$ . ✓ Why this step? With Top-2 the sum is $2$ , not $1$ — this is the classic slip.
Products: $0.8 \cdot 0.4 = 0.32$ ; $0.6 \cdot 0.3 = 0.18$ ; $0.4 \cdot 0.2 = 0.08$ ; $0.2 \cdot 0.1 = 0.02$ . Sum $= 0.60$ .
$L_{aux} = 0.01 \cdot 4 \cdot 0.60 = 0.024$ .

Verify: $\sum_{i} P_{i} = 0.4 + 0.3 + 0.2 + 0.1 = 1$ ✓ (probabilities always sum to 1 regardless of $K$ ). The perfectly-balanced target here would be $α K = 0.01 \cdot 2 = 0.02$ ; our $0.024 > 0.02$ signals mild imbalance. ✓

Worked example Example 4 — Cell D: a dead expert (

$f_{i} = 0$ ) Setup: $N = 3$ , $K = 1$ . Expert 3 is completely starved: $f_{3} = 0$ . Loads $f = (0.7, 0.3, 0)$ , probabilities $P = (0.5, 0.3, 0.2)$ , $α = 0.01$ .

Forecast: Expert 3 still has $P_{3} = 0.2 > 0$ . Does it contribute to the loss? Does it receive a gradient?

Term for expert 3: $f_{3} P_{3} = 0 \cdot 0.2 = 0$ . Why this step? A zero factor annihilates the whole product — a starved expert adds nothing to the loss sum.
Its gradient: $\partial L / \partial P_{3} = α N f_{3} = 0.01 \cdot 3 \cdot 0 = 0$ . Why this step? The gradient is proportional to load, and load is zero — so the balancing loss gives expert 3 no push at all.
Live terms: $f_{1} P_{1} = 0.7 \cdot 0.5 = 0.35$ ; $f_{2} P_{2} = 0.3 \cdot 0.3 = 0.09$ . Sum $= 0.44$ .
$L_{aux} = 0.01 \cdot 3 \cdot 0.44 = 0.0132$ .

Verify: $L_{aux} = 0.0132$ . The subtle lesson: the aux loss does not directly rescue a dead expert (its gradient is $0$ ). Rescue is indirect — because $\sum_{i} P_{i} = 1$ , pushing $P_{1}, P_{2}$ down for the busy experts leaves probability mass that flows to $P_{3}$ . This is why the note calls it a feedback controller, not a symmetric objective. ✓

Worked example Example 5 — Cell E: which probabilities rise, which fall?

Setup: $N = 4$ , $K = 1$ , $α = 0.01$ . Current state $f = (0.6, 0.3, 0.1, 0.0)$ . We do not need $P$ to answer the gradient-direction question.

Forecast: rank the four experts from " $P$ pushed down hardest" to " $P$ pushed down least."

Gradient on each $P_{i}$ : $\frac{\partial L}{\partial P _{i}} = α N f_{i}$ . Why this step? From the parent — $f_{i}$ is a frozen constant, so differentiating $α N f_{i} P_{i}$ w.r.t. $P_{i}$ leaves $α N f_{i}$ .
Plug in: $g_{1} = 0.01 \cdot 4 \cdot 0.6 = 0.024$ ; $g_{2} = 0.012$ ; $g_{3} = 0.004$ ; $g_{4} = 0.0$ . Why this step? Each gradient is just load × $α N$ .
Interpret sign: gradient descent moves $P_{i}$ against a positive gradient, i.e. downward. Larger positive gradient → larger downward push. Why this step? $P_{i}$ enters the loss with a positive coefficient, so minimizing lowers the busiest experts' probabilities.

Verify: Ranking of downward push: expert 1 ( $0.024$ ) > expert 2 ( $0.012$ ) > expert 3 ( $0.004$ ) > expert 4 ( $0$ ). The overloaded expert 1 is throttled hardest and the empty expert 4 is left alone — the controller behaves exactly as designed. Because softmax renormalizes, the mass shed by 1 and 2 lands mostly on 3 and 4. ✓

Worked example Example 6 — Cell F: capacity factor and token dropping

Setup: $B = 12$ tokens, $N = 4$ experts, capacity factor $C = 1.0$ , Top-1. After routing, the counts are: expert 1 → 6 tokens, expert 2 → 3, expert 3 → 2, expert 4 → 1. How many tokens get dropped? (See Expert Capacity and Token Dropping.)

Forecast: which expert overflows, and how many of its tokens spill?

Expert capacity $= C \cdot \frac{B}{N} = 1.0 \cdot \frac{12}{4} = 3$ tokens each. Why this step? Capacity is the hard ceiling per expert; $C = 1$ means exactly the fair share $B / N$ .
Compare each expert's assignment to the ceiling $3$ :
- Expert 1: assigned 6, keeps 3, drops $6 - 3 = 3$ .
- Expert 2: assigned 3, keeps 3, drops 0.
- Expert 3: assigned 2, drops 0. Expert 4: assigned 1, drops 0. Why this step? Any token beyond capacity bypasses the expert as a residual.
Total dropped $= 3$ . Fraction dropped $= 3/12 = 0.25$ . Why this step? Dropped-token rate is the health metric for capacity settings.

Verify: Processed tokens $= 12 - 3 = 9$ ; expert loads after capping $= (3, 3, 2, 1)$ , summing to $9$ ✓. A quarter of the batch was dropped — a strong back-pressure signal telling the router to stop over-feeding expert 1. Raising $C$ to $1.5$ gives capacity $⌈ 1.5 \cdot 3 ⌉ = 5$ , dropping only $6 - 5 = 1$ token. Trade-off: higher $C$ = fewer drops but more compute/memory. ✓

Worked example Example 7 — Cell G: the restaurant word problem

Setup (the parent's chef story, made numeric): 8 chefs (experts), 100 customers (tokens) in a batch, Top-1. The gate sends 60 customers to chef A, 25 to chef B, 15 to chef C, and 0 to the other 5 chefs. The gate's average probabilities are $P_{A} = 0.55$ , $P_{B} = 0.25$ , $P_{C} = 0.15$ , and $0.01$ each for the 5 idle chefs. $α = 0.01$ , $N = 8$ .

Forecast: will this restaurant's aux loss be closer to the balanced $α K = 0.01$ , or several times bigger?

Loads: $f_{A} = 60/100 = 0.6$ , $f_{B} = 0.25$ , $f_{C} = 0.15$ , and $f = 0$ for the 5 idle chefs. Why this step? "Customers sent" ÷ batch size gives the dispatch fraction.
Products for busy chefs: $0.6 \cdot 0.55 = 0.33$ ; $0.25 \cdot 0.25 = 0.0625$ ; $0.15 \cdot 0.15 = 0.0225$ . Why this step? The 5 idle chefs have $f = 0$ → zero product (Cell D again).
Sum $= 0.33 + 0.0625 + 0.0225 = 0.415$ .
$L_{aux} = 0.01 \cdot 8 \cdot 0.415 = 0.0332$ .

Verify: $L_{aux} = 0.0332$ vs balanced $α K = 0.01$ → about 3.3× larger. Sanity on probabilities: $0.55 + 0.25 + 0.15 + 5 (0.01) = 1.00$ ✓. Business reading: chef A is over-worked, five chefs sit idle — the loss is signalling exactly the waste the parent's intuition described. ✓

Worked example Example 8 — Cell H: the "double-

$α$ " exam trap Setup: A student computes $L_{aux} = 0.024$ (from Ex 3) and task loss $L_{task} = 2.5$ . They write the total as $L_{task} + α L_{aux}$ with $α = 0.01$ . What total do they get, what is the correct total, and why do they differ?

Forecast: is the student's number too big or too small, and by roughly what factor in the balancing term?

Student's (wrong) total: $2.5 + 0.01 \cdot 0.024 = 2.5 + 0.00024 = 2.50024$ . Why this step? Just following their formula literally.
Correct total: $L_{total} = L_{task} + L_{aux} = 2.5 + 0.024 = 2.524$ . Why this step? $α$ is already inside $L_{aux} = α N \sum_{i} f_{i} P_{i}$ ; multiplying again applies it twice.
Compare balancing contributions: correct $0.024$ vs student $0.00024$ — the student's is $100 \times$ ( $= 1/ α$ ) too small. Why this step? Multiplying by $α = 0.01$ a second time shrinks it by $100$ .

Verify: $0.024/0.00024 = 100 = 1/ α$ ✓. See the parent's "Don't double-count $\alpha$" callout — this is exactly the mistake it warns against. ✓

Recall Rapid self-test

Balanced fixed-point loss for $K = 1$ , any $N$ , $α = 0.01$ ? ::: $α K = 0.01$ (independent of $N$ ). With Top-2, what does $\sum_{i} f_{i}$ equal? ::: $K = 2$ , not 1. Gradient on $P_{i}$ from the aux loss? ::: $α N f_{i}$ — proportional to current load. Does a dead expert ( $f_{i} = 0$ ) get a balancing gradient? ::: No — its gradient is $0$ ; it is rescued indirectly via softmax renormalization. Expert capacity with $C = 1$ , $B = 12$ , $N = 4$ ? ::: $3$ tokens per expert.

Related deep dives: Top-K Routing in MoE, GShard, MoE Training Dynamics, Gradient Routing in MoE, Mixture of Experts (MoE) Architecture.