[KongchangAI]
· 3 min read· 1,820 words

Understanding MoE Routing: Three Core Problems and Two Fixes

Understanding MoE Routing: Three Core Problems and Two Fixes

A 32-expert tiny model's three-round experiment reveals the inherent tension and trade-offs in MoE routing.

Through three rounds of controlled experiments on English, Python, and arithmetic datasets using a tiny 32-expert character model, the article exposes MoE routing's core conflicts: Token Choice without balancing collapses (one expert takes 57% of tokens), load-balancing loss smooths distribution but adds tuning overhead, and Expert Choice achieves perfect expert balance yet causes token dropping—while oddly yielding the lowest eval loss. The inherent tension between expert balance and token coverage means every routing strategy is a trade-off.

MoE Routing Explained

What Is MoE Routing and Why It Matters So Much

The core idea behind Mixture of Experts (MoE) architectures isn't complicated: replace a single feed-forward layer with multiple "expert" sub-networks, then use a routing network to decide which experts should process each token. This sparse activation strategy lets you dramatically scale up model parameters while compute only grows linearly—one of the key paths to making today's large models cheaper and more efficient.

But routing mechanisms are far less elegant than they appear. A researcher who has spent years working on MoE and sparsity recently used a set of elegant controlled experiments to expose three core problems with routing: expert load imbalance, uneven token distribution, and the inherent trade-offs of routing strategies themselves. He trained a tiny character-level model with just 32 expert layers, running three rounds each on English text, Python code, and arithmetic datasets. The results are highly illuminating.

How Token Choice and Expert Choice Differ

Token Choice is the most intuitive routing approach: each token passes through the routing network and selects the top-k experts (usually k=1 or 2) with the highest scores. This naturally guarantees that every token gets processed by at least one expert, but it can't control how many tokens each expert receives—which is the root cause of routing collapse. Expert Choice flips the perspective: each expert picks a fixed number of top-k tokens based on its own preferences. Since each expert's capacity is predetermined, load balancing at the expert level is mathematically guaranteed—but this creates the opposite problem: some tokens may not be selected by any expert (token dropping). The two strategies also differ in their communication patterns during distributed training: Token Choice requires all-to-all communication to send tokens to the devices hosting the relevant experts, while Expert Choice has experts pull tokens toward themselves—a more regular communication pattern but with weaker guarantees on token coverage.

Experimental Design: 32 Experts, Three Training Rounds

The clever part of the experiment is variable control. The model itself is tiny, but the expert layer is fixed at 32—enough to observe routing behavior without being dragged down by training costs. The only difference across the three rounds is the routing strategy:

  • Round 1: Token Choice routing, no load-balancing loss
  • Round 2: Token Choice routing, with load-balancing loss
  • Round 3: Expert Choice routing, where experts select tokens

These three configurations correspond to the most typical approaches in MoE routing design, and the experimental results clearly show the cost of each.

When Routing Collapses: 57% of Tokens Flow to a Single Expert

The first round's results are an almost textbook case of "routing collapse." Early in training, a single expert swallowed 57% of tokens, while 17 out of 32 experts were virtually ignored. This means more than half of the compute was wasted on redundant experts, and the parameters of the remaining experts were never adequately trained.

This phenomenon isn't rare in MoE training. The routing network and expert networks start in a "cold start" state, and certain experts gain more tokens due to tiny initialization advantages, which in turn trains them more thoroughly, creating a positive feedback loop. The end result is severely wasted model capacity—you pay the parameter storage cost of 32 experts but only get the effective capacity of a few.

Load imbalance doesn't just reduce parameter efficiency—it also leads to training instability. Inactive experts receive zero gradients, which can cause communication bottlenecks and computational dead spots in distributed training.

The Mathematical Intuition Behind the Positive Feedback Loop

Routing collapse is essentially a "rich get richer" dynamic. Suppose the routing network's initial weights are near-random, and some expert e gets a slightly higher routing probability on the first batch of tokens due to a small initialization advantage. Because this expert is activated more often, its gradients update more frequently, and its parameters adapt faster to the current data distribution, earning even higher probability in subsequent routing decisions. This loop is mathematically similar to the temperature effect in softmax—as the variance of routing logits increases, the softmax output tends toward a one-hot distribution, i.e., winner-takes-all. Notably, the speed of collapse is closely related to learning rate, batch size, and the initialization scale of the routing network. Too large a learning rate or too large an initial variance in routing logits both accelerate collapse. This is why many MoE implementations add noise to routing logits (e.g., Noisy Top-K Gating) or use a smaller routing learning rate to slow this process down.

Fix One: Load-Balancing Loss

The second round added a load-balancing loss—the most commonly used fix in mainstream MoE implementations today. The idea is to add a penalty term to the training objective that encourages even distribution of tokens across experts.

The effect was immediate: token distribution became roughly uniform, and every expert got a workload. Routing collapse was effectively curbed, and model capacity was fully utilized.

But load-balancing loss isn't free. It essentially introduces a constraint unrelated to language modeling into the optimization objective, creating tension with the main task goal. Tuning the loss coefficient requires experience—too weak and it can't prevent collapse; too strong and it may force the model to assign tokens to inappropriate experts, hurting model quality.

