Load balancing in MoE: Visual walkthrough — AI-ML deep dive

Before any symbols, the cast of characters:

Step 1 — The router turns one token into a bar chart of scores

WHAT. Take a single token, call its vector $x$ . The router multiplies it by a learned matrix $W_{g}$ and squashes the result with softmax:

g_{i} (x) = makes them sum to 1 softmax (raw scores per expert W_{g} x)_{i}

$g_{i} (x)$ — the router's probability that expert $i$ is the right one for this token. A number between 0 and 1.
The subscript $i$ runs over the $N$ experts, so one token produces $N$ of these numbers.
softmax guarantees $\sum_{i = 1}^{N} g_{i} (x) = 1$ : the whole bar chart's heights add to exactly one, like slices of a pie.

WHY softmax and not just the raw scores? Because we want a probability — a comparable, normalized preference — not an arbitrary magnitude. Softmax also keeps everything positive and differentiable, which we will need in Step 5.

PICTURE. One token in, one bar chart out. The tallest bar is where the router wants to send it.

Step 2 — Dispatch: keep only the top bar(s), throw a hard switch

WHAT. In Top-K routing we keep only the $K$ tallest bars (Switch Transformer uses $K = 1$ ). That is a hard decision — a token either goes to an expert or it does not.

1 [expert i \in TopK for token b] = {10 if expert i is chosen otherwise

$1 [\dots]$ — the indicator: a light switch that reads 1 when the statement inside is true, else 0.
$b$ — the index of which token in the batch (token number 1, 2, 3, …).

WHY split "score" from "dispatch"? This is the crux of the whole derivation. The score $g_{i}$ is a smooth dial. The dispatch $1 [\dots]$ is a hard on/off switch. A smooth dial has a gradient; a hard switch does not. Remember these two objects — they become $P_{i}$ and $f_{i}$ in a moment.

PICTURE. The smooth bar chart on the left; the flattened 0/1 "who actually got the token" on the right.

Step 3 — Two summaries over the whole batch: $f_{i}$ (hard load) and $P_{i}$ (soft desire)

WHAT. One token is not enough to judge balance — imbalance is a batch phenomenon. So we average over all $B$ tokens in the batch, in two different ways.

The hard load — what fraction of tokens actually landed on expert $i$ :

f_{i} = \frac{1}{B} b = 1 \sum B 1 if token b went to i 1 [expert i \in TopK for token b]

The soft desire — the router's average probability for expert $i$ :

P_{i} = \frac{1}{B} b = 1 \sum B router’s score for token b g_{i} (x_{b})

$B$ — number of tokens in the batch; dividing by it turns a count into a fraction.
$f_{i}$ is built from switches → non-differentiable, a plain count.
$P_{i}$ is built from dials → differentiable through $W_{g}$ .

WHY two? $f_{i}$ is what actually stresses the hardware (real tokens, real compute). $P_{i}$ is the only quantity we can push with gradient descent. We will need one to measure imbalance and the other to fix it.

Two useful sanity constants fall out immediately:

i = 1 \sum N f_{i} = K, i = 1 \sum N P_{i} = 1.

Each token contributes $K$ dispatches (so the fractions add to $K$ ), and each token's softmax scores add to 1.

PICTURE. Two bar charts side by side over a whole batch: blue $f_{i}$ (counts) and orange $P_{i}$ (probabilities).

Recall Which one carries the gradient?

$f_{i}$ or $P_{i}$ — which is differentiable? ::: $P_{i}$ (built from smooth softmax scores). $f_{i}$ is a hard count, gradient zero.

Step 4 — Watch collapse happen: the feedback loop we must break

WHAT. With no correction, the router drifts to a degenerate state: almost all tokens pile onto one or two experts.

WHY it happens (the runaway loop). Random init makes expert 1 slightly better at some common pattern → router sends it more tokens → expert 1 gets more gradient and improves → router sends it even more → experts 2…N starve. This is positive feedback, and it is exactly what we must counter.

PICTURE. An arrow-loop: "more tokens → more gradient → better expert → more tokens." Beside it, the collapsed bar chart with $f_{1} \approx 0.9$ towering over flat neighbours.

Step 5 — Build the corrective loss: multiply hard load by soft desire

WHAT. We want a single number that is big when the router keeps favouring already-busy experts, and whose gradient pushes probability off those experts. The Switch Transformer answer:

L_{aux} = α N i = 1 \sum N how busy, = const f_{i} \cdot how wanted, has gradient P_{i}

$α$ — a small weight (typically $0.01$ ) so this nudge never drowns the real task loss.
$N$ — rescales the sum so the balanced value is a clean constant (we compute it in Step 6).
$f_{i} \cdot P_{i}$ — the key product: the measured load $f_{i}$ acts as a coefficient on the pushable probability $P_{i}$ .

WHY multiply, and why treat $f_{i}$ as constant? Differentiate. Only $P_{i}$ has a gradient, so $f_{i}$ is frozen (detached) during backprop:

\frac{\partial L _{aux}}{\partial P _{i}} = α N f_{i} .

Read that literally: the push on expert $i$ 's probability is proportional to how overloaded it currently is. Busy experts (big $f_{i}$ ) get their $P_{i}$ pushed down hardest. And because $\sum_{i} P_{i} = 1$ is fixed by softmax, dropping $P_{i}$ for busy experts automatically raises it for idle ones. The loss is a feedback controller, cancelling the runaway loop of Step 4.

Why not $\sum_{i} f_{i}^{2}$ or $\sum_{i} P_{i}^{2}$ alone?

$\sum_{i} f_{i}^{2}$ — non-differentiable (it is all hard counts). No gradient, useless for training.
$\sum_{i} P_{i}^{2}$ — differentiable, but only tames the router's confidence; it never sees the actual dispatched load that stresses hardware.
$f_{i} P_{i}$ — couples the measured imbalance (constant $f_{i}$ ) to the pushable probability ( $P_{i}$ ). Best of both.

PICTURE. For each expert, a blue $f_{i}$ bar and an orange $P_{i}$ bar; the gradient arrow on $P_{i}$ is tall exactly where $f_{i}$ is tall — a downward push scaled by the blue height.

Step 6 — The balanced fixed point: what the loss settles to

WHAT. Suppose the controller wins: the router equalizes probability ( $P_{i} \approx 1/ N$ ) and dispatch equalizes ( $f_{i} \approx K / N$ ). Plug in:

L_{aux} = α N i = 1 \sum N \frac{K}{N} \cdot \frac{1}{N} = α N \cdot N \cdot \frac{K}{N ^{2}} = α K .

Each of the $N$ terms equals $\frac{K}{N} \cdot \frac{1}{N}$ ; there are $N$ of them; the outer $α N$ scales it.
The result $α K$ is a small constant — the loss floor. Any imbalance pushes $L_{aux}$ above $α K$ .

WHY this matters. It tells you the loss's job is done when it flatlines near $α K$ , not near zero. The extra factor $N$ in the formula exists precisely to make this floor tidy.

PICTURE. A gauge: needle at $α K$ when balanced (green), swinging higher (red) as one expert hogs load.

Step 7 — Worked numbers: 4 experts, 8 tokens, Top-1

WHAT. From the parent's Example 1. Dispatch: tokens 1–4 → E1, tokens 5,6 → E2, token 7 → E3, token 8 → E4. So:

f_{1} = \frac{4}{8} = 0.5, f_{2} = \frac{2}{8} = 0.25, f_{3} = \frac{1}{8} = 0.125, f_{4} = \frac{1}{8} = 0.125

(check: $\sum f_{i} = 1 = K$ .) Average router probabilities:

P_{1} = 0.6, P_{2} = 0.2, P_{3} = 0.1, P_{4} = 0.1 (\sum P_{i} = 1) .

Now the loss with $α = 0.01$ , $N = 4$ :

L_{aux} = 0.01 \cdot 4 \cdot (0.5 \cdot 0.6 + 0.25 \cdot 0.2 + 0.125 \cdot 0.1 + 0.125 \cdot 0.1)

= 0.04 \cdot (0.30 + 0.05 + 0.0125 + 0.0125) = 0.04 \cdot 0.375 = 0.015 .

The balanced floor here would be $α K = 0.01 \cdot 1 = 0.01$ .
We measured $0.015 > 0.01$ → the batch is imbalanced, and the loss correctly reports it.

WHY show the number? So you see the floor ( $0.01$ ) and an imbalanced value ( $0.015$ ) side by side — the gap is the signal the router learns from.

PICTURE. The two bar charts for this exact example, the per-term products, and the two dashed lines: floor $0.01$ vs actual $0.015$ .

Recall Recompute the floor

For $N = 4$ , $K = 1$ , $α = 0.01$ , what is the perfectly-balanced $L_{aux}$ ? ::: $α K = 0.01 \times 1 = 0.01$ .

Step 8 — The hard backstop: capacity and token dropping

WHAT. The soft loss nudges; sometimes we also want a wall. Each expert gets a hard seat count:

Expert Capacity = C \cdot \frac{B}{N}

$\frac{B}{N}$ — the fair share of tokens per expert.
$C$ — the capacity factor (typically $1.0$ – $1.5$ ), a little slack above the fair share.
Tokens that arrive after an expert is full are dropped — they skip the expert and pass through as a residual.

WHY a second mechanism? The aux loss is a soft incentive; dropping is a hard limit. Dropping creates back-pressure: overload expert 1 → its tokens get dropped → the model does worse on them → gradients discourage over-routing to expert 1. See Expert Capacity and Token Dropping for the full treatment.

Degenerate cases to keep straight:

$C$ too low (e.g. $0.5$ ): many tokens dropped even when routing is reasonable → accuracy suffers.
$C$ too high (huge): the wall never binds → you are back to relying on the soft loss alone, and you waste memory.
$α = 0$ : no soft push at all → collapse returns (Step 4).
$α$ too large: balance dominates the task loss → the router spreads tokens randomly, ignoring which expert is actually best. A dumb but perfectly balanced model.

PICTURE. Four buckets each with a capacity line at $C \cdot B / N$ ; expert 1 overflows and the spilled tokens are drawn as red residual arrows going around it.

The one-picture summary

The whole pipeline on one canvas: token → router bar chart ( $g_{i}$ ) → hard dispatch → the two batch summaries $f_{i}$ (blue, constant) and $P_{i}$ (orange, differentiable) → their product summed and scaled to $α N \sum_{i} f_{i} P_{i}$ → gradient $α N f_{i}$ pushing busy experts' $P_{i}$ down → back toward the floor $α K$ .

Recall Feynman retelling — say it like you'd explain it to a friend

A little router gives every token a bar chart of how much it likes each expert. We keep the tallest bar and actually send the token there. Over a whole batch we make two summaries: a blue one, $f_{i}$ — how busy each expert really got (just a headcount, can't be nudged), and an orange one, $P_{i}$ — how much the router wanted each expert on average (this one we can nudge). We build a penalty by multiplying them expert-by-expert and adding up: $α N \sum_{i} f_{i} P_{i}$ . The magic is in the calculus: since only orange can move, the pressure on each expert's want-level equals its blue busy-ness, $α N f_{i}$ . So the busiest experts get their probability shoved down hardest, and because all probabilities must add to one, the quiet experts get lifted for free. When everything evens out, the penalty rests at a tiny floor, $α K$ ; any lopsidedness lifts it above that, and the gradient keeps pushing until it settles back. If we also want a hard wall, we give every expert a seat limit $C \cdot B / N$ and drop the overflow — a physical backstop when the gentle nudge isn't enough.

See also: Switch Transformer · GShard · MoE Training Dynamics · Gradient Routing in MoE