MoE Router Collapse: Why 90% of Tokens Hit 2 Experts

The 8-Expert Model That Only Uses 2

You train a Mixture of Experts model with 8 experts, expecting distributed specialization. After a few thousand steps, you check the routing statistics and find 87% of tokens going to experts 0 and 3. The other six experts? Basically decorative.

This is router collapse, and it’s one of the most frustrating failure modes in MoE training. Your model has 8x the parameters but uses a fraction of them. The paper that first systematically addressed this — Shazeer et al.’s “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” (2017) — remains the foundational reference. You can read it here.

The core insight is deceptively simple: without explicit load balancing, routers learn to send everything to whichever experts happen to perform slightly better early in training. Those experts get more gradient signal, improve faster, and attract even more tokens. It’s a rich-get-richer dynamic that starves most experts of training signal entirely.

What Actually Causes Router Collapse

The gating mechanism in MoE looks innocuous:

$$G(x) = \text{Softmax}(W_g \cdot x)$$

where $W_g$ projects the input $x$ to a vector of expert scores. You pick the top-$k$ experts (usually $k=1$ or $k=2$) and route the token there. The problem emerges from how gradients flow.

When expert $i$ processes more tokens, it receives more gradient updates. Its weights improve. Next batch, the router notices expert $i$ produces lower loss on average. The router weights $W_g$ shift to favor expert $i$ even more. Within a few hundred steps, you’ve got a runaway feedback loop.

The mathematical framing helps here. Let $f_i$ be the fraction of tokens routed to expert $i$. The expected gradient magnitude for expert $i$ scales with $f_i$. If we denote the initial routing distribution as roughly uniform ($f_i \approx 1/N$ for $N$ experts), any small perturbation $\epsilon$ gets amplified:

$$f_i^{(t+1)} \propto f_i^{(t)} + \alpha \cdot \nabla_{W_g} \mathcal{L}$$

The gradient $\nabla_{W_g} \mathcal{L}$ is dominated by the experts that actually processed tokens. Experts with low $f_i$ contribute almost nothing to this gradient, so their router weights stagnate.

The Auxiliary Load Balancing Loss

Shazeer et al.’s fix is elegant: add a loss term that penalizes uneven load distribution.

$$\mathcal{L}_{aux} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i$$

where $f_i$ is the fraction of tokens actually routed to expert $i$, and $P_i$ is the average router probability assigned to expert $i$ across the batch. The coefficient $\alpha$ (typically 0.01 to 0.1) controls how aggressively you enforce balance.

This loss has a clever property: it’s minimized when $f_i = P_i = 1/N$ for all experts. Any deviation from uniform distribution increases the loss. The multiplication of $f_i \cdot P_i$ means experts that receive both high routing probability AND high actual traffic get penalized most heavily.

Here’s what the implementation looks like in practice:

import torch
import torch.nn.functional as F

def compute_load_balance_loss(router_logits, expert_mask, num_experts):
    """
    Compute auxiliary load balancing loss for MoE routing.

    Args:
        router_logits: (batch_size, seq_len, num_experts) raw logits from router
        expert_mask: (batch_size, seq_len, num_experts) binary mask of selected experts
        num_experts: int
    """
    # Router probabilities (what the router "wanted" to do)
    router_probs = F.softmax(router_logits, dim=-1)

    # P_i: average probability assigned to each expert
    # Shape: (num_experts,)
    expert_probs = router_probs.mean(dim=[0, 1])

    # f_i: fraction of tokens actually routed to each expert
    # This uses the hard assignment from top-k selection
    tokens_per_expert = expert_mask.float().sum(dim=[0, 1])
    total_tokens = expert_mask.shape[0] * expert_mask.shape[1]
    expert_fraction = tokens_per_expert / total_tokens

    # The aux loss
    aux_loss = num_experts * (expert_frobs * expert_fraction).sum()

    return aux_loss

Wait, there’s a typo — expert_frobs should be expert_probs. I’ve left it because this is exactly the kind of bug that ships to production at 11pm. The corrected version:

    aux_loss = num_experts * (expert_probs * expert_fraction).sum()

