DL0130 MoE Router Collapse and Auxiliary Loss

Router collapse is the failure mode in which the learned router of a sparse MoE layer stops using most of its experts and concentrates almost all token assignments on a small subset, so a layer you paid experts’ worth of memory for behaves like a much smaller model. The cause is a positive feedback loop: an expert that happens to win slightly more tokens early in training receives more gradient signal, becomes genuinely better on those tokens, so the router raises its logit further, while starved experts never see enough tokens to become useful. Under token-choice top-k the visible symptoms are a skewed load histogram, a large dropped-token fraction once the hot experts hit their capacity buffer, and a validation loss that tracks a dense model of the active-parameter size rather than the total-parameter size. Production training stacks suppress it with two extra loss terms rather than with a new routing algorithm: an auxiliary load-balancing loss that pushes router probability mass away from overloaded experts, and a router z-loss that penalizes the squared log-sum-exp of the router logits so the logits cannot grow without bound and destabilize the softmax in low precision. Both are differentiable surrogates for quantities that are not: the actual assignment counts are piecewise constant, and the actual numerical blow-up is a hardware property, so each loss attacks a proxy the gradient can reach.

(1) Collapse Is Self-Reinforcing: a small early advantage in router logits compounds through the gradient the winning expert receives, which is why collapse typically happens in the first few thousand steps and is nearly irreversible afterwards.
(2) The Assignment Is Non-Differentiable: the load fraction comes from a top-k selection and has zero gradient almost everywhere, so the balance loss must pair it with the differentiable mean gate probability .
(3) Balance Loss Is A Normalized Dot Product: equals 1 under a uniform assignment and grows toward under total collapse, giving a scale-free objective independent of expert count.
(4) Its Gradient Is Load-Proportional: the derivative with respect to is exactly , so probability mass is pushed down in proportion to how overloaded each expert already is.
(5) Z-Loss Bounds The Logits: penalizing the squared log-partition function caps , which matters because the relative round-off in grows linearly with in bfloat16 and can flip top-k selection.
(6) Coefficients Are A Quality Trade: and are the common settings, because a large buys perfect balance by actively fighting the language-modeling objective.

Figure 1: Router collapse is a closed loop, not a single bad step: the logit gap, the dispatch skew, and the gradient imbalance each amplify the next. The two stabilizers cut the loop at different points, with the z-loss constraining the logit magnitudes and the balance loss constraining the dispatch distribution.

The derivation of the balance loss starts from what you actually want to penalize, namely the variance of the per-expert token counts, and then asks which part of that quantity carries a gradient. The counts themselves come from a top-k over the router logits, so they are a step function of the parameters and give nothing to backpropagation. The fix used by GShard and simplified by Switch Transformers is to pair the non-differentiable load vector with the differentiable importance vector , the mean softmax probability per expert, and minimize their inner product. Treating as a constant, the objective is linear in with coefficient , so each step lowers the router’s probability for exactly the experts that were overloaded in the current batch, and because is recomputed every step this becomes a self-correcting controller whose fixed point is the uniform assignment. The factor is a normalization choice: it makes the minimum value 1 regardless of how many experts you have, so the same transfers from an 8-expert layer to a 256-expert layer. The z-loss has an entirely different motivation: it is a numerical guard, derived from the observation that a logit stored with relative precision produces an absolute error of about , and exponentiation converts that absolute error into a relative error of the same size, so large logits make the softmax and therefore the selected expert set unreliable.

Mathematical Formulation:

Where:

• is the hidden state of token , the router matrix, and the router logits before any selection.
• is the number of tokens the statistics are aggregated over (micro-batch, device batch, or global batch), the expert count, and the set of experts selected for token .
• is the load, the fraction of tokens dispatched to expert ; it is piecewise constant in the parameters and therefore contributes no gradient.
• is the importance, the mean router probability assigned to expert ; it is smooth, so all of the balance-loss gradient flows through it and into .
• exactly when for all , and approaches when one expert takes every token, so the value is directly readable as an imbalance factor.
• is the log-partition function; squaring it penalizes large logits in either direction and, by the sandwich bound, keeps within of a small target.
• and are the balance and z-loss coefficients, commonly and ; is the ordinary next-token cross-entropy.

Figure 2: The two losses guard different quantities. Left: without a balance term the maximum load fraction runs from the uniform to near 1.0 within a few thousand steps, while holds it close to uniform. Right: the relative round-off in grows linearly with , which is why bfloat16 routers become unreliable at moderate logits and the z-loss keeps near .