Mixture-of-Experts (MoE) architecture: Worked examples

Everything here uses only three tools. Let us fix their meaning once, in plain words, before any symbol appears again.

Definition Notation used on this whole page

Fix these once — every example reuses them.

$x$ — the token representation: the vector of numbers the model currently holds for one token (word-piece). It is the input to the MoE layer.
$E_{i} (x)$ — the output of expert $i$ when fed $x$ . Each expert $E_{i}$ is its own small feed-forward network; $E_{i} (x)$ is the vector it returns.
$y$ — the output of the MoE layer for this token: the gate-weighted blend of the chosen experts' outputs, $y = \sum_{i \in T} g_{i} E_{i} (x)$ . This $y$ is what the next layer receives.
$N$ — total number of experts. $k$ — how many experts run per token (Top- $k$ ). $h_{i}, p_{i}, g_{i}$ — logit, softmax probability, and final gate for expert $i$ . $T$ — the set of Top- $k$ indices that survive.
$α$ — small weight on the load-balance loss. $f_{i}$ — fraction of a batch's tokens sent to expert $i$ . $P_{i}$ — average router probability for expert $i$ over the batch.

The scenario matrix

Every example below is tagged with the cell it covers.

Cell	Case class	What could break the naive intuition
A	Standard Top-2 route	nothing tricky — the baseline
B	Exact tie in logits ( $k = 1$ )	who wins Top- $k$ ? deterministic tie-break
B2	Tie straddling the $k^{th}$ cutoff	tie splits across the keep/drop boundary
C	Negative & zero logits	softmax still works; $e^{0} = 1$ , $e^{< 0} \in (0, 1)$
D	All logits equal (degenerate)	uniform routing, maximal entropy
E	Top-1 ( $k = 1$ )	the gate still matters (gradient!)
F	Capacity overflow → dropped tokens	residual carries the token, output changes
G	Load-balance loss, balanced vs collapsed	$L_{aux}$ min vs worst case
H	Params vs FLOPs word problem	capacity $\neq =$ compute
I	Exam twist: change $k$ , watch FLOPs & output	which quantity moves, which doesn't

We will hit A–I across ten examples.

Before the examples, one picture of the whole pipeline — every example is a walk along this same arrow chain, so keep it in mind. Cyan boxes are the fixed machinery; amber is where a token's fate is decided.

Figure — Mixture-of-Experts (MoE) architecture

Figure — the routing pipeline shared by every example: logits $h$ → softmax $p$ → Top- $k$ keeps some, zeros the rest → re-normalise survivors → gate-weighted blend $y$ . Amber marks the two decision points (which experts survive, whether the expert has room).

Cell A — the baseline route

The bar picture below is the whole example in one glance: the tall grey bars are the softmax probabilities $p$ ; the two cyan bars are the survivors Top-2 keeps; the two faded bars are zeroed. Watch how re-normalising stretches the two cyan bars up so they fill the height $1$ between them.

Figure — Cell A: softmax probabilities (grey), the Top-2 survivors (cyan), and the same two after re-normalising to sum to 1 (amber outline). Experts 1 and 4 are faded — they contribute nothing.

Worked example A · Standard Top-2, distinct logits

$N = 4$ experts, $k = 2$ . Logits $h = [1.0, 3.0, 2.0, 0.0]$ . Find the two gate values and the output formula.

Forecast: guess which two experts win, and roughly what fraction of the output the top one owns — half? two-thirds? more?

Exponentiate. $e^{h} = [2.718, 20.086, 7.389, 1.000]$ , sum $= 31.193$ . Why this step? Softmax needs positive, fast-growing scores so a clearly-better logit dominates.
Normalise (softmax). $p = [0.0871, 0.6440, 0.2369, 0.0321]$ . Why? Now the four numbers are comparable weights summing to $1$ (grey bars in the figure).
Pick Top-2. Largest are expert $2$ ( $0.6440$ ) and expert $3$ ( $0.2369$ ). Experts $1, 4$ get gate $0$ . Why? Only the $k = 2$ chosen experts will run — that is where the FLOP saving lives (cyan bars).
Re-normalise survivors. Denominator $= 0.6440 + 0.2369 = 0.8809$ . $g_{2} = 0.6440/0.8809 = 0.7311$ , $g_{3} = 0.2369/0.8809 = 0.2689$ . Why? Used gates must sum to $1$ so $y$ is a proper weighted blend (amber-outlined bars, stretched taller than the cyan ones).
Output. $y = 0.7311 E_{2} (x) + 0.2689 E_{3} (x)$ .

Verify: $g_{2} + g_{3} = 0.7311 + 0.2689 = 1.0$ ✓. Only $2$ experts evaluated ✓. Top expert owns ~73%, more than two-thirds — the softmax amplified the gap.

Cell B — an exact tie ( $k = 1$ )

Worked example B · Two logits are identical

$N = 3$ , $k = 1$ . Logits $h = [2.0, 2.0, 1.0]$ . Who runs?

Forecast: if two experts tie for first and $k = 1$ , does the layer crash, pick both, or pick one? Guess the rule.

Softmax. $e^{h} = [7.389, 7.389, 2.718]$ , sum $= 17.496$ . $p = [0.4223, 0.4223, 0.1554]$ . Why? Ties in logits stay ties in probability — softmax is order-preserving.
Top-1 with a tie. We need exactly one winner. Frameworks break ties by lowest index, so expert $1$ wins. Why this step? Routing must be deterministic so the same token always trains the same expert; "lowest index" is the standard, reproducible convention.
Re-normalise. Single survivor → $g_{1} = 0.4223/0.4223 = 1.0$ . Why? One expert must own the whole output.
Output. $y = 1.0 \cdot E_{1} (x)$ .

Verify: $g_{1} = 1$ , one expert only ✓. The gate is $1$ here, but note it still carries gradient — see Cell E for why that matters even at $k = 1$ .

Cell B2 — a tie straddling the cutoff

Cell B had the tie at the top with $k = 1$ . The nastier case is a tie sitting exactly on the keep/drop boundary: more experts are tied for the last surviving slot than there are slots left. The figure shows this — two equal-height bars share rank 2, but Top-2 has only one slot left after taking rank 1.

Figure — Cell B2: logits $[3, 2, 2, 1]$ , $k = 2$ . Rank-1 (cyan) is taken outright; experts 2 and 3 tie for the single remaining slot (amber-striped). The lowest-index tie-break keeps expert 2, drops expert 3.

Worked example B2 · Tie on the

$k^{th}$ boundary $N = 4$ , $k = 2$ . Logits $h = [3.0, 2.0, 2.0, 1.0]$ . Experts 2 and 3 tie. Which two run?

Forecast: Top-2 wants two winners, but three experts (2, 3 — and clearly 1) are candidates for the top with a tie at rank 2. Guess who is dropped.

Softmax. $e^{h} = [20.086, 7.389, 7.389, 2.718]$ , sum $= 37.582$ . $p = [0.5345, 0.1966, 0.1966, 0.0723]$ . Why? The tie in logits becomes an exact tie in probability for experts 2 and 3.
Fill slots by rank. Slot 1 → expert $1$ (clear top, $0.5345$ ). One slot left. Experts $2$ and $3$ are tied for it. Why this step? Top- $k$ fills highest first; the conflict only appears at the last slot.
Break the boundary tie by lowest index. Expert $2$ takes the last slot; expert $3$ is dropped even though its probability equals expert $2$ 's. Why this rule? We must keep exactly $k = 2$ experts and the choice must be reproducible; "lowest index wins" is the deterministic convention (same rule as Cell B, now applied at the cutoff, not the top).
Re-normalise survivors ${1, 2}$ . Denom $= 0.5345 + 0.1966 = 0.7311$ . $g_{1} = 0.7311/0.7311$ … let us compute: $g_{1} = 0.5345/0.7311 = 0.7311$ , $g_{2} = 0.1966/0.7311 = 0.2689$ . Why? Survivor gates must sum to $1$ .
Output. $y = 0.7311 E_{1} (x) + 0.2689 E_{2} (x)$ . Expert 3, despite tying, contributes nothing.

Verify: $g_{1} + g_{2} = 0.7311 + 0.2689 = 1.0$ ✓; exactly $2$ experts run ✓; the dropped expert 3 had identical probability to the kept expert 2 — the tie-break, not the score, decided it.

Cell C — negative and zero logits

Worked example C · Signs and a zero

$N = 4$ , $k = 2$ . Logits $h = [- 1.0, 0.0, 2.0, - 3.0]$ . Do negatives cause negative probabilities?

Forecast: a logit of $- 3$ — does that expert get a negative weight, or just a tiny positive one?

Exponentiate — watch the signs. $e^{- 1} = 0.3679$ , $e^{0} = 1.0$ , $e^{2} = 7.389$ , $e^{- 3} = 0.0498$ . Why? $e^{x}$ is always positive: $e^{0} = 1$ exactly, $e^{< 0}$ falls into $(0, 1)$ , never below $0$ . This is why softmax uses exponentials — no negative shares can ever appear.
Sum & normalise. Sum $= 0.3679 + 1.0 + 7.389 + 0.0498 = 8.807$ . $p = [0.0418, 0.1136, 0.8390, 0.00565]$ .
Top-2. Experts $3$ ( $0.8390$ ) and $2$ ( $0.1136$ ). Why? Highest two probabilities; the negative-logit experts $1, 4$ get gate $0$ .
Re-normalise. Denom $= 0.8390 + 0.1136 = 0.9526$ . $g_{3} = 0.8807$ , $g_{2} = 0.1193$ . Why? The two survivors' gates must sum to $1$ so the output is a proper convex blend — the negative-logit experts already dropped out at step 3.
Output. $y = 0.8807 E_{3} (x) + 0.1193 E_{2} (x)$ .

Verify: every $p_{i} > 0$ despite negative logits ✓; $g_{2} + g_{3} = 1.0$ ✓. The $- 3$ expert got $0.00565$ , near-zero but never negative.

Cell D — the degenerate case: all logits equal

Worked example D · Uniform routing (maximum uncertainty)

$N = 8$ experts, $k = 2$ . All logits equal: $h = [c, c, \dots, c]$ for any constant $c$ . What does the router do?

Forecast: if the router has no preference, what probability does each expert get, and does the value of $c$ matter?

Softmax of equal logits. $p_{i} = e^{c} / (8 e^{c}) = 1/8 = 0.125$ for every $i$ . Why? The shared $e^{c}$ cancels top and bottom — softmax is shift-invariant, so only differences between logits matter, never their absolute level.
Top- $k$ under a full tie. With $k = 2$ , every $p_{i}$ ties, so Top- $k$ falls back to lowest indices and picks experts $1, 2$ . Why? Deterministic tie-break again (Cell B logic), now applied to all $8$ .
Re-normalise. $g_{1} = g_{2} = 0.125/0.25 = 0.5$ .
Meaning. This is maximum-entropy routing — the router is maximally unsure. It is also exactly the state the load-balance loss (Cell G) wants on average across a batch.

Verify: $8 \times 0.125 = 1.0$ ✓; $g_{1} + g_{2} = 1.0$ ✓; answer independent of $c$ ✓ (shift-invariance).

Cell E — Top-1, why the gate still bites

Worked example E ·

$k = 1$ but the gate is not just "1" $N = 3$ , $k = 1$ . Logits $h = [0.5, 2.0, 1.0]$ . Common myth: " $k = 1$ means gate $= 1$ , so it does nothing."

Forecast: with only one expert chosen, is the output just $E_{2} (x)$ , or is it scaled by something?

Softmax. $e^{h} = [1.6487, 7.389, 2.718]$ , sum $= 11.756$ . $p = [0.1403, 0.6286, 0.2312]$ .
Top-1. Expert $2$ wins ( $0.6286$ ).
Two conventions for the gate:
- Re-normalised (as parent note): denom $= 0.6286$ , so $g_{2} = 1.0$ , output $= E_{2} (x)$ .
- Un-normalised (used by many $k = 1$ routers, e.g. Switch): keep the raw softmax value $g_{2} = 0.6286$ , output $= 0.6286 E_{2} (x)$ . Why does the un-normalised version exist? Because that $0.6286$ carries gradient to the router. If you always force the gate to $1$ , the router logit for the winner gets no gradient signal from the main loss, and it cannot learn how confident to be. See Softmax and Gating Functions.
Output (Switch-style): $y = 0.6286 E_{2} (x)$ .

