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.