The router of §13.1 has a problem that sounds fatal. It makes a hard, discrete decision (these two experts, not the other six), and a hard decision has no gradient. Nudge a router weight by a tiny amount and, almost always, the same two experts win, so the derivative of "which experts" with respect to the router is zero almost everywhere. How does the router learn anything?

And a second problem waits on the hardware side. A GPU wants to know the shape of a matrix product before it runs it. Routing decides, token by token, how many rows each expert's product will have, and nothing stops the router from sending half the batch to one expert. This section solves both problems. It derives how gradients reach the router, defines the gating functions real models use, introduces expert capacity and the token dropping that comes with it, describes the alternative where experts choose tokens, and builds the sort, gather, and scatter implementation that every fast MoE kernel is organized around.

How a token chooses experts

For a batch we flatten all sequences into $T$ tokens, $\mathbf{X} \in \R^{T\times d}$ with row $\mathbf{x}_t$. The router is a matrix $\mathbf{W}_r \in \R^{d\times E}$ for $E$ experts, and it produces one logit per token and expert,

Each token keeps the set $\mathcal{K}_t$ of its $k$ largest logits. Then the logits become gates $g_{t,i}$, the weights in $\mathbf{y}_t = \sum_{i\in\mathcal{K}_t} g_{t,i}\operatorname{FFN}_i(\mathbf{x}_t)$. Three conventions are in use, and they are not interchangeable.

• Top-k, then softmax. $g_{t,i} = e^{r_{t,i}}/\sum_{j\in\mathcal{K}_t} e^{r_{t,j}}$ for $i \in \mathcal{K}_t$. Shazeer et al. (2017) and Mixtral use this. The gates of each token sum to 1.
• Softmax, then top-k. $\mathbf{p}_t = \operatorname{softmax}(\mathbf{r}_t)$ over all $E$ experts and $g_{t,i} = p_{t,i}$ for $i\in\mathcal{K}_t$, without renormalizing. The Switch Transformer (with $k = 1$) and DeepSeek-V2 use this. The gates of a token sum to less than 1.
• Sigmoid, then top-k, then normalize. $s_{t,i} = \sigma(r_{t,i})$, and the chosen $s$ values are divided by their sum. DeepSeek-V3 uses this, with a twist for load balancing that §13.3 and §13.5 explain.

The first two are closer than they look. Renormalizing the full softmax over the chosen set gives back the first convention exactly, because the full partition function cancels,

In plain words, "softmax over the chosen logits" and "softmax over everything, then renormalize the chosen ones" are the same function. The only real choice is whether to renormalize.

That choice matters most at $k = 1$. With renormalization, a single chosen expert always gets gate $1$, so the output $\mathbf{y}_t = \operatorname{FFN}_i(\mathbf{x}_t)$ does not depend on the router's logits at all, and the router receives no gradient from the language-modeling loss. Shazeer et al. conjectured that routing to at least two experts was needed for useful router gradients. Fedus et al. (2022) showed that $k = 1$ works if the gate is the unnormalized probability $p_{t,i}$, because then the output is scaled by how confident the router was, and that scale has a gradient.

import numpy as np

def top_k_router(logits, k):
    idx = np.argsort(-logits, axis=1)[:, :k]
    p = np.exp(logits) / np.exp(logits).sum(1, keepdims=True)   # softmax over ALL experts
    return idx, np.take_along_axis(p, idx, axis=1)

How the router gets a gradient

The choice of experts has no gradient, but the gates do, and that is enough. Fix one token and drop its index. Let the chosen set be $\mathcal{K}$, the gates $g_i = e^{r_i}/\sum_{j\in\mathcal{K}} e^{r_j}$, the expert outputs $\mathbf{f}_i = \operatorname{FFN}_i(\mathbf{x})$, and the upstream gradient $\mathbf{G} = \partial\mathcal{L}/\partial\mathbf{y}$, a row vector the size of the output. Since $\mathbf{y} = \sum_{i\in\mathcal{K}} g_i\mathbf{f}_i$, each gate's gradient is a dot product,

