Train an MoE layer with nothing but the language-modeling loss and watch the router. For a few dozen steps the tokens spread out over the experts. Then one expert pulls ahead by a little, by luck of initialization, so it receives a few more tokens, so it trains a little more, so it gets a little better, so the router sends it even more. A few dozen steps later it handles almost everything, and the other experts sit idle. You have paid for eight experts and you are training one.

This is expert collapse, and it follows directly from the router gradient of §13.2, which only rewards experts the router already chose. This section explains why it happens, derives the auxiliary loss that the Switch Transformer used to prevent it and proves where that loss is minimized, adds the router z-loss that keeps logits from blowing up, builds the auxiliary-loss-free bias method of DeepSeek-V3, and ends with what to monitor during a real run.

Why routers collapse

Recall the router gradient for one token with chosen set $\mathcal{K}$, gates $g_i$, and $a_i = \partial\mathcal{L}/\partial g_i$,

Two facts combine into a feedback loop. An expert that is chosen more often receives more gradient updates, so it improves faster and its outputs become more useful, which makes $-a_j$ larger and pushes its logit up. And an expert that is never chosen receives no router gradient at all, so nothing in the task loss can ever tell the router it would have been useful. Rich experts get richer, and poor ones never get a chance. Shazeer et al. (2017) observed exactly this, a gating network that converges to large weights for the same few experts, an imbalance that reinforces itself, and the DeepSeek-V3 report calls the end state routing collapse.

Collapse costs more than wasted parameters. A collapsed expert overflows its capacity, so tokens are dropped (§13.2). In expert-parallel training, the GPU holding it does most of the work while the others wait (§13.4). And the model degenerates toward a dense model the size of one expert, with all the machinery of an MoE and none of the benefit.

Play with the simulation below before reading on. It trains a toy router whose experts improve with the tokens they receive, which is the feedback loop above in miniature, and lets you compare three ways of fighting it.

In the toy, a small auxiliary loss only delays the collapse, because the task signal is strong, and a larger coefficient is needed. Bias balancing keeps the load even from the start, as long as its step $\gamma$ is large enough to keep up with the pull, and it does so without the router's weights ever being pushed toward balance, which is the property DeepSeek wanted. We will see both reasons in the math.

The Switch auxiliary loss

Fedus et al. (2022) add a single term to the loss, for a batch of $T$ tokens routed among $E$ experts,

where $\mathbf{p}_t = \operatorname{softmax}(\mathbf{r}_t)$ is the router's full softmax over all $E$ experts, $f_i$ is the fraction of token assignments that went to expert $i$ (the Switch paper has $k = 1$, and dividing by $k$ makes $\sum_i f_i = 1$ for any $k$), $P_i$ is the average probability the router gave expert $i$, and $\alpha$ is a small coefficient, $10^{-2}$ in Switch and ST-MoE. In plain words, multiply each expert's share of the work by its share of the router's belief, and add those up.

The two factors play different roles. The count $f_i$ is a step function of the logits, so its gradient is zero almost everywhere, and the loss treats it as a constant. All of the gradient flows through $P_i$, which is smooth. That is the trick. The non-differentiable thing we want to balance (the actual load) sets the weights, and the differentiable proxy (the probabilities) is what moves.

Where the loss is minimized

Start with the value at perfect balance. If every expert gets $f_i = 1/E$ and the router's average probabilities are also uniform, $P_i = 1/E$, then

The factor $E$ in front is there so this value does not depend on the number of experts. At full collapse with $k = 1$, every token goes to expert 1 with probability near 1, so $f_1 = P_1 = 1$ and the loss is $\alpha E$. With top-$k$ collapsed onto $k$ experts that share the probability evenly, $f_i = P_i = 1/k$ on those $k$, and the loss is $\alpha E k/k^2 = \alpha E/k$.

Why is uniform the minimum? Look at the loss in the regime it is designed for, where the router's probabilities match its actual routing, $P_i = f_i$. Then $\mathcal{L}_{\text{aux}} = \alpha E\sum_i f_i^2$, and with $\sum_i f_i = 1$, the Cauchy-Schwarz inequality gives

with equality exactly when all $f_i$ are equal. In plain words, the sum of squared shares is smallest when the shares are equal, so the loss bottoms out at uniform routing and grows as load concentrates.

The gradient tells you how it gets there. Differentiate through $P_i$ only, using the softmax Jacobian $\partial p_{t,i}/\partial r_{t,j} = p_{t,i}(\delta_{ij} - p_{t,j})$,

In plain words, for every token, the loss lowers the logit of each expert whose load is above the token's probability-weighted average load and raises the logits of experts whose load is below it, in proportion to how much probability the token gave them. Overloaded experts become less attractive to every token at once.

