Load balancing in MoE

What Is Load Balancing in MoE?

Why do we need it?

Without balancing: The gating network (router) can converge to a degenerate solution where it sends most tokens to a small subset of experts. The remaining experts receive almost no gradient signal and fail to learn.
With balancing: All experts are trained on diverse data, develop specialized skills, and the model uses its full parameter budget efficiently.

The Load Balancing Problem from First Principles

Step 1: How Routing Works

Recall the MoE gating function. For token $\mathbf{x}$ , we compute gating scores for $N$ experts:

g_i(\mathbf{x}) = \text{softmax}(\mathbf{W}_g \mathbf{x})_i

where $g_i(\mathbf{x})$ is the probability of routing to expert $i$ . In Top-K routing (typically $K=2$ ), we pick the $K$ experts with highest $g_i$ and normalize:

\tilde{g}_i(\mathbf{x}) = \begin{cases} \frac{g_i(\mathbf{x})}{\sum_{j \in \text{TopK}} g_j(\mathbf{x})} & \text{if } i \in \text{TopK} \\ 0 & \text{otherwise} \end{cases}

The output is:

y = \sum_{i \in \text{TopK}} \tilde{g}_i(\mathbf{x}) \, E_i(\mathbf{x})

The gating network $\mathbf{W}_g$ is learned via gradient descent on the main task loss. Early in training, random initialization might make expert 1 slightly better at some common pattern. The router sends more tokens to expert 1, it gets more gradient, improves faster, attracting even more tokens—a feedback loop. Experts 2-8 starve.

Step 3: Measuring Load Imbalance

Define the fraction of tokens routed to expert $i$ over a batch of $B$ tokens:

f_i = \frac{1}{B} \sum_{b=1}^{B} \mathbb{1}[\text{expert } i \in \text{TopK for token } b]

Perfectly balanced: $f_i = K/N$ for all $i$ (if Top-2 out of 8 experts, each expert should see $2/8 = 25\%$ of tokens on average).

Collapsed: $f_1 \approx 0.9$ , $f_2 \approx 0.1$ , $f_3 \ldots f_8 \approx 0$ .

Load Balancing Loss: Derivation from Scratch

The standard approach comes from the Switch Transformer (and the related GShard) paper. The key insight: we want two things balanced simultaneously, and we need a term that is (a) minimized by a uniform distribution and (b) differentiable through the router weights.

The Two Quantities We Track

Define, over a batch of $B$ tokens with $N$ experts:

$f_i$ : fraction of tokens actually dispatched to expert $i$ (this is a hard count, non-differentiable):

f_i = \frac{1}{B} \sum_{b=1}^{B} \mathbb{1}[\text{expert } i \in \text{TopK for token } b]

$P_i$ : average routing probability the router assigns to expert $i$ (this is differentiable through $\mathbf{W}_g$ ):

P_i = \frac{1}{B} \sum_{b=1}^{B} g_i(\mathbf{x}_b)

The Switch Transformer Auxiliary Loss

\boxed{\;\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i\;}

where $\alpha$ is a hyperparameter (typically $0.01$ ).

Formula Why This EXACT Form? (Corrected)

This is a carefully engineered loss, and its correctness depends on subtle facts. It is not simply "minimize a dot product":

Fact 1 — Constraints matter. By construction $\sum_i f_i = K$ (each token picks $K$ experts) and $\sum_i P_i = 1$ (softmax sums to one). The bare dot product $\sum_i f_i P_i$ under only these sum constraints would be minimized by putting all $f$ -mass on the smallest- $P$ expert (a corner solution), not by uniformity. So the "minimized at uniform" story for the raw dot product is false.

Fact 2 — Only $P_i$ carries gradients. $f_i$ is a hard indicator count → its gradient is zero (it is treated as a constant during backprop). So the effective differentiable loss is: $\mathcal{L}_{\text{aux}} = \alpha N \sum_i \underbrace{f_i}_{\text{const}} \cdot P_i \quad\Rightarrow\quad \frac{\partial \mathcal{L}_{\text{aux}}}{\partial P_i} = \alpha N f_i.$ The gradient on each $P_i$ is proportional to the current measured load $f_i$ .

Fact 3 — This is why it works. The router is pushed to lower the probability $P_i$ of exactly those experts that are currently overloaded (large $f_i$ ). Because $\sum_i P_i = 1$ (softmax), lowering $P_i$ for overloaded experts automatically raises probability for under-used ones. This is a feedback controller, not a symmetric anti-correlation objective. The product $f_i P_i$ is used precisely so that the magnitude of the corrective push on $P_i$ scales with how overloaded expert $i$ is.