Why the Coefficient α Matters More Than You’d Think

The load balancing coefficient $\alpha$ is surprisingly sensitive. Too low (< 0.001), and collapse still happens — just slower. Too high (> 0.5), and you’re forcing uniform routing at the expense of actual task performance. The model routes evenly but learns nothing.

I’ve seen teams set $\alpha = 0.1$ because “the Switch Transformer paper used it” without realizing that value was tuned for their specific architecture and scale. Fedus et al. (2021) in the Switch Transformer paper actually report experimenting with values from 0.001 to 1.0 and finding the sweet spot varies with model size.

Here’s a diagnostic I find useful:

def diagnose_routing_health(router_probs, top_k=2):
    """
    Check if your router is collapsing.

    Returns a dict with diagnostic metrics.
    """
    # router_probs: (batch, seq, num_experts)
    num_experts = router_probs.shape[-1]

    # Flatten to (total_tokens, num_experts)
    probs_flat = router_probs.view(-1, num_experts)

    # Get top-k expert indices per token
    _, top_indices = probs_flat.topk(top_k, dim=-1)

    # Count how many unique experts are used across entire batch
    unique_experts = torch.unique(top_indices).numel()

    # Expert usage entropy (should be ~log(num_experts) for uniform)
    usage_counts = torch.bincount(top_indices.flatten(), minlength=num_experts).float()
    usage_dist = usage_counts / usage_counts.sum()
    entropy = -(usage_dist * torch.log(usage_dist + 1e-10)).sum()
    max_entropy = torch.log(torch.tensor(float(num_experts)))

    # The "collapse ratio" - 1.0 means uniform, 0.0 means one expert dominates
    collapse_ratio = entropy / max_entropy

    return {
        'unique_experts_used': unique_experts,
        'expected_experts': num_experts,
        'usage_entropy': entropy.item(),
        'collapse_ratio': collapse_ratio.item(),  # <0.7 is concerning, <0.5 is collapsed
        'top_expert_fraction': usage_dist.max().item(),
    }

Run this every few hundred steps. If collapse_ratio drops below 0.7, increase $\alpha$. If it stays pegged at 1.0 but validation loss plateaus, you might be over-regularizing.

The Capacity Factor Trick

Load balancing loss alone doesn’t prevent a subtler failure: buffer overflow. Each expert has a fixed capacity — say, it can process at most $C \cdot B / N$ tokens per batch, where $B$ is batch size, $N$ is number of experts, and $C$ is the capacity factor (typically 1.0 to 2.0).

If more tokens want to route to expert $i$ than its capacity allows, excess tokens get dropped or sent to a fallback. This creates a weird gradient situation: dropped tokens contribute nothing to that expert’s training.

The capacity factor $C$ controls the buffer size. Setting $C = 1.0$ means each expert can handle exactly its “fair share” — fine if routing is perfectly balanced. In practice, you need $C > 1.0$ to handle variance. But higher $C$ means more memory per expert.

$$\text{Expert Capacity} = \left\lceil \frac{C \cdot B \cdot S}{N} \right\rceil$$

where $S$ is sequence length. At $C = 1.25$, you’re allocating 25% extra buffer. Switch Transformer uses $C = 1.25$ for top-1 routing. For top-2 routing, you often need $C = 2.0$ or higher because tokens hit two experts each.

My best guess on why this trips people up: the capacity calculation is done at the expert level, not globally. If you’re debugging OOM errors in MoE training, check whether your capacity factor math accounts for the actual sequence lengths in your data.

Noise Injection: The Underrated Fix

Shazeer et al. proposed adding tunable noise to router logits during training:

$$G(x) = \text{Softmax}(W_g \cdot x + \epsilon \cdot \text{Softplus}(W_{noise} \cdot x) \cdot \mathcal{N}(0, 1))$$

The noise magnitude is learned (via $W_{noise}$), which lets the model adapt exploration-exploitation tradeoff during training. Early in training, higher noise encourages exploration of different experts. Later, the model can learn to reduce noise for confident routing decisions.

This is weirdly underused in practice. Most implementations I’ve seen just use the auxiliary loss and call it a day. But noise injection helps specifically in the early training phase when collapse is most likely to start.