Older losses have the same goal. Shazeer et al. (2017) penalized the squared coefficient of variation of each expert's total gate weight (its importance), plus a smooth estimate of its load built from the noisy gating's normal CDF. GShard's loss is $\frac1E\sum_e \frac{c_e}{S}m_e$, where $c_e/S$ is the fraction of the group's $S$ tokens whose first choice is expert $e$ and $m_e$ is its mean gate, so up to the coefficient it is the Switch loss divided by $E^2$. DeepSeek writes $f_i$ scaled by $E/(kT)$ so that it equals 1 at balance, which moves the factor $E$ inside $f$ and leaves the loss unchanged.

import numpy as np

def switch_aux_loss(R, idx, alpha):
    T, E = R.shape
    p = np.exp(R) / np.exp(R).sum(0, keepdims=True)        # placeholder: softmax over the wrong axis
    f = np.bincount(idx.ravel()) / T
    return alpha * (f * p.mean(0)).sum()

How large should $\alpha$ be? Too small and collapse wins, as in the widget. Too large and the balancing gradient fights the language-modeling gradient on every token, which costs quality. Switch swept $\alpha$ from $10^{-1}$ to $10^{-5}$ and chose $10^{-2}$ as large enough to balance without overwhelming the main loss. That tension, a hyperparameter whose right value depends on the run, is what motivated the bias method below.

The router z-loss

Balancing fixes where tokens go. A second failure is about how confident the router is. Router logits tend to grow during training, and large logits are dangerous in low precision. A bfloat16 number has 8 significant bits (§8.5), so near 100 consecutive representable values are 0.5 apart, and near 1,000 they are 4 apart. Two experts whose logits differ by 0.3 become a tie, or a gap of 0.5, depending on rounding. The exponentials inside the softmax amplify such roundoff, and Zoph et al. (2022) found that this was a source of training instability in large MoE models.

The probabilities depend only on the gaps between logits, yet at height 1,000 bfloat16 has erased the gaps entirely. ST-MoE's fix is the router z-loss,

with $\alpha_z = 10^{-3}$, the value Zoph et al. chose by a sweep. In plain words, penalize the size of each token's log-partition function, which keeps the logits near zero. It is the same leash that PaLM's z-loss puts on the output softmax (§8.8), applied to the router.

Why the log-sum-exp and not the logits themselves? Because adding a constant $c$ to all of a token's logits leaves the softmax unchanged but shifts $\operatorname{lse}$ by exactly $c$. So the z-loss can pull the overall level of the logits toward zero without touching the routing. It also penalizes a large spread, since $\operatorname{lse}(\mathbf{r}) \ge \max_j r_j$. For example, logits $(10, 9, 8, 0)$ have $\operatorname{lse} = 10.41$ and a squared value of 108.3, while the same logits shifted down by 8, $(2, 1, 0, -8)$, route identically with $\operatorname{lse} = 2.41$ and a squared value of 5.8. The gradient is simple,

because $\partial\operatorname{lse}/\partial r_j = p_j$. Each logit is pushed down in proportion to its probability and to how large the partition function is. In ST-MoE's stability study, the z-loss made all three runs stable and slightly improved quality, where alternatives such as update clipping also stabilized training but badly hurt quality.

import numpy as np

def router_z_loss(R, alpha_z):
    T = R.shape[0]
    lse = np.log(np.exp(R).sum(1))                  # overflows for large logits
    return float(alpha_z * (lse ** 2).mean()), np.zeros_like(R)

Balancing without a loss

Every auxiliary loss has the same weakness. Its gradient flows into the router weights, the same weights that the language-modeling loss is training, and on every token the two gradients disagree to some degree. Wang et al. (2024) call this interference. A large $\alpha$ balances well and hurts perplexity, and a small one hurts less and balances badly.

Their alternative changes the routing decision without changing the gradient. Give every expert a bias $b_i$ that is added to its score only for choosing the top $k$, not for the gate, and adjust the biases by a fixed rule after each step. In DeepSeek-V3's form, with sigmoid affinities $s_{t,i} = \sigma(r_{t,i})$,

and after each training step, with $c_i$ the number of tokens expert $i$ received and $\bar c$ the mean over experts,

In plain words, an overloaded expert's bias goes down by $\gamma$ and an underloaded expert's goes up by $\gamma$, which nudges the next batch's selections toward balance, while the gates, and therefore the gradients, still come from the unbiased scores.

The bias is not a parameter of the model in the gradient sense. Nothing backpropagates into it, so balancing creates no interference with the main loss. It acts like a price in a market. Busy experts get more expensive and idle ones cheaper until supply and demand match. DeepSeek-V3 used $\gamma = 0.001$ for the first 14.3T training tokens and $\gamma = 0$ for the last 500B, and reports dropping no tokens at all, in training or inference.