Fact 4 — Why $f_i \cdot P_i$ and not $\sum f_i^2$ or $\sum P_i^2$ alone? Using $\sum_i f_i^2$ alone is non-differentiable (no gradient). Using $\sum_i P_i^2$ alone regularizes the router's confidence but ignores the actual dispatched load (which is what causes hardware/compute imbalance). The cross term $f_i P_i$ couples the measured hard load (as a scaling constant) to the differentiable probability, giving a gradient that directly corrects the observed imbalance. It is minimized (given fixed $f$ ) when the router shifts probability mass away from heavily loaded experts.

Key takeaways (corrected mental model):

The loss does not magically become minimal at $f_i = P_i = 1/N$ as an unconstrained dot product. Its purpose is achieved dynamically: gradients scale with load $f_i$ and push $P_i$ down for busy experts.
At the training fixed point, if the router equalizes probabilities ( $P_i \approx 1/N$ ) and dispatch equalizes ( $f_i \approx K/N$ ), the loss settles to $\alpha N \cdot N \cdot (K/N)(1/N) = \alpha K$ , a small constant. Imbalance raises it above this.

Total Training Loss

\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \mathcal{L}_{\text{aux}}

The router now optimizes a trade-off: route tokens to the best expert for the task (minimize $\mathcal{L}_{\text{task}}$ ) while keeping load balanced (minimize $\mathcal{L}_{\text{aux}}$ by not over-concentrating probability on already-busy experts).

Alternative: Capacity Factor and Token Dropping

Some MoE implementations (like Switch Transformer) also use a capacity factor $C$ :

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

How it works:

If $C=1.0$ , each expert can handle exactly $B/N$ tokens (perfect balance).
If an expert's assignments exceed capacity, excess tokens are dropped (their representation passes through as a residual, bypassing the expert).
Dropping creates back-pressure: if expert 1 is overloaded, tokens get dropped, the model performs worse on those tokens, and gradients discourage over-routing to expert 1.

Why this helps:

Hard constraint: physically prevents one expert from taking all tokens.
Encourages the router to explore under-utilized experts.
Trade-off: dropping tokens can hurt performance if $C$ is too low. Typical values: $C \in [1.0, 1.5]$ .

Worked example Example 1: Computing the Switch Auxiliary Loss

Setup: 4 experts, batch of 8 tokens, Top-1 routing. Routing outcomes (hard dispatch):

Tokens 1,2,3,4 → Expert 1
Tokens 5,6 → Expert 2
Token 7 → Expert 3
Token 8 → Expert 4

Compute $f_i$ (fraction dispatched):

f_1 = 4/8 = 0.5, \quad f_2 = 2/8 = 0.25, \quad f_3 = 1/8 = 0.125, \quad f_4 = 1/8 = 0.125

(Note $\sum_i f_i = 1 = K$ since $K=1$ .)

Average gating probabilities (softmax, so $\sum_i P_i = 1$ ):

P_1 = 0.6, \quad P_2 = 0.2, \quad P_3 = 0.1, \quad P_4 = 0.1

Auxiliary loss ( $\alpha = 0.01$ , $N=4$ ):

