Sparse routing and gating networks: Visual walkthrough

This page rebuilds the central result of Sparse routing and gating networks — the sparse Mixture-of-Experts forward pass — one picture at a time. We assume you have never seen a "gate", a "softmax", or an "expert". Every symbol is earned before it is used.

Step 1 — What is an "input" and what is an "expert"?

WHAT. An input is just a list of numbers. When a Transformer (see Transformer Architecture) reads the word "chlorine", it hands the next layer a vector — a fixed-length list of $d$ numbers. We call it $x$ .

Two more everyday quantities will show up later, so let us name them now:

An expert is a small neural network that eats one $x$ and returns a transformed vector of the same size. Think of it as one specialist doctor: give it a patient $x$ , it gives back an opinion $E_{i} (x)$ .

Figure — Sparse routing and gating networks

In the picture, the orange dot on the left is $x$ . The violet boxes are the $N$ experts. Right now none are lit — we have not decided who to call. That decision is the whole game.

Step 2 — Scoring the experts (why a linear map?)

WHAT. Before we can pick experts, we need a number for each one saying "how relevant are you to this $x$ ?" We produce $N$ such numbers, called logits, with one matrix multiply.

WHY a linear map and not something fancier? We want the cheapest possible judge. A single learned matrix $W_{g}$ turns the $d$ -number input into $N$ scores in one shot, and it is differentiable so training can improve it. Anything heavier would eat the compute savings we are trying to win.

Look at the bars: expert 3 scored high (long magenta bar), expert 1 scored negative. These are not yet probabilities — they can be huge, tiny, or negative. We fix that next.

Step 3 — Turning scores into a probability with softmax (why softmax?)

WHAT. We convert the raw logits $s$ into $N$ numbers that are all positive and add up to $1$ — a probability distribution over experts. That conversion is the softmax.

