Auxiliary Loss & Load Balancing.

An MoE only saves compute if tokens spread across experts. Left alone, a router does the opposite — it piles tokens onto a handful of experts and starves the rest. This article is about the force that stops that: the load-balancing auxiliary loss. What it is, the exact formula, why it works, why it quietly damages the model, and the bias trick DeepSeek-V3 used to delete it.

The problem: routing collapse

The router is trained jointly with the rest of the model. Early in training one expert gets, by chance, slightly better at something. The router sends it more tokens. More tokens means more gradient, so it improves faster, so the router sends it even more. Rich-get-richer. Within a few thousand steps a small clique of experts handles almost everything and the others are dead — near-zero load, never trained.

This is routing collapse, and it defeats the entire point of MoE. You provisioned 256 experts of capacity; you're paying to store all of them; and 16 are doing the work. Worse, dead experts can't recover — no tokens means no gradient means no improvement means no reason for the router to ever pick them.

MoE has no natural pressure toward balance. The default dynamics actively favour collapse. Balance has to be imposed by an explicit objective — that objective is the auxiliary loss.

The auxiliary loss, exactly

Define it over a batch of T tokens and N experts, with top-k routing. Two quantities per expert i:

f_i = (1/T) * (number of tokens that selected expert i)     # fraction of dispatch
P_i = (1/T) * sum over tokens t of  g_t,i                    # mean router probability

  where g_t,i = softmax(router_logits_t)_i  is the gate prob
  the router assigned token t to expert i.

f_i is the actual load — what share of tokens landed on expert i. P_i is the soft load — the average probability mass the router put on expert i, whether or not it was selected. The Switch Transformer auxiliary loss is their dot product, scaled:

L_aux = alpha * N * sum over i of  ( f_i * P_i )

Add it to the language-modelling loss: L_total = L_LM + L_aux. The coefficient alpha (Switch used 0.01) sets how hard balance is enforced.

Why this particular product?

The subtlety: f_i is a count, produced by an argmax/top-k — it has no usable gradient. You can't backprop through "how many tokens chose expert i." But P_i is a smooth softmax average — fully differentiable. The product f_i · P_i treats f_i as a constant weight and pushes gradients only through P_i:

d L_aux / d P_i  ∝  alpha * N * f_i

Read that: the gradient pressure to lower the router probability for expert i is proportional to how overloaded i already is. Overloaded experts (large f_i) get their logits pushed down hardest; starved experts barely get touched, so their relative probability rises. The loss is minimised when load is uniform — every f_i = P_i = 1/N, giving L_aux = alpha · N · N · (1/N)² = alpha, a constant floor. Any imbalance pushes it above that floor.

Tuning alpha is a tightrope

The coefficient is the whole problem in miniature:

• alpha too small — balancing is too weak, collapse wins, experts die. You wasted capacity.
• alpha too large — balancing dominates the LM signal. The router is forced to spread tokens evenly even when that means sending a token to an expert that handles it worse. Balance improves; model quality drops.

There's no setting that's free. Even at a good alpha, you're trading some language-modelling quality for balance every single step. Switch settled on 0.01 empirically; it's model- and data-dependent and brittle.

DeepSeek-V2: three balancing losses at once

At scale, "balance" isn't one thing. DeepSeek-V2 ran three auxiliary losses simultaneously, each targeting a different granularity:

• Expert-level balance — the classic f_i · P_i loss, so no individual expert collapses.
• Device-level balance — experts are sharded across GPUs; you also need each device's group of experts to get even load, or one GPU bottlenecks the step. Same form, grouped by device.
• Communication balance — expert parallelism dispatches tokens across the network (all-to-all). This loss caps how much any device sends/receives so the interconnect doesn't choke.

Three coefficients to tune, all interacting. Functional, but a lot of brittle knobs — and every one of them is still leaking balancing gradient into the model.

The real cost: gradient interference