The Specific Form of Load-Balancing Loss

The most common form of load-balancing loss is the auxiliary loss, exemplified by Switch Transformer: for each expert, compute the fraction of tokens it's selected for, f_i, and its average routing probability, p_i, then take L_balance = α · N · Σ(f_i · p_i), where N is the number of experts and α is the loss coefficient. When all experts are used uniformly, f_i = p_i = 1/N, so Σ(f_i · p_i) = 1/N, and multiplying by N gives 1; if one expert monopolizes traffic, this value increases significantly, producing a penalty. This loss is added to the main language modeling loss for joint optimization. The coefficient α is typically around 0.01, but the specific value depends heavily on model scale and data distribution. There's also a "router z-loss" used to stabilize the scale of routing logits and prevent unbounded growth. Although these auxiliary losses are effective, they optimize the routing distribution itself rather than task performance, so they compete for gradients with the main objective.

Fix Two: Expert Choice Routing

The third round adopted the Expert Choice strategy: instead of tokens choosing experts, each expert selects the tokens it most wants to process. This achieves perfect load balancing—each expert processes exactly the same number of tokens.

But the cost is equally obvious. Since each expert selects independently, some tokens may be ignored by all experts (processed by none), while others are processed redundantly by multiple experts. This token-level imbalance—"some feast, some starve"—doesn't exist in Token Choice.

Interestingly, despite the token coverage problem, Expert Choice achieved the lowest evaluation loss across all three rounds. This suggests that perfect expert load balancing may be more important than perfect token coverage—or that current evaluation metrics aren't sensitive enough to penalize ignored tokens.

The Trade-Offs and Takeaways

Comparing the two fixes reveals a clear trade-off:

DimensionLoad-Balancing LossExpert Choice
Expert LoadRoughly uniformCompletely uniform
Token CoverageFull coverageSome tokens ignored or duplicated
Implementation ComplexityRequires tuning loss coefficientSimpler logic
Evaluation PerformanceGoodBest in this experiment

This result offers several takeaways for MoE practitioners. Load-balancing loss is a "soft constraint" that preserves Token Choice's coverage advantage but introduces hyperparameter tuning overhead. Expert Choice is a "hard constraint" that achieves perfect balance at the expert level but shifts the imbalance to the token level.

Which approach to choose depends on the priorities of your specific scenario. If token coverage is critical to the task (e.g., every input position needs to be adequately processed), load-balancing loss may be safer. If expert computational efficiency is the primary bottleneck, Expert Choice's strict balancing is more appealing.

Evaluation Metric Selection Bias

The result that Expert Choice achieved the lowest evaluation loss needs to be understood in the context of the evaluation metric. Standard language modeling evaluation loss (e.g., cross-entropy) measures the model's prediction quality for each token. Under Expert Choice, tokens processed redundantly by multiple experts get richer representations and more accurate predictions; ignored tokens may have worse predictions, but if they're a small fraction of the data or inherently easy to predict, their effect on pulling down the overall average loss is limited. In other words, cross-entropy is a token-averaged metric that doesn't additionally penalize the structural flaw of "some tokens not being processed at all." If you instead used the worst-token loss or token-type-stratified evaluation metrics, the conclusion might reverse. This suggests that MoE researchers comparing routing strategies shouldn't just look at a single average evaluation loss, but should pay attention to the tail behavior of the loss distribution and task-specific coverage requirements.

Beyond Either/Or

The value of this experiment lies not just in showing the effects of two fixes, but in using a minimalist setup to expose the core tension of MoE routing: there's an inherent tension between expert balance and token coverage, and any routing strategy must make a choice within this tension.

The community is also exploring more hybrid strategies, such as adding token coverage constraints on top of Expert Choice, or using learnable routing temperatures to dynamically adjust the degree of balancing. These directions all point to the same goal: letting MoE's sparse activation truly deliver on its parameter efficiency advantage, rather than being dragged down by routing problems.

For teams currently using or planning to use MoE architectures, the recommendation is to monitor expert load distribution and token coverage rate from early in training. Once routing collapse occurs, it's very difficult to fix by continuing training later.

One-line Summary

A three-round controlled experiment with a 32-expert tiny model reveals the inherent tension between expert balance and token coverage in MoE routing, and the trade-offs of two fixes.

Paragraph Summary

Using a tiny character-level model with just 32 expert layers, the article runs three rounds of controlled experiments on English text, Python code, and arithmetic datasets, systematically revealing the core contradictions of MoE routing. Round 1's Token Choice routing without load balancing suffers severe collapse—a single expert swallows 57% of tokens while 17 experts sit nearly idle. Round 2 adds load-balancing loss, making token distribution roughly uniform but introducing hyperparameter tuning overhead and tension with the main task objective. Round 3 adopts Expert Choice routing, achieving perfect expert load balancing but causing some tokens to be ignored or redundantly processed—yet surprisingly achieving the lowest evaluation loss. The article points out that an inherent tension exists between expert balance and token coverage, that any routing strategy is a trade-off between the two, and recommends monitoring expert load distribution and token coverage rate early in training.

#MoE routing#Mixture of Experts#load balancing loss#Expert Choice routing#Token Choice routing
Share:

Related articles