Verify: winner is expert $2$ ✓; un-normalised gate $= 0.6286$ ✓. Even at $k = 1$ the gate is a learnable multiplier, not a no-op.

Cell F — capacity overflow and dropped tokens

First, one new quantity.

The figure below draws this literally: each expert is a fixed-height bin of height $C$ . Cyan fill = tokens that fit and get processed; the dashed white line is the capacity ceiling $C$ ; amber above the line = overflow tokens that are dropped and pushed sideways by an amber arrow onto the residual path. Read it bin by bin as you follow the steps.

Figure — Cell F: four experts as fixed-capacity bins. Expert E1 receives 200 tokens but its bin only holds C=128; the 72 tokens above the dashed ceiling (amber) are dropped and carried forward unchanged by the residual connection. Experts E2–E4 stay under C and drop nothing, yet their unused slots cannot be lent to E1.

Worked example F · Tokens dropped past capacity

Batch $= 512$ tokens, $N = 4$ experts, capacity factor $= 1.0$ (no slack). The router (Top-1) sends counts $[200, 120, 92, 100]$ to experts $[1, 2, 3, 4]$ . How many tokens are dropped and what happens to them?

Forecast: which experts overflow, and where do the overflow tokens go — deleted, re-routed, or passed through unchanged?

Capacity per expert. $C = 1.0 \times \frac{512}{4} = 128$ tokens. Why this formula? Each expert gets an equal fixed slot count so memory is pre-allocated and fixed regardless of routing luck; here the capacity factor $1.0$ gives exactly the fair share.
Compare each count to $C$ .
- Expert 1: $200 > 128$ → keeps $128$ , drops $200 - 128 = 72$ .
- Expert 2: $120 \leq 128$ → keeps all, drops $0$ .
- Expert 3: $92 \leq 128$ → drops $0$ .
- Expert 4: $100 \leq 128$ → drops $0$ . Why? Only over-subscribed experts overflow; under-full experts waste slots but drop nothing.
Total dropped $= 72$ .
Where do the $72$ go? They skip the MoE layer and pass forward via the residual connection unchanged — the layer's output for those tokens is just their input $x$ . Why? A skip keeps the network well-defined even when an expert is full; the token is not deleted, just not transformed here. See Conditional Computation.

Verify: drops $= max (200 - 128, 0) + max (120 - 128, 0) + max (92 - 128, 0) + max (100 - 128, 0) = 72 + 0 + 0 + 0 = 72$ ✓. Note experts $2, 3, 4$ together wasted $(128 - 120) + (128 - 92) + (128 - 100) = 72$ slots — same number, but slots are not transferable across experts.

Cell G — the load-balance loss, best vs worst

The bowl figure gives the intuition the algebra hides: plot $\sum_{i} f_{i} P_{i}$ as we slide from perfectly uniform usage toward collapse. It is a bowl with its lowest point exactly at uniform — any move away from the centre only raises the penalty.

Figure — Cell G: the load-balance term $N \sum f_{i} P_{i}$ traced as usage moves from uniform (centre, value 1.0) to full collapse (edge, value N). The curve is a bowl bottoming out at uniform — the visual proof that uniform is the unique minimum.

Worked example G ·

$L_{aux}$ balanced vs collapsed Recall $L_{aux} = α N \sum_{i = 1}^{N} f_{i} P_{i}$ , where $f_{i}$ = fraction of tokens sent to expert $i$ , $P_{i}$ = mean router probability for expert $i$ . Take $N = 4$ , $α = 1$ for clarity. Compute the loss for two extremes.

Case (i) — perfectly balanced: $f = P = [\frac{1}{4}, \frac{1}{4}, \frac{1}{4}, \frac{1}{4}]$ . Case (ii) — collapsed: one expert hogs everything, $f = P = [1, 0, 0, 0]$ .

