Load Balancing › Expert Collapse

The most severe manifestation of load imbalance is expert collapse, a failure mode where the model converges to using only one or two experts for nearly all tokens. It breaks the sparse-computation promise of MoE. When collapse occurs, the model has effectively reverted to a dense architecture, paying the memory costs of multiple experts while receiving the computational benefits of only one or two. The ability to scale model capacity without proportionally scaling computation is lost.

Expert collapse is one of those failure modes that seems surprising at first but makes perfect sense once you understand the underlying dynamics. The architecture is designed to be flexible, and this flexibility, without constraints, allows it to find a degenerate solution where the gating network simply routes everything to the best-performing expert and ignores the rest. From the optimizer's perspective, this is a perfectly valid strategy: if one expert can handle all the tokens adequately, why maintain eight? The problem is that this strategy violates the whole point of having eight experts in the first place.

Expert collapse occurs when the rich-get-richer dynamics run to their logical extreme. The process typically unfolds through a predictable sequence of stages, each building on the previous one:

1. Early training: The gating network's random initialization creates slight expert preferences. These initial biases are typically small and might seem insignificant, but they provide the seed from which imbalance can grow.
2. Divergence: Preferred experts receive more tokens and more gradients, improving faster. The gap between experts begins to widen, though it may still be subtle enough to escape notice in standard training metrics.
3. Reinforcement: The gating network learns to rely more heavily on better experts. At this stage, the feedback loop is firmly established and accelerating. The router's preferences become increasingly pronounced.
4. Saturation: A small subset of experts handles nearly all tokens. The model has effectively reorganized itself around a handful of dominant experts, with others receiving only occasional, often random, token assignments.
5. Collapse: Non-preferred experts stop receiving training signal entirely. The process reaches its endpoint, where certain experts are completely excluded from the model's functioning.

Once collapsed, recovery is nearly impossible through normal training. The ignored experts become stale: their weights remain at initialization values or drift randomly without meaningful gradients. The gating network has learned to avoid them, so they never get the chance to improve. This creates a stable but pathological equilibrium where the model is stuck with reduced capacity and no path back to balanced utilization.

Think of expert collapse as analogous to a market monopoly. Once a dominant player captures enough market share, the remaining competitors cannot attract enough business to invest in improvements, which causes them to fall further behind, which causes them to lose even more business. The end state, where one player controls everything, is self-sustaining even if it is inefficient for the overall system. Breaking out of this equilibrium requires external intervention, just as MoE training requires explicit balancing mechanisms.

Collapse can be detected by monitoring expert utilization during training. A healthy MoE layer shows all experts receiving tokens in roughly equal proportions, with natural fluctuations based on the content of each batch. A collapsing layer shows increasing concentration in a small number of experts, with the trend persisting across batches rather than averaging out over time.

The warning signs include:

• Monotonically increasing utilization for a subset of experts across consecutive batches
• Near-zero routing probabilities for some experts that persist and worsen over time
• Decreasing entropy in the router's softmax outputs, measured as a trend rather than momentary fluctuation
• Sharp drops in validation performance as the model loses access to specialized expert knowledge
• Increasing gradient norms for popular experts and decreasing norms for unpopular ones

The important distinction is between natural variation and systematic trend. In any given batch, some experts will naturally receive more tokens than others simply because the batch content favors certain topics. This per-batch variation is healthy and expected. Collapse is characterized by trends that persist and amplify across many batches. Monitoring tools should compute rolling averages over many steps to distinguish signal from noise.

Catching collapse early is important because intervention becomes harder as the process advances. In the divergence stage, simple adjustments to the auxiliary loss coefficient or learning rate may be sufficient. By the saturation stage, the model may require restarting from an earlier checkpoint with modified hyperparameters. At full collapse, recovery through continued training is practically impossible.

The softmax operation in the gating network contributes to collapse in a basic way that makes this failure mode particularly difficult to avoid. Recall from the gating networks chapter that router scores pass through a softmax to produce routing probabilities.

To understand this quantitatively, consider how softmax transforms a set of scores. The softmax function takes a vector of real-valued scores and converts them into a probability distribution, where all values are positive and sum to one:

$$g_i = \frac{e^{s_i}}{\sum_{j=1}^{E} e^{s_j}}$$

where:

• $g_i$: the routing probability assigned to expert $i$
• $s_i$: the raw score (logit) for expert $i$
• $E$: the total number of experts
• $e^{s_i}$: the exponential function applied to the score, so a positive value
• $\sum_{j=1}^{E} e^{s_j}$: the normalizing sum over all $E$ experts

Why does this formula make sense as a routing mechanism? Notice that the exponential function is always positive (no matter how negative the score), ensures that all outputs are non-negative, and that dividing by the sum guarantees outputs sum to one. It is a clean way to turn arbitrary real-valued scores into valid probabilities.

This formula reveals why collapse is self-reinforcing. The exponential function at the heart of softmax has a powerful sharpening effect on probability distributions. Small differences in raw scores $s_i$ get amplified by the exponential, and larger differences get amplified even more. To understand this intuitively, consider that if one score increases by just 1 unit, its contribution to the numerator increases by a factor of approximately 2.7 (since $e^1 \approx 2.718$). This multiplicative effect means that even modest advantages in raw scores translate into substantial advantages in routing probability.

If expert $E_1$ consistently has slightly higher scores than others, the softmax concentrates probability mass on $E_1$, reducing training signal to other experts. As the gap grows over successive training steps, the softmax's sharpening effect accelerates collapse. What starts as a small preference can rapidly snowball into complete dominance. This is why monitoring metrics like router entropy, which we discuss shortly, provides valuable early warning signals before collapse becomes irreversible.