MoE Load-Balancing Loss

A mixture-of-experts router (see mixture-of-experts) starts training with roughly random preferences over experts. That symmetry is unstable: an expert that gets picked slightly more often in the first few steps receives more gradient updates, becomes marginally better at whatever tokens it sees, and is therefore picked even more often next time. Left alone, this positive feedback collapses the router onto a small subset of experts within a few hundred steps, wasting the rest of the model's parameters entirely. The load-balancing loss is the auxiliary term that breaks the feedback loop.

The original formulation

Shazeer et al., 2017 introduced the first version of this idea for sparsely-gated mixture-of-experts layers: an importance loss that penalises the coefficient of variation of the total router weight received by each expert across a batch, pushing every expert toward receiving a comparable share of total gate mass. The intuition it establishes, penalise imbalance in how much weight each expert is receiving in aggregate, is the one every later variant refines.

The Switch Transformer simplification

Fedus, Zoph and Shazeer, 2021 replace this with a formulation designed for top-1 routing at large scale:

f_i = (1/T) * count(tokens routed to expert i)      # fraction dispatched, hard
P_i = (1/T) * sum_t router_prob(token_t, expert i)   # fraction of router mass, soft
aux_loss = alpha * N * sum_{i=1}^{N} f_i * P_i

T is the number of tokens in the batch, N is the number of experts, and alpha is a small coefficient (0.01 in the Switch Transformer paper). The construction is deliberate: f_i is the actual dispatch fraction, computed from the discrete argmax routing decision, so it is not differentiable; P_i is the router's own softmax probability mass for expert i, averaged over the batch, and it is differentiable. Multiplying a non-differentiable count by a differentiable probability gives the optimiser a usable gradient (through P_i) that still tracks the real, hard dispatch imbalance (through f_i). The product f_i * P_i is minimised, for fixed total mass, when both are uniform at 1/N, and the N * prefactor normalises the loss so its scale does not change as you add more experts.

Why the loss doesn't guarantee balance, and what backstops it

The auxiliary loss is a soft penalty added to the main task loss; it discourages imbalance but places no hard ceiling on how many tokens any one expert can receive in a given batch. Production MoE implementations add a second, architectural backstop: a fixed expert capacity, the maximum number of tokens an expert will process per batch, usually expressed as a capacity factor (for example, 1.25x the perfectly-uniform share). Tokens routed to an expert that is already full are dropped: their expert computation is skipped, typically falling back to passing the token through via the residual connection unchanged. Token dropping is a real cost, a dropped token gets no expert-FFN transformation at that layer at all, so a well-tuned load-balancing loss and a sane capacity factor work together: the loss keeps imbalance small on average, and the capacity factor bounds the worst case cheaply rather than requiring a hard, non-differentiable balancing constraint in the loss itself.

The coefficient is a real trade-off

alpha trades task quality against balance directly. Too small, and the router can still drift toward collapse, especially early in training when gradients are noisy. Too large, and the auxiliary loss dominates the update, actively pushing tokens toward experts that are a worse fit for them purely to satisfy balance, which shows up as a measurable regression on the primary task loss. Published values cluster in the 0.01-range for the aggregate load-balancing term, but this is a hyperparameter tuned per model, not a constant.

When it falls down

• Balance is measured per batch, not globally. A router can be well balanced within each batch while still specialising heavily by topic across the corpus (some experts see mostly code, others mostly prose), which is often desirable specialisation, not a bug the loss should suppress.
• Small batches make the balance signal noisy. f_i and P_i are batch averages; at small batch size, the realised routing fraction per expert is a noisy estimate, and the loss can chase noise rather than genuine imbalance.
• Interacts with router logit scale. A router whose logits have drifted to extreme values (see z-loss-logit-regularisation) produces a near-saturated, low-entropy P_i, which starves the load-balancing loss of the soft gradient it needs; z-loss and the load-balancing loss are usually tuned together, not independently.
• Capacity factor is a silent quality knob. Raising it reduces token dropping but increases compute and memory (experts must be provisioned for the larger buffer), so it is as much a cost decision as a quality one.

Further reading

• Shazeer et al., 2017, Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer - the original importance/load-balancing loss for MoE.
• Fedus, Zoph and Shazeer, 2021, Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity - the simplified top-1 auxiliary loss, capacity factor, and token dropping in production-scale training.