class NoisyTopKRouter(torch.nn.Module):
    def __init__(self, hidden_dim, num_experts, top_k=2, noise_std=0.1):
        super().__init__()
        self.gate = torch.nn.Linear(hidden_dim, num_experts, bias=False)
        self.noise_std = noise_std
        self.top_k = top_k

    def forward(self, hidden_states):
        # hidden_states: (batch, seq, hidden_dim)
        logits = self.gate(hidden_states)  # (batch, seq, num_experts)

        if self.training:
            # Add noise during training only
            noise = torch.randn_like(logits) * self.noise_std
            logits = logits + noise

        probs = F.softmax(logits, dim=-1)
        top_k_probs, top_k_indices = probs.topk(self.top_k, dim=-1)

        # Normalize top-k probs to sum to 1
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        return top_k_probs, top_k_indices, probs

The Surprising Ablation: Auxiliary Loss Weight Annealing

One thing that surprised me reading through follow-up work: annealing the auxiliary loss coefficient helps significantly. Start with high $\alpha$ (say 0.1) to force balanced routing early, then decay it to a lower value (0.01) as training progresses.

The intuition is that you need strong regularization to prevent collapse in the first few thousand steps, but once experts have differentiated, you want the router to specialize. Keeping $\alpha$ high throughout training forces artificial uniformity that hurts final performance.

I haven’t tested this systematically myself, but Fedus et al. (2021) mention this in passing, and subsequent work like BASE Layers (Lewis et al., 2021) takes it further with capacity-aware routing schemes.

Expert Dropout: A Simple Regularization That Works

A technique that doesn’t get enough attention: randomly dropping experts during training. Not dropping their outputs — actually removing them from the routing candidates.

def expert_dropout_routing(router_logits, num_experts, drop_prob=0.1):
    """
    Randomly mask out some experts before routing.
    Forces the router to not over-rely on any single expert.
    """
    if not self.training:
        return F.softmax(router_logits, dim=-1)

    # Create random mask (1 = keep, 0 = drop)
    mask = (torch.rand(num_experts, device=router_logits.device) > drop_prob).float()

    # Ensure at least 2 experts survive
    if mask.sum() < 2:
        # Force-keep two random experts
        indices = torch.randperm(num_experts)[:2]
        mask[indices] = 1.0

    # Apply mask by setting dropped expert logits to -inf
    masked_logits = router_logits + (1 - mask) * (-1e9)

    return F.softmax(masked_logits, dim=-1)

This is brutal but effective. By randomly removing experts from consideration, you force the router to develop fallback paths. The resulting model is more robust and shows better expert utilization in my experience.

What the Paper Doesn’t Tell You About Hyperparameter Sensitivity

The original MoE paper and most follow-ups report results with carefully tuned hyperparameters. What they don’t emphasize: the interaction between load balancing coefficient, learning rate, and batch size is finicky.

Higher learning rates accelerate collapse. The router weights update faster, so the rich-get-richer dynamic kicks in sooner. If you’re seeing collapse despite using auxiliary loss, try reducing your router’s learning rate specifically (many implementations let you set a separate LR for the gating network).

Smaller batch sizes also hurt. With fewer tokens per batch, the estimate of $f_i$ (actual routing fraction) is noisier. This noise can actually help prevent collapse early on, but it makes the auxiliary loss signal less reliable. Batch sizes below 64 tokens/expert seem to cause instability.

Comparing Against GShard’s Random Token Dropping

The GShard paper (Lepikhin et al., 2020) took a different approach to capacity overflow: instead of dropping tokens deterministically when an expert is full, they randomly selected which tokens to drop. This prevents the bias where “late” tokens in a sequence always get dropped.

But this creates its own problem. Random dropping means the gradient signal is stochastic in a way that’s hard to control. GShard compensates with larger batch sizes and more aggressive load balancing.

If I were choosing between approaches:
– Switch Transformer’s top-1 routing with auxiliary loss: simpler, works well at scale, but you lose the mixture effect of top-2
– GShard’s top-2 with random dropping: better theoretical capacity but more hyperparameter tuning
– BASE Layers’ learned routing: most sophisticated, but requires significantly more engineering

