MoE Routing.

The router is the brain of an MoE layer. It decides, for every token at every layer, which experts run. Get it right and experts specialise usefully; get it wrong and the model either collapses onto a few experts (5.1's failure mode) or routes incoherently. This article unpacks how routing actually works — the gating math, the choices that matter, and the non-differentiability problem hiding inside top-k.

The basic gate

The router is one small linear layer mapping a token's hidden state to a score per expert, followed by a selection step:

h        = token hidden state, shape (d_model,)
logits   = W_g · h                  # (N,) one logit per expert, W_g: (N, d_model)
gates    = softmax(logits)          # (N,) probabilities summing to 1
top_k    = argtopk(gates, k)        # indices of k highest
y        = Σ_{i in top_k} gates_i · Expert_i(h)

W_g is tiny — N × d_model, negligible next to the experts themselves. That a single linear layer governs a model with hundreds of experts is part of MoE's elegance and part of its fragility: the whole routing policy is one matrix.

Softmax before or after top-k?

A subtle but real design choice: do you softmax over all N logits and then pick top-k, or pick top-k logits and softmax only those? It changes the gate magnitudes.

# Option A — softmax over all, then select (gates don't sum to 1 after selection)
gates = softmax(logits)[top_k]            # e.g. [0.34, 0.12] — sums to 0.46

# Option B — select, then softmax over the k chosen (gates sum to 1)
gates = softmax(logits[top_k])            # e.g. [0.74, 0.26] — sums to 1.0

Option B (renormalise over the chosen experts) keeps the combined output at a consistent scale regardless of how confident the router was, which interacts more cleanly with the residual stream and normalisation. DeepSeek-V2/V3 use a softmax-then-normalise variant. The choice affects gradient flow and the effective scale of the MoE output — small detail, real consequences.

Why top-k breaks differentiability

Here's the catch that shapes everything. argtopk — picking the k largest — is a discrete, non-differentiable operation. You can't backprop through "which experts were chosen": the selection is a hard step function of the logits, derivative zero almost everywhere. So how does the router learn anything?

The answer is the gate weights. The chosen experts' outputs are multiplied by their gate values gates_i, and those are differentiable. Gradient flows back through the gate value into W_g: if expert i produced a useful output, the loss pushes gates_i up, which raises logits_i, making expert i more likely to be chosen next time. The selection itself is non-differentiable, but the weighting of selected experts carries the learning signal. The router learns "how much to trust an expert I picked," and that indirectly shapes "which to pick."

Top-k routing learns through the gate weights of the experts it actually used — never through the ones it didn't. An expert that's never selected gets no gradient and never improves. That's exactly the starvation dynamic behind routing collapse, and why explicit balancing (5.4–5.6) is mandatory.

Noisy top-k gating

Shazeer's 2017 router added tunable Gaussian noise to the logits before selection:

logits = W_g · h + StandardNormal() · softplus(W_noise · h)
top_k  = argtopk(logits, k)

Two reasons. First, exploration: noise occasionally bumps a non-favourite expert into the top-k, giving starved experts a chance to receive gradient and improve — a direct counter to collapse. Second, it smooths the load distribution by making selection stochastic rather than deterministic. Modern large MoEs (Switch, DeepSeek) often drop noise in favour of explicit balancing losses or bias terms, but the idea — inject exploration so every expert gets some traffic — survives.

Token-choice vs expert-choice routing

Two opposite framings of the assignment problem:

• Token-choice (standard, DeepSeek): each token picks its top-k experts. Simple, but nothing guarantees experts receive equal load — a popular expert can be swamped while others starve. Needs explicit balancing.
• Expert-choice (Zhou et al. 2022): flip it — each expert picks its top-capacity tokens. Load is balanced by construction (every expert takes exactly its quota), but a token might be chosen by zero experts (dropped) or many (uneven compute per token).

Token-choice matches autoregressive decoding better (you can't see future tokens to do expert-choice at inference), so it dominates production LLMs. The price is that load balancing must be solved separately — the subject of 5.4–5.6.

Worked routing example

4 experts, top-2, softmax-then-select-then-renormalise:

token h → logits = [2.0, 0.5, 1.5, -1.0]

step 1: select top-2 logits      → experts {0, 2}, logits [2.0, 1.5]
step 2: softmax over chosen       → [exp(2.0), exp(1.5)] / sum
                                   = [7.39, 4.48] / 11.87
                                   = [0.623, 0.377]
step 3: output = 0.623·E0(h) + 0.377·E2(h)

Experts 1 and 3 do nothing for this token and receive no gradient from it. Run this over a batch and tally how many tokens hit each expert — that tally is the load, and keeping it even is the balancing problem.

Routing in DeepSeek specifically

DeepSeek's router has a few distinctive choices that 5.6 expands on:

• Fine-grained experts — many small experts (256 in V3) instead of few large ones, so each token's top-k draws from a richer menu and specialisation is finer.
• Shared experts — a few experts every token always uses (no routing), capturing common knowledge so routed experts can specialise on the rest.
• Sigmoid gate + bias balancing — V3 scores experts with a sigmoid and adds a per-expert bias term, adjusted during training to equalise load without an auxiliary loss. This is the auxiliary-loss-free method of 5.6 — the router's selection uses logits+bias, but the gate values used for the weighted sum exclude the bias, so balancing doesn't distort the learned combination.

What this sets up

Routing gives the mechanism; it doesn't give balance. Because gradient only reaches selected experts and token-choice guarantees nothing about load, an unmanaged router collapses. Article 5.3 first looks at what the router's choices actually mean — what experts specialise in — then 5.4 starts the balancing story with the classic auxiliary loss, 5.5 covers capacity limits and token dropping, and 5.6 shows DeepSeek's bias-based scheme that fixes balance without the auxiliary loss's downsides.

References

• Shazeer et al. (2017), Sparsely-Gated MoE — noisy top-k gating. arXiv:1701.06538
• Fedus et al. (2021), Switch Transformers — top-1 routing. arXiv:2101.03961
• Zhou et al. (2022), Mixture-of-Experts with Expert Choice Routing. arXiv:2202.09368
• DeepSeek-AI (2024), DeepSeek-V3 — sigmoid gate + bias balancing. arXiv:2412.19437