Why Your Mixture-of-Experts Model Silently Drops Tokens

A Mixture-of-Experts (MoE) transformer can have 8x the parameters of a dense model and cost roughly the same per token. The catch nobody puts in the marketing slide: some of your tokens get routed to an expert that's already full, and the model computes nothing for them. They fall through on the residual connection and skip the FFN entirely. This is called token dropping, and it's not a bug — it's the load-balancing mechanism doing exactly what it was told.

If you've ever fine-tuned an MoE, or wondered why a sparse model's quality is weirdly sensitive to batch composition, this is the reason. Let's go through the mechanism: top-k routing, expert capacity, and the auxiliary load-balancing loss that ties them together.

TL;DR

• MoE replaces each dense FFN with N expert FFNs plus a router. A softmax router picks the top-k experts per token (top-1 in Switch Transformer, top-2 in Mixtral 8x7B). Only those experts run, so FLOPs stay near a dense model's while parameter count multiplies.
• Experts have a fixed capacity per batch. Capacity = capacity_factor * tokens_per_batch / num_experts. Tokens beyond an expert's slot count are dropped and pass through the residual unchanged.
• Routing is learned and naturally collapses. Without intervention, the router sends most tokens to a few favored experts, so the rest sit idle and the popular ones overflow.
• The auxiliary load-balancing loss fixes collapse by penalizing the product of "fraction of tokens routed to expert i" and "mean router probability for expert i," pushing the distribution toward uniform.
• Token dropping is a training/inference-throughput trade-off, tuned via capacity factor. Higher capacity = fewer drops, more memory and compute waste.

What is a Mixture-of-Experts layer actually doing?

An MoE layer swaps one feed-forward network for N independent FFNs (the experts) plus a small gating network (the router). For each token, the router produces a distribution over experts, keeps the top-k, and sends the token only to those. The outputs are combined weighted by the router's gate values.

The point is decoupling parameters from compute. Mixtral 8x7B has 8 experts per layer and routes top-2, so it holds ~47B parameters but activates ~13B per token. You pay dense-model FLOPs for a much larger parameter budget. That's the whole pitch of sparse models, and it's why the frontier labs are widely believed to run MoE architectures even though they don't confirm specifics.

Here's the core routing logic, stripped down:

import torch
import torch.nn.functional as F

def moe_route(x, router_w, num_experts, k=2, capacity_factor=1.25):
    # x: [tokens, d_model]
    T = x.shape[0]
    logits = x @ router_w                      # [T, num_experts]
    probs = F.softmax(logits, dim=-1)          # gate distribution

    topk_vals, topk_idx = probs.topk(k, dim=-1)  # [T, k]
    # renormalize the kept gates so they sum to 1 per token
    topk_vals = topk_vals / topk_vals.sum(dim=-1, keepdim=True)

    # each expert can hold this many tokens this batch
    capacity = int(capacity_factor * T / num_experts)

    dispatched = torch.zeros(T, dtype=torch.bool)
    counts = torch.zeros(num_experts, dtype=torch.long)
    for t in range(T):                          # illustrative; real impl is vectorized
        for slot in range(k):
            e = topk_idx[t, slot].item()
            if counts[e] < capacity:
                counts[e] += 1
                dispatched[t] = True
            # else: this (token, expert) assignment is DROPPED
    return topk_idx, topk_vals, dispatched, counts

The counts[e] < capacity check is where tokens die. When an expert fills its slots, every later token that wanted it gets no compute from that expert. If all of a token's top-k experts are full, that token is fully dropped and only its residual survives.

Why do experts overflow if routing is learned?

Because a freshly initialized router has no reason to spread load evenly, and gradient descent actively makes it worse. Early in training, one expert gets marginally better at some frequent pattern, so the router sends it more tokens, so it gets more gradient signal, so it gets better still. This is a rich-get-richer feedback loop called expert collapse: a handful of experts absorb most traffic while the rest stay undertrained and effectively dead.

Collapse is catastrophic for a sparse model. If 3 of 8 experts handle 90% of tokens, you've paid for 8x the parameters but you're getting the representational capacity of ~3 experts, and those three constantly overflow their capacity and drop tokens. You get the memory cost of MoE with the quality of a small dense model.

How does the auxiliary load-balancing loss stop expert collapse?

It adds a differentiable penalty that's minimized when tokens are spread uniformly across experts. The GShard/Switch Transformer formulation, per MoE layer, is:

f_i = (1/T) * sum_t  1{ argmax router(x_t) == i }     # fraction of tokens routed to expert i
P_i = (1/T) * sum_t  router(x_t)_i                      # mean router probability for expert i

L_aux = alpha * N * sum_i ( f_i * P_i )