The evidence is quantitative. Wang et al. trained 1B and 3B MoE models with 64 routed experts, 2 shared, and 6 active, and measured imbalance with $\operatorname{MaxVio} = (\max_i c_i - \bar c)/\bar c$, how far the busiest expert is above the fair share. With the auxiliary loss, the 1B model reached perplexity 9.56 with a MaxVio of 0.72 over the validation set. Bias balancing reached 9.50 with a MaxVio of 0.04. At 3B the numbers were 7.97 and 0.52 versus 7.92 and 0.04. Better balance and better perplexity at once is the signature of removing interference. DeepSeek-V3's own ablations, at 15.7B and 228.7B total parameters, show the auxiliary-loss-free model ahead on most benchmarks, though not all (at the larger size its MMLU was 67.2 against 68.3).

DeepSeek-V3 does keep one tiny auxiliary loss, a sequence-wise balance term with $\alpha = 10^{-4}$, to stop any single sequence from piling onto a few experts. The paper also compares where balance is enforced. Requiring balance within every sequence is stricter than requiring it over a whole batch, and it fights specialization (a sequence of code should be allowed to use the code-heavy experts). In their 1B ablation, a sequence-wise auxiliary loss reached validation loss 2.258, while bias balancing and a batch-wise auxiliary loss both reached 2.253, and the bias-balanced model showed clearer domain specialization of experts.

The same question hides inside every ordinary auxiliary loss. Frameworks usually compute $f_i$ and $P_i$ within each micro-batch on each GPU, and a micro-batch of a few long sequences is close to the strict sequence-wise case. Qiu et al. (2025) synchronized the counts $f_i$ across all data-parallel ranks, so that the loss balances the global batch instead, and reported lower perplexity and more domain-specialized experts at no extra communication worth mentioning (one vector of $E$ numbers per layer).

import numpy as np

def bias_route_and_update(S, b, k, gamma):
    adj = S + b
    idx = np.argsort(-adj, axis=1, kind="stable")[:, :k]
    g = np.take_along_axis(adj, idx, 1)                   # placeholder: gates use the biased scores
    g = g / g.sum(1, keepdims=True)
    c = np.bincount(idx.ravel(), minlength=S.shape[1])

Expert collapse, and its cousins

"Collapse" covers a few related failures, and it helps to name them separately.

• Load collapse. A few experts take most tokens, as above. Balancing losses and biases are the cure.
• Dead experts. An expert that stops being chosen gets no gradient, its weights freeze, and it can never recover through the task loss alone. A balancing mechanism is the only thing that brings it back, by raising its score until it is tried again.
• Representation collapse. Chi et al. (2022) argue that the router's training pulls token representations toward clusters around the experts' router vectors (the columns of $\mathbf{W}_r$), which reduces what the representation can express. Their fix scores tokens against experts on a low-dimensional hypersphere.
• Redundant experts. Experts are all used but learn nearly the same function, which wastes capacity without any imbalance. This is one motivation for the shared experts of DeepSeekMoE (§13.5), which take the common knowledge so routed experts do not each have to store it.

Several training choices help prevent collapse before balancing has to. Switch initializes weights with one tenth of the usual scale, and computes the router in float32 even when the rest of the model runs in bfloat16 (§13.6). A router whose logits start near zero routes almost uniformly at first, which gives every expert some tokens and some gradient.

Monitoring expert usage

In a real run you watch routing the way you watch the loss. These are the numbers worth logging per MoE layer, over a window of steps.

The load entropy $H(\mathbf{f}) = -\sum_i f_i\log f_i$ divided by $\log E$ is 1 for perfect balance and 0 for total collapse. Different layers often behave differently, so log per layer, not just an average. And measure on held-out data too, since routing on a validation set can differ from routing on the training stream (Wang et al. report MaxVio both ways).

import numpy as np

def routing_stats(idx, E, cap):
    c = np.bincount(np.asarray(idx).ravel(), minlength=E)
    f = c / c.sum()
    return {"load": f, "maxvio": 0.0, "entropy": 1.0, "drop": 0.0}

In PyTorch

In PyTorch the auxiliary loss is three lines, and autograd confirms the gradient we derived, flowing only through $P_i$ because $f_i$ comes from integer counts. The z-loss uses torch.logsumexp, which is stable.

Read more

• Fedus, Zoph & Shazeer (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Section 2.2 defines the auxiliary loss, and the appendix compares router noise options.
• Zoph et al. (2022). ST-MoE: Designing Stable and Transferable Sparse Expert Models. The router z-loss, a systematic stability study, and practical recommendations.
• Wang et al. (2024). Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts. The bias method, MaxVio, and the interference argument.
• DeepSeek-AI (2024). DeepSeek-V3 Technical Report. Bias balancing at scale, the sequence-wise loss, and the batch versus sequence ablation.
• Shazeer et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. The importance and load losses, in Section 4 and Appendix A.
• Chi et al. (2022). On the Representation Collapse of Sparse Mixture of Experts. How routing can collapse representations, and a hypersphere router that avoids it.
• Qiu et al. (2025). Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models. Why the batch over which balance is measured matters, and global-batch balancing.

Watch next

▶

▶