Load Balancing Across Experts

"They keep sending everyone to me. I am flattered, I am overloaded, and seven of my colleagues have not seen a single token all morning."

A Popular Expert Holding Up the Whole Batch

The previous section moved tokens to their chosen experts with an all-to-all exchange, taking the routing decision as given. Now we confront the decision itself. In a dense layer every device does an identical share of the work, so balance is automatic; nobody chooses who computes. A mixture of experts breaks that symmetry on purpose: the gate sends each token to a small number of experts, and the experts are scattered across devices for the memory reasons that motivated expert parallelism in the first place. The gate is now, whether or not its designers think of it this way, a load balancer for the cluster. When it does that job badly, the consequences land on the hardware immediately and on model quality more slowly, and the two failures reinforce each other.

The trouble is that nothing in the basic gating objective asks for balance. The gate is trained to route each token to whichever experts most reduce the loss, and real data is not uniform: some patterns are common, some are rare, and a few experts quickly become specialists for the common patterns. Left alone, the router discovers that sending most tokens to a handful of experts is locally optimal, a self-reinforcing collapse sometimes called routing collapse. Those experts improve because they see the most data, which makes the gate prefer them even more, while the starved experts stay weak and the gate learns to avoid them further. The result is a model that has paid for many experts but trained only a few.

1. Why Imbalance Is a Distributed-Systems Problem Beginner

Expert parallelism places a disjoint set of experts on each device. When a batch arrives, the gate's choices induce a load on each device equal to the number of tokens routed to the experts it holds. The forward pass through the expert layer cannot finish until every device has processed its share, so the wall-clock cost of the layer is set by the busiest device, not the average one. This is exactly the straggler effect introduced in Chapter 2: a synchronous step proceeds at the speed of its slowest participant, and one overloaded device makes everyone else wait at the next all-to-all barrier. Adding more experts and more devices does not help if the gate keeps funnelling tokens to the same few; you have bought parallel hardware and then serialized it through a hot spot.

The mechanism that produces the hot spot is also familiar. A gate that concentrates tokens on popular experts is producing a skewed partition of the batch, the same pathology that wrecks a Spark shuffle when one key dominates, studied in Chapter 7. There, one popular join key sends a disproportionate share of rows to a single reducer; here, one popular expert receives a disproportionate share of tokens on a single device. The cause differs (a learned router rather than a data distribution) but the cure rhymes: detect the skew, and reshape the assignment so no single partition carries far more than its fair share.

2. The Auxiliary Load-Balancing Loss Intermediate

The classic fix, introduced with the sparsely gated mixture of experts and carried into Switch Transformer, is to add a term to the training objective that penalizes imbalance directly. For a batch of $T$ tokens and $E$ experts, define $f_e$ as the fraction of tokens dispatched to expert $e$ under the (hard, top-$k$) routing, and $P_e$ as the mean gate probability assigned to expert $e$ across the batch,

where $g_e(x_t)$ is the softmax gate probability of expert $e$ for token $t$. The auxiliary loss is the scaled dot product of these two vectors,

Two facts make this the right object. First, the dot product $\sum_e f_e P_e$ is minimized, subject to each vector summing to one, when both are uniform at $1/E$, so driving it down pushes usage toward the even split we want; the factor $E$ keeps the target value near a constant as $E$ changes, and $\alpha$ (a small weight such as $10^{-2}$) sets how hard balance competes with the task loss. Second, the gradient is well behaved: the hard counts $f_e$ are not differentiable, but $P_e$ is, so the gradient flows through the gate probabilities and gently lowers the logits of experts that are already popular ($f_e$ large) while raising the rest. The term nudges the router toward fairness without dictating any single token's destination, which is what lets the experts still specialize. Section 17.7 picks up the companion control, capacity factors and token dropping, that bounds the damage when balance is imperfect.

The code below trains a tiny gate over synthetic tokens that genuinely cluster toward a few popular experts, first with no balancing term and then with the auxiliary loss switched on. It reports per-expert usage, the load on the busiest of four devices (two experts each), and the max-to-mean ratio that quantifies the skew.

import numpy as np

rng = np.random.default_rng(7)
N, d, E = 6000, 16, 8          # tokens, feature dim, experts

# Tokens that genuinely cluster toward experts 0 and 1, so an unregularized
# gate piles onto them: a learned, self-reinforcing popularity skew.
pop = rng.standard_normal((E, d)); pop[0] *= 3.0; pop[1] *= 2.2
assign = rng.choice(E, size=N, p=[.40, .22, .10, .08, .07, .06, .04, .03])
X = pop[assign] + 0.6 * rng.standard_normal((N, d))

def softmax(z):
    z = z - z.max(axis=1, keepdims=True); e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)

def usage_fraction(W):                      # top-1 routing, fraction per expert
    return np.bincount(np.argmax(softmax(X @ W), axis=1), minlength=E) / N

def train(aux_weight, steps=800, lr=0.5):
    W = 0.01 * np.random.default_rng(0).standard_normal((d, E))   # fixed init
    for _ in range(steps):
        g = softmax(X @ W)
        # Task pull: a rich-get-richer signal that nudges each token toward the
        # expert already most confident for it, so popular experts grow more
        # popular and an unregularized gate tends to collapse onto a few of them.
        hard = np.zeros((N, E)); hard[np.arange(N), np.argmax(g, axis=1)] = 1.0
        grad_task = X.T @ (g - hard) / N
        # Auxiliary balance loss aux = E * ||P||^2, the smooth surrogate for
        # E * sum_e f_e P_e, minimized at uniform usage P_e = 1/E. Its gradient
        # flows through the softmax mean P_e = mean_t g and pushes mass off hot experts.
        P = g.mean(axis=0)
        v = 2.0 * E * P
        grad_aux = X.T @ (g * (v - (g @ v)[:, None])) / N
        W -= lr * (grad_task + aux_weight * grad_aux)
    return W

