The Routing Collapse Problem

In the previous chapter we built a Mixture-of-Experts block, sliced the experts fine, and added a small always-on shared pool. The promise was intoxicating: hundreds of experts, only a handful active per token, ten times the parameters at the same active compute. There is one problem nobody warned you about, and it is the single most common reason MoE training runs are quietly abandoned. The router collapses. A handful of experts win every token; the rest go dark, never receive gradient, and become dead weight that consumes memory and bandwidth while contributing nothing. By the time you notice, you have wasted a week of H800 time.

The collapse, in one sentence: a top-$k$ router with no balancing pressure is a positive-feedback system. Whatever expert wins a token early gets more gradient, becomes more attractive, wins more tokens, gets more gradient. The dynamic is exponential, not gradual.

The Real Problem

Picture the first 1000 steps of an MoE training run with 64 experts and top-2 routing. At step 0 the router weights are random — every expert has roughly an equal chance of being picked. By step 50, tiny accidents have nudged the router: expert 17 happened to win a few extra tokens this minute, nothing dramatic. By step 500, expert 17 is handling 8% of all tokens. By step 1000, expert 17 and three friends are handling 70% of the corpus and the other 60 experts have effectively stopped learning. The loss curve still goes down — slowly — because the four winners are fitting harder. Most of the model is dead.

This is not a numerical glitch. It is the natural behaviour of an unbiased top-$k$ router. The mechanism is the same one that produces winner-take-all markets, viral cascades, and the rich-get-richer dynamic of preferential attachment. MoE happens to wire it directly into the compute graph of your trillion-parameter model.

What the failure looks like in production

The Feedback Loop That Eats the Router

The cleanest way to see why collapse is inevitable is to trace the feedback loop one cycle at a time. Pretend there are two experts, A and B, and at step 0 the router has a microscopic preference for A — say A wins 51% of tokens, B wins 49%.

1. A gets more tokens. More tokens means more gradient signal reaching A's weights — A learns faster than B this step.
2. A becomes more useful. Because A has learned more, A contributes more usefully to the loss when it is picked. The router's gradient now flows in the direction that picks A more strongly: the router weight for A grows.
3. A becomes more attractive. A larger router weight for A means a larger softmax score for A on the next batch — A wins even more tokens than 51%.
4. Repeat. Each cycle of the loop amplifies the asymmetry. What started as 51/49 becomes 60/40, then 80/20, then 99/1 — and B is functionally dead.

There is no opposing force. Nothing in the standard top-$k$ router penalises an expert for winning too much. Every feedback path points the same way. Without an intervention, this runs to completion.

The hospital analogy, broken

In Chapter 5 we modelled MoE as a hospital with a triage desk. Routing collapse is what happens when the triage nurse keeps sending every patient to the cardiologist because the cardiologist saw the most patients yesterday and is therefore "most experienced." The dermatologist sees no patients, learns nothing new, and the gap to cardiology widens. By the end of the year the hospital has one extraordinarily skilled cardiologist and a building full of forgotten specialists. The hospital's budget bought nine expert salaries and got the output of one.

The Mathematics of Collapse

Let $E$ denote the number of experts, $k$ the top-k value, and $N$ the number of tokens in a batch. Define the load of expert $i$ as the number of tokens dispatched to it:

$\ell_i = \sum_{t=1}^{N} \mathbb{1}\big[ i \in \mathrm{TopK}(W_r h_t, k) \big]$

where $h_t$ is the hidden state of token $t$ and $W_r \in \mathbb{R}^{E \times d}$ is the router. The fair share is $\bar{\ell} = N k / E$. The standard imbalance metric is the coefficient of variation:

$\mathrm{CoV}(\ell) = \frac{\sigma(\ell)}{\mu(\ell)} = \frac{\sqrt{\tfrac{1}{E}\sum_i (\ell_i - \bar{\ell})^2}}{\bar{\ell}}$

Perfect balance gives $\mathrm{CoV} = 0$. A single-expert monopoly gives $\mathrm{CoV} = \sqrt{E - 1}$ — for $E = 256$, that is about 16. The practical alarm threshold in production MoE training is $\mathrm{CoV} > 0.5$.

Why the dynamic is exponential, not linear

Let $p_i$ be the probability that a random token picks expert $i$ in its top-$k$. At step $s$, the gradient update to the router weight for expert $i$ is approximately proportional to $p_i \cdot g_i$ where $g_i$ is the per-token usefulness gradient. To a first approximation the bias evolves as:

$b_i^{(s+1)} = b_i^{(s)} + \eta \cdot p_i^{(s)} \cdot \bar{g}$

Because top-k is monotone in the bias and softmax is monotone in the score, $p_i^{(s+1)}$ is an increasing function of $b_i^{(s+1)}$. Substitute: experts with larger $p_i$ grow their bias faster, which grows their $p_i$ on the next step. This is a discrete-time analogue of $\dot{x} = \alpha x$ — pure exponential amplification. Small initial asymmetries blow up; they do not average out.

Manual Numerical Walkthrough

Let us follow the feedback loop with explicit numbers for a four-expert, top-1 router over six steps. Every value is computed by hand.

Visualizing a Routing Collapse

Below is the four-step dynamic from the walkthrough, except generalized to 8 experts and top-2 routing, with 64 tokens per step. Press Play with Balancing OFF and watch the load distribution evolve. Within 20 to 40 steps the bars on the right go red — those are dead experts. The CoV trace at the bottom climbs past the alarm threshold of 0.5 within a few seconds of real time. Then toggle Balancing ON, reset, and play again — same initial conditions, completely different outcome.

Three behaviours to watch for. First, the collapse is non-monotonic in the short term — individual steps can bounce around — but the trend is unmistakable within a few dozen steps. Real training runs see the same shape over thousands of steps. Second, once an expert goes red (dead), it stays dead. There is no recovery path without external intervention. Third, the Top-2 Share metric is the earliest reliable warning: it climbs measurably before any individual expert dies, giving you time to act if your training infra is watching.

Plain Python: Simulating the Collapse

Here is the same dynamic in NumPy, stripped to the bone. No router weight matrix, no token embeddings, no transformer — just the feedback loop in isolation. If you can read this 30-line program you understand the entire problem the rest of this chapter is built to solve.

After step 0 with the seeded RNG, load may look like [9, 11, 16, 7, 14, 21, 30, 20]. Already unequal — and the bias has not even updated yet.

If expert 6 was visited 30 times (fair = 16), its bias grows by 0.10 · (30 - 16)/16 ≈ 0.088. Next step it scores higher for every token by 0.088 — and noise of ±0.30 is not enough to dislodge it.

1import numpy as np
2
3rng = np.random.default_rng(0)
4N_EXPERTS, K, BATCH = 8, 2, 64
5STEPS, LR = 60, 0.10
6
7# Router bias per expert. Tiny random init — every expert is "equal".
8bias = (rng.random(N_EXPERTS) - 0.5) * 0.5
9cumulative_load = np.zeros(N_EXPERTS)
10
11for step in range(STEPS):
12    load = np.zeros(N_EXPERTS, dtype=int)
13    # Each token gets noisy logits = bias + per-token affinity noise.
14    noise = (rng.random((BATCH, N_EXPERTS)) - 0.5) * 0.6
15    scores = bias[None, :] + noise                       # (BATCH, N_EXPERTS)
16    topk = np.argpartition(-scores, K, axis=1)[:, :K]    # (BATCH, K)
17    for row in topk:
18        for e in row:
19            load[e] += 1
20
21    # Positive feedback: experts that won this step grow more attractive
22    # because they received more gradient and their router weight updates
23    # in the direction that scores their winning tokens higher. We model
24    # that compactly as: bias_e += LR * surplus / fair_share.
25    fair = (BATCH * K) / N_EXPERTS
26    bias += LR * (load - fair) / fair
27    cumulative_load += load
28
29cov = cumulative_load.std() / cumulative_load.mean()
30top2_share = np.sort(cumulative_load)[-2:].sum() / cumulative_load.sum()
31dead = int((cumulative_load < 0.2 * fair * STEPS).sum())
32print(f"CoV={cov:.2f}  top-2 share={top2_share:.0%}  dead={dead}/{N_EXPERTS}")

Run this snippet at home and you will see CoV land somewhere in the 1.0 to 1.5 range after 60 steps with $E = 8$, top-2 share well over 50%, and two or three dead experts. The exact numbers depend on the seed; the qualitative outcome does not. Change $E$ to 64 and $\mathrm{BATCH}$ to 256 and you get a richer version of the same collapse — more visibly skewed, more dead experts in absolute terms.

The minimal counterexample. Replace the bias update with $b_i \leftarrow b_i - 0.06 \cdot (\ell_i - \bar{\ell})/\bar{\ell}$ — flipping the sign so winners get pushed down. Re-run. The collapse evaporates instantly. CoV stays under 0.2 forever. That sign flip is the entire idea behind DeepSeek's auxiliary-loss-free balancing in section 6.3 — a single line of code that prevents days of wasted training.

PyTorch: Detecting Collapse in Real Training

Before we fix the collapse in sections 6.2 to 6.4, we need a way to detect it inside a real training loop. The simulator above made the dynamic visible because we instrumented it from the start; production MoE training runs are silent unless you wire diagnostics yourself. The following function is the minimum viable instrumentation — call it once per step, log every value, alert when any one crosses a threshold.

Healthy run: 0.10 – 0.35. Drifting: 0.35 – 0.60. Collapse: > 1.0.

