Worked examples — Sparse routing and gating networks

This page takes the machinery from the parent note on sparse routing and runs the numbers — by hand — through every situation a Mixture-of-Experts layer can hit. If a symbol shows up, we re-earn it here so you never have to flip back.

Before any example, let us define the symbols we will use constantly, from zero.

The exponential $e^s$ just means "$2.71828\dots$ raised to the power $s$." When $s$ is bigger, $e^s$ is bigger; when $s=0$, $e^0=1$; when $s$ is very negative, $e^s$ shrinks toward 0 but never hits it. That "never hits 0" is exactly why softmax gives every expert a tiny but non-zero share — important later.

---

The scenario matrix

Every worked example below is tagged with the cell it covers. Together they fill the whole grid.

#	Cell class	What makes it tricky
C1	Plain positive logits	baseline softmax + top-2, nothing weird
C2	Negative & mixed-sign logits	does softmax still work when scores go below 0?
C3	Zero / tied logits (degenerate)	ties in the argmax, uniform output
C4	Limiting value: one huge logit	softmax → one-hot, the "confident router", overflow-safe form
C5	Renormalization after top-k	why the survivors get re-scaled and by how much
C6	Compute / FLOP accounting	$N$ vs $k$ — the sublinear-scaling payoff
C7	Load balancing (Switch aux loss)	imbalanced batch, compute $f_i,P_i,{\mathcal{L}}_{\text{aux}}$, CV
C8	Expert Choice capacity	experts pick tokens; solve for capacity $C$
C9	Word problem (real-world)	translate a scenario into routing arithmetic
C10	Exam twist	temperature-scaled logits change the winner

---

Worked examples

Example 1 — Plain positive logits (cell C1)

Forecast: guess now — expert 1 wins (biggest logit), but is its probability above 60%? Write your guess.

1. Exponentiate each logit. Why this step? Softmax needs positive weights so the outputs can be probabilities; $e^s$ turns any real score into a positive number. $e^2=7.389,e^1=2.718,e^0=1.000$
2. Sum them. Why? We need a common denominator $Z$ so the three outputs add to exactly 1. $Z=7.389+2.718+1.000=11.107$
3. Divide. Why? Dividing each positive numerator by their sum $Z$ is exactly the softmax formula $G(x{)}_i=e^{s_i}/Z$, guaranteeing the shares are fractions of one whole. $G(x)=\left(\frac{7.389}{11.107},\frac{2.718}{11.107},\frac{1.000}{11.107}\right)=(0.665,0.245,0.090)$
4. Top-1 (argmax). Why? Switch routing keeps only the single biggest gate to run just one expert. Winner $=$ expert 1.

Verify: $0.665+0.245+0.090=1.000$ ✓ (probabilities sum to 1). Expert 1's share is $0.665>0.60$, so the forecast "above 60%" is correct.

---

Example 2 — Negative and mixed-sign logits (cell C2)

Forecast: does a negative logit produce a negative probability? (No — predict why.)

1. Exponentiate. Why? $e^s$ is positive for every real $s$, including negatives — that is the whole point of using the exponential. $e^{-1}=0.368,e^0=1.000,e^3=20.086$
2. Sum. Why? The same normalizing denominator $Z$ that will force the outputs to add to 1. $Z=0.368+1.000+20.086=21.454$
3. Divide. Why? Applying the softmax formula $G(x{)}_i=e^{s_i}/Z$ converts the positive numerators into shares of one whole. $G(x)=(0.0172,0.0466,0.9363)$

Verify: all three are strictly positive despite $s_1=-1$ being negative ✓, and $0.0172+0.0466+0.9363=1.0001\approx 1$ (rounding) ✓. Takeaway: negative logits are fine — they just map to small probabilities, never negative ones.

---

Example 3 — Zero / tied logits, the degenerate case (cell C3)

Forecast: with a perfect tie, what fraction should each expert get?

1. Exponentiate. Why? Same softmax rule applies. $e^0=1$ for all four.
2. Sum and divide (softmax). Why? Building the denominator $Z$ and dividing is the softmax step that turns equal logits into equal probabilities. Every numerator is 1, so $Z=4$ and $G(x)=\left(\frac{1}{4},\frac{1}{4},\frac{1}{4},\frac{1}{4}\right)=(0.25,0.25,0.25,0.25).$
3. Top-1 under a tie. Why care? argmax must break the tie somehow, and which index it returns is framework-dependent (some return the lowest index, some are unspecified). You cannot rely on a particular tie-break. This is the degenerate router: it carries no information, so training must push logits apart.

Verify: $4\times 0.25=1.000$ ✓. Uniform is the maximum-entropy (least confident) gate.

---

Example 4 — Limiting value: one dominant logit (cell C4)

Forecast: roughly what probability does expert 1 grab — 0.9? 0.99? more?

1. Subtract the max, then exponentiate. Why? From the intuition callout above, subtracting $m={\max }_j{s}_j=10$ before $e^{(\cdot )}$ gives the identical answer but never overflows. $s-m=(0,-10,-10)$, so $e^0=1,e^{-10}=0.0000454,e^{-10}=0.0000454$.
2. Sum. Why? The denominator $Z$; here it is $1+0.0000454+0.0000454=1.0000908$, dominated by the top term.
3. Divide. Why? The softmax step $G(x{)}_i=e^{s_i-m}/Z$; because one term swamps $Z$, that expert grabs almost all the mass. $G(x)=(0.999909,0.0000454,0.0000454).$

Verify: sum $=1.0000$ ✓. Expert 1 holds $>0.9999$. Limit statement: as $s_1\to \infty$ with others fixed, $G(x)\to (1,0,0)$ — the router becomes deterministic, and top-1 vs weighted output become identical. This is the "confident router" regime a trained MoE lives in. (You can confirm the plain form $e^{10}/(e^{10}+2)$ gives the same $0.999909$ — the max-subtraction only changes the arithmetic, not the answer.)

---

Example 5 — Renormalization after top-k (cell C5)

Forecast: after we drop experts 3 and 4, do the surviving weights still add to 1? If not, what do we fix?

1. Select top-2. Why? We only pay to run $k=2$ experts, so we keep the two biggest gates. Survivors: experts 1 and 2 with gates $0.50,0.30$.
2. Note the leak. Why? Because zeroing experts 3 and 4 discards their probability, we must check the survivors still form a whole. $0.50+0.30=0.80\ne 1$: we threw away $0.20$ of the probability mass, which would silently scale the activation down by 20%.
3. Renormalize. Why? Dividing each survivor by their sum $0.80$ restores a proper weighted average that sums to 1, undoing the leak from step 2. $G_1^{\prime }=\frac{0.50}{0.80}=0.625,G_2^{\prime }=\frac{0.30}{0.80}=\mathrm{0.375.}$
4. Combine. Why? The MoE output is the gate-weighted mix of the kept expert outputs, $y={\sum }_{i\in \text{TopK}}{G}_i^{\prime }{E}_i(x)$. $y=0.625\times 10+0.375\times 4=6.25+1.5=\mathrm{7.75.}$

Verify: renormalized weights $0.625+0.375=1.000$ ✓. The answer $7.75$ lies between $E_2(x)=4$ and $E_1(x)=10$ ✓ (a weighted average must sit between its inputs). Without renormalization we'd have gotten $0.5\cdot 10+0.3\cdot 4=6.2$, artificially small — the sanity check catches the mistake.

---

Example 6 — Compute accounting: $N$ vs $k$ (cell C6)

Forecast: we made the model 128× bigger in parameters — did per-token compute go up 128×, 2×, or barely at all?

1. Dense cost per token. Why? We need a baseline to compare against; the two FFN matrices (up- and down-projection, introduced above) dominate. $2d{d}_{\text{ffn}}=2\cdot 512\cdot 2048=2,097,152$ multiply-adds. Call this unit $D$.
2. MoE active expert cost. Why? A token runs through exactly $k=2$ experts, each the same size as the dense FFN, so we multiply $D$ by $k$. Active $=k\cdot D=2D=4,194,304$.
3. Gating cost. Why? We must also pay to compute the scores $s=W_{g}x$ — one matrix-vector product of $N\cdot d=128\cdot 512=65,536$ multiply-adds. Compare to $2D\approx 4.2$M — the gate is $\approx 1.5\%$ of expert cost, negligible.
4. Ratio. Why? Dividing MoE active cost by dense cost tells us the real compute overhead we pay. Active MoE / dense $=2D/D=2$ (plus the tiny gate). Compute scales with $k$, not $N$.
5. Parameters. Why the punchline? Because storing $N$ full experts is what actually grows memory. Total expert params $=N\times (\text{one FFN})=128\times$ a dense FFN. Parameters grew ~128×; active compute grew ~2×. That gap is the sublinear scaling the parent note promised.

Verify: $2\cdot 512\cdot 2048=2,097,152$ ✓; active $=2\times$ that $=4,194,304$ ✓; gate $=128\times 512=65,536$ ✓; gate/active $=65536/4194304=0.0156=1.56\%$ ✓.

---

Example 7 — Load balancing: Switch aux loss on an imbalanced batch (cell C7)

Forecast: the loss is minimized at value $1$ for perfect balance. Will our imbalanced batch score above or below 1?

1. Fractions $f_i$. Why? $f_i$ = share of tokens dispatched to expert $i$ (the "hard" count); we need it as the detached weight in the loss. $f_1=6/8=0.75,f_2=2/8=0.25$.
2. Aux loss. Why this formula? Gradients flow only through the soft $P_i$; multiplying by the detached $f_i$ penalizes probability piled on already-overloaded experts. ${\mathcal{L}}_{\text{aux}}=1\cdot 2\cdot (0.75\cdot 0.7+0.25\cdot 0.3)=2\cdot (0.525+0.075)=2\cdot 0.600=\mathrm{1.200.}$
3. Ideal minimum. Why 1? At perfectly uniform routing every $f_i=1/N$ and every $P_i=1/N$, so each product is $\frac{1}{N}\cdot \frac{1}{N}=\frac{1}{N^2}$, and there are $N$ of them: ${\sum }_i{f}_i{P}_i=N\cdot \frac{1}{N^2}=\frac{1}{N}$. Multiplying by the front factor $\alpha N$ gives $\alpha N\cdot \frac{1}{N}=\alpha =1$, the floor. Our $1.200>1$ ✓ (imbalance costs extra), matching the forecast.
4. CV of load. Why a separate metric? CV is the monitoring number (not differentiable), so we watch it even though we don't optimize it. Loads are $6,2$; mean $\mu =4$; deviations $\pm 2$; $\sigma =\sqrt{\frac{(6-4{)}^2+(2-4{)}^2}{2}}=\sqrt{4}=2$; $\text{CV}=\sigma /\mu =2/4=0.5$.
5. Balanced check. Why? To confirm the loss really bottoms out at the uniform routing we claimed. $f=P=(0.5,0.5)$: $\mathcal{L}=2(0.5\cdot 0.5+0.5\cdot 0.5)=2\cdot 0.5=1.000$, and loads $4,4\Rightarrow \sigma =0\Rightarrow \text{CV}=0$.

Verify: imbalanced loss $1.200$ ✓; balanced loss $1.000$ ✓ (the floor); balanced $\text{CV}=0$ ✓; imbalanced $\text{CV}=0.5$ ✓. Larger CV and larger aux loss agree that the first batch is worse — two independent diagnostics point the same way.

---

Example 8 — Expert Choice capacity (cell C8)

Forecast: in Expert-Choice, every expert processes the same number of tokens by construction. Guess whether $C$ is a fraction or a whole number for these values.

1. Total slots. Why? Each of $B$ tokens claims $k$ expert-slots, and we must count them all to share fairly. Total slots $=B\cdot k=1024\cdot 2=2048$.
2. Slots per expert. Why divide by $N$? Spreading the slots evenly across all $N$ experts is exactly what "balanced" means. $C=\frac{Bk}{N}=\frac{2048}{64}=32$.
3. $k=1$ case. Why redo it? To show capacity scales linearly with the target experts-per-token. $C=\frac{1024\cdot 1}{64}=16$.
4. Why perfectly balanced? Why no aux loss? Each expert selects its own top-$C$ tokens by affinity score $G(x{)}_i$, always filling exactly $C$ slots — imbalance is impossible by construction, so no ${\mathcal{L}}_{\text{aux}}$ is needed (contrast Example 7). See Conditional Computation.

Verify: $1024\cdot 2/64=32$ ✓; $1024\cdot 1/64=16$ ✓; both whole numbers ✓. Every expert handles exactly $C$ tokens $\Rightarrow$ CV of load $=0$ automatically.

---

Example 9 — Word problem, real-world routing (cell C9)

Forecast: guess the saving factor before computing — is it near 2×, 50×, or 100×?

