Sparse routing and gating networks: Foundations

This page builds every symbol, word, and picture that the parent topic leans on. If the parent wrote $W_{g} x$ , $softmax$ , $TopK$ , $1 [\cdot]$ , or "load", and never stopped to explain it — we stop here. Read top to bottom; each block only uses symbols already earned above it.

0. The picture the whole topic lives inside

Figure — Sparse routing and gating networks

Look at the figure. On the left an input arrow enters a small box (the gating network). That box shines a spotlight on just 2 of the 8 tall boxes (the experts). Only the lit experts do work; the dark ones stay asleep. Their two outputs are blended and leave on the right.

Everything below is just: what is that arrow, what is that box, and how does the spotlight get decided.

1. A vector — the input $x$

Plain words: a row of $d$ dials, each set to some value.
Picture: an arrow in space, OR simpler — a vertical stack of $d$ little number-boxes (see figure s02, leftmost column).
Why the topic needs it: every "token" flowing through a Transformer arrives as such a list of $d$ numbers. The gating network must read this list to make its decision.

The symbol $R$ means "all the real numbers" (any decimal, positive or negative). $R^{d}$ means " $d$ of them stacked". So $x \in R^{d}$ reads: " $x$ lives in $d$ -dimensional real space."

2. A linear layer — $W x + b$

The parent writes $logits = W_{g} x + b_{g}$ . Let us earn every piece.

What "dotted" means — the dot product:

$(row i) \cdot x = \sum_{j = 1}^{d} W_{ij} x_{j}$

Plain words: multiply matching entries and add them all up — one number.
Why THIS tool? We need to turn a length- $d$ input into $N$ scores (one score per expert). A dot product is the simplest way to ask "how much does this input agree with this direction?" — big when input and row point the same way, negative when opposite. Each row of $W_{g}$ is a "taste profile" for one expert; the dot product measures how well the input matches that taste.

The figure shows the length- $d$ input becoming $N = 4$ scores. The coral row of $W$ combines with the input to make score #2; a different row makes each other score.

$logits_{i} = dot product j = 1 \sum d (W_{g})_{ij} x_{j} + (b_{g})_{i}$

Recall Why is a plain dot-product-plus-bias called "linear"?

Because doubling the input doubles the score (before the bias), and adding two inputs adds their scores. No curves, no thresholds — a straight-line relationship. ::: A straight-line relationship: scaling or adding inputs scales or adds outputs; the graph has no bends.

The word logits just means "raw scores, not yet turned into probabilities." They can be any size, positive or negative.

3. Turning scores into a choice — the softmax

The parent writes $G (x)_{i} = \frac{exp ( logits _{i} )}{\sum _{j} exp ( logits _{j} )}$ . Two new symbols: $exp$ and $\sum$ .

Why THIS tool (why not just use the scores directly)? We want probabilities: numbers that are all positive and add to 1, so we can read them as "chance of picking this expert." Raw logits can be negative — you can't have a negative chance. $exp$ fixes both problems: it makes every score positive, and it magnifies gaps (a slightly bigger logit becomes a much bigger weight), giving a confident-but-soft choice.

The figure shows the same 4 logits before and after softmax. Notice: the biggest logit (coral) grabs the biggest share, the negative logit still gets a small positive share, and the four output bars stack to height 1.

Recall Why must softmax outputs sum to 1?

Because the denominator is the sum of every numerator, so adding all the fractions puts the same total on top and bottom. ::: The denominator equals the sum of all numerators, so the fractions add to (total)/(total) = 1.

Because $exp$ is smooth and the division is smooth, softmax is differentiable — that word is the next stop.

4. Differentiable, gradients, and backprop

The parent says softmax is chosen because it is "differentiable (enables backprop)." Let's ground that.

Picture: stand on a hillside (the loss surface). The gradient is the arrow pointing straight uphill; training walks the opposite way, downhill, to reduce error. See figure s04.

Why the topic cares: neural networks learn by backpropagation — sending "which way to nudge" signals backward through every operation. An operation with a jump or corner (like a hard "pick the top one") has no well-defined nudge there, so no signal passes. This is exactly why the parent's top- $k$ step is troublesome and why the auxiliary loss routes gradients through the smooth $P_{i}$ instead of the jumpy $f_{i}$ . Hold that thought for §7.

