Backprop through
a MoE router

The layer, exactly

One token at a time. A token is a vector x ∈ ℝ^d. The layer holds E experts and activates k of them.

Everything hinges on one structural fact: the router's output is a scalar multiplier on each expert's output vector. The router does not touch the token content path directly. It only scales. So the only way loss information can reach the router is by asking, for each expert, "how much did scaling you up help?"

The worked example

Every number on this page comes from this one setting. d=2, E=4, k=2 — small enough to verify by hand, large enough to show the non-selected experts.

Where backprop enters

Nothing special happens here. The MoE block sits inside a transformer, typically after a residual add and a norm, so ȳ = ∂L/∂y is handed down by whatever is above. In our example ȳ = [0.4, −0.3].

The one thing worth noting: because y is a sum over selected experts, this same ȳ gets broadcast unchanged to every selected branch. Addition is a gradient fan-out. All the differentiation between experts happens in the next step.

The gate gradient is a dot product

This is the step that makes a router trainable at all, and it is worth staring at.

The output is y = Σ_e g_e E_e(x). Differentiate with respect to the scalar g_e. The product rule on a scalar-times-vector gives a vector, and the chain rule contracts it against ȳ:

Read it as an alignment test. −ȳ is the direction the output wants to move to reduce loss. If expert e's output vector points that way, the dot product ȳ · E_e(x) is negative, the gradient is negative, and descent raises that gate. If the expert points the wrong way, the gate gets lowered.

So expert 2 here is actively unhelpful for this token — its output has a large component along the loss-increasing direction — and the router will learn to send this token elsewhere. Expert 4 is mildly helpful and gets promoted. Credit assignment for routing is just cosine alignment between an expert's answer and the direction the loss wants to go, scaled by both magnitudes:

Two consequences fall straight out of that formula, and both bite in practice:

• An expert that outputs near-zero gets a near-zero gate gradient. A dead expert is invisible to the router, which is one half of the dead-expert failure loop.
• Expert output norm is a confound. An expert with a large output norm produces a large-magnitude gate gradient regardless of whether it is right. Norms drifting apart across experts silently reweights routing pressure. This is part of why MoE blocks are sensitive to what normalization sits after the expert.

The other branch out of y is the ordinary one — into the expert body:

Note the gate appears here too, as a plain scale factor. An expert with a small gate learns slowly, because every gradient reaching its weights is multiplied by that gate. This is the second half of the dead-expert loop: low gate → small weight updates → the expert stays generic → its outputs stay unaligned → gate stays low. Nothing in the main loss breaks this cycle, which is exactly why the auxiliary loss exists.

Top-k has no useful derivative

Between the gate and the probability sits a hard selection. This is the part that makes MoE training qualitatively different from training a dense layer.

TopK maps a probability vector to an index set. Perturb p slightly and, almost everywhere, the index set does not move at all — the function is piecewise constant, so its Jacobian is exactly zero. At the measure-zero set where two probabilities cross, it jumps, and the derivative does not exist.

So we do not differentiate through the selection. We differentiate through the values of the gates that were selected, and treat the mask as a constant for this backward pass:

Some caveats on the "no gradient" framing, because they matter:

• Variant A leaks a little. If the softmax is taken over all E logits before selection, unselected experts still get logit gradient through the softmax denominator. See the next section — the number is not zero.
• Variant B does not. If you top-k first and softmax over only the k survivors (Mixtral-style), unselected logits get exactly zero, full stop.
• The loss is still continuous in the parameters. When the ranking flips, the gate value at the crossover is equal on both sides, so y does not jump — only its slope does. Training is piecewise smooth, not discontinuous, which is why plain SGD works at all here.

The softmax Jacobian spreads the signal

Two sparse numbers go in. Four dense numbers come out. This is where a competition between experts gets created out of what was independent per-expert feedback.

The softmax Jacobian is the standard one:

s̄ is the probability-weighted mean of the incoming signal — it is the baseline. Every logit is scored against it. A logit only rises if its expert beat the current average expert.

Note also that expert 4 ends up with a negative logit gradient of larger magnitude than the unselected experts, which is right: it was the one genuinely useful expert here, so it gets promoted hardest. And expert 2, whose gate signal was strongly positive, is the only logit pushed down. The zero-sum structure is the competition: the router cannot raise one expert without lowering the others.

Into the router weights, and back into the token

The router is a plain linear map, so the last step is the textbook one. Write δ_h = ∂L/∂h.

Every row is a scaled copy of the token itself. Descent moves w₂ away from x and moves w₄ toward it — the router is literally storing a prototype per expert, and each step nudges prototypes toward the tokens they handled well. That is the whole learning rule.

Gradient to the token

x was consumed twice — once by the router, once by each running expert — so its gradient is a sum over both paths:

Forgetting the router path is a classic bug in hand-written fused MoE kernels. It is small in magnitude but it is the only thing that lets earlier layers shape their outputs to be routable — without it, layers below get no signal about how to make their representations easy to dispatch.

Expert weights

