**Linwei Li, Hongye Jin, Binxuan Huang, Xiaotian Han, Xin Liu**

<aside>

In this post, we show how router $z$-loss — originally introduced as a numerical-stability fix — is widely mis-interpreted as a MoE load-balancing / stability tool, and how, at large coefficients, it actively drives catastrophic last-layer expert collapse during sparse MoE training. We also uncover a hidden second side of MoE load balancing — under-specialization in early layers — and show how the same $z$-loss that breaks the last layer accidentally helps the first ones. Along the way, we introduce a practical new diagnostic, $\text{MaxVio}_{\text{mass}}$ — a combine-weight-share variant of the standard MaxVio that uses the same formula but aggregates per-expert combine mass instead of token counts, and that catches the silent failure mode the count-based MaxVio is blind to.

</aside>

This study is based on a series of experiments with ~50B-parameter MoE models with 1B activated parameters. We believe the insights are more likely to transfer to modern ultra-sparse larger-scale MoE training.

1. Causes of expert collapse

When people talk about expert collapse in a Mixture-of-Experts (MoE) model, the picture is usually the same: a routing pathology in which “the model consistently selects only a few experts, hindering sufficient training of the other experts” (Wang et al., 2024 — DeepSeek aux-loss-free, restating the classical observation from Shazeer et al., 2017). The standard cure follows directly from the standard diagnosis: if collapse means concentration, then push back on concentration — with an auxiliary load-balancing loss, with an auxiliary-loss-free expert bias, with $z$-loss to keep logits bounded, and so on.

The dominant metric for “is the router collapsing?” is MaxVio, introduced alongside the aux-loss-free method (Wang et al., 2024, §4.1):

$$ \text{MaxVio} = \frac{\max_i \text{Load}_i - \overline{\text{Load}}}{\overline{\text{Load}}} $$

where $\text{Load}_i$ is the number of tokens routed to expert $i$ and $\overline{\text{Load}}$ is the mean load. It is a one-sided, token-count statistic that asks one question: is there an expert that gets far more tokens than the average? Every closely related metric on the standard MoE dashboard — the aux loss $\alpha \sum_i f_i P_i$ (Fedus et al., 2021, Lepikhin et al., 2020), router entropy $H(p)$, coefficient of variation of load, token drop rate — is a different way of asking the same question.

That question about load-balancing is real, but it is not the whole question.

Two failure modes, not one

In this post we make a sharper claim, drawn from a series of MoE diagnostic studies we ran on 512-expert, ultra-sparse, 50B MoEs (details in later sections). Expert collapse, on closer inspection, decomposes into at least two distinct failure modes that look nothing alike under the hood:

1. Over-specialization (the classical mode). A small set of experts wins the routing competition and absorbs a disproportionate share of tokens. The other experts are starved of gradient signal and never learn. This is what MaxVio sees, what $\sum_i f_i P_i$ aux-loss penalizes, and what aux-loss-free bias corrects.
2. Under-specialization (the silent mode). Token counts can remain perfectly balanced: every expert receives its fair share of tokens, MaxVio stays near zero, and the load-balancing dashboard looks healthy. Yet the experts may still fail to develop meaningful functional differentiation. Their parameters may appear dissimilar because of random initialization and noisy early updates, but this dissimilarity does not translate into stable routing preferences. Instead, multiple experts respond to similar token populations, produce weak or low-attribution outputs, and are often dominated by the shared experts or the residual stream. In this regime, capacity is not wasted through starvation, but through redundancy.

These two failure modes are not uniformly distributed across depth. Empirically, with AdamW optimizer, early layers are more prone to under-specialization (shallow routers are weak, hidden states are not yet diverse enough to give the router a strong signal to specialize on), while later layers are more prone to over-specialization (the rich get richer until a minority of selective experts dominate the combine-weight share). The depth split is one of the cleanest signatures in our experiments, and shows up again whenever we look at any depth-resolved metric.

Evidence: per-token routing geometry across depth

The figure below is from a per-token routing dump on a 512-expert, top-8 softmax-routed MoE checkpoint at step 40k steps at 16B tokens with only seq-wise aux loss.

Figure 1. Per-token routing geometry across depth on a 512-expert, top-8 softmax-routed MoE at step 40k (bsz=512 ctx=8192), trained with no router $z$-loss and only a sequence-wise aux load-balance loss ($\alpha_{\text{seq}} = 10^{-3}$).

Top-k margin (left) = $r_{\text{top-}1} - r_{\text{top-}k}$ per token — the raw-logit gap between the strongest expert and the weakest of the 8 selected. It measures how confidently the router picks its top-k. A small margin means the top-1 and the 8th-ranked expert are nearly tied: the router has no strong preference, the chosen 8 are almost interchangeable with neighboring candidates, and any expert in that pool is functionally redundant. A large margin means the router clearly distinguishes its preferred experts from the rest.

MaxVio_count vs MaxVio_mass (right) — these two curves use the same formula, $(\max - \text{mean}) / \text{mean}$, applied to two different per-expert quantities: