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.