Spot the two config choices that invite router collapse in this MoE training script.
Click any words you think contain an error. Click again to unmark.
Two bugs: init_std 1.0 gives the router large arbitrary logits that pick favorites from step one, and aux_loss_coef 0.0 removes the balancing pressure that would correct the skew.
Think of a brand new manager assigning tickets to a team of 8 specialists. Bug one: on day one, before knowing anyone, the manager already has loud, confident favorites, purely by accident. Bug two: nobody in the company ever tells the manager to spread the work around. The favorites get all the practice, so they genuinely become better, so the manager trusts them even more. Within months, two people do everything while six sit idle, and the team is no better than a much smaller one. The fix mirrors the bugs: start the manager with no opinions (tiny initial preferences), and add a standing rule that rewards spreading work evenly. In MoE training those are a near-zero router init and a load-balancing penalty in the loss.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Router collapse is the signature failure mode of sparse models, and what makes it interesting is that it is not a bug in the code. Every component works as written. The failure lives in the dynamics: routing decisions determine which experts learn, and what experts learn determines future routing decisions. Close that loop with the wrong constants and the system eats itself.
This question hands you a six-line config and asks which two lines arm the loop. Answering well requires more than pattern matching on suspicious zeros; one of the zeros in this file is perfectly fine. You need a causal model of collapse: what seeds it, what amplifies it, and what normally holds it in check.
The sections below build that model, walk through each config line against it, establish healthy defaults, and finish with the monitoring practices that catch collapse in the first few hundred steps, when it is still cheap to fix.
The feedback loop that defines collapse
A routed MoE layer trains two coupled systems at once. The experts learn to transform the tokens they receive. The router learns to score experts per token, and tokens flow to the top scorers. Coupling is the trap: experts only improve on traffic they actually see, and the router only favors experts that look good.
Run the loop forward from a slight imbalance. Expert 3 receives a few more tokens than average in early steps. It accumulates more gradient signal, so its outputs improve slightly faster. The router, scoring outcomes, nudges expert 3's logits upward, sending it more traffic. Each cycle widens the gap. Meanwhile the starved experts barely change from initialization, so they keep losing every comparison.
The end state is stable and useless: two or three experts handle nearly all tokens, the rest are dead weight occupying VRAM, and your 8-expert model behaves like a dense model a fraction of its size. Stability is what makes collapse vicious; the system is not oscillating toward recovery, it has found an attractor.
Every anti-collapse mechanism targets one of two variables: how big the initial imbalance is, and how strongly the system restores toward uniform. That is the lens for reading this config.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- Switch Transformer popularized the auxiliary balancing loss with coefficient near 0.01 and small router init as defaults for stable sparse training.
- ST-MoE added router z-loss after observing that large router logits destabilize training, the same magnitude problem this init_std creates.
What an interviewer would ask next. Try answering before peeking at the approach.
QWhy does softmax make a large init so much worse than a linear scorer would?
Exponentiation turns additive logit gaps into multiplicative probability ratios; consider what a 2-logit gap does to expert shares.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Flagging jitter_noise: 0.0 as the bug. Exploration noise is a nice to have; the actual collapse drivers are the oversized router init and the disabled balancing loss.
60 second bullets to scan on the way to the call.
The rich get richer feedback loop, step by step
Why large router init means confident routing at step zero
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.