WHY softmax and not "just divide by the sum"? Plain division breaks on negative scores (probabilities can't be negative). Softmax first sends every score through $exp (\cdot)$ — the exponential — which is always positive and preserves order (bigger logit → bigger weight). It is also smooth, so gradients can flow back into $W_{g}$ . That is exactly the tool for "make positive, keep ranking, stay differentiable".

Same experts as Step 2, but now the bars are all positive, colored by height, and their total area is exactly one full block. Expert 3 grabbed most of the probability; the negative-logit experts got tiny slivers, not zero. That "not zero" is a problem for compute — fix in Step 4.

Step 4 — Sparsifying: keep only the Top- $k$ (why throw the rest away?)

WHAT. We keep only the $k$ largest gate values and set every other gate to $0$ . Here $k$ is a tiny fixed number like $1$ or $2$ ; $N$ can be $128$ or $2000$ .

WHY? A doctor's opinion costs money to compute — running an expert is the expensive part. Softmax gave every expert a nonzero weight, so naively we'd run all $N$ . But most weights are near-zero noise. By hard-cutting to the top $k$ , we run only $k$ experts and skip $N - k$ of them entirely. This is the heart of Conditional Computation: spend compute only where it matters.

With $k = 2$ , only the two tallest bars survive (bright), the rest are grayed to $0$ . The two survivors are re-scaled so their heights add back to a full block. Only these two experts will actually run.

Common mistake Degenerate Top-

$k$ : ties and too-few experts Two corner cases the picture hides:

Ties for the $k$ -th place. If two experts have exactly equal gate values right at the cutoff, "the top $k$ " is ambiguous. In practice we break ties deterministically — e.g. by lowest expert index — so the choice is reproducible. With real-valued logits exact ties are vanishingly rare, but the tie-break must exist or two runs could disagree.
Fewer than $k$ usable experts. If $N < k$ , or a capacity limit (Step 6) has already filled every other expert, a token may end up with fewer than $k$ live experts. Then Top- $k$ simply keeps whatever remains and the renormalization denominator $\sum_{j \in TopK} G (x)_{j}$ runs over that smaller set — the sum-to-one guarantee still holds. In the extreme where no expert is available, the token is dropped (passed through unchanged via a residual), which is exactly why capacity limits are set with slack.

Step 5 — Combining the chosen experts into the output

WHAT. Run each surviving expert on $x$ , weight its output by its (renormalized) gate value, and add them up. That sum is the layer's output $y$ .

WHY a weighted sum? The gate value $G^{'} (x)_{i}$ already encodes "how much do I trust expert $i$ for this $x$ ". Multiplying each expert's answer by its trust and summing is the natural blend — a confident expert dominates, a marginal one contributes a whisper.

Two bright arrows (the two chosen experts) flow into the sum node; the gray experts send nothing. Their outputs, scaled by the little gate weights, merge into the single orange output $y$ .

The compute payoff, in one line. A token pays for $k$ experts, each an FFN of hidden width $d_{ffn}$ (Step 1) and hence cost $O (d \cdot d_{ffn})$ : $active cost = O (k \cdot d \cdot d_{ffn}) + tiny gate O (N \cdot d)$

$k$ — experts actually run (fixed, small).
$d$ — input width; $d_{ffn}$ — expert hidden width, the expensive dimension.
$N \cdot d$ — the cost of the gate matrix multiply, tiny next to the expert FFNs.

The cost depends on $k$ (fixed), not on $N$ . So we can grow $N$ — and total parameters — almost for free. This is the sublinear scaling MoE is famous for.

Step 6 — Backprop through a hard choice (why Top- $k$ needs a trick)

WHAT. Top- $k$ makes a discrete decision — expert in or out. There is no smooth slope on "should this expert have been in the set?", so the ordinary chain rule cannot push a gradient through the selection itself.

WHY it matters. Consider the two kinds of experts after a forward pass:

Selected experts do get a gradient: their output $E_{i} (x)$ is multiplied by the live, differentiable weight $G^{'} (x)_{i}$ , so both the expert and that gate value are trained normally.
Non-selected experts get zero gradient. The gate learns nothing about whether skipping them was wise, because their contribution was hard-cut to $0$ before any loss was computed.

PICTURE + FIX. The standard workarounds:

Keep the soft gate weight on the backward path. Because the surviving weight $G^{'} (x)_{i}$ is a real, differentiable number, the gate does receive a usable gradient for the experts it did pick — the selection is discrete but the magnitude is smooth. This alone gives the gate a learning signal.
Straight-through estimator. Treat the hard 0/1 mask as the identity during the backward pass: forward uses the discrete Top- $k$ , but backward pretends the mask was the smooth softmax, letting a little gradient reach even near-miss experts.
Load-balancing loss (below). An auxiliary term that adds gradient pressure spreading tokens across all experts, so no expert is starved of data long enough to fossilize.

The figure contrasts the solid gradient path through a selected expert (bright arrow back to the gate) with the dead, dashed path of a skipped expert — and shows the straight-through shortcut that revives a trickle of signal.

Step 7 — The degenerate case: collapse, and the balancing loss

WHAT. What if the gate learns "Expert 3 is best for everything"? Then Top- $k$ always picks expert 3, the other $N - 1$ experts never run, never learn, and your $128$ -expert model behaves like a $1$ -expert model. This is collapse, and it must be handled or the whole idea fails.

WHY it happens. Early in training one expert is randomly slightly better; the gate prefers it; that expert gets all the data and improves more; the gate prefers it even harder. A runaway feedback loop — a Long-tail Distribution over experts where one expert eats everything.

THE FIX — a differentiable balancing loss. Define two per-expert quantities over a batch of $B$ tokens:

The Switch Transformer auxiliary loss adds: $L_{aux} = α N \sum_{i = 1}^{N} f_{i} \cdot P_{i}$

$α$ — a hyperparameter (a knob you set, typically around $1 0^{- 2}$ ) that trades off balancing pressure against the main task loss: too small and experts collapse, too large and the model routes uniformly at random and stops specializing.
$f_{i}$ is treated as a detached constant — gradients flow only through the differentiable $P_{i}$ .
The factor $N$ keeps the loss $O (1)$ for any expert count.

Common mistake "The loss is minimized only at perfectly uniform routing"

That over-sells it. With $f$ held fixed, $\sum_{i} f_{i} P_{i}$ is a linear function of $P$ over the simplex $\sum_{i} P_{i} = 1$ , so it is actually minimized by dumping all of $P$ onto the single smallest- $f_{i}$ expert, and if $f$ happens to be uniform the term is constant in $P$ . So $L_{aux}$ is not a function whose lone minimizer is uniform routing.

What it really does is provide a gradient direction each step: $\partial L_{aux} / \partial P_{i} \propto f_{i}$ , so the gate is pushed to lower probability on the currently overloaded experts (large $f_{i}$ ) and raise it on underloaded ones. Because $f$ is recomputed every batch, this is a coupled feedback loop whose stable fixed point — where the pushes cancel — is the balanced state $f_{i} \approx P_{i} \approx 1/ N$ , at which the loss takes the value $1$ . That fixed point, not a static global minimum, is why the loss works.

Left bars: collapsed load — one expert towers, the rest starve. Right bars: after repeated aux-loss nudges, load spreads toward the dashed uniform line $1/ N$ .

The one-picture summary

The whole pipeline in one frame: $x \to$ linear gate $\to$ softmax $\to$ Top- $k$ cut (a discrete step needing a backprop trick) $\to$ run only the chosen experts $\to$ weighted sum $= y$ , with the balance loss watching over expert usage.

Recall Feynman retelling — say it in plain words

An input arrives as a list of numbers. A cheap little judge (one matrix) scores every expert. We squash those scores through softmax so they become positive and sum to one — a vote share per expert. We keep only the two biggest votes and throw the rest away, because running an expert costs money and we only want the best two. We re-scale those two votes so they add to one, run just those two experts on the input, weight each answer by its vote, and add them up — that sum is the output. The keeping-the-top-two step is a hard yes/no with no slope, so to train the judge we lean on the smooth vote sizes it did assign (plus a straight-through shortcut) to carry gradient back. The only remaining danger is one expert hogging all the votes and the others starving; a balancing loss keeps nudging probability off overused experts each batch until the work settles into an even spread. Because we always run just two experts no matter how many exist, we can pile up thousands of experts almost for free.

Recall Quick self-test

What does softmax fix that plain division cannot? ::: Plain division breaks on negative logits; softmax runs every score through $exp$ first, making them all positive while keeping their order, and stays differentiable. Why does active compute grow with $k$ and not $N$ ? ::: We only run the top- $k$ experts, so cost is $O (k d d_{ffn})$ where $d_{ffn}$ is the expert hidden width; the extra $N - k$ experts are never computed, only the tiny $O (N d)$ gate scales with $N$ . Why can't ordinary backprop pass through Top- $k$ , and what do we do? ::: Top- $k$ is a discrete select — no slope on membership. We rely on the smooth surviving gate weights for selected experts and use a straight-through estimator to leak a little gradient to near-miss ones. Is $L_{aux} = α N \sum_{i} f_{i} P_{i}$ minimized only at uniform routing? ::: No — with $f$ fixed it is linear in $P$ . It supplies a per-step gradient $\propto f_{i}$ pushing probability off overloaded experts; the balanced state $f_{i} \approx P_{i} \approx 1/ N$ (loss value $1$ ) is its stable fixed point, not a static minimizer. What is $α$ ? ::: A hyperparameter weighting the balance loss against the task loss — too small lets experts collapse, too large forces near-random routing.

See also: Model Parallelism (how the $N$ experts get spread across devices), Knowledge Distillation (compressing a huge MoE into a dense student), Neural Architecture Search, Attention Mechanisms, Batch Normalization.