Routing and Gating

"They told me I was free to choose any expert I wanted. Then they told me seven of them were full, the eighth was on another continent, and my gradient was due in 40 milliseconds."

A Token Awaiting Its Assignment

In Section 17.2 we assembled the mixture-of-experts layer: a bank of $E$ independent feed-forward networks (the experts) plus a gate that activates only a few of them per token. The layer's promise, capacity that grows with $E$ while per-token compute stays fixed, rests entirely on the gate making good, cheap decisions. This section is about that gate. It is the most delicate component in the whole sparse design, because a single small matrix decides, token by token, where computation happens and therefore where data must move. A poorly behaved gate does not merely waste capacity; it can route nearly every token to the same one or two experts, leaving the rest idle and the model effectively dense and tiny. Understanding why that happens, and the routing variants invented to prevent it, is the prerequisite for everything that follows: load balancing (Section 17.6), capacity limits (Section 17.7), and the all-to-all shuffle that physically delivers tokens to experts on other machines (Section 17.5).

1. The Gate Is a Tiny Learned Softmax Beginner

The gate, also called the router, is almost embarrassingly small. For an MoE layer with $E$ experts operating on token representations of dimension $d$, the gate is a single linear map $W_g \in \mathbb{R}^{d \times E}$, often with no bias and no nonlinearity of its own. Given a token's hidden vector $x \in \mathbb{R}^d$, the gate produces one score per expert and turns those scores into a distribution:

The vector $p$ assigns a probability to each expert, but the layer does not run all $E$ experts; that would forfeit the entire point of sparsity. Instead it selects the top $k$ experts by score (typically $k = 1$ or $k = 2$), routes the token only to those, and combines their outputs weighted by the corresponding gate probabilities. Writing $\mathcal{T}_k(h)$ for the index set of the $k$ largest entries of $h$, the layer output for token $x$ is

where $E_i$ is the $i$-th expert network. The renormalization in the denominator keeps the surviving gate weights summing to one, so an unrouted expert contributes nothing and the kept experts share full responsibility for the token. This is the standard "noisy top-$k$ gating" recipe from the work that revived MoE for deep learning (Shazeer et al., 2017), minus the noise, which we add in Section 4. The whole router is a few thousand parameters bolted onto a layer that may hold billions, yet it governs the behavior of all of them.

2. Why the Hard Selection Hurts: Gradients Through a Discrete Choice Intermediate

Softmax is smooth and differentiable; the top-$k$ that follows it is neither. Selecting the $k$ largest scores is a discrete, piecewise-constant operation: nudge a token's hidden vector by an infinitesimal amount and the chosen expert set usually does not change at all, then suddenly flips when one score overtakes another. The selection itself has a gradient of zero almost everywhere and is undefined at the ties. This is the central awkwardness of routing, and it shapes how the gate learns.

The practical resolution is that gradients flow only through the experts that were actually chosen, and they reach the gate only through the continuous weights $p_i$ that multiply those experts' outputs. Concretely, the backward pass treats the index set $\mathcal{T}_k(h)$ as a fixed, non-differentiable mask for the current step, while the gate probabilities $p_i$ of the selected experts remain live numbers in the computation graph. So the router does receive a learning signal: if expert $i$ produced a useful output for this token, the loss pushes $p_i$ up, which pushes the corresponding logit $h_i$ up, which makes the gate more likely to select expert $i$ for similar tokens next time. What the gate never directly learns is the counterfactual, how an unchosen expert would have performed, because that expert was never run. A token only ever gets feedback about the roads it took, never the roads it skipped. This one-sidedness is exactly what drives the instabilities in Section 5, and it is why exploration (Section 4) has to be injected deliberately rather than emerging on its own.

3. Two Ways to Match Tokens and Experts Intermediate

So far we have described token-choice routing: each token looks at its gate scores and picks its own top-$k$ experts. This is the standard, the formulation used by the GShard and Switch Transformer line of work, and it has an intuitive appeal, every token goes where it most wants to go. Its weakness is equally intuitive. Because each token chooses independently and the gate is shared, nothing stops a large fraction of tokens from all preferring the same expert. The load on each expert is an emergent statistic of the gate, not a controlled quantity, and gates left to themselves tend to concentrate. An expert that receives twice its fair share of tokens becomes the slowest stage in the layer, and on a real cluster it is a single overloaded device that the other devices wait for.

