MoE Load Balancing and Gradient Interference

Whether the balancing signal enters the loss function or the routing score is what decides the training quality of a MoE model.

Top-K routing and routing collapse

A sparse Mixture-of-Experts (MoE) layer splits a single feed-forward block into $N$ experts. A gate function scores every expert for each token, activates only the top $K$, and skips the computation for the rest. The design keeps the total parameter count while cutting the per-token compute.

The trouble is that this selection reinforces itself. When one expert keeps getting picked, only that expert receives gradient and improves, and the improvement makes it get picked even more often. The remaining experts are left out of training, which ends in routing collapse, where a small subset does all the work.

A second cost appears on the training infrastructure side. Distributed training caps how many tokens each expert may process, and tokens routed to an overloaded expert are dropped once that capacity is exceeded. Load skew is therefore a quality problem and a wasted-compute problem at the same time.

One term needs separating here. Load balancing in this context means distributing tokens across experts, and it shares only the name with L4/L7 network load balancing in front of a server. The thing being balanced differs, and so does the way the signal is applied.

The interference gradient from auxiliary loss

The standard remedy since Switch Transformer is an auxiliary loss. It multiplies the selection frequency $f_i$ of expert $i$ by the average gating score $p_i$ assigned to that expert. The products are summed over experts and added to the main loss with a coefficient $\alpha$.

$\mathcal{L}_{aux} = \alpha \, N_E \sum_i f_i p_i$

The value grows as frequency and score concentrate on the same expert. Reducing this term therefore pushes toward an even distribution, and the single coefficient $\alpha$ sets how hard that push is.

The problem is that one coefficient holds two objectives at once.

The root cause is that the auxiliary loss enters the gradient computation of the routing scores directly. Backpropagation cannot separate the signal coming from next-token prediction and the signal coming from load balance, so it sends them through together. Wang, Gao, Zhao, Sun, and Dai (arXiv:2408.15664) call this mixed-in component the interference gradient.

Changing only the selection with a bias

Loss-Free Balancing, proposed in the same paper, removes the loss term and moves the point of intervention. It keeps one scalar bias $b_i$ per expert and adds it to the routing score only when choosing the top $K$. Once the selection is made, the output is weighted with the original score, without the bias.

# every step (or every batch)
for expert i in 1..N:
    e_i = c_avg - c_i          # average load minus the actual load of expert i
    b_i = b_i + u * sign(e_i)  # u: update rate (0.001 recommended in the paper)
 
# score used only for the top-K selection
g_i,t = s_i,t + b_i            # s_i,t: original routing score
 
# the weighted output uses s_i,t, without the bias

Because the bias touches selection alone, the gradient flowing into the model parameters stays intact. An expert carrying less than the average load gets a higher bias and a better chance of being picked on the next step, and an overloaded expert moves the other way. The balancing signal shifts the selection boundary instead of passing through the loss.

The update is also cheap. It adds one value per expert on each step, so it costs $O(N)$ per batch, less than the $O(N \times T)$ of an auxiliary loss that scales with the token count $T$.

The controlled experiment in the paper reads two metrics together. MaxVio is the normalized deviation between the most loaded expert and the average load, and lower is better balanced.

At both scales perplexity dropped while MaxVio fell by more than 90%. The relationship in which balance had to be bought with quality disappeared entirely.

Why the sign update converges

The original paper verified only experimentally that a rule defined by a single sign does not diverge. Han and Zhong (arXiv:2512.03915) formalize token-to-expert assignment as an integer program (IP) and identify what the rule actually is. The objective is $\max \sum_{i,k} \gamma_{ik} x_{ik}$. The constraints are that each token goes to exactly $K$ experts and each expert takes at most $L$ tokens.

Relaxing the expert capacity constraint with a dual variable $p_k$ gives the following Lagrangian.

$\mathcal{L}(x,p) = \sum_{i,k}(\gamma_{ik}+p_k)x_{ik} - L\sum_k p_k$

Here $\gamma_{ik}+p_k$ sits in exactly the same position as a routing score with the bias added. The dual variable is updated by a rule that raises it when the load falls short of capacity and lowers it when the load exceeds capacity.

$p_k^{(n+1)} \leftarrow p_k^{(n)} + \varepsilon_k^{(n)}(L - A_k^{(n)}), \quad \varepsilon_k^{(n)} = \frac{u}{|L - A_k^{(n)}|}$

$A_k^{(n)}$ is the number of tokens expert $k$ received on step $n$. The step size cancels the absolute value of the deviation, so only the sign and $u$ remain in the update. That is exactly the sign update from DeepSeek.

The same paper draws three guarantees out of this correspondence.

The stochastic online setting relies on the strong convexity of the dual objective. Shrinking the step size to $u/n$ yields a proven $O(\log T)$ bound on expected regret. The trick found empirically turns out, in hindsight, to be standard dual ascent on an assignment problem.

Micro-batch balance and expert specialization

The unit at which load is measured matters as much as where the balancing signal is injected. Qiu et al. (arXiv:2501.11873) point out that most MoE training frameworks compute the load balancing loss per micro-batch for the convenience of parallel execution. A micro-batch in a large model holds very few sequences.

With few sequences, that loss effectively enforces sequence-level balance. Even the tokens of a sequence filled entirely with code have to spread evenly across all experts for the loss to go down. Specialization, where experts diverge by domain, is blocked at exactly this point.

The fix is to widen the measurement unit. Synchronizing expert selection frequencies across several micro-batches and computing the loss per global batch lifts the balance requirement to the level of the whole batch. Since only a batch mixing many domains has to come out even, an individual sequence is free to concentrate on particular experts. Selection frequency is a single vector with one element per expert, so the synchronization traffic is close to nothing.

In experiments from 3.4B (0.6B active) to 43B (6.6B active), the global-batch approach improved both pre-training perplexity and downstream results. The gain grew as the balance batch size went from 2 to 128.

Trade-offs across the three approaches

Putting the three approaches side by side exposes the selection criteria.

The two axes are not exclusive. Whether the signal comes from a bias or a loss term, and whether load is measured per micro-batch or per global batch, are separate choices. Recent open source MoE models converge on using the loss-free bias update together with global-batch statistics.

The boundary of these methods is equally clear. Everything above concerns balancing during training while keeping the gradient clean. A deployed model runs into the hot expert problem again, where requests pile onto certain experts according to the prompt distribution of real service traffic. Skew at inference time belongs to serving frameworks such as vLLM and TensorRT-LLM, through token-level expert grouping and all-to-all communication optimization.

Summary

The point of the loss-free approach is to shift the focus of load balancing from the loss function to the routing selection. An auxiliary loss mixes the balancing signal into backpropagation, so the single coefficient $\alpha$ trades balance against quality. The bias update touches only the top-K selection and leaves the gradient alone.

In the 3B/200B token experiment by Wang et al., perplexity fell from 7.97 to 7.92 and MaxVio from 0.52 to 0.04 at the same time. The formalization by Han and Zhong shows that this sign update is equivalent to dual ascent on the token-to-expert assignment problem, giving the heuristic a convergence argument. The unit at which load is measured remains a separate choice, and expert skew at serving time is still a problem for the serving framework to handle.