5. Picking the best few — $ar g max$ and TopK

Picture: sort the bars from §3 tallest-first, keep the front $k$ , throw the rest to zero. This is a staircase decision — nudging a logit a hair usually changes nothing, until it crosses a neighbour and the choice snaps. That snap is the non-differentiable corner from §4.

Recall Why is TopK not differentiable?

Because its output stays flat as you nudge inputs, then jumps abruptly when the ranking flips — flat-then-jump has no smooth slope. ::: The selected set is constant under tiny nudges and jumps discontinuously when two scores swap order; no smooth gradient exists.

6. The indicator $1 [\cdot]$ and counting load

The parent defines $Load_{i} = \sum_{x} 1 [i \in TopK (G (x))]$ .

Why the topic needs it: to count how many tokens landed on expert $i$ , we tally a 1 each time that expert was chosen. Summing $1 [i \in TopK]$ over the whole batch is literally "count the visits."
Note the trap: $1 [\cdot]$ is a hard yes/no — it inherits the staircase problem. That is why load counting is a monitoring tool, not something you can backprop through directly.

7. Fractions $f_{i}$ , mean-probability $P_{i}$ , and why they differ

Now every symbol in the parent's balance loss $L_{aux} = α N \sum_{i} f_{i} P_{i}$ is buildable:

Symbol	Plain words	Smooth?	Built from
$f_{i}$	fraction of tokens that actually got sent to expert $i$	❌ staircase	§5 TopK + §6 indicator
$P_{i}$	average softmax probability the gate gave expert $i$	✅ smooth	§3 softmax
$α$	small dial (a chosen constant) scaling the penalty	—	just a number
$N$	number of experts	—	a count

Why two versions of "how busy is expert $i$ "? Because we need one to measure reality ( $f_{i}$ , the true count) and one to push the gate ( $P_{i}$ , the smooth thing gradients can travel through, from §4). The product $f_{i} P_{i}$ multiplies "how overloaded it is" by "the knob we can turn," so training turns down the knob on the crowded experts.

Recall Why can gradients flow through

$P_{i}$ but not $f_{i}$ ? $P_{i}$ is an average of smooth softmax outputs (§3–4); $f_{i}$ counts TopK hits, a staircase with no slope (§5). ::: $P_{i}$ averages differentiable softmax values so it has a slope; $f_{i}$ counts discrete TopK selections which jump, so it has no usable gradient.

8. Big-O — reading the cost claims

The parent writes costs like $O (k \cdot d \cdot d_{ffn})$ .

Why it matters here: the headline of the whole topic is that active compute is $O (k \cdot d \cdot d_{ffn})$ — it depends on $k$ (a tiny constant like 1 or 2), not on $N$ , the total expert count. So you can pile on experts (grow total parameters) while per-token cost barely moves. That is conditional computation in one line. Here $d_{ffn}$ is just the hidden width of each expert's feed-forward network — another size number.

Prerequisite map

Equipment checklist

Cover each answer and see if you can state it before revealing.

x \in R^{d}

means

an ordered list of

d

real numbers — the input token, one number per feature.

W \in R^{N \times d}

a grid of numbers with

N

rows and

d

columns; each row is one expert's "taste profile."

The dot product

(row) \cdot x

computes

multiply matching entries and sum them — a single number measuring agreement between the row and the input.

exp (z)

always returns

a positive number; it grows fast and turns any logit into a positive weight.

Softmax guarantees its outputs

are each between 0 and 1 and sum to exactly 1 — a valid probability distribution.

"Differentiable" means

smooth, no jumps or corners, so you can always ask which way and how much to nudge (the gradient exists).

ar g max_{i}

returns

the index (position) of the largest entry, not its value.

TopK is not differentiable because

its selection stays flat then snaps when the ranking flips — a staircase with no slope.

1 [condition]

equals

1 if the condition is true, 0 if false — used to count expert visits.

f_{i}

versus

P_{i}

f_{i}

is the true fraction of tokens routed to expert

i

(staircase, no gradient);

P_{i}

is the mean softmax probability (smooth, carries gradients).

O (k \cdot d \cdot d_{ffn})

tells us

active compute scales with

k

(tiny constant), not with

N

— so total parameters can grow almost for free.