Foundations — Load balancing in MoE

Before you can read a single load-balancing formula, you need to be fluent in the pieces those formulas are built from. This page introduces each symbol and idea from absolute zero, in an order where every new thing only uses things already defined. Nothing here assumes you have seen Mixture of Experts (MoE) Architecture before — we build it.

0. The scene: what an MoE layer physically is

Look at the figure. A single input arrives on the left. Instead of one big network processing it, we have $N$ stations. A little box at the bottom — the router — decides which stations the input actually visits. Only the chosen stations do work; the rest stay idle for this input.

Why bother? Because you get the knowledge of a huge model (all $N$ experts' parameters) while only paying the compute of the few you actually use. That trade is the whole reason MoE exists — see Mixture of Experts (MoE) Architecture.

1. The token $\mathbf{x}$ — the thing being routed

The picture of a vector: an arrow / a column of numbers, e.g. $\mathbf{x}=[0.2,-1.1,0.7,\dots ]$. Each number is one feature the earlier layers extracted.

Why we need it: the router looks at $\mathbf{x}$ and must decide "which experts suit this token?" So $\mathbf{x}$ is the question; the routing decision is the answer.

2. Counting symbols: $N$, $B$, $K$, $i$, $b$

These are plain counters. Meet them before any formula uses them.

The picture: a grid, rows = experts $i$, columns = tokens $b$. A cell is lit if token $b$ went to expert $i$. This grid is the routing decision, and everything about balancing is about the shape of the lit cells — see Top-K Routing in MoE.

3. Softmax — turning scores into a choice

The router first produces a raw score for each expert. Scores can be any numbers (negative, huge, tiny). We need to turn them into something interpretable: probabilities that are all positive and add up to $1$. The tool that does exactly this is softmax.

Why this tool and not, say, "just divide each score by the total"? Two reasons, both visible in the figure:

1. Raw scores can be negative; plain division would give negative "probabilities," which are meaningless. The $e^{(\cdot )}$ (the exponential — a function that turns any real number into a positive one) fixes that: $e^{\text{anything}}>0$.
2. The denominator ${\sum }_j{e}^{s_j}$ (the sum over all experts) forces the outputs to add to exactly $1$. That is what "${\sum }_i{P}_i=1$" means later — it is not magic, it is baked into softmax.

4. Gating scores $g_i(\mathbf{x})$ — the router's opinion

Break the pieces:

• ${\mathbf{W}}_g$ = the router's weight matrix. Multiplying ${\mathbf{W}}_g\mathbf{x}$ turns the token into $N$ raw scores (one per expert). This is the only thing the router learns.
• Wrapping softmax around it makes those scores into probabilities: $g_1,\dots ,g_N$, all positive, summing to $1$.

Picture: for one token, a bar chart of $N$ bars whose heights add to $1$. A tall bar = "the router really wants this expert."

Why we need $g_i$: it is soft (a smooth number), so gradients can flow through it during training. Hold onto that word "soft" — it is the whole reason $P_i$ exists later.

5. Top-K and the indicator $1[\cdot ]$ — the hard choice

The router has $N$ soft opinions $g_i$, but a token only visits $K$ experts. So we keep the $K$ biggest $g_i$ and throw the rest away. That is Top-K routing.

So $1[\text{expert }i\in \text{TopK for token }b]$ reads: "$1$ if token $b$ actually went to expert $i$, else $0$." That is exactly a lit cell in the grid from Section 2.

The tilde $\stackrel{~}{}$ just means "the cleaned-up version." The output blends the chosen experts: $y={\sum }_{i\in \text{TopK}}{\tilde{g}}_i(\mathbf{x})E_i(\mathbf{x})$ where $E_i(\mathbf{x})$ is expert $i$'s output for the token.

6. Two ways to measure "how busy is expert $i$": $f_i$ and $P_i$

Here is the payoff of the soft/hard distinction. We describe an expert's load two ways.

Look at the two bar charts. $f_i$ (left, teal) counts actual visits — jagged, because it comes from hard switches. $P_i$ (right, orange) is the router's average intention — smooth, because it comes from softmax. They usually look similar, but they are not the same object, and the load-balancing loss deliberately multiplies one by the other. That is the entire trick of the Switch Transformer auxiliary loss you meet in D2.

7. Two useful sum facts (memorise the pictures)

Picture for ${\sum }_i{f}_i=K$: each token is a coin worth $K$ "visits"; across the batch of $B$ tokens the visits get spread over experts, and when you total the fractions they must come back to $K$. Nothing is created or lost. These two facts are the constraints that make the load-balancing loss behave the way it does.

8. Capacity $C$ — the physical shelf size

Picture: a shelf holding at most that many tokens. If more tokens are routed to expert $i$ than the shelf holds, the overflow is dropped (skips the expert, passes through unchanged). That overflow idea belongs to Expert Capacity and Token Dropping. It is the hard safety net that complements the soft auxiliary loss.

Prerequisite map

Every arrow means "you must understand the source before the target makes sense." The two paths into the loss — the soft $P_i$ path and the hard $f_i$ path — are exactly the two quantities the Switch Transformer and GShard losses combine.

Equipment checklist

Cover the right side; answer, then reveal.