\mathcal{L}_{\text{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.3 + 0.05 + 0.0125 + 0.0125) = 0.04 \cdot 0.375 = 0.015

Why this step? The gradient w.r.t. $P_i$ is $\alpha N f_i$ . For expert 1: $0.01\cdot4\cdot0.5 = 0.02$ ; for expert 3: $0.01\cdot4\cdot0.125 = 0.005$ . The router is pushed to lower $P_1$ four times harder than $P_3$ — exactly proportional to how overloaded each expert is. Softmax then redistributes that probability toward under-used experts.

Worked example Example 2: Effect of Balancing on Expert Utilization

Without the auxiliary loss:

After 1000 steps: Expert 1 sees 80% of tokens, experts 2-8 see 2-3% each.
Experts 2-8 have near-random weights (no training signal).
Effective model capacity: 1/8 of total parameters.

With auxiliary loss ( $\alpha = 0.01$ ):

After 1000 steps: Each expert sees 10-15% of tokens (roughly uniform).
All experts develop distinct specializations (e.g., syntax, math, dialogue, ...).
Effective capacity: ~8x larger, better performance at same parameter count.

Why this step? The gradient $\propto f_i$ prevents the feedback loop: even if expert 1 is initially better, its large $f_1$ produces a large downward push on $P_1$ , freeing probability mass so other experts get dispatched and finally receive learning signal.

Common Hyperparameters

Hyperparameter	Typical Value	Purpose
$\alpha$ (aux loss weight)	0.01	Trade-off between task loss and balance
Capacity factor $C$	1.0 - 1.5	Buffer for load variation
Top-K	1 or 2	Number of experts per token

Why $\alpha = 0.01$ ? Empirically tuned. Too high: router flattens probabilities to satisfy balance and ignores which expert is best (hurts accuracy). Too low: load collapse returns. $0.01$ is a robust sweet spot.

Common mistake Mistake 1: Setting

$\alpha$ Too High Wrong idea: "More load balancing is always better, so I'll set $\alpha = 1.0$ ." Why it feels right: You really want to avoid expert collapse, so a strong penalty seems safe.

Reality: With $\alpha = 1.0$ , the auxiliary loss dominates. Because the gradient $\alpha N f_i$ is huge, the router flattens $P_i$ toward uniform regardless of which expert is actually best for each token. Result: all experts become mediocre generalists, performance drops below a dense baseline. You've traded expert collapse for expert homogeneity.

The fix: Use $\alpha \in [0.001, 0.01]$ . The task loss should dominate, with the auxiliary loss as a gentle regularizer. Monitor both $\mathcal{L}_{\text{task}}$ and expert utilization; if utilization is still imbalanced, raise $\alpha$ gradually.

Common mistake Mistake 2: Thinking

$\sum_i f_i P_i$ Is a Symmetric "Anti-Correlation" Objective Wrong idea: "The loss makes $f_i$ and $P_i$ anti-correlated, so it's minimized at the uniform point $f_i=P_i=1/N$ ." Why it feels right: Products of two "balanced" quantities feel like they should be smallest when everything is equal.

Reality: Under the true constraints $\sum_i f_i = K$ and $\sum_i P_i = 1$ , the raw dot product $\sum_i f_i P_i$ is actually minimized at a corner — dumping all $P$ -mass on the single smallest- $f$ expert — not at uniformity. Moreover, $f_i$ is non-differentiable; gradients flow only through $P_i$ . So the loss is not a symmetric anti-correlation objective at all. It is a feedback controller: it treats $f_i$ as a constant weight on $P_i$ and pushes probability away from currently-overloaded experts. Balance emerges dynamically over training, not by static minimization.

The fix: Reason about the gradient $\partial\mathcal{L}_{\text{aux}}/\partial P_i = \alpha N f_i$ , not about "where the function is minimal." The corrective push on each expert is proportional to its measured load.

When Load Balancing Isn't Enough

Very small batch sizes: With $B=16$ and $N=8$ experts (Top-2), stochastic imbalance is unavoidable. Solution: accumulate stats over micro-batches, or increase $B$ .
Highly skewed data: If 90% is English and 10% is code, an early-specialized expert may dominate. Solution: data-aware balancing or accept mild imbalance if task performance improves.
Sparse expert gradients: Experts fed by few tokens get noisy updates. Solution: increase capacity factor $C$ or use gradient accumulation.

Recall Explain to a 12-Year-Old

Imagine a group project with 8 friends and a team leader (the router) who hands out tasks. Left alone, the leader might dump ALL the drawing on Alex just because Alex started well. Alex gets overloaded and the other 7 never learn to draw.

Load balancing is like the teacher watching a tally sheet of who got how many tasks. If Alex's tally is huge, the teacher gives the leader a nudge proportional to Alex's tally: "stop preferring Alex so much." The nudge is bigger for whoever is most overloaded. Since the leader's preferences must add up to 100%, pulling preference away from Alex automatically gives more to the quiet friends — so they finally get tasks and learn.

The clever bit: the tally itself (how many tasks each got) is just a number the teacher reads — the teacher only changes the leader's preferences, and changes them harder for the busiest person. That's exactly how the MoE auxiliary loss works.

Connections

Mixture of Experts (MoE) Architecture: The base architecture; load balancing is essential for MoE at scale.
Top-K Routing in MoE: The routing mechanism that creates the load imbalance problem.
Switch Transformer: Source of the $\alpha N \sum f_i P_i$ auxiliary loss and capacity factor.
GShard: Related importance/load auxiliary losses for large sharded MoE.
Expert Capacity and Token Dropping: The hard-constraint alternative to the soft auxiliary loss.
MoE Training Dynamics: How MoE models converge, including balancing effects.
Gradient Routing in MoE: How gradients flow through the router (only via $P_i$ ).

#flashcards/ai-ml

What is the core problem that load balancing solves in MoE?

Without load balancing, a few experts receive most tokens and dominate training, while other experts collapse into uselessness, wasting parameters and compute. Load balancing ensures all experts receive roughly equal training signal and develop diverse skills.

State the Switch Transformer auxiliary load-balancing loss.

\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i

, where

f_i

is the fraction of tokens dispatched to expert

i

(non-differentiable count),

P_i

is the average router softmax probability for expert

i

(differentiable),

N

is the number of experts, and

\alpha \approx 0.01

. Note

\alpha

is already inside, so total loss is

\mathcal{L}_{\text{task}} + \mathcal{L}_{\text{aux}}

(not

\mathcal{L}_{\text{task}} + \alpha\mathcal{L}_{\text{aux}}

Which term in the auxiliary loss carries gradients, and what is that gradient?

Only

P_i

carries gradients (softmax is differentiable).

f_i

is a hard dispatch count and is treated as a constant. The gradient is

\partial\mathcal{L}_{\text{aux}}/\partial P_i = \alpha N f_i

, i.e., proportional to the measured load of expert

i

Why is it wrong to say

\sum_i f_i P_i

is minimized at the uniform point

f_i=P_i=1/N

Under the true constraints

\sum_i f_i = K

and

\sum_i P_i = 1

, the dot product

\sum_i f_i P_i

is minimized at a corner — concentrating all

P

-mass on the smallest-

f

expert — not at uniformity. The loss works not by being minimal at uniform, but as a feedback controller: its gradient

\alpha N f_i

pushes probability away from overloaded experts, and softmax redistributes it.

Why use the cross term

f_i P_i

rather than

\sum f_i^2

\sum P_i^2

alone?

\sum f_i^2

alone is non-differentiable (no gradient from the hard count).

\sum P_i^2

alone regularizes only router confidence and ignores the actual dispatched load that causes real compute imbalance. The cross term couples the measured hard load

f_i

(as a constant scaling) to the differentiable

P_i

, so the corrective gradient on each

P_i

scales directly with observed overload.

What does "perfectly balanced" mean in Top-K routing?

Each expert receives

f_i = K/N

of tokens on average. For Top-2 out of 8 experts, each should see

2/8 = 25\%

What is the capacity factor

C

and how does it enforce balance?

Expert capacity

= C \cdot B/N

. Tokens assigned beyond capacity are dropped (pass through as residual). Overloaded experts drop tokens, hurting performance, so gradients discourage over-routing. Typical

C \in [1.0, 1.5]

What happens if

\alpha

is set too high (e.g., 1.0)?

The auxiliary gradient

\alpha N f_i

becomes huge, flattening router probabilities toward uniform regardless of which expert is best per token. Experts become mediocre generalists and performance can fall below a dense base

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, MoE yaani Mixture of Experts mein hume ek problem aati hai jise samajhna bahut zaroori hai. Socho ek restaurant hai jismein 8 expert chefs hain, lekin ek gating system hai jo 90% customers ko sirf 2 chefs ke paas bhej deta hai. Wo 2 chefs overworked ho jaate hain aur baaki 6 chefs khaali baithe rehte hain. Yahi hoti hai "expert collapse" — router ek feedback loop mein fas jaata hai, jahan jo expert shuruaat mein thoda better hota hai, usko zyada tokens milte hain, wo aur fast improve karta hai, aur aur tokens attract karta hai. Baaki experts starve ho jaate hain, unko gradient signal hi nahi milta, aur wo useless reh jaate hain.

Ab load balancing yahi problem solve karta hai. Iska core idea simple hai — hum chahte hain ki saare experts ko roughly equal training signal mile taaki har expert apni-apni diverse aur complementary skills develop kare. Iske liye ek auxiliary loss add ki jaati hai, jaise Switch Transformer wali $\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_i f_i \cdot P_i$ . Yahan $f_i$ batata hai ki kitne fraction tokens actually expert $i$ ko gaye (ye hard count hai, non-differentiable), aur $P_i$ average routing probability hai (jo differentiable hai, isiliye gradient iske through flow karta hai). Perfectly balanced situation wo hoti hai jab har expert ko $K/N$ fraction tokens milte hain — jaise Top-2 out of 8 mein har expert ko 25% tokens.

Ye matter kyun karta hai? Kyunki agar load balance nahi hoga to aap apne poore parameter budget ka waste kar rahe ho — aapne 8 experts banaye lekin sirf 2 hi kaam kar rahe hain, matlab compute aur memory sab bekaar. Load balancing ensure karti hai ki model apni full capacity efficiently use kare, saare experts specialize karein, aur scaling ka jo faayda hum MoE se lena chahte the wo actually mile. Isiliye jab bhi tum large efficient models padho, ye concept core hai — bina iske MoE ka poora structure hi collapse kar jaata hai.

What Is Load Balancing in MoE?

The Load Balancing Problem from First Principles

Step 1: How Routing Works

Step 2: Why Collapse Happens

Step 3: Measuring Load Imbalance

Load Balancing Loss: Derivation from Scratch

The Two Quantities We Track

The Switch Transformer Auxiliary Loss

Total Training Loss

Alternative: Capacity Factor and Token Dropping

Common Hyperparameters

When Load Balancing Isn't Enough

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Scaling & Efficient Architectures

Connections