The gates are a softmax over the chosen logits, whose Jacobian (§0.2) is $\partial g_i/\partial r_j = g_i(\delta_{ij} - g_j)$ for $i, j \in \mathcal{K}$. A logit outside $\mathcal{K}$ does not appear in any gate, so as long as a small change does not alter the selection, its derivative is zero. The chain rule gives, for $j\in\mathcal{K}$,

In plain words, among the experts it tried, the router turns up the ones whose outputs pointed more downhill than the gate-weighted average, and turns down the others. It learns nothing at all about the experts it did not try.

That second sentence is the heart of routing. The router is in the position of a gambler who only sees the payoff of the machines it plays. If one expert happens to be chosen often early on, it trains more, gets better, gets chosen even more, and the others never get a chance to show what they could do. This is why routers are given noise, auxiliary losses, or bias corrections to keep exploring, the subject of §13.3. The gradient into $\mathbf{W}_r$ follows as usual, $\partial\mathcal{L}/\partial\mathbf{W}_r = \mathbf{X}^\top\,\partial\mathcal{L}/\partial\mathbf{R}$, and the gradient into each expert's weights is the expert's ordinary backward pass with its upstream gradient scaled by the gate, $g_i\mathbf{G}$.

The two agree, and the four unchosen experts get exactly zero.

The unnormalized convention changes this picture in one instructive way. With $g_i = p_i$ taken from the full softmax over all $E$ experts, every logit appears in every gate through the partition function, and the Jacobian $\partial p_i/\partial r_j = p_i(\delta_{ij} - p_j)$ holds for all $j$, so

An unchosen expert now receives a gradient, but it is $-p_j\sum_{i\in\mathcal{K}}p_ia_i$, the same signal for every unchosen expert scaled only by its probability. It says how much the chosen experts helped, not whether the skipped ones would have. When the chosen experts helped ($\sum_{i\in\mathcal{K}}p_ia_i < 0$), gradient descent lowers every unchosen logit, which sharpens the router toward the choice it already made. In plain words, the full softmax lets the router learn how confident to be, which is what keeps $k = 1$ trainable, but it still learns nothing about the quality of experts it did not try.

import numpy as np

def router_logit_grad(r, F_out, G, k):
    sel = np.argsort(-r)[:k]
    out = np.zeros_like(r)
    out[sel] = F_out[sel] @ G                     # placeholder: forgot the softmax Jacobian
    return out

Noisy top-k gating

Shazeer et al. (2017) added noise to the logits before choosing, with a learned, input-dependent scale,

with $\mathbf{W}_{\text{noise}}\in\R^{d\times E}$ and $\operatorname{softplus}(z) = \log(1 + e^z)$ to keep the scale positive. The top $k$ of $\tilde{\mathbf{r}}_t$ are chosen and softmaxed. In plain words, the router sometimes tries an expert just outside its current top $k$, so that expert gets a chance to prove itself, and the model can learn how much exploration each token needs.

The noise had a second job. It makes "is expert $i$ in the top $k$?" a random event with a smooth probability, computable with the normal CDF, so Shazeer et al. could write a differentiable estimate of each expert's load and penalize imbalance (§13.3). Later Transformer MoEs kept the exploration but made it lighter. Switch Transformers compared several options in an appendix and settled on input jitter, multiplying the router's input by random factors near 1, which scored slightly better than plain argmax and clearly better than sampling from the softmax. Balancing itself moved to auxiliary losses and bias terms. The cell shows the exploration effect on one token whose third-best expert is just behind the second.

Without noise the token would use experts 0 and 1 every time. With noise, expert 2 is in the top two about half the time and receives gradient on those steps, while experts 3 and 4, far behind, are almost never tried. The noise explores near the boundary, not at random.

Expert capacity

Now the hardware problem. On an accelerator each expert processes its tokens as one matrix product, and the simplest fast implementation gives every expert a buffer of fixed size, its capacity. With $T$ tokens, $k$ choices each, and $E$ experts, a perfectly balanced router would give each expert $kT/E$ tokens. Capacity is that fair share times a slack factor, the capacity factor $\text{CF}$,

