Core ML

Expert Routing and Load Balancing

The router in a mixture-of-experts model scores each token against all experts and picks the top-k — but without a load-balancing loss, all tokens pile onto the same few popular experts, wasting most of the expert capacity.

Expert routing computes a score for each expert from the token representation, selects the top-k by score, then applies an auxiliary load-balancing loss to prevent all tokens from routing to the same few experts

Why Does This Exist?

Mixture of experts gives you N expert FFNs and routes each token to k of them. The router — a small learned linear layer followed by softmax — picks which k experts each token uses. Simple, and it works well enough in a toy setting.

At scale it fails in a specific way: expert collapse. Left to its own gradients, the router quickly discovers that a few experts tend to do well, starts routing more tokens to them, which makes those experts train more, which makes them even better, which routes still more tokens to them. Within a few thousand steps, you have a model where 90% of tokens route to 2 or 3 experts out of 64. The other experts are essentially untrained and wasted. Total parameter count is 64× the dense baseline; useful parameter count is barely 3×.

Routing isn't just a pick-the-best problem — it's a scheduling problem. You need to distribute work across experts while still routing each token to the experts best suited for it. That tension is the entire substance of this page.

Think of It Like This

A restaurant with 8 chefs where customers keep asking for the same 2

A restaurant hires 8 chefs, each specialising in a different cuisine. But the menu only shows the ratings customers gave last week, and last week 2 chefs got the best ratings. Now every customer requests those 2 chefs. Those 2 are overwhelmed, the other 6 are idle, and the customers who get turned away from their first choice get slow service. If the restaurant guaranteed each chef a roughly equal share of the load, all 8 would develop, the food would be faster, and the two popular chefs wouldn't be the bottleneck.

Load balancing in MoE is that guaranteed share — enforced by a loss term, not by the customers.

How It Actually Works

The routing computation

For a token with hidden state $x \in \mathbb{R}^{d_{model}}$ , the router computes a score for each of the $N$ experts:

g_i = \text{softmax}(W_r \, x)_i \quad i = 1, \ldots, N

The top-k indices by $g_i$ are selected; those k experts run and their outputs are weighted by the corresponding $g_i$ values.

Why top-k is used instead of sampling

The softmax scores form a proper probability distribution, so you could sample k experts rather than taking the top-k. Sampling adds variance to routing during training, which some papers argue helps avoid collapse. In practice, top-k routing is more stable and predictable, and combined with an auxiliary loss it handles collapse without sampling noise. Top-k is the default in Switch Transformer and Mixtral.

The auxiliary load-balancing loss

Switch Transformer introduced an auxiliary loss that penalises uneven expert utilisation:

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

where $f_i$ is the fraction of tokens dispatched to expert $i$ in the current batch, and $P_i$ is the average router probability assigned to expert $i$ over the batch. If routing is perfectly uniform, $f_i = P_i = 1/N$ for all $i$ , and the loss is minimised. When tokens cluster on a few experts, the product $f_i \cdot P_i$ is large for the popular experts and small for the unpopular ones, increasing the loss.

The coefficient $\alpha$ (typically 0.01) controls how strongly the auxiliary loss pushes toward uniformity. Too large and the auxiliary loss dominates, experts get perfectly balanced but the model ignores token-expert affinity. Too small and collapse re-emerges.

The capacity factor

Even with load balancing, any batch may have more tokens wanting a particular expert than it can process in parallel. The capacity factor $C$ sets a hard limit: expert $i$ can receive at most $C \times (\text{batch\_size} / N)$ tokens per batch. Tokens routed to an over-capacity expert are dropped — their output is zero — or sent to their second-choice expert (token dropping vs. expert choice routing). $C=1.0$ means strict balance; $C=1.25$ or $C=2.0$ allows some slack at the cost of extra memory.

Show Me the Code

Computing the auxiliary load-balancing loss and demonstrating how it penalises skewed routing.

1import numpy as np2
3rng = np.random.default_rng(21)4n_tokens, n_experts, k = 16, 8, 25alpha = 0.016
7def aux_loss(router_probs: np.ndarray) -> float:8    """Switch Transformer auxiliary load-balancing loss."""9    # router_probs: (n_tokens, n_experts) — softmax scores10    top_k_idx = np.argsort(router_probs, axis=-1)[:, -k:]11    # f_i: fraction of tokens routed to each expert12    f = np.zeros(n_experts)13    for idx in top_k_idx.flatten():14        f[idx] += 115    f /= n_tokens16    # P_i: mean router probability per expert17    P = router_probs.mean(0)18    return alpha * n_experts * float(np.dot(f, P))19
20# Balanced routing: uniform-ish probs21probs_balanced = np.ones((n_tokens, n_experts)) / n_experts22probs_balanced += rng.standard_normal((n_tokens, n_experts)) * 0.0523probs_balanced = np.abs(probs_balanced)24probs_balanced /= probs_balanced.sum(-1, keepdims=True)25
26# Collapsed routing: all weight on first 2 experts27probs_collapsed = np.zeros((n_tokens, n_experts))28probs_collapsed[:, :2] = 0.4529probs_collapsed[:, 2:] = 0.05 / 630probs_collapsed /= probs_collapsed.sum(-1, keepdims=True)31
32loss_balanced = aux_loss(probs_balanced)33loss_collapsed = aux_loss(probs_collapsed)34
35print(f"Balanced routing  aux loss: {loss_balanced:.5f}")36print(f"Collapsed routing aux loss: {loss_collapsed:.5f}")37print(f"Collapse penalty: {loss_collapsed / loss_balanced:.1f}× higher")38# Balanced routing  aux loss: ~0.0013039# Collapsed routing aux loss: ~0.0050040# Collapse penalty: ~3-4× higher

The collapsed routing incurs a meaningfully larger loss, which backpropagates into the router weights and pushes it toward more even assignment.

Watch Out For

Setting alpha too high or too low

Too high an alpha and the auxiliary loss dominates training — the router learns to balance perfectly regardless of which expert is best for a token, degrading MoE quality. Too low and collapse re-emerges even with the loss. Typical values are 0.001 to 0.01; the right value is usually found by monitoring the fraction of tokens per expert during early training and adjusting if any expert's fraction falls below 1/(2N) or climbs above 3/N.

Ignoring token dropping at capacity

With a capacity factor of 1.0, tokens routed to an over-capacity expert are silently dropped — their contribution is zero. In a model with many MoE layers, systematic token dropping can cause certain token types to lose meaningful representation. Monitor the drop rate per layer during training; if it consistently exceeds 1–2% in a layer, either raise the capacity factor or reduce the number of experts in that layer.

The Quick Version

The router computes a softmax score for each expert and selects the top-k by score; those k experts run and their outputs are summed weighted by score.
Expert collapse (most tokens routing to few experts) is the default failure mode without intervention — the router's gradients reinforce popular choices.
An auxiliary load-balancing loss penalises skewed routing by penalising the product of expert-usage fraction and average router probability.
The capacity factor limits tokens-per-expert per batch, dropping or redirecting overflow — a necessary guard against hotspot experts.
Coefficient $\alpha$ controls the load-vs-quality trade-off; 0.01 is the common default; monitor per-expert routing fractions during training.