Mixture-of-Experts with Expert Choice Routing

Sparse mixture-of-experts (MoE) layers promise a sweet deal: grow the parameter count while keeping the activated compute per token roughly fixed. In practice that deal often breaks at the router. When every token independently picks its top-$k$ experts, some experts overflow their capacity, others starve, and training papers reach for an auxiliary load-balancing loss.

This paper’s move is almost embarrassingly simple. Instead of tokens choosing experts, experts choose tokens. Each expert independently takes its top-$k$ highest-affinity tokens. Capacity is uniform by construction. Tokens that matter can still be claimed by many experts; easy tokens can be claimed by few. The affinity matrix is the same object Switch and GShard already compute — only the TopK axis flips.

What this paper explains

Standard MoE FFN layers replace a dense feed-forward block with $e$ expert FFNs plus a gating network. For token representations $X \in \mathbb{R}^{n \times d}$, the gate produces affinities between tokens and experts. Prior work (Switch top-1, GShard top-2) does token choice: each of the $n$ tokens selects $k$ experts.

Zhou et al. study expert choice: each of the $e$ experts selects $k$ tokens, with

where $c$ is a capacity factor. Setting $c = 2$ matches the activation budget of GShard top-2 gating so comparisons stay fair. The method returns three matrices used to gather tokens, weight them, and scatter expert outputs back — same systems skeleton as token-choice MoE, different index selection.

Prior limits

• Load imbalance. Independent token decisions pile traffic onto popular experts. Overloaded experts drop tokens or stretch step latency; under-used experts waste capacity.
• Auxiliary losses. Switch and GShard encourage balance with an extra loss term. That term is a second objective fighting the primary likelihood.
• Uniform compute per token. Token choice gives every token exactly $k$ experts. The paper argues important tokens should be allowed more compute and simple tokens less — at the same global FLOP budget.
• Under-specialization. Chronically under-filled experts never see enough diverse traffic to specialize.

The mechanism

Affinity, then transpose the TopK. Compute a token–expert score matrix $A \in \mathbb{R}^{n \times e}$ (same softmax-style gate as prior MoE work). Token choice runs TopK along the expert axis for each token. Expert choice runs TopK along the token axis for each expert:

1. For expert $i$, keep the $k$ tokens with largest $A_{\cdot, i}$.
2. Build index / gate / one-hot gather tensors $(I, G, P)$ so each expert receives exactly $k$ tokens.
3. Run expert FFNs in parallel at fixed batch size $k$.
4. Scatter and combine weighted expert outputs back into the token stream.

Because each expert’s bucket size is exactly $k$, load balance is perfect without an auxiliary loss. Because buckets are chosen independently, the number of experts covering a given token is variable — often zero for some tokens at a layer (the paper mitigates this by interleaving dense Transformer layers) and larger than $c$ for others.

Capacity factor. $c$ is the average number of experts per token if assignments were uniform. $c = 2$ is the paper’s default (GShard-matched). $c = 1$ matches Switch top-1 activation cost and still beats top-1 gating in their pre-training curves; even $c = 0.5$ remains competitive with top-1 in the reported plots.

Algorithm walkthrough

1. Choose expert count $e$ and capacity factor $c$; set $k = (n/e)\,c$.
2. Compute affinities $A$ with the learned gate on $X$.
3. For each expert, TopK over tokens → indices $I$ and gates $G$.
4. Gather tokens with $P$, apply expert $i$’s FFN to its $k$ tokens, weight by $G$.
5. Scatter/combine into the layer output; continue the Transformer stack.
6. Optional capped variant: after expert choice, limit how many experts a single token may keep (their EC-CAP2), trading a little quality for a hard per-token fan-in.

The opening scene contrasts a token-choice pile-up (one hot expert overflowing) with expert-choice buckets that stay full and even. The slider above lets you change $c$ and watch bucket size $k$ and the implied average experts-per-token move together.

What to notice when reading

• Figure 1 in the paper is the whole thesis: arrows reverse from “token → experts” to “expert → tokens.”
• They drop the balancing auxiliary loss — balance is structural, not regularized.
• Interleaving MoE and dense layers matters; some tokens can be skipped by a given MoE layer.
• Autoregressive generation is called out as a limitation of the training expert-choice implementation (selection uses a full sequence view). Do not silently assume decode-time routing is identical.

Results and evidence

Numbers below are from the paper’s reported experiments (NeurIPS 2022 version / Parallel extract of the PDF). Treat them as the authors’ measurements under their setup, not universal constants.

Limitations

Stated or clearly implied by the authors:

• The expert-choice selection they study looks at a batch/sequence of tokens at once; auto-regressive decoding needs a different online assignment story.
• Skipping tokens at MoE layers can hurt unless dense layers are interleaved.
• Results are on their T5-style pre-train / fine-tune stack; exact GLUE numbers will not transfer unchanged to every modern decoder-only MoE.
• Perfect expert load balance is not the same as perfect token coverage — some tokens may receive fewer experts than others by design.

How to read the paper

1. Abstract + §1 — token-choice failure modes.
2. §3.1–3.2 — pitfalls, then expert-choice equations and Figure 1.
3. §3.3–3.4 — optional bipartite constraints / capping; architecture table.
4. §4 training curves (Fig. 2–4) and Tables 2–4 for downstream numbers.
5. Final limitations paragraph on autoregressive generation.

Knowledge check

Keep reading

1. Original paper (arXiv:2202.09368) — Figure 1, §3, Tables 2–4.
2. Switch Transformers (token-choice top-1) — Fanout Daily 2026-07-30-switch-transformers and arXiv:2101.03961.
3. GShard (top-2 token choice) — arXiv:2006.16668.

Sources

• Zhou et al., Mixture-of-Experts with Expert Choice Routing, NeurIPS 2022 — arXiv:2202.09368
• Fedus et al., Switch Transformers, 2021 — arXiv:2101.03961
• Lepikhin et al., GShard, 2020 — arXiv:2006.16668