For a first MoE implementation, I’d start with top-1 routing and auxiliary loss. It’s the most forgiving.

Practical Debugging Checklist

When your MoE model underperforms:

1. Check routing entropy every 500 steps. Collapse usually starts in the first 10% of training.
2. Log expert utilization histograms. You want them roughly uniform, not bimodal.
3. Watch for capacity overflow. If >5% of tokens are being dropped, increase capacity factor.
4. Verify your auxiliary loss is actually backpropagating. I’ve seen codebases where the aux loss was computed but never added to the total loss.
5. Compare router gradients across experts. Dead experts have near-zero router gradients.

And when you’re debugging this at 2am because your training run collapsed after 50k steps, a bag of dark chocolate espresso beans will be more helpful than another arxiv paper.

When Router Collapse Isn’t Actually the Problem

Sometimes what looks like collapse is actually correct behavior. If your task genuinely only requires two types of computation, the router finding two experts is optimal. The other experts are redundant.

This happens most often with:
– Small datasets where there isn’t enough diversity to justify many experts
– Tasks with strong inductive biases (some NLP tasks really do cluster into a few patterns)
– Over-parameterized MoE where you have more experts than the task needs

Before debugging collapse, ask: does this task actually benefit from $N$ experts? Sometimes the answer is no, and you should use a smaller MoE or a dense model.

The Token-Choice vs Expert-Choice Question

Recent work like Expert Choice routing (Zhou et al., 2022) flips the paradigm: instead of tokens choosing experts, experts choose tokens. Each expert picks its top-$k$ tokens from the batch.

This elegantly solves load balancing — every expert processes exactly the same number of tokens by construction. But it introduces new problems: some tokens might not be chosen by any expert, or might be chosen by all of them. The practical implications for sequence modeling (where you need to maintain causality) are still being worked out.

I haven’t used Expert Choice in production, so take this with appropriate skepticism. The preliminary results look promising for encoder-only models, less clear for autoregressive generation.

FAQ

Q: How do I know if my MoE model has collapsed during training?
Monitor the routing entropy and expert utilization histogram. If entropy drops below 70% of maximum ($\ln(N)$ for $N$ experts) or if fewer than half your experts receive more than 5% of tokens, you’re likely seeing collapse. Catch it early — collapse that starts at step 1000 is easier to fix than collapse entrenched at step 50000.

Q: Can I fix router collapse by just restarting training with different random seed?
Sometimes, but it’s treating symptoms rather than cause. Different initialization might delay collapse, but if your auxiliary loss coefficient is too low or your learning rate is too high, collapse will recur. Fix the hyperparameters, not the seed. That said, if you’re 90% through training and just noticed collapse, restarting with new seed AND adjusted hyperparameters is often faster than trying to recover.

Q: Should I use top-1 or top-2 routing for my first MoE implementation?
Start with top-1. It’s simpler, uses less memory (each token hits only one expert), and the load balancing dynamics are easier to debug. Top-2 theoretically gives better quality by mixing expert outputs, but the capacity management is trickier. Graduate to top-2 once you’ve confirmed your top-1 routing is stable.

References

• Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017. https://arxiv.org/abs/1701.06538

• Fedus, W., Zoph, B., & Shazeer, N. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR. https://arxiv.org/abs/2101.03961

• Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., … & Chen, Z. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ICLR 2021. https://arxiv.org/abs/2006.16668

• Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., … & Chi, E. (2022). Mixture-of-Experts with Expert Choice Routing. NeurIPS 2022. https://arxiv.org/abs/2202.09368

If you’re building MoE models, start with the auxiliary loss at $\alpha = 0.01$, use top-1 routing, and monitor routing entropy religiously. Only add complexity (noise injection, expert dropout, capacity tuning) when you see specific problems. The failure mode you’re most likely to hit is collapse, and now you know exactly what to look for.

What I’m still curious about: whether Expert Choice routing will become the default for encoder models, and whether someone will figure out how to make it work cleanly for autoregressive generation. The load balancing problem feels “solved” for training, but inference-time routing efficiency is still wide open.