This is where sparse MoE gets its efficiency in the backward pass too: only k of E expert weight matrices receive any gradient for a given token. It is also why MoE optimizer states behave oddly — with Adam, an expert's second-moment estimate keeps decaying while it sits idle, so it takes an artificially large step the moment it is picked again.

Variant A vs Variant B: same forward, different gradients

The order of softmax and top-k is a one-line implementation choice with a real effect on what the router can learn.

Variant B's derivation, on the same numbers. With Z = p₂+p₄ = 0.69978:

With k=2 the two gradients are exact negatives — the competition is now strictly between the two chosen experts and completely sealed off from the rest. Note the magnitude is also ~20% larger than Variant A's, because renormalization removes the Z < 1 attenuation.

Which is better is genuinely contested. Variant B gives cleaner, larger, better-scaled gradients and a k-invariant output magnitude. Variant A gives the router a weak global exploration term for free. In practice Variant B has won on recent open models, and they compensate with stronger auxiliary balancing — or, increasingly, with bias-based balancing that sidesteps the aux gradient entirely.

Auxiliary losses, and why they exist

The main loss cannot fix imbalance, because it only ever sees the experts that were already picked. So a second gradient is injected directly into h.

Load-balancing loss · step 1: name the thing you want to minimize

Fix one MoE layer and one batch of T tokens. Two different numbers describe "how much did expert e get used", and keeping them apart is the entire trick.

Imbalance is a property of f, not of P. You can have beautifully uniform probabilities and still overflow one expert's buffer, because top-k takes a hard slice through those probabilities. So the quantity to minimize is a second moment of f, and it has a very readable normalization:

That identity I = 1 + E²·Var(f) is the whole justification for the shape of the loss. Minimizing Σ_e f_e² is minimizing the variance of expert load — just written in a form with no subtraction in it, so it survives being differentiated.

Step 2: make it differentiable by swapping one factor

I(f) is the target, and its gradient is identically zero. The Switch Transformer move is to write Σ f_e² as Σ f_e·f_e and replace one of the two factors with its smooth counterpart:

Three conventions differ between papers, and all three cause real bugs. The k in the denominator: Switch used k=1, so a lot of code writes f_e = (1/T)·#{tokens} and silently stops summing to 1 when someone raises k. Where the E lives: DeepSeek folds it into the load term, so their formula reads α Σ f_eP_e with f_e = (E/kT)·#{…} — the same loss, a different-looking equation. What "the batch" means: computing f from one micro-batch on one device balances load only within that shard, so large runs all-reduce the counts across the expert-parallel group before forming the loss. Some models instead compute it per sequence, which is a much stronger constraint and a different training signal.

Step 3: read the number

Because of the E factor, L_aux/α sits on a fixed scale no matter how many experts you have, which makes it the one MoE diagnostic worth putting on a dashboard. Log that, not the raw loss: 1.0 is uniform, up to ~1.3 is healthy, past ~2 a real fraction of your parameters is idle, and a value pinned near E means the router has collapsed and the run will not recover on its own.

It is computed per layer, on every batch, and added to whatever the model is actually training on:

Why a loss that's linear in P balances anything

This is the part the formula hides, and it is the real answer to "how does that expression cause balance". Hold f fixed, as the definition demands, and L_aux is a linear function of P. A linear function on the probability simplex is minimized at a vertex — so the loss, taken literally at one instant, wants the router to send everything to whichever expert is currently emptiest. That is not balance. That is the opposite collapse.

Balance comes from the fact that f is re-measured from scratch every step. Each expert's coefficient is its own current load, so the push always points away from whoever is full and toward whoever is empty, and it reverses the instant the ordering does. It is a repulsive force with a moving reference point, not a bowl with a minimum at the bottom. Uniform load is the fixed point of the dynamics, not the minimum of the loss — which is exactly why it can oscillate, and why α behaves more like a damping constant than like a regularization strength.

The gradient it produces

Once the definitions are pinned down this part is trivial, because f is a constant:

A flat per-expert constant is not yet a balancing force — it becomes one only after passing through the softmax Jacobian from the ∂L/∂h section, which scores each expert's load against the probability-weighted average load. Above average, logits go down; below average, up.

Worked, with T = 8, E = 4, α = 0.01, counts [4, 3, 1, 0]:

This term reaches every expert on every token, including ones that never fire. That is the property the main loss lacks, and it is the whole point. Note it is also the one term that is genuinely batch-coupled — a token's router gradient depends on how the other tokens in the batch were routed, which means changing your batch size or your expert-parallel sharding changes the router's gradient even with everything else fixed.

Router z-loss

A numerical-stability term, not a balancing one. It penalizes the log-partition function, keeping logits small so that the exponentials stay in range in bf16:

Every entry is positive, so every logit is pushed down together — this is the one router term that does not sum to zero across experts. It shifts the whole logit vector without changing the softmax much, which is exactly what you want from a regularizer that is supposed to be scale-only.

Four things that change the gradient in a real system

Capacity overflow zeroes the gradient too

