MoE Capacity Factor: Why Mixture-of-Experts Drops Your Tokens

Your Mixture-of-Experts fine-tune looks clean in eval and degrades in production, and the degradation tracks concurrency rather than input difficulty. Nothing throws. Logits look normal. The usual suspects — sampling params, quantization, prompt drift — all check out. The culprit is often the MoE capacity factor: a fixed-size buffer per expert that silently discards tokens when the router sends too many of them to the same place. A dropped token doesn't error. It just skips the feed-forward block at that layer and continues as if nothing happened.

Key takeaways

• Every MoE layer allocates each expert a fixed buffer: capacity = capacity_factor × tokens × top_k / num_experts. Tokens that arrive after the buffer fills are dropped — they bypass the expert MLP and pass through on the residual stream alone.
• Routers self-reinforce ("rich get richer"), so load imbalance is the default state, not an anomaly. Auxiliary load-balancing loss and router z-loss exist to fight it.
• Drop rate depends on the token distribution in the current batch, which is why the failure appears under production traffic mixes and not in single-sequence eval.
• Fixes, cheapest first: raise capacity_factor, switch to a dropless block-sparse kernel (MegaBlocks-style grouped GEMM), or use auxiliary-loss-free balancing (DeepSeek-V3's per-expert routing bias) so you stop paying a quality tax for balance.
• Always instrument drop rate per layer per expert. If you don't log it, you cannot see it.

What is the MoE capacity factor, exactly?

The MoE capacity factor is a multiplier on the expected number of tokens per expert, used to size a static buffer.

In a top-k MoE layer with E experts and T tokens in the batch, each token picks k experts, so there are T·k assignments to distribute. If routing were perfectly uniform, each expert would receive T·k/E of them. Capacity is that ideal times a slack factor:

capacity = capacity_factor * T * k / E

Enter fullscreen mode Exit fullscreen mode

With capacity_factor = 1.0, an expert can take exactly its fair share and not one token more. At 1.25 it tolerates 25% overload. Switch Transformer swept this range and landed in roughly 1.0–2.0 depending on the setup.

Why a fixed size at all? Because expert dispatch is an all-to-all communication followed by a batched matmul, and both want static shapes. Variable-length expert batches mean dynamic allocation, recompiled kernels on TPU/XLA, and ragged collectives. A fixed buffer makes the whole thing a single [E, capacity, d_model] tensor you can scatter into and gather from with one permutation. The cost of that convenience is that overflow has to go somewhere, and "somewhere" is the floor.

What actually happens to a dropped token?

It skips the expert MLP for that layer and keeps only its residual.

Concretely, the layer computes x + Σ_j g_j · E_j(x) over the token's selected experts. When an assignment is dropped, its gate g_j is zeroed. If all k assignments are dropped, the sum is empty and the layer degenerates to x + 0 — an identity. The token traverses a 100-plus-layer network with one of its transformer blocks reduced to a no-op.

Two details make this worse than it sounds.

First, the drop is position-dependent and order-dependent. Buffers fill in arrival order, so a token's fate depends on how many tokens ahead of it in the flattened batch picked the same expert. Move the same sequence to a different position in a batch and it may route identically but survive differently. (This is one mechanism, among several, behind batch-size-dependent output changes in MoE serving.)

Second, most implementations do not renormalize the gates over surviving experts. If a top-2 token loses one of its two experts, the layer output is the single surviving expert scaled by its raw gate, which may be 0.3. You get a 0.3× MLP contribution, not a rescaled full one. That's a quiet attenuation applied unevenly across your batch.

In an encoder classifying sentiment, a few percent of dropped tokens is survivable. In autoregressive decode, the damage lands on whichever token you happen to be generating, and one degraded hidden state produces a wrong token that then conditions everything after it.

Why does the router overload one expert in the first place?

Because routing is learned, and learned routing has a positive feedback loop.

The router is typically a single linear layer: logits = x @ W_g, softmax, take top-k. An expert that receives more tokens gets more gradient, becomes better at those tokens, so the router raises its score for them, so it receives even more tokens. Left alone, a large fraction of experts collapse into near-dead weight while a handful absorb the traffic. This is why every serious MoE recipe carries a balancing mechanism.

The classic one is the Switch/GShard auxiliary load-balancing loss:

L_aux = α · E · Σ_i f_i · P_i

Enter fullscreen mode Exit fullscreen mode

where f_i is the fraction of tokens dispatched to expert i and P_i is the mean router probability for expert i over the batch. f_i comes from a hard argmax and carries no gradient; P_i does. The product penalizes assigning high probability mass to already-overloaded experts. α is usually small — 0.01 is a common default — because this is a regularizer fighting your loss function, and turning it up buys balance by trading away modeling quality.

ST-MoE added the router z-loss, (1/T)·Σ_t (logsumexp_i z_t,i)², typically with a coefficient around 1e-3. It penalizes large router logits. The point is numerical: big logits round badly in bf16/fp16, and small rounding differences flip a top-k decision, which changes the expert entirely. Routers are discontinuous functions of their inputs, so precision noise becomes routing noise.

How do I measure the drop rate in my MoE layer?

Count it directly in the dispatch path. Here is a top-k router with capacity, written so the drop bookkeeping is visible:

import torch
import torch.nn.functional as F

def route_with_capacity(x, w_gate, num_experts, top_k=2, capacity_factor=1.25):
    """x: [T, d_model] flattened tokens. Returns gates, indices, keep-mask, stats."""
    T = x.shape[0]
    logits = x @ w_gate                                   # [T, E]
    probs = F.softmax(logits, dim=-1, dtype=torch.float32)
    topk_p, topk_i = probs.topk(top_k, dim=-1)            # [T, k]

    capacity = int(capacity_factor * T * top_k / num_experts)

    keep = torch.zeros_like(topk_i, dtype=torch.bool)
    offsets = torch.zeros(num_experts, dtype=torch.long, device=x.device)

    # Slot 0 fills buffers before slot 1. This priority ordering is not cosmetic:
    # under overload it decides whether a token keeps its best expert or its worst.
    for slot in range(top_k):
        e = topk_i[:, slot]                               # [T]
        onehot = F.one_hot(e, num_experts)                # [T, E]
        rank = (onehot.cumsum(dim=0) * onehot).sum(-1) - 1
        pos = rank + offsets[e]
        keep[:, slot] = pos < capacity
        offsets = offsets + onehot.sum(dim=0)

    gates = topk_p * keep                                 # dropped assignments -> gate 0

    stats = {
        "assign_drop_rate": (~keep).float().mean().item(),
        "token_fully_dropped": (~keep).all(dim=-1).float().mean().item(),
        "max_expert_load": offsets.max().item() / max(capacity, 1),
    }
    return gates, topk_i, keep, stats

Enter fullscreen mode Exit fullscreen mode

Log all three numbers per layer. assign_drop_rate is the headline. token_fully_dropped is the one that correlates with visible quality loss — those tokens got an identity layer. max_expert_load above 1.0 tells you how far over capacity your worst expert ran, which is what you'd need to raise capacity_factor to in order to eliminate drops.

The matching auxiliary loss:

def load_balance_loss(probs, topk_i, num_experts, alpha=0.01):
    mask = F.one_hot(topk_i, num_experts).sum(dim=1).clamp(max=1).float()  # [T, E]
    f = mask.mean(dim=0)      # dispatch fraction, no gradient
    P = probs.mean(dim=0)     # mean routing prob, differentiable
    return alpha * num_experts * torch.sum(f * P)

Enter fullscreen mode Exit fullscreen mode

Why does token dropping show up in production but not in eval?

Because capacity is computed over the tokens in the current batch, and the routing distribution of a batch depends on what's in it.

A router trained on mixed data learns experts that specialize — some drift toward code and structured syntax, others toward prose, others toward particular languages or numeric content. Eval at batch size 1 on a short prompt produces a small T, so capacity = cf · T · k / E may round down to a handful of slots per expert, but the token count is small enough that imbalance rarely overflows. Now serve a batch of twenty concurrent requests that are all code completions. Every token in the batch prefers the same two or three experts. Those experts blow through capacity; the rest of the buffer sits empty.

Long-context prefill is the same failure with one request. A 100k-token prefill flattens into a huge T of highly homogeneous tokens — a single log file, a single codebase — and homogeneous tokens route homogeneously.

So the drop rate is a function of traffic composition, which is exactly the variable your eval harness holds constant.

Should I just go dropless?

Usually yes for training, and it's the default in most modern MoE inference stacks.

Dropless MoE (MegaBlocks-style) reformulates the expert computation as a block-sparse matmul, so each expert processes exactly the tokens routed to it with no padding and no truncation. There's no capacity, so there are no drops. You pay for it with variable memory (a badly imbalanced batch means one enormous expert GEMM) and a dependence on grouped-GEMM kernels rather than plain batched matmul. Modern serving engines lean on grouped GEMM for MoE anyway, so the marginal cost is small.

That said, dropless does not fix imbalance — it converts a quality problem into a latency problem. A saturated expert becomes a straggler, and with expert parallelism, one slow rank stalls the all-to-all for everyone. You still want the router balanced.

The cleanest current answer to balancing is DeepSeek-V3's auxiliary-loss-free scheme: keep a per-expert bias b_i that is added to the routing logits for selection only, while the gate value comes from the unbiased score. After each step, nudge each bias by a fixed amount based on whether that expert was over- or under-loaded:

# selection is biased; the gate value is not
sel = (logits + bias).topk(top_k, dim=-1).indices
gates = probs.gather(-1, sel)

with torch.no_grad():
    load = F.one_hot(sel, num_experts).sum(dim=(0, 1)).float()
    bias += gamma * torch.sign(load.mean() - load)   # gamma: small update speed

Enter fullscreen mode Exit fullscreen mode

This is a controller, not a loss term. It applies zero interference gradient to the model, which is the whole point — the auxiliary loss buys balance by corrupting the objective, and at large expert counts that tax is measurable.

What should I actually check?

Run this order:

1. Instrument first. Log per-layer assign_drop_rate and token_fully_dropped on real traffic, not eval batches. If drops are zero, the problem is elsewhere; stop here.
2. Check whether gates are renormalized after drops. Unnormalized surviving gates attenuate the MLP contribution unevenly.
3. Raise capacity_factor to your observed max_expert_load as a diagnostic. If quality recovers, drops were the cause — now decide whether to pay the memory or go dropless.
4. Check inference capacity separately from training capacity. They're frequently different config values, and a model trained dropless but served with capacity_factor = 1.0 will drop tokens it never learned to survive.
5. Watch expert-parallel stragglers once dropless. Balance still matters; it just shows up in your p99 instead of your accuracy.

The short answer

A Mixture-of-Experts layer drops tokens because expert dispatch needs static tensor shapes, so each expert gets a fixed buffer sized by the MoE capacity factor — capacity_factor × tokens × top_k / num_experts — and any assignment arriving after that buffer fills is discarded, leaving the token to skip the feed-forward block entirely and pass through on the residual alone. Routers self-reinforce toward imbalance, so overflow is the normal state rather than an edge case, and because capacity is computed over the current batch, the drop rate tracks traffic composition — which is why the failure surfaces under production concurrency and homogeneous long-context prefills but never in single-sequence eval. Fix it by instrumenting drop rate per layer, then either raising the capacity factor, moving to a dropless block-sparse kernel, or adopting auxiliary-loss-free routing bias so you get balance without taxing the training objective.