The load-balancing loss that makes MoE work, read closely

Make a neural network bigger and it can memorise more — but every extra parameter is also extra multiply-adds on every example. That linear coupling of capacity to cost is the wall large models keep hitting. Conditional computation is the escape hatch: what if only a small, input-dependent slice of the network fired for each example? Then you could grow the parameter count almost for free.

A Mixture-of-Experts layer is the cleanest realisation of that idea. Replace one big feed-forward block with K smaller ones — the experts — and add a tiny gating network that looks at the input and emits a weight for each expert. The layer's output is the weighted combination of the experts it picked. If the gate only ever turns on, say, 2 of 2048 experts per token, you pay for 2 experts' worth of compute while owning 2048 experts' worth of parameters. That is how the paper reaches 137 billion parameters at “only minor losses in computational efficiency,” and the same trick sits inside modern sparse LLMs.

The forward pass is small enough to read in one sitting:

import math

def softmax(z):
    m = max(z)
    e = [math.exp(v - m) for v in z]         # shift for numerical stability
    s = sum(e)
    return [v / s for v in e]

def gate(Wg, bg, x):                         # a tiny router: which experts, how much
    logits = [dot(Wg[k], x) + bg[k] for k in range(K)]
    return softmax(logits)                   # dense weights that sum to 1

def expert(We, be, k, x):                    # expert k is just a linear map here
    return dot(We[k], x) + be[k]

def predict(model, x):
    g = gate(model.Wg, model.bg, x)
    f = [expert(model.We, model.be, k, x) for k in range(K)]
    yhat = sum(g[k] * f[k] for k in range(K))  # mixture: weighted sum of experts
    return yhat, g

A quick tour: the gate turns the input into a probability vector over experts with a softmax; each expert maps the input to an output (here just a line — in a real model it is a two-layer MLP); the prediction is Σ g_k · f_k. That is the whole mechanism. Our demo uses dense gating (every expert gets some weight) rather than the paper's sparse top-k, because the pathology we are chasing shows up with the plain softmax too, and dense gating keeps the gradients exact and easy to test. Sparsity is the compute story; what follows is the correctness story, and they are separable.

Here is the sentence most summaries skip. The authors write: “the gating network tends to converge to a state where it always produces large weights for the same few experts. This imbalance is self-reinforcing, as the favored experts are trained more rapidly and thus are selected even more by the gating network.” Read that twice — it is the reason the layer needs saving from itself.

The feedback loop is baked into the gradient. When you differentiate the mixture, an expert's weight update is scaled by g_k, the gate mass it received. An expert the gate ignores gets a near-zero gradient, so it never improves, so the gate keeps ignoring it. A tiny early lead — from nothing more than random initialisation — snowballs into monopoly. In a real MoE the stakes are brutal: the experts you paid to store sit frozen at their initial noise while a handful do all the work, and effective capacity silently drops by an order of magnitude.

To see collapse cost something we need a task where you genuinely need several experts — otherwise one expert covering everything is a fine solution and balance is cosmetic. So we borrow the site's favourite hard shape: XOR. Points fall in the plane [−1, 1]², and the target is y = sign(x₁)·sign(x₂) plus a little noise — +1 in two opposite quadrants, −1 in the other two. No single line can fit it; but four experts, one per quadrant, fit it beautifully — if the gate learns to send each quadrant to its own expert. The gate is linear, and four linear scores are exactly enough to carve the plane into four territories.

Now the balancing loss itself. The paper defines each expert's importance as the total gate mass it collected over a batch, then penalises the squared coefficient of variation of those importances — one clean scalar that is zero when every expert is used equally and grows as usage skews:

def cv(xs):                                  # coefficient of variation = std / mean
    mean = sum(xs) / len(xs)
    if mean == 0: return 0.0
    var = sum((v - mean) ** 2 for v in xs) / len(xs)
    return math.sqrt(var) / mean

# Importance(expert k) = how much total gate mass expert k received over the batch.
# Shazeer et al., eq. (importance loss):  L_importance = w * CV(Importance)^2
def balancing_loss(gates, w):
    importance = [sum(g[k] for g in gates) for k in range(K)]
    return w * cv(importance) ** 2           # 0 when every expert is used equally

# The gradient this adds to each gate logit pushes mass OFF over-used experts and ONTO
# starved ones — c_k is positive for experts above the average, negative below it:
#     c_k = (2 * w * K / N**2) * (Importance[k] - N / K)

The site runs the TypeScript above (the Python and C++ are line-for-line translations). The comment on the last lines is the crux: minimising CV² adds a term to each gate logit that drains weight from over-subscribed experts and hands it to starved ones — the exact counter-force to the collapse spiral. Everything below is computed by these functions; nothing is drawn by hand.

The experiment is a controlled A/B. For each of 16 seeds we build one dataset and one random initialisation, then train two MoEs from that identical start: one with the balancing loss off, one with it on. The only thing that differs is w_importance. First, one representative seed's training history — the coefficient of variation of expert usage as it learns:

The endpoint of that collapse is stark. Here is where each of the four experts' gate mass ended up, same seed:

And because each expert owns a region of the input, we can draw the collapse directly. Colour each point of the plane by the expert the gate would route it to:

“Sparse MoE gives you 1000× the parameters at nearly constant compute” is the line everyone repeats, and it is true. But it is a promise about the architecture that the architecture cannot keep on its own. Wire up the experts and the gate, train the obvious loss, and the gate quietly hands the whole job to a clique while the parameters you paid for rot. The averages, over all 16 seeds, are not subtle:

• We demoed dense gating, not top-k. The real layer keeps only the top few experts per token — that is where the compute savings come from. We used dense softmax to isolate the balancing dynamic with exact gradients; as the §02 caveat notes, sparsity makes collapse worse, so our demo understates the problem rather than inventing it.
• Importance is not the whole story. The paper pairs the importance loss with a second load loss that equalises the number of examples routed to each expert (importance can be balanced while one expert still gets a flood of low-weight traffic). Production systems add a hard capacity factor and drop overflow tokens. We implemented the importance term because it is the one that cleanly kills the collapse; the rest is plumbing for real hardware.
• Our experts are lines. Real experts are MLPs with their own nonlinearity; our linear experts keep the math auditable and the XOR task honestly needs four of them. The collapse dynamic is a property of the gate, not the expert, so it transfers.

And the paper's own results — on real tasks, not a toy — are worth stating as theirs. Applied between stacked LSTMs on machine translation, the MoE beat a strong baseline while costing less to run:

• Switch to top-k and watch experts die permanently. Replace dense gating with top-1 routing and re-run Fig 3: starved experts should get zero gradient and freeze at their init, so no amount of extra training revives them without the loss.
• Sweep the loss weight. Slide w_importance from 0 upward and plot effective-experts-used against it — there is a knee where the router unlocks, and past it a regime where forcing perfect balance starts to fight the task and nudge MSE back up.
• Add more experts than quadrants. With K = 8 on a 4-region task, the honest outcome is redundancy — do the spare experts specialise into sub-regions, or does the balancing loss just spread traffic thinner? That is the over-provisioning question every real MoE faces.

Mixture-of-Experts is one of those ideas whose headline (“free parameters!”) is genuinely true and also completely hides where the difficulty lives. The experts are easy. The gate is a small, greedy thing that, left to its own devices, will happily strangle its own capacity — and the fix is a single auxiliary term you could write on a napkin. Go turn up the depth on the neural-net page and picture a router bolted in front of a rack of them, its one job to keep any single expert from eating the whole meal.