In plain words, each expert agrees to take at most $\text{CF}$ times its fair share of tokens. (Implementations differ in rounding, and GShard computes it per group of tokens rather than per batch.)

A token that arrives at a full expert is dropped for that expert. It is not lost. Its contribution from that expert is zero, and because the MoE layer sits inside a residual connection, the token's vector passes to the next layer unchanged, as if the FFN had output zero. If a token's other choices had room, it still gets their (now smaller) sum. Every drop is computation the layer skipped for that token, so the capacity factor is a real trade-off. Larger $\text{CF}$ drops fewer tokens but pads buffers with empty slots that cost memory, compute, and communication. Switch Transformers trained with $\text{CF} = 1.0$ or $1.25$, found the lower factors gave the better speed and quality trade-off, and report that typically fewer than 1% of tokens were dropped. ST-MoE (Zoph et al. 2022) used 1.25 in training and 2.0 at evaluation.

Who gets dropped when an expert overflows? The usual rule is arrival order, the position of the token in the flattened batch, with first choices served before second choices. GShard places every token's first choice, in token order, and only then the second choices into whatever room is left. Riquelme et al. (2021) proposed batch prioritized routing, which serves tokens with larger gates first, so the tokens the router is most confident about are the last to be dropped.

A worked example

Eight tokens, $E = 4$ experts, $k = 1$, $\text{CF} = 1$, so $\operatorname{cap} = \lceil 8/4\rceil = 2$. The router sends tokens 0 to 7 to experts $(0, 1, 0, 0, 2, 0, 1, 3)$. Expert 0 receives tokens 0 and 2 into its two slots, and tokens 3 and 5 find it full, so they are dropped. Experts 1, 2, and 3 receive two, one, and one tokens, all accepted, leaving two empty slots. Two of eight tokens (25%) are dropped and two of eight slots are wasted padding, even though only one expert was overloaded. The widget below runs the same bookkeeping on 24 tokens, with a router you can skew toward one expert.

Even a perfectly fair router drops tokens, because randomness alone makes some experts busier than others. If each token picks one of $E$ experts uniformly and independently, an expert's load is binomial with mean $T/E$ and standard deviation about $\sqrt{T/E}$, and the expected fraction dropped is

The fluctuations shrink relative to the mean as $T$ grows, so big routing groups drop less at the same capacity factor.

A uniform random router at $\text{CF} = 1$ still drops 13% of a 64-token group and 1.6% of a 4,096-token group. Real routers are less balanced than uniform, so their drop rates are higher, which is why balancing (§13.3) and capacity go together.

import numpy as np

def dispatch(idx, E, cap):
    idx = np.asarray(idx)
    T, k = idx.shape
    pos = np.full((T, k), -1)
    counts = np.zeros(E, dtype=int)
    for t in range(T):                       # placeholder: token-major order, no capacity check

Experts choose tokens

Token-choice routing lets every token pick its favorite experts and then deals with the crowding. Zhou et al. (2022) turned it around. In expert-choice routing every expert picks its favorite tokens. Compute the token-to-expert affinities $\mathbf{S} = \operatorname{softmax}(\mathbf{X}\mathbf{W}_r) \in \R^{T\times E}$ (a softmax over experts for each token), and let each expert $i$ take the $\operatorname{cap}$ tokens with the largest scores in column $i$,