for label, aw in [("no aux loss ", 0.0), ("with aux loss", 1.0)]:
    u = usage_fraction(train(aw))
    dev = u.reshape(4, 2).sum(axis=1)                  # 4 devices, 2 experts each
    print(f"{label}: usage% = [{', '.join(f'{100*x:4.1f}' for x in u)}]")
    print(f"               busiest device = {100*dev.max():4.1f}%  (ideal 25.0%),"
          f"  max/mean expert = {u.max()/u.mean():.2f}x")

no aux loss : usage% = [ 0.0, 55.4,  0.0,  0.0, 31.1, 13.5,  0.0,  0.0]
               busiest device = 55.4%  (ideal 25.0%),  max/mean expert = 4.43x
with aux loss: usage% = [13.1, 14.1, 13.5, 16.5, 10.3,  7.5, 10.4, 14.6]
               busiest device = 30.0%  (ideal 25.0%),  max/mean expert = 1.32x

The improvement is decisive here, and the honest lesson of the demo is why: left alone, the rich-get-richer pull lets a few experts swallow the batch and starves the rest into dead weight, exactly the routing collapse this section warns about. The auxiliary loss counters that feedback and brings the idle experts back into play. On real data the gate also carries genuine specialization pressure, so a practical $\alpha$ trades a little of that specialization for balance rather than flattening usage all the way; pushing $\alpha$ higher buys more balance at a rising cost to model quality, which is the central tension of the whole topic.

3. Balanced by Construction, and Balanced Without a Loss Advanced

The auxiliary loss treats imbalance as something to penalize after the fact. Two other strategies attack it more directly. The first is to change who does the choosing. In the token-choice routing of Section 17.3, each token picks its top-$k$ experts, and nothing stops a popular expert from being everyone's pick. Expert-choice routing inverts the selection: each expert picks the top-$T k / E$ tokens it wants from the batch. Because every expert selects exactly the same number of tokens, the load is uniform by construction, with no auxiliary loss required and no device able to become a hot spot. The cost is that a token may be chosen by several experts or by none, so expert-choice trades guaranteed device balance for variable per-token capacity, a different knob on the same trade-off that Section 17.7 formalizes.

The second strategy keeps token choice but removes the auxiliary loss entirely. The auxiliary loss has a known side effect: its gradient perturbs the task objective, and an $\alpha$ large enough to enforce balance can measurably hurt quality, an interference the literature calls the balance-versus-specialization tension. The auxiliary-loss-free approach replaces the loss term with a per-expert bias $b_e$ added to the routing logits only for the top-$k$ selection. After each step the biases are nudged by a simple control rule: lower $b_e$ for experts that were overloaded, raise it for experts that were starved,

where $\bar f = 1/E$ is the target fraction and $\gamma$ is a small update rate. The bias steers the routing toward balance, but because it is added only for selection and not to the gate value used in the weighted combination, it never appears in the task gradient, so it balances load without distorting what the experts learn. This is the mechanism DeepSeek-V3 uses at scale, and it is the current state of the art for keeping a large mixture of experts balanced while letting the experts specialize freely.

import torch, torch.nn.functional as F

def load_balance_loss(gate_logits, expert_index, num_experts):
    # gate_logits: (T, E) router scores; expert_index: (T,) the top-1 pick per token.
    P = F.softmax(gate_logits, dim=-1).mean(dim=0)                 # mean prob per expert
    f = torch.bincount(expert_index, minlength=num_experts).float() / gate_logits.size(0)
    return num_experts * torch.dot(f, P)                          # alpha applied by caller

# In the training step the layer hands this back; you just add it to the task loss:
# loss = task_loss + alpha * aux_loss          # alpha ~ 1e-2 in Switch Transformer

4. The Balance-Versus-Specialization Tension Advanced

Perfect balance and perfect specialization pull in opposite directions, and a good mixture of experts lives at a deliberate compromise between them. Push balance to its limit and you would route tokens to experts almost at random, guaranteeing a flat load but destroying the very specialization that makes a sparse model worth more than a smaller dense one; the experts would all learn the same average function. Push specialization to its limit and the gate collapses onto a few experts, which trains well but utilizes the cluster terribly and starves most of the model. The whole engineering art is to enforce just enough balance that no device becomes a straggler, while leaving the gate free to send genuinely different tokens to genuinely different experts.

This is why the field has drifted from heavy auxiliary losses toward lighter-touch controllers. The auxiliary loss enforces balance by adding a force to the loss landscape that every token feels, which inevitably bends the task objective. The bias-update approach enforces balance by adjusting only the selection threshold, leaving the gradient that shapes the experts untouched, so it buys balance at a lower price in specialization. Either way, balance is never free, and the right operating point depends on how skewed your data is and how tight your hardware budget is. The performance models of Chapter 3 give you the language to put a number on the throughput cost of a given imbalance, and that number is what tells you how hard to push balance for a given cluster.