Forecast: which case gives the smaller loss — balanced or collapsed? (The loss is a penalty, so smaller is better.)

Balanced sum. $\sum f_{i} P_{i} = 4 \times (\frac{1}{4} \cdot \frac{1}{4}) = 4 \times \frac{1}{16} = \frac{1}{4}$ . Loss $= 1 \times 4 \times \frac{1}{4} = 1.0$ . Why the factor $N$ ? It rescales so the minimum loss is a clean constant $1$ regardless of how many experts — the penalty scale is $N$ -independent.
Collapsed sum. $\sum f_{i} P_{i} = 1 \cdot 1 + 0 + 0 + 0 = 1$ . Loss $= 1 \times 4 \times 1 = 4.0$ .
Compare. Balanced $= 1.0 <$ collapsed $= 4.0$ . So the loss punishes collapse and is minimised at uniform usage. Why use $P_{i}$ (soft) not just $f_{i}$ (hard count)? $f_{i}$ comes from a hard argmax, which has zero gradient; multiplying by the smooth $P_{i}$ lets gradients flow back and actually push the router toward balance. See Load Balancing and Auxiliary Losses.
Why uniform is the unique minimum (brief proof). Tie $f_{i} = P_{i} = q_{i}$ with $\sum q_{i} = 1$ . Then $\sum q_{i}^{2}$ is being minimised. By the identity $\sum q_{i}^{2} = \frac{1}{N} + \sum_{i} (q_{i} - \frac{1}{N})^{2}$ , the sum equals its floor $\frac{1}{N}$ only when every $q_{i} = \frac{1}{N}$ ; any spread adds the strictly-positive squared-deviation term. That is the bowl in the figure — one lowest point, walls rising in every direction. Why this matters? It guarantees the loss cannot be fooled: no lopsided distribution ties the uniform score.

Verify: balanced loss $= 1.0$ , collapsed loss $= 4.0$ , and $1.0 < 4.0$ ✓. The uniform value $\frac{1}{4}$ is indeed the minimum of $\sum f_{i} P_{i}$ (i.e. $\sum q_{i}^{2}$ ) under $\sum q_{i} = 1$ .

Cell H — a word problem: capacity is not compute

Worked example H · The hospital budget

A model has $N = 128$ experts of $10$ M params each; $k = 2$ ; shared trunk (attention + router + embeddings) $= 200$ M params. Answer: (a) total params, (b) active params per token, (c) the per-token FLOP cost, (d) the knowledge-equivalent dense model.

Forecast: guess the ratio between the "knows-as-much-as" model and the "costs-as-much-as" model. 2×? 10×? more?

Total (dense) params. $128 \times 10 M + 200 M = 1280 M + 200 M = 1480 M$ . Why? Every expert is stored and trained; capacity scales with $N$ . See Sparse vs Dense Models.
Active params per token. $2 \times 10 M + 200 M = 20 M + 200 M = 220 M$ . Why? Only $k = 2$ experts run; the number of weights touched scales with $k$ , not $N$ — this is Conditional Computation.
Per-token FLOPs. A dense feed-forward path costs roughly $2 \times (params touched)$ FLOPs per token (one multiply + one add per weight). So the MoE token cost $\approx 2 \times 220 M = 440 M$ FLOPs — driven by the $220$ M active params, not the $1480$ M stored ones. Why is compute $\neq =$ storage? Storage counts weights that merely exist in memory ( $1480$ M); FLOPs count only weights actually multiplied this forward pass ( $220$ M). The $1260$ M idle expert params sit in memory contributing $0$ FLOPs for this token. See Scaling Laws.
Knowledge-equivalent dense model $\approx 1480$ M params (a dense model that stores as much would need all $1480$ M and would also compute all of them).
Ratio. $1480/220 \approx 6.7 \times$ more stored knowledge for the same per-token compute a $220$ M dense model would spend.