1. Active FLOP/s. Why $k$, not $N$? Only $k=2$ experts fire per token, so we multiply the token rate by 2 experts and the per-expert cost. $50000\frac{\text{tok}}{\text{s}}\times 2\text{experts}\times 4\text{MFLOP}=400000\text{MFLOP/s}=0.0004\text{PFLOP/s}.$
2. Dense-equivalent FLOP/s. Why $N$? A dense model would hit all 100 experts per token, so we multiply by 100 instead of 2. $50000\times 100\times 4=20,000,000\text{MFLOP/s}=0.02\text{PFLOP/s}.$
3. Saving factor. Why the ratio? Dividing dense by active isolates exactly how many times more work the naive model does. $\frac{N}{k}=\frac{100}{2}=50$. Sparsity does the same routing job for $1/50$ of the dense-all cost.

Verify: active $=50000\cdot 2\cdot 4=400000$ MFLOP/s $=0.0004$ PFLOP/s ✓; dense $=50000\cdot 100\cdot 4=2\times 10^7$ MFLOP/s $=0.02$ PFLOP/s ✓; ratio $=2\times 10^7/4\times 10^5=50$ ✓ — a 50× saving, matching $N/k$. (Recall $1$ PFLOP/s $=10^9$ MFLOP/s.) Related idea: shrinking a big model differently, see Knowledge Distillation.

---

Example 10 — Exam twist: temperature changes the winner (cell C10)

Forecast: does lowering $T$ change which expert wins, or only how sure the router is?

Temperature $T$ is a positive knob that divides every logit before softmax: small $T$ makes the gate sharper (more one-hot), large $T$ flatter. Because it divides, the order of the logits is preserved for any $T>0$.

1. $T=1$ (no change). Why start here? It is the ordinary softmax baseline to compare against. $e^3=20.086,e^{2.5}=12.182,Z=32.268$, so $G(x)=(0.6225,0.3775)$. Winner: expert 1.
2. $T=0.1$ (sharpen). Why? Dividing the logits by $0.1$ multiplies them by 10, blowing the gap up and pushing softmax toward one-hot. $s/T=(30,25)$. These are large, so we use the max-subtracted form (subtract $m=30$): $e^0=1,e^{-5}=0.006738$, $Z=1.006738$, giving $G(x)\approx (0.99331,0.00669).$ Winner: still expert 1, now near-certain.
3. Why the winner is fixed for $T>0$. Why? Dividing all logits by the same positive number is monotonic — it can't reorder them — so argmax is invariant to positive temperature scaling.
4. The trap: $T<0$. Why forbid it? Dividing by a negative $T$ flips every sign, so the smallest logit becomes the largest and expert 2 would win — a negative temperature silently inverts your router. (Contrast this with the normalization tricks in Batch Normalization, which rescale but never invert.)

Verify: $T=1$: $0.6225+0.3775=1.000$ ✓ and $0.6225>0.5$ so expert 1 wins ✓. $T=0.1$: $0.99331+0.00669=1.000$ ✓, expert 1 still wins ✓ but sharper. Order preserved across $T=1\to 0.1$ ✓; sign-flip at $T<0$ makes expert 2 win ✓.

---

Recall

Recall Why does softmax never output a negative probability, even for negative logits?

Because $e^s>0$ for every real $s$; a positive numerator over a positive sum is always in $(0,1)$. ::: Exponentials are strictly positive.

Recall Why subtract the max logit before exponentiating in softmax?

To prevent numerical overflow on large logits; subtracting $m={\max }_j{s}_j$ leaves the fraction identical (multiply top and bottom by $e^{-m}$) but caps the largest exponent at $e^0=1$. ::: Overflow safety, exact same answer.

Recall After top-

$k$, why renormalize the surviving gates? Dropping experts removes probability mass, so survivors no longer sum to 1; the output would be scaled down. Dividing by the survivors' sum restores a proper weighted average. ::: To keep $y$ a true convex combination.

Recall In the Switch aux loss

$\alpha N{\sum }_i{f}_i{P}_i$, which factor carries the gradient and why? Only $P_i$ (mean gate probability) is differentiable; $f_i$ is a detached hard count. So gradients push down probability on overloaded experts. ::: $P_i$ carries the gradient; $f_i$ is detached.

Recall Why does Expert-Choice routing need no auxiliary loss?

Each expert fills exactly $C=Bk/N$ slots by picking its own top-$C$ tokens, so load balance is structural, not learned. ::: Balance is enforced by construction.

See also: Transformer Architecture, Model Parallelism, Attention Mechanisms, Neural Architecture Search, Long-tail Distribution.