Routers, Experts, and Load Balancing

Lesson 2 of 3 in Mixture of Experts: Sparse by Design.

The router is almost embarrassingly small. It is one learned matrix of shape d_model × N: multiply the token’s hidden state by it and you get N raw scores — router Logits, one per expert. A Softmax turns them into probabilities, the top k experts win, and their outputs are combined, weighted by those probabilities (usually renormalized over just the winners). A billion-parameter routing decision made by a matrix that might be a few thousand times smaller than the experts it commands.

Two properties matter. First, routing is per token and per layer: the same token can visit different experts at every MoE layer, and two adjacent tokens in the same sentence can take completely different paths. Second, the router is trained jointly with the experts by ordinary gradient descent — nobody labels which expert should get which token. The division of labor is emergent.

One token through one MoE layer (k = 2, N = 8)

  1. Token hidden state

    The d_model vector arriving from the attention sublayer.

  2. Router: linear scores over 8 experts

    One matrix multiply: hidden state × (d_model × N) matrix → N logits.

  3. Top-k selection

    Softmax the logits, keep the k highest-probability experts. Everything below the cut gets no compute.

  4. Expert 3 (FFN)

    An ordinary feed-forward network — one of eight identical-shaped copies.

  5. Expert 7 (FFN)

    The second winner. Runs in parallel with Expert 3.

  6. Six unselected experts: no compute

    They stay resident in memory — the next token may need them — but contribute nothing to this token.

  7. Weighted sum by router probability

    output = w₃ · E₃(x) + w₇ · E₇(x), with w₃ and w₇ renormalized to sum to 1.

  8. Add back to the residual stream

    Exactly where a dense FFN’s output would have been added.

Joint training hides a trap. Early in training, some experts — by pure chance — get slightly better at the tokens they happen to receive. The router notices and sends them more tokens. More tokens mean more gradient updates, so those experts improve further, so the router prefers them more strongly. It is a rich-get-richer loop, and left alone it converges on a degenerate model: a handful of overworked experts doing everything while the rest never learn anything useful. The literature calls this expert collapse (or routing collapse), and it has been the central engineering obstacle of MoE since Shazeer et al. (2017) documented it.

The auxiliary loss, precisely

The Switch Transformer formulation (Fedus, Zoph & Shazeer, 2021) is the one to know because of its simplicity. For a batch of T tokens routed over N experts, compute two vectors: f_i, the fraction of tokens actually dispatched to expert i, and P_i, the mean router probability the softmax assigned to expert i across the batch. The auxiliary loss is

L_aux = α · N · Σᵢ fᵢ · Pᵢ

added to the ordinary language-modeling loss. The dot product Σ fᵢ · Pᵢ is minimized when both distributions are uniform — every expert getting 1/N of the tokens and 1/N of the probability mass — so the loss pushes toward balance. The factor N normalizes the scale so the loss sits at α under perfect balance regardless of expert count.

One subtlety explains the two-term design: f_i comes from a hard top-k choice, so it is not differentiable — no gradient flows through it. P_i is a softmax output and is differentiable. Multiplying them lets the gradient flow through P_i while f_i supplies the signal about where tokens actually went. The coefficient α is kept small (the Switch authors used on the order of 10⁻²) — large enough to prevent collapse, small enough not to distract from predicting the next token.

Two older companions from Shazeer et al. (2017) still echo through modern recipes. Noisy top-k gating adds learned Gaussian noise to router logits before selection, so near-tied experts trade tokens stochastically and undertrained experts keep getting chances. And capacity factors cap how many tokens an expert may accept per batch: capacity ≈ (tokens per batch ÷ N) × capacity factor. Tokens that overflow a full expert are typically dropped — they skip the FFN entirely and ride the residual stream unchanged to the next layer. A capacity factor above 1.0 buys slack at the price of wasted, padded compute; tuning it is a genuine trade-off, not a formality.

Later work has explored dropping the auxiliary loss in favor of other balancing mechanisms, but the underlying constraint never goes away: hard routing decisions do not backpropagate, so every MoE recipe needs some mechanism that keeps all N experts in the game.

Interactive checkpoint quiz (1 questions) — open this page in a browser to take it.