1import torch
2import torch.nn.functional as F
3
4@torch.no_grad()
5def routing_diagnostics(router_logits: torch.Tensor, k: int,
6                        n_experts: int) -> dict[str, float]:
7    """
8    router_logits: (N, E)  — N = batch*seq, E = number of experts
9    Returns diagnostics every MoE training loop should log.
10    """
11    # 1. Pick top-k experts per token.
12    topk_w, topk_idx = router_logits.topk(k, dim=-1)        # (N, k)
13
14    # 2. Count how many tokens each expert received this batch.
15    flat = topk_idx.reshape(-1)                              # (N*k,)
16    load = torch.bincount(flat, minlength=n_experts).float() # (E,)
17
18    # 3. Imbalance metrics.
19    fair = load.sum() / n_experts
20    cov = load.std(unbiased=False) / load.mean().clamp(min=1.0)
21    dead = (load < 0.2 * fair).sum()
22    top2_share = load.topk(2).values.sum() / load.sum()
23
24    # 4. Router entropy — low entropy means routing has crystallised.
25    p = F.softmax(router_logits, dim=-1).mean(dim=0)         # (E,)
26    entropy = -(p * (p + 1e-12).log()).sum()
27    max_entropy = torch.log(torch.tensor(float(n_experts)))
28
29    return {
30        "load_cov": cov.item(),
31        "dead_experts": int(dead.item()),
32        "top2_share": top2_share.item(),
33        "norm_entropy": (entropy / max_entropy).item(),
34    }
35
36# Inside the training loop:
37#   logits = router(hidden_states)      # (B*T, E)
38#   stats = routing_diagnostics(logits, k=2, n_experts=64)
39#   if stats["load_cov"] > 0.5:
40#       wandb.alert(title="MoE routing collapse imminent",
41#                   text=str(stats))

Three of the four metrics are derived from the load vector and one is derived from the router probabilities directly. Together they cover the failure space:

What Changes at Massive Scale

Everything in this section was demonstrated with 8 experts and 64 tokens per step. Real DeepSeek-V3 runs with 256 routed experts, top-8 routing, 4M tokens per batch, sharded across 2000+ H800s. The collapse dynamic does not get gentler at scale. It gets sharper and more expensive.

Why scale makes it worse

1. More experts means more dead-expert capacity at risk. With 8 experts, losing 3 to collapse is 37.5% of the model. With 256 experts, the same fractional collapse loses 96 experts — 96 × the FFN parameters of the dense baseline, wasted.
2. Sparser routing amplifies the feedback. Top-8 of 256 is a 3% activation ratio. Each expert sees, on average, 3% of the batch. Variance in that 3% is high — random fluctuations in early training are large enough to seed asymmetries that the feedback loop then amplifies.
3. Cross-device traffic patterns become catastrophic. Experts are sharded across nodes. A collapsed router sends all tokens to a few nodes — the all-to-all becomes an all-to-few, saturating a tiny slice of the network while the rest of the cluster idles. GPU utilization drops from 60% to 15% as collapse advances. You see this in NCCL profiles before you see it in the loss curve.
4. Recovery is logistically impossible. At 64 experts locally you might try resetting the router and warm-starting from a checkpoint. At 256 experts across 2000 GPUs, restarting means hours of checkpoint loading and rewinding the dataloader; doing it more than once per training run is unacceptable.

The DeepSeek-V3 evidence

The DeepSeek-V3 technical report shows the auxiliary-loss-free balancer holding norm_entropy above 0.95 and load CoV below 0.15 for the full 14.8T-token pretraining run. Earlier MoE attempts in the open literature — GShard, Switch Transformer, ST-MoE — all required a non-trivial auxiliary balance loss in their training objective. The progression of approaches in sections 6.2 through 6.4 mirrors the historical progression of the field: we knew about collapse for years before we knew the right way to prevent it.

Engineering Reality and Gotchas

Several practical lessons recur across MoE training postmortems:

1. Collapse can hide behind a falling loss curve. The surviving experts continue to learn and the loss continues to drop — slowly. Teams without routing diagnostics frequently ship models with 20% dead experts and never know. Always log the four metrics above; cheaper than wasted compute by a factor of millions.
2. The first 1000 steps are critical. Once an asymmetry establishes itself in the early steps it is very hard to undo. Some training infrastructures apply a router warmup — uniform routing for the first $N$ steps, then gradually fade in the learned router — specifically to let every expert collect a baseline gradient signal before the feedback loop activates.
3. Random init does not save you. A common misconception is that initialising the router weights uniformly gives every expert a fair start. It does, technically — but as the simulator showed, even a fair start collapses within tens of steps. The asymmetry is generated by the stochasticity of the data stream, not by the init.
4. Recoverability is myth. Once an expert has missed thousands of gradient steps, its parameters are not just stale — they are actively wrong relative to where the rest of the model has moved. Reviving it requires far more compute than preventing its death cost.
5. Fine-grained MoE (chapter 5) makes the stakes higher. When each expert is small, you can afford more of them — which is why DeepSeek scaled to 256 experts. But more experts means more surface area for collapse, and a dead small expert is still a dead expert. The fine-grained gains of chapter 5 are conditional on solving the collapse problem this chapter addresses.

Carry one sentence forward into the next section: top-$k$ routing without a balancing pressure is a positive-feedback loop, and positive-feedback loops always collapse to a degenerate equilibrium. Everything else in this chapter is engineering around that single dynamical fact.