A sparse Mixture-of-Experts layer only saves compute if the router spreads tokens across experts evenly. Classic token-choice routing — each token picks its favourite experts — leaves that balance to luck and props it up with an auxiliary loss, yet still overloads popular experts and drops tokens. Expert-choice routing (Zhou et al., 2022) flips the selection on its head: instead of tokens choosing experts, each expert reaches into the batch and picks the top-k tokens it wants. That single transposition changes everything downstream — load balance becomes exact by construction, the auxiliary loss disappears, and compute per token becomes adaptive. The catch is subtle: the selection is a global operation over the token axis, which sits awkwardly with the strict left-to-right causality of autoregressive decoding. This piece works through the transposed selection, its math, why the balance is guaranteed, and exactly where the causality caveat bites.
The routing problem: who picks whom
A Mixture-of-Experts layer replaces one dense feed-forward block with E parallel expert FFNs and a small router that sends each token to only a few of them. The promise is a large parameter count at a small active compute cost — but only if the router keeps every expert roughly equally busy. If half the tokens pile onto expert 3, that expert becomes a throughput bottleneck while the rest sit idle.
There are two ways to frame the assignment. In token-choice routing, the token is the agent: it looks at all experts and selects its top few. In expert-choice routing, the expert is the agent: it looks at all tokens and selects the ones it will process. Both start from the same affinity scores between tokens and experts — they simply run the top-k along a different axis of the same matrix. That choice of axis is the entire story, because it decides who is guaranteed a fixed workload and who is not.
Token-choice and its load-balance headache
The dominant scheme — used by Switch Transformer, GShard, Mixtral and most production MoEs — is token-choice. For each token you compute affinity scores over experts, s = softmax(x W_g) with W_g: [d, E], then route the token to its top-k experts (often k = 1 or 2). Every token gets exactly k experts — simple and causal-friendly.
The trouble is that nothing constrains how many tokens land on any given expert. Token preferences are data-dependent and lumpy, so some experts are swamped and others starve. Because each expert runs on fixed-size hardware buffers, every expert has a capacity; tokens beyond it are dropped (skipped via the residual). To fight this, token-choice models add an auxiliary load-balancing loss that nudges the router toward uniform usage. It helps, but it is a soft penalty fighting the main objective — balance stays approximate, the loss weight is fiddly to tune, and dropped tokens still happen.
The transposed selection: experts pick tokens
Expert-choice routing keeps the same affinity computation but reverses the selection. Build the score matrix over the whole group of tokens at once, S = softmax(X W_g) with X: [n, d] and S: [n, E], so S[i, e] is how much token i and expert e like each other. Token-choice takes top-k along the expert axis (each row). Expert-choice takes top-k along the token axis (each column).
Concretely, each expert e scans its column S[:, e] across all n tokens and keeps the k highest-scoring ones — the tokens it most wants. Because every expert independently claims exactly k tokens, each processes a fixed, identical amount of work regardless of the data. The router is no longer a marketplace where tokens compete for scarce slots; it is a set of experts each filling a fixed quota — a guaranteed, non-negotiable batch size.
The math: a gating matrix and top-k over tokens
Following the paper, let the per-expert quota be k. From the score matrix S: [n, E] compute, for every expert (i.e. along the token axis), the top-k tokens and their gate values:
S = softmax(X · W_g) # [n, E] token-expert affinity
G, I = TopK( Sᵀ , k ) # over token axis, per expert
G: [E, k] gate weights I: [E, k] token indices
P = onehot(I) # [E, k, n] permutation/gather
X_in[e] = P[e] · X # [k, d] tokens routed to expert e
Y[e] = G[e] ⊙ Expert_e( X_in[e] ) # scale outputs by gate
out = scatter_add(Y, I) # back to [n, d]The key line is TopK applied to the transposed scores Sᵀ: [E, n]. Each expert yields exactly k (index, gate) pairs, so I has shape [E, k] — a fixed E×k assignment. The gathered inputs are dense [E, k, d] tensors ideal for batched matmuls, and the final scatter_add sums each token’s expert outputs back into its residual slot.
Guaranteed load balance, no auxiliary loss
The headline property falls straight out of the shapes. Every expert selects exactly k tokens, so every expert does exactly k FFN evaluations. Utilisation is perfectly uniform by construction — not encouraged by a penalty, but forced by the selection rule. There is nothing left for a load-balancing loss to fix, so expert-choice models drop it entirely, which removes a hyperparameter and a term that used to tug against the language-modelling objective.
Because the per-expert count is fixed at k, there is also no notion of an expert overflowing: no token is ever dropped because an expert ran out of capacity. (Tokens can still be un-selected — see below — but that is a routing decision, not a hardware overflow.) Empirically, this stable balance lets expert-choice reach a target quality in fewer steps than a comparably sized token-choice model, precisely because no capacity is wasted on idle experts and no gradient signal is spent policing balance.
Variable tokens-per-token: adaptive compute
Token-choice gives every token the same treatment: exactly k experts, full stop. Expert-choice makes no such promise about tokens. Because experts pick independently, a single token can be chosen by many experts, by one, or by none at all. The number of experts a token receives is now variable and data-dependent.
This is a feature, not a bug. It lets the model spend more compute on the tokens many experts find valuable — rare, ambiguous or pivotal tokens attract several — while cheap, predictable tokens (a trailing space, an obvious continuation) may be picked by zero experts and pass through on the residual. This is adaptive computation: the FLOP budget flows where it helps most rather than spreading uniformly. The aggregate is still fixed — total selections equal E × k — but its distribution across tokens is learned. The mild risk is that a token receiving zero experts gets no MoE signal that step, so architectures keep a shared/dense path or residual so no token is stranded.
A worked example
Take a group of n = 1024 tokens, E = 8 experts, and a capacity factor c = 2. The capacity factor sets the average number of experts per token, so total selections are c × n = 2048. Split evenly across the eight experts, each expert’s quota is:
k = c · n / E = 2 · 1024 / 8 = 256 tokens per expert
total routed = E · k = 8 · 256 = 2048 = c · n ✓So every expert runs its FFN on exactly 256 tokens — a clean [8, 256, d] batch — regardless of how the 1024 tokens actually distribute their preferences. Compare token-choice at k = 2: total routing is also 2048, but the per-expert load might be 500 on the busiest expert and 40 on the quietest — forcing you to either over-provision or cap capacity and drop ~250 tokens. Same total work, wildly different worst case. Expert-choice turns that ragged histogram flat, which is what makes it hardware-friendly on both GPUs and CPU SLM inference.
The capacity factor and the compute budget
The capacity factor c is the knob that sets the compute-quality trade-off. With c = 1 the average token gets one expert and total MoE FLOPs match a single dense FFN pass; c = 2 spends twice that but lets popular tokens gather more experts. Raising c raises each expert’s quota k = c·n/E linearly, and with it both compute and coverage — fewer tokens end up selected by zero experts.
Two shapes matter. First, cost is governed by E × k expert-token evaluations, independent of how lumpy preferences are, so you can size buffers exactly with no capacity-drop slack. Second, the quota k scales with the group size n: a larger routing group gives each expert a richer pool and sharper selectivity — which, as the next section shows, is exactly the source of its main limitation.
The causality caveat for autoregressive decoding
Here is the sharp edge. To pick its top-k tokens, an expert must rank a whole group of tokens against one another — the selection for token i depends on the scores of other tokens in the group, including tokens that come after it in the sequence. That is a global, non-causal operation over the token axis. In a decoder-only language model trained to predict the next token, letting the routing of position i depend on positions > i is a causality violation: it leaks future information into the present.
It gets worse at inference. Autoregressive decoding produces one token at a time; when you generate token i the future tokens do not exist yet, so the top-k-over-tokens cannot be computed as it was at training. This is why expert-choice fits encoders and full, fixed sequences — where the whole group is visible at once — but needs care for generative decoders. Practical remedies restrict selection to a causal window, route over the batch dimension instead of time, or fall back to token-choice at decode time. Ignore the caveat and you get a model that quietly cheats in training and cannot reproduce its own routing when it generates.
Practical implications and pitfalls
For CPU-bound SLM inference the appeal is concrete: uniform per-expert batches mean predictable memory and cache behaviour with no wasted padding, and dropping the auxiliary loss simplifies the recipe. If your deployment is an encoder (retrieval, classification, embedding) or you score full sequences, expert-choice is often the cleaner, better-balanced option and worth reaching for first.
The pitfalls track its assumptions. It presumes a group of tokens visible together, so any autoregressive-generation path must handle the causality problem explicitly — do not port an encoder recipe to a decoder unchanged. Balance is guaranteed only across the routing group, not within a single sequence, so with tiny groups the selectivity weakens. And because some tokens may receive zero experts, keep a residual or shared-expert path. Used within those bounds, expert-choice buys exact load balance and adaptive compute for free; used blindly on a generative decoder, it is a correctness bug waiting to surface at inference.