The elegant part is why it uses the product of two quantities. f_i is the hard token count — non-differentiable because of the argmax. P_i is the soft mean probability — fully differentiable. Multiplying them gives you gradients (through P_i) that are scaled by how overloaded each expert already is (via f_i). If expert 3 is swamped, its f_3 is large, so the loss pushes hard to lower P_3 — the router learns to send fewer tokens there.

Under the constraints sum_i f_i = 1 and sum_i P_i = 1, the sum sum_i f_i P_i is minimized at the uniform point f_i = P_i = 1/N, where the loss equals exactly 1 after the N scaling. So the target is a nice round number, and alpha (typically small, on the order of 1e-2) sets how aggressively balancing competes with the language-modeling loss.

def load_balancing_loss(probs, topk_idx, num_experts, alpha=1e-2):
    T = probs.shape[0]
    # f_i: fraction of tokens whose top-1 choice is expert i
    top1 = topk_idx[:, 0]
    f = torch.bincount(top1, minlength=num_experts).float() / T
    # P_i: mean gate probability mass on expert i
    P = probs.mean(dim=0)
    return alpha * num_experts * torch.sum(f * P)

Set alpha too low and you get collapse and heavy token dropping. Set it too high and you flatten the router into near-uniform noise — it stops specializing, which defeats the purpose of having experts at all. It's a genuine tension, not a knob you can max out.

What is expert capacity and how do I tune it?

Capacity is the hard cap on tokens per expert per batch, set by the capacity factor:

capacity = capacity_factor * (tokens_per_batch / num_experts)

A capacity_factor of 1.0 means each expert gets exactly its fair share of slots — perfectly efficient only if routing is perfectly balanced, which it never is, so you drop a lot. Training runs typically use 1.25 to 2.0 to absorb the imbalance the aux loss can't fully remove. Inference sometimes pushes higher, or uses a "drop-and-pad" vs. "no-drop" mode.

The trade-off is concrete:

• Higher capacity factor → fewer dropped tokens, better quality, but each expert's buffer is padded out to capacity, wasting compute and memory on empty slots. Throughput drops.
• Lower capacity factor → tighter buffers, higher throughput, but more tokens skip their expert. Quality degrades, and it degrades unevenly — tokens later in the batch's routing order are the ones that hit full experts.

That last point causes a subtle, real failure mode: MoE quality can depend on batch composition. Pack a batch with many similar sequences and they'll all fight for the same experts, blowing capacity and dropping tokens that a more diverse batch would have served fine. If you see mysterious eval variance on a sparse model, check your batching before you blame the weights.

Why does training add a router z-loss on top?

Because the aux loss balances which expert gets picked, but does nothing about the magnitude of the router logits — and large logits make routing numerically unstable and prone to sudden, discontinuous flips. ST-MoE introduced the router z-loss to keep logits small:

L_z = (1/T) * sum_t ( log sum_j exp( logit_{t,j} ) )^2

It penalizes the log-sum-exp of the router logits, squared. Large logits get squashed, the softmax stays in a well-conditioned range, and training doesn't diverge when a single logit spikes. In practice you carry three losses: the main cross-entropy, the load-balancing aux loss, and this z-loss, each with its own small coefficient.

What does this mean if I'm fine-tuning or serving MoE?

Three things worth internalizing.

First, don't strip the aux loss during fine-tuning. It's tempting to drop it because your task-specific dataset is small and you want all the gradient going to the LM loss. But a narrow fine-tuning distribution is exactly where routing re-collapses fastest — a few experts capture your domain and the rest atrophy. Keep balancing on, even at a reduced coefficient.

Second, treat capacity factor as a first-class serving parameter. If your production traffic is more homogeneous than your training mix (say, all customer-support queries), you'll drop more tokens at the same capacity factor than your eval suggested. Measure the actual drop rate, don't assume it.

Third, watch for dead experts. Log per-expert token counts. If some experts see near-zero traffic after training, you're carrying parameters you paid for and don't use, and you can sometimes prune or merge them.

The direct answer

A Mixture-of-Experts model silently drops tokens because each expert has a fixed capacity per batch — capacity_factor * tokens / num_experts — and any token routed to an expert that has already filled its slots gets no FFN computation and passes through unchanged on the residual connection. This capacity limit exists to bound memory and keep expert compute balanced. The router that decides where tokens go tends to collapse onto a few favored experts, causing them to overflow, so MoE training adds an auxiliary load-balancing loss (penalizing the product of token-fraction and mean gate probability per expert) and a router z-loss (bounding logit magnitude) to spread load toward uniform. Token dropping is therefore a tunable throughput-versus-quality trade-off, not a defect: raise the capacity factor to drop fewer tokens at the cost of wasted compute, lower it for speed at the cost of quality — and remember that dropping is worst when a batch's tokens all want the same experts.