Verify: total $= 1480$ M ✓; active $= 220$ M ✓; per-token FLOPs $= 2 \times 220 M = 440$ M ✓; ratio $1480/220 = 6.727 \dots$ ✓ (~6.7×). Capacity ( $N$ -driven, $1480$ M stored) grew $6.7 \times$ but compute ( $k$ -driven, $440$ M FLOPs) stayed fixed.

Cell I — exam twist: change $k$ , what moves?

The paired bar figure shows the same softmax vector twice — once with the $k = 2$ survivors highlighted, once with the $k = 3$ set. Notice the grey $p$ bars are identical between panels; only which bars are lit and how tall the re-normalised gates rise changes.

Figure — Cell I: identical softmax $p$ in both panels. Left ( $k = 2$ ) keeps experts 1,2 and re-normalises them taller; right ( $k = 3$ ) keeps 1,2,3, so each survivor gate is shorter (shared over a bigger denominator).

Worked example I · Same logits, flip

$k$ from 2 to 3 $N = 4$ , logits $h = [3.0, 2.0, 1.0, 0.0]$ (same softmax for both parts). We ran $k = 2$ ; the exam now asks $k = 3$ . What changes: the softmax $p$ ? the gates $g$ ? the FLOPs? the set of runners?

Forecast: name every quantity that changes and every one that stays the same before computing.

Softmax (identical for both). $e^{h} = [20.086, 7.389, 2.718, 1.000]$ , sum $= 31.193$ . $p = [0.6440, 0.2369, 0.0871, 0.0321]$ . Why unchanged? Softmax is computed over all logits before Top- $k$ — it never sees $k$ .
$k = 2$ gates. Survivors $1, 2$ ; denom $= 0.6440 + 0.2369 = 0.8809$ ; $g_{1} = 0.7311, g_{2} = 0.2689$ .
$k = 3$ gates. Survivors $1, 2, 3$ ; denom $= 0.6440 + 0.2369 + 0.0871 = 0.9680$ ; $g_{1} = 0.6653, g_{2} = 0.2447, g_{3} = 0.0900$ . Why did $g_{1}$ shrink? A third expert now shares the normalised budget, so every surviving gate is divided by a larger denominator.
How FLOPs scale with $k$ . Per-token expert FLOPs $\approx k \times (FLOPs of one expert)$ — each running expert adds the same fixed cost, so the total is linear in $k$ . Going $k : 2 \to 3$ multiplies expert FLOPs by $\frac{3}{2} = 1.5$ (a $50%$ increase). The router and shared trunk cost is unchanged.
What moved: the runner set (2 → 3 experts), the gate values, and the expert FLOPs (×1.5). What stayed: the softmax $p$ , the logits, all stored params.

Verify: $k = 2$ : $g_{1} + g_{2} = 1.0$ ✓. $k = 3$ : $g_{1} + g_{2} + g_{3} = 1.0$ ✓, and $g_{1} (k = 3) = 0.6653 < 0.7311 = g_{1} (k = 2)$ ✓; FLOP ratio $3/2 = 1.5$ ✓. Softmax vector identical in both ✓.

Recall Quick self-test

All four logits equal — what is each softmax probability for $N = 5$ ? ::: $1/5 = 0.2$ each (shift-invariance, value of the constant irrelevant). A $- 4$ logit — can its softmax probability be negative? ::: No; $e^{- 4} \approx 0.018 > 0$ . Softmax is always positive. Logits $[3, 2, 2, 1]$ , $k = 2$ — who is dropped? ::: The tie at rank 2 is broken by lowest index: expert 2 kept, expert 3 dropped. Batch 800, $N = 5$ , capacity factor 1.0, one expert gets 200 tokens — how many dropped there? ::: $C = 160$ , dropped $= 200 - 160 = 40$ . At $k = 1$ , why keep the raw gate instead of forcing it to 1? ::: So the gate carries gradient to the router — it learns confidence, not just choice. Does raising $k$ change the softmax $p$ ? ::: No — softmax runs over all logits before Top- $k$ ; only the survivor set, gates, and FLOPs change.

Back to the parent MoE note · related: Feed-Forward Networks (Transformer FFN), Model Parallelism & Expert Parallelism.