where $c$ plays the role of the capacity factor (the paper's default is $c = 2$, matching the compute of top-2 token choice). In plain words, each expert fills its buffer with the tokens that want it most.

Three consequences follow. Load is perfectly balanced by construction, since every expert processes exactly $\operatorname{cap}$ tokens, so no balancing loss is needed and no buffer slot is wasted. Tokens receive a variable number of experts, from zero to $E$, which spends more compute on tokens many experts find relevant. Zhou et al. report that most tokens in their analysis went to one or two experts, 23% to three or four, and about 3% to more than four, and that expert choice improved training convergence by more than 2x over Switch and GShard routing at the same compute. And tokens can go unprocessed, if no expert ranks them in its top $\operatorname{cap}$.

The catch is causality. Whether token $t$ is chosen depends on how the other tokens in the batch score, including tokens that come later in the same sequence. In training a decoder that leaks information about the future into the present (by counting the possible routing patterns, Wang et al. 2024 show that a 9-layer model with 16 experts and two per token could leak more than 50 bits per token, enough for a token to learn which token comes next), and at inference the future tokens do not exist yet, so the routing cannot be reproduced token by token. Zhou et al. note the method may not apply directly to autoregressive generation. That is why expert choice is used in encoders (Komatsuzaki et al. use it in the T5 encoder and ViT, §13.5), while decoder-only LLMs route by token choice.

A third design enforces balance exactly from the token side. BASE layers (Lewis et al. 2021) treat routing during training as a linear assignment problem, giving every expert exactly $T/E$ tokens of the batch while maximizing the total affinity, and route each token greedily to its best expert at inference. The training assignment has the same dependence on the rest of the batch as expert choice. Clark et al.'s S-BASE router (§13.5) approximates the assignment with a few Sinkhorn iterations, which rescale the rows and columns of the affinity matrix until every expert's column sums to its share.

import numpy as np

def expert_choice(S, cap):
    T, E = S.shape
    idx = np.argsort(-S, axis=1)[:, :1]           # placeholder: this is token choice, top-1
    return idx, np.take_along_axis(S, idx, 1), np.ones(T, dtype=int)

Sort, gather, compute, scatter

A loop over tokens is hopeless on a GPU, and even the loop over experts in §13.1 hides a cost, because each nonzero call scans the whole batch. Fast MoE layers reorganize the batch so that each expert's tokens sit next to each other in memory, run each expert on its contiguous block, and then put the results back. Here are the steps, with shapes, for $T$ tokens and $k$ choices each.

1. Route. Get idx and gates, both $(T, k)$. Flatten to $Tk$ assignments, with expert = idx.ravel() and token = repeat(arange(T), k).
2. Sort. order = argsort(expert, stable) groups the assignments by expert. counts = bincount(expert, minlength=E) gives each expert's number of rows, and offsets = [0, cumsum(counts)] gives where each expert's block starts.
3. Gather. Xs = X[token[order]], shape $(Tk, d)$. This is the permutation of tokens into expert order. Each token appears $k$ times.
4. Compute. Expert $i$ processes rows offsets[i] to offsets[i+1] of Xs. On a GPU these $E$ products of different heights run as one grouped GEMM (§13.4).
5. Scatter. Multiply each output row by its gate and add it back to its token's row, out[token[order]] += gates[order] * Ys. Because a token appears $k$ times, this must be a scatter-add, not an assignment.

The backward pass mirrors it. The gradient of a gather is a scatter-add, and the gradient of a scatter-add is a gather, so autograd handles both. With capacity, step 2 also discards assignments whose slot is $-1$ (or pads each block to $\operatorname{cap}$ rows), which is what turns this into the fixed-shape version.

import numpy as np

def moe_forward(X, W_r, experts, k):
    raise NotImplementedError

In PyTorch

The same algorithm in PyTorch uses argsort, bincount, split for the per-expert blocks, and index_add_ for the scatter. Autograd differentiates through the gather and scatter, and the router gets its gradient through gates, exactly as derived above.

Read more

• Shazeer et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. Section 2.1 defines noisy top-k gating.
• Lepikhin et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. Top-2 gating with capacity, groups, and the first-choices-first dispatch order, in Algorithm 1.
• Fedus, Zoph & Shazeer (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Top-1 routing, the capacity factor, and its effect on speed and quality.
• Zhou et al. (2022). Mixture-of-Experts with Expert Choice Routing. Experts pick tokens, and what that buys and costs.
• Riquelme et al. (2021). Scaling Vision with Sparse Mixture of Experts. Batch prioritized routing, which drops the least confident tokens first.
• Lewis et al. (2021). BASE Layers: Simplifying Training of Large, Sparse Models. Routing as a balanced assignment problem, with no balancing loss and no dropped tokens.
• Gale et al. (2022). MegaBlocks: Efficient Sparse Training with Mixture-of-Experts. Dropless MoE with block-sparse kernels, the endpoint of the sort and gather approach.

Watch next

▶

▶