Step back from tuning and look at what the auxiliary loss is. It's a second objective bolted onto the model whose only goal is even load — a goal the user does not care about. Its gradient and the language-modelling gradient point in different directions, and they're summed. Every step, some of the update is spent making load even rather than making predictions better.

Concretely: a token whose best expert is already busy gets nudged toward a worse expert, because the aux loss wants that busy expert relieved. Multiply across billions of tokens and you get a measurable quality tax. The aux loss fixes collapse by actively corrupting the routing the model would otherwise learn.

The auxiliary loss is a hack that works: it trades model quality for load balance, every step, by design. DeepSeek's question for V3 was — can we get the balance without paying the quality tax?

DeepSeek-V3: auxiliary-loss-free balancing

The answer is yes, and it's elegant. The key realisation: the aux loss conflates two jobs — deciding which expert handles a token (selection) and how much each expert's output is weighted (the gate value). Collapse is a selection problem. So fix selection directly, and leave the gradients alone.

Add a per-expert bias term b_i used only for the top-k selection, not for the gate weight:

select top-k experts by:     s_i + b_i        # biased scores, selection only
weight chosen expert by:     g_i = softmax(s_i)  # original score, NO bias

  s_i   = router affinity (sigmoid/softmax of router logit)
  b_i   = balancing bias, NOT a learned parameter
  g_i   = the gate value actually multiplying the expert output

The bias tilts which experts win the top-k, but the output weighting uses the clean, unbiased score — so no balancing signal enters the gradient. Then update b_i with a dead-simple rule after each step, watching the realised load:

for each expert i:
    if load_i > target:   b_i = b_i - gamma     # overloaded → make it less attractive
    if load_i < target:   b_i = b_i + gamma     # underloaded → make it more attractive

  gamma = small fixed step (bias update speed)
  target = T*k / N   (even share of the k*T dispatch slots)

Overloaded experts get their bias nudged down, so on the next step they win fewer top-k slots; starved experts get nudged up and win more. It's a feedback controller on selection. No second loss, no gradient interference — the LM gradient stays pure, and balance is handled by a non-gradient control loop running alongside.

A tiny safety net

DeepSeek-V3 keeps one very small sequence-wise auxiliary loss with a minuscule coefficient — not to drive balance (the bias does that), but to discourage extreme imbalance within a single sequence, which the batch-level bias controller can miss. It's a guardrail, not the mechanism. The heavy lifting is the bias.

Why this matters

• Quality. Removing the balancing gradient removes the per-step quality tax. DeepSeek reports better performance at equal balance versus the aux-loss approach.
• Fewer knobs. One gamma and a target replace three brittle, interacting alphas.
• Stability. A control loop on realised load is more robust than hoping a gradient coefficient holds across a long run.

The lesson generalises: when a regulariser fights your real objective, ask whether the thing it controls actually needs to be in the gradient. Often — as with expert selection — it doesn't, and a cheap control loop does the job cleaner.

Next, 5.5 looks at the other half of balancing in practice: capacity factor — what happens to tokens when an expert's buffer is full, why dropped tokens exist, and how to size the buffers.

References

• Shazeer et al. (2017), Outrageously Large Neural Networks — original sparsely-gated MoE and load-balancing loss. arXiv:1701.06538
• Fedus et al. (2021), Switch Transformers — the f_i · P_i auxiliary-loss formulation used here. arXiv:2101.03961
• Lepikhin et al. (2020), GShard — load balancing at scale with expert parallelism. arXiv:2006.16668
• DeepSeek-AI (2024), DeepSeek-V2 — expert/device/communication balancing losses. arXiv:2405.04434
• Wang et al. (2024), Auxiliary-Loss-Free Load Balancing — the bias-based scheme. arXiv:2408.15664
• DeepSeek-AI (2024), DeepSeek-V3 — aux-loss-free balancing in a production model. arXiv:2412.19437