Expert-choice routing (Zhou et al., 2022) inverts the selection to fix this at the source. Instead of every token picking its experts, every expert picks its tokens. Each expert is given a fixed capacity $C$ (a number of token slots) and selects the $C$ tokens that score it highest. Because every expert takes exactly $C$ tokens, the per-expert load is perfectly uniform by construction; there is no balancing loss to tune and no possibility of an overloaded expert. The cost is symmetric to token-choice's: a token that no expert ranks within its top $C$ receives no expert at all and is dropped (passed through by a residual connection), while a token loved by many experts can be processed several times. Token-choice can overload an expert but never starves a token; expert-choice can starve a token but never overloads an expert. The demo in Section 6 measures exactly this trade.

4. Noisy Top-k Gating: Buying Exploration on Purpose Intermediate

Recall from Section 2 that the gate only ever learns about experts it actually used. Early in training this is dangerous: whichever experts happen to score slightly higher at initialization get all the tokens, get all the gradient, improve fastest, and score even higher, a rich-get-richer loop that can lock the layer onto a handful of experts before the others ever run. The original fix, and still a common one, is to add learned noise to the gate scores before the top-$k$ selection:

The noise term perturbs the ranking just enough that an expert which is currently second-best sometimes wins the slot, runs on the token, and produces a gradient that lets the gate discover whether it was actually any good. A second learned matrix $W_{\text{noise}}$ controls how much noise each token injects, so the model can be exploratory where it is uncertain and decisive where it is confident. This is direct exploration bought at the price of a little routing randomness, and it matters most in the opening phase of training when the gate's preferences are still arbitrary. As training proceeds the useful experts genuinely pull ahead and the noise becomes a minor jitter rather than a coin flip. The mechanism is the gating analogue of the exploration-exploitation balance that recurs throughout distributed learning, and it is worth keeping in mind that without it, a top-$k$ gate has no built-in reason to ever try an expert it has not already favored.

5. The Two Failure Modes: Collapse and Fluctuation Advanced

Two instabilities haunt every router, and naming them precisely is half the cure. The first is representation collapse (also called routing collapse): the gate learns to send nearly all tokens to a small subset of experts. The remaining experts, starved of tokens, never receive gradient, never improve, and the model's effective capacity shrinks to the few experts in use, defeating the entire purpose of the sparse layer. Collapse is the equilibrium that the rich-get-richer loop of Section 4 reaches if nothing opposes it. The standard countermeasure is an auxiliary load-balancing loss that penalizes uneven expert utilization, which we develop in full in Section 17.6; expert-choice routing sidesteps collapse entirely by making balance structural rather than learned.

The second failure mode is router fluctuation: the expert assigned to a given token keeps changing from step to step as the gate weights wobble during training. Each time a token's preferred expert flips, the token's representation is suddenly shaped by a different expert with different parameters, which destabilizes the experts downstream (they keep seeing a shifting distribution of inputs) and slows convergence. Fluctuation is most acute when several experts have nearly equal gate scores for a token, so a tiny weight update tips the top-$k$ from one to another. Modern training recipes damp it with router z-losses that keep gate logits small and well-separated, and with careful learning-rate schedules for the gate. Collapse is too much agreement (everything to a few experts); fluctuation is too little (assignments that will not settle). A healthy router lives between them, and keeping it there is an active subject of research.

6. Demo: Token-Choice Skews, Expert-Choice Balances Intermediate

The cleanest way to feel the difference between the two routing regimes is to run both on the same batch and look at the resulting per-expert load. The code below builds a deliberately biased gate (two experts made intrinsically more attractive, the collapse pressure of Section 5), then routes 4096 tokens two ways: token-choice top-$1$, where each token picks its single best expert, and expert-choice with capacity $C = T/E$, where each expert claims its top $C$ tokens. It reports, for each scheme, the per-expert load fractions, the coefficient of variation (standard deviation over mean, zero when perfectly balanced), and for expert-choice the fraction of tokens left unserved. Everything is plain NumPy; there is no framework and no GPU, so the routing logic is fully visible.

import numpy as np

rng = np.random.default_rng(7)
T, D, E, k = 4096, 64, 8, 1          # tokens, model dim, experts, top-k

# A batch of token representations and a learned gate (router) weight matrix.
X = rng.standard_normal((T, D))
# Bias the gate so a few experts are intrinsically more attractive: this is the
# representation-collapse pressure the section warns about.
W_gate = rng.standard_normal((D, E))
W_gate[:, 0] += 1.8                  # expert 0 made artificially appealing
W_gate[:, 3] += 1.1                  # expert 3 too

def softmax(z, axis):
    z = z - z.max(axis=axis, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

scores = X @ W_gate                  # (T, E) raw gate logits
probs = softmax(scores, axis=1)      # per-token distribution over experts

# ---- Token-choice top-k routing: each token picks its own top-k experts. ----
topk = np.argsort(-scores, axis=1)[:, :k]      # (T, k) chosen experts per token
tc_load = np.bincount(topk.reshape(-1), minlength=E)

# ---- Expert-choice routing: each expert picks its top-C tokens. ----
# Capacity per expert so total picks match token-choice: C = T*k / E.
C = (T * k) // E
ec_load = np.zeros(E, dtype=int)
covered = np.zeros(T, dtype=bool)
for e in range(E):
    order = np.argsort(-scores[:, e])          # tokens this expert wants most
    chosen = order[:C]
    ec_load[e] = len(chosen)
    covered[chosen] = True

def stats(load):
    load = np.asarray(load, dtype=float)
    frac = load / load.sum()
    cv = load.std() / load.mean()              # coefficient of variation
    return frac, cv, load.max() / load.mean()

tc_frac, tc_cv, tc_max = stats(tc_load)
ec_frac, ec_cv, ec_max = stats(ec_load)

np.set_printoptions(precision=3, suppress=True)
print("tokens T, experts E, top-k :", T, E, k)
print("expert-choice capacity C   :", C, "tokens per expert")
print()
print("token-choice load (counts):", tc_load)
print("token-choice load (frac)  :", tc_frac)
print(f"token-choice  CV={tc_cv:.3f}  max/mean={tc_max:.3f}  busiest expert={tc_frac.max():.1%} of tokens")
print()
print("expert-choice load (counts):", ec_load)
print("expert-choice load (frac)  :", ec_frac)
print(f"expert-choice CV={ec_cv:.3f}  max/mean={ec_max:.3f}  busiest expert={ec_frac.max():.1%} of tokens")
print()
print(f"tokens left unserved by expert-choice : {int((~covered).sum())} of {T} ({(~covered).mean():.1%})")

tokens T, experts E, top-k : 4096 8 1
expert-choice capacity C   : 512 tokens per expert

token-choice load (counts): [872 387 469 548 343 517 600 360]
token-choice load (frac)  : [0.213 0.094 0.115 0.134 0.084 0.126 0.146 0.088]
token-choice  CV=0.315  max/mean=1.703  busiest expert=21.3% of tokens

expert-choice load (counts): [512 512 512 512 512 512 512 512]
expert-choice load (frac)  : [0.125 0.125 0.125 0.125 0.125 0.125 0.125 0.125]
expert-choice CV=0.000  max/mean=1.000  busiest expert=12.5% of tokens

tokens left unserved by expert-choice : 1476 of 4096 (36.0%)

The numbers make the trade concrete. Token-choice respects every token's preference and serves all 4096, but its busiest expert does 1.7 times the work of an average one, and on a real cluster that single expert is the device every other device waits for at the all-to-all barrier. Expert-choice gives perfect load balance for free, no auxiliary loss and no tuning, but at this capacity it drops more than a third of the tokens. Neither is strictly better; the choice depends on whether your bottleneck is straggling devices (favor expert-choice) or token coverage and simplicity (favor token-choice with a balancing loss). This is the seam where routing meets the systems layer, and it is why the next section moves from "which expert" to "which machine".

# Megatron-Core / Tutel-style top-k router; the heavy lifting is one config + one call.
import torch
import torch.nn.functional as F

class TopKRouter(torch.nn.Module):
    def __init__(self, d_model, num_experts, k):
        super().__init__()
        self.gate = torch.nn.Linear(d_model, num_experts, bias=False)
        self.k = k
    def forward(self, x):                     # x: (tokens, d_model)
        logits = self.gate(x)                 # (tokens, num_experts)
        probs = F.softmax(logits, dim=-1)
        weights, experts = probs.topk(self.k, dim=-1)   # hard top-k select
        weights = weights / weights.sum(dim=-1, keepdim=True)  # renormalize
        return experts, weights               # who runs, and with what weight

7. From Routing Decision to Network Hop Beginner

Everything in this section has treated routing as a question of which experts compute on a token. In a distributed MoE that question is inseparable from a second one: where does the token have to go? When the $E$ experts are sharded across many devices, as they must be once the layer is too large for one accelerator, the gate's top-$k$ output is a list of destination devices, and the token's hidden vector must be physically shipped to each chosen expert's device, processed there, and shipped back to be combined. This is the all-to-all communication pattern, the third member of the collective family alongside the all-reduce of Chapter 15 and the reduce-scatter of Chapter 16, and it is the subject of Section 17.5. The reason load balance matters so much, the reason this section spent so long on collapse and capacity, is that an unbalanced router turns this all-to-all into a traffic jam: the device holding the popular expert receives a flood of tokens while others sit idle, and the whole layer runs at the speed of its busiest hop. Routing quality and communication cost are the same problem viewed from two ends.