Each expert has a fixed buffer of capacity = capacity_factor · k · T / E slots. Tokens past that are dropped: their contribution to y is zeroed, so on the way back ∂L/∂g_e = ȳ · 0 = 0 for that token-expert pair. A dropped token teaches the router nothing about the expert that rejected it. The residual connection is the only reason the token survives at all. Worth knowing: this makes the drop pattern a silent, batch-order-dependent source of gradient noise.

Combine is the transpose of dispatch

In expert-parallel training, forward is: build a dispatch mask, all-to-all the tokens to their expert's device, run, all-to-all back, combine with gates. The backward pass is the exact mirror — the backward of the combine all-to-all is a dispatch all-to-all. If you write a custom kernel, the dispatch/combine pair must be exact transposes or your router gradient will be silently wrong in a way that still trains, just worse.

Noisy top-k makes load itself differentiable

The original sparsely-gated MoE added learned noise before selection: h_e = (W_gx)_e + ε·softplus((W_nx)_e), ε ∼ 𝒩(0,1). With noise, the probability that expert e lands in the top-k is a smooth function of the parameters — computable in closed form from the Gaussian CDF — so you can put a load-balancing loss on that smooth quantity and get a real gradient into W_n, rather than relying on the f·P surrogate. Most large-scale systems dropped this for simplicity; a light input jitter is the vestigial remnant.

Loss-free balancing removes the aux gradient entirely

The newest approach (DeepSeek-V3) keeps a per-expert bias b_e used only for the top-k comparison, never for the gate value. After each step it is updated by a plain rule outside of autograd — decrement overloaded experts, increment idle ones. Since b never enters the gate, it contributes no gradient, and the router's gradient is purely task-driven again. This directly targets the tension in the callout above: balance without paying for it in the task gradient.

The backward pass, written out

Autograd handles all of this for you. Writing it manually once is still the fastest way to be sure you know where each term comes from — and it is what you need if you fuse the layer.

# ---------- forward (Variant A) ----------
h = x @ Wr.T                                  # [T,E]
p = softmax(h, dim=-1)                        # [T,E]
topv, topi = p.topk(k, dim=-1)                # hard selection
g = zeros_like(p).scatter(-1, topi, topv)     # [T,E], k nonzeros per row
y = sum(g[:, e:e+1] * expert[e](x) for e in range(E))

# ---------- backward, given gy = dL/dy  [T,d] ----------
dg = zeros_like(g)                            # [T,E]
dx = zeros_like(x)

for e in range(E):
    m = g[:, e] != 0                          # tokens routed to expert e
    if not m.any(): continue                  # no gradient at all — skipped
    Ee = expert[e](x[m])                      # [n_e, d]  (recompute or stash)

    dg[m, e] = (gy[m] * Ee).sum(-1)           # dot product — the gate gradient
    dEe      = g[m, e:e+1] * gy[m]            # gate scales the expert gradient
    dtheta[e], dx_e = expert[e].backward(dEe)
    dx[m]   += dx_e                           # expert path into the token

dp  = dg                                      # top-k mask: already zero elsewhere
dh  = p * (dp - (dp * p).sum(-1, keepdim=True))     # softmax Jacobian

# ---- router-only terms, added straight into dh ----
f   = (g != 0).float().mean(0)                # load fractions — NO gradient
c   = alpha * E * f / T                       # dL_aux/dp, constant per expert
dh += p * (c - (c * p).sum(-1, keepdim=True))
dh += cz * 2 * logsumexp(h, -1, keepdim=True) * p   # z-loss

dWr = dh.T @ x                                # [E,d]
dx += dh @ Wr                                 # router path — easy to forget

Two checks that catch most mistakes:

• Zero-sum invariant. Before the z-loss line, dh.sum(-1) must be zero to floating-point tolerance for every token. Any softmax-backward bug breaks this immediately. The z-loss term deliberately breaks it — check before adding it.
• Finite differences on the boundary. Run gradcheck in float64 on a two-expert toy, but choose inputs away from ties. Near a top-k crossing, finite differences will disagree with the analytic gradient, and that disagreement is real, not a bug — it is the discontinuity from Figure 3.

Nine things to remember

• 01
  The router's gradient is manufactured, not received. ∂L/∂g_e = ȳ · E_e(x) — a dot product between the loss direction and the expert's answer.
• 02
  Top-k contributes zero Jacobian. You differentiate the selected gate values, with the mask held constant.
• 03
  The gate scales the expert's own gradient, so low-gate experts learn slowly — the second half of the collapse loop.
• 04
  The softmax Jacobian turns independent per-expert feedback into a zero-sum competition around the baseline s̄.
• 05
  In Variant A, unselected experts get −p_js̄ through the normalizer. In Variant B they get exactly zero.
• 06
  The balancing loss is the only term that reaches every expert on every token. It is also the only one that couples tokens within a batch.
• 07
  Load f is a constant in the graph. Only P carries gradient. Wrapping f in a stop-gradient is not optional.
• 08
  ∂L/∂x has two paths. The router path W_rᵀδ_h is small and easy to drop in a fused kernel — and it is what teaches lower layers to be routable.
• 09
  Dropped tokens produce zero gate gradient. Capacity factor is a gradient hyperparameter, not just a memory one.