Mixture-of-Experts (MoE) architecture: Exercises

Before we start, one reminder of every symbol you will meet, in plain words:

Definition Symbol glossary (all earned before use)

$d$ ::: the length of a token's vector — how many numbers describe one word-piece (the embedding dimension).
$N$ ::: total number of experts stored in the layer.
$k$ ::: how many experts each token is actually sent to (Top- $k$ ).
$x$ ::: one token's vector (a list of $d$ numbers describing that word-piece).
$W_{g}$ ::: the router's learnable weight matrix ( $N$ rows, one "preference vector" per expert); it turns a token $x$ into $N$ scores via $W_{g} x$ .
$E_{i}$ ::: expert number $i$ — itself a small feed-forward network; $E_{i} (x)$ is the vector it outputs when fed token $x$ .
$h_{i}$ ::: the router's raw score ("logit") for expert $i$ — bigger means "more relevant".
$p_{i}$ ::: that score after softmax, now a probability (all $p_{i}$ add to 1).
$g_{i}$ ::: the final gate weight after keeping only Top- $k$ and re-normalizing.
$T$ ::: the Top- $k$ index set — the set of indices of the $k$ highest-scoring experts (the ones that actually run).
$f_{i}$ ::: fraction of a batch's tokens that landed on expert $i$ .
$P_{i}$ ::: average router probability given to expert $i$ over the batch.
$C$ ::: expert capacity — the max tokens one expert may accept per batch.
$α$ ::: the load-balancing loss coefficient — a small positive number (typically $\approx 0.01$ ) that scales how strongly the auxiliary balance penalty $L_{aux}$ nudges the router, without overpowering the main training objective.

Level 1 — Recognition

L1.1

For an MoE layer with $N = 32$ experts and Top- $k$ with $k = 2$ : how many experts are stored, and how many do computation for a single token?

Recall Solution

Stored $= N = 32$ . Every expert physically exists and is trained. Compute for one token $= k = 2$ . Only the two highest-scoring experts run. This is the whole MoE idea: capacity (32) is decoupled from per-token compute (2). See Sparse vs Dense Models and Conditional Computation.

L1.2

Which quantity does the FLOP (compute) cost per token scale with — $N$ or $k$ ? State it in one line.

Recall Solution

Compute per token scales with $k$ (active experts), not $N$ (total experts). Growing $N$ grows knowledge/memory; only growing $k$ grows per-token compute.

L1.3

Name each of the four stages the router performs, in order.

Recall Solution

Score — dot product $h = W_{g} x$ , where $W_{g}$ is the router's learnable weight matrix and each output number $h_{i}$ is one score per expert.
Softmax — turn scores into probabilities $p_{i}$ that sum to 1.
Top- $k$ + re-normalize — keep the $k$ largest (their indices form the set $T$ ), zero the rest, rescale so used gates sum to 1.
Combine — weighted sum $y = \sum_{i \in T} g_{i} E_{i} (x)$ , where $E_{i} (x)$ is the output of expert $i$ (a small feed-forward network) on token $x$ , and $T$ is the Top- $k$ index set. Mnemonic: Score → Softmax → Top-k → Combine (one word per stage, matching the four steps above).

Level 2 — Application

L2.1

A model has $N = 128$ experts, each expert holds $8$ M parameters, and $k = 1$ . Shared (attention + embeddings + router) $= 60$ M params. Compute: (a) total expert params, (b) active expert params per token, (c) total params, (d) active params per token.

Recall Solution

(a) Total expert params $= 128 \times 8 M = 1024 M = 1.024 B$ . (b) Active expert params $= 1 \times 8 M = 8 M$ . (c) Total params $= 1024 M + 60 M = 1084 M$ . (d) Active params/token $= 8 M + 60 M = 68 M$ . The model knows like a 1.08B model but computes like a 68M one — a $\approx 16 \times$ gap between capacity and cost. This gap is why Scaling Laws treat MoE so favourably.

L2.2

Router logits for a token are $h = [1.0, 3.0, 2.0]$ over $N = 3$ experts, $k = 2$ . Find the softmax probabilities, the Top-2 gates, and write the output $y$ .

Recall Solution

Softmax. $e^{h} = [e^{1}, e^{3}, e^{2}] = [2.71828, 20.0855, 7.38906]$ , sum $= 30.1928$ . $p = [0.09003, 0.66524, 0.24473]$ . Top-2. Largest two are expert 2 ( $0.66524$ ) and expert 3 ( $0.24473$ ). Expert 1 gets gate $0$ . Re-normalize. Denominator $= 0.66524 + 0.24473 = 0.90997$ . $g_{2} = 0.66524/0.90997 = 0.73106$ , $g_{3} = 0.24473/0.90997 = 0.26894$ . Output. $y = 0.73106 E_{2} (x) + 0.26894 E_{3} (x)$ . Expert 1 is never evaluated. (Check $g_{2} + g_{3} = 1$ : yes, so this is a proper convex blend.)

Figure — Mixture-of-Experts (MoE) architecture

L2.3

Batch $= 2048$ tokens, $N = 16$ experts, capacity factor $= 1.5$ . Compute the expert capacity $C$ . If expert 5 receives $240$ tokens, how many are dropped?

Recall Solution

$C = 1.5 \times \frac{2048}{16} = 1.5 \times 128 = 192 tokens .$ Expert 5 got $240$ , exceeding $C = 192$ , so $240 - 192 = 48$ tokens are dropped — they skip the expert and pass through the residual unchanged. Higher capacity factor would reduce drops at the cost of more memory/compute.

Level 3 — Analysis

L3.1

An 8-expert layer trains without any auxiliary loss. After a while, one expert receives $90%$ of tokens, three share the rest, and four are dead. Explain the mechanism that produced this, and name the loss that prevents it.

Recall Solution

Mechanism — "rich get richer." Early on one expert is slightly better, so the router raises its probability, so it sees more tokens, so it trains faster, so it gets even more tokens. This positive feedback loop is routing collapse: a few experts hog everything, the rest never receive gradients and stay useless. Prevention. The auxiliary load-balancing loss $L_{aux} = α N \sum_{i = 1}^{N} f_{i} P_{i},$ minimized when every $f_{i} = P_{i} = 1/ N$ (uniform usage). It nudges the router to spread tokens. Here $α$ is the small coefficient (typically $\approx 0.01$ ) that keeps this a gentle nudge rather than the main objective.

L3.2

Compute $L_{aux}$ for two batches with $N = 4$ , $α = 1$ . Assume $f_{i} = P_{i}$ for simplicity. (a) Balanced: $f = [0.25, 0.25, 0.25, 0.25]$ . (b) Collapsed: $f = [0.85, 0.05, 0.05, 0.05]$ . Which is larger, and why is that the right direction for a penalty?

Recall Solution

Formula: $L = α N \sum_{i} f_{i} P_{i} = 4 \sum_{i} f_{i}^{2}$ (since $f_{i} = P_{i}$ ). (a) $\sum f_{i}^{2} = 4 \times 0.2 5^{2} = 4 \times 0.0625 = 0.25$ . So $L_{a} = 4 \times 0.25 = 1.0$ . (b) $\sum f_{i}^{2} = 0.8 5^{2} + 3 (0.0 5^{2}) = 0.7225 + 0.0075 = 0.73$ . So $L_{b} = 4 \times 0.73 = 2.92$ . Collapsed is larger ( $2.92 > 1.0$ ). Correct direction: the imbalanced batch is penalised more, pushing the router back toward uniform. The balanced case hits the theoretical minimum of $L = 1$ (equals $α$ when perfectly uniform, since $N \cdot N \cdot (1/ N)^{2} = 1$ ).

L3.3

Why does the load-balancing loss use the soft probability $P_{i}$ rather than only the hard dispatch fraction $f_{i}$ ? What would break if you used $f_{i}$ alone?

Recall Solution

$f_{i}$ comes from a hard Top- $k$ /argmax decision — "which expert did we actually pick" — which is a step function. Its derivative with respect to router weights is zero almost everywhere and undefined at the jumps, so no gradient flows and the router can't be corrected. $P_{i}$ is the softmax probability: smooth and differentiable. Multiplying $f_{i}$ (which counts where tokens went) by $P_{i}$ (which is learnable) gives a term whose gradient actually reaches $W_{g}$ , letting the optimizer reduce imbalance. See Softmax and Gating Functions.

Level 4 — Synthesis

L4.1

You must design a MoE layer with a compute budget equal to a dense FFN with $2 \times$ the FFN width (so $\approx 2$ expert-FFNs of compute per token), but you want maximum knowledge capacity. Each expert equals one standard FFN in size. Choose $k$ , then choose $N$ to hit a 32× capacity-to-compute ratio. Justify.

Recall Solution

Compute budget = 2 expert-FFNs per token $\Rightarrow$ set $k = 2$ (each active expert costs one FFN). Capacity-to-compute ratio $= N / k$ . We want this $= 32$ . $N = 32 \times k = 32 \times 2 = 64 experts .$ So: $N = 64$ , $k = 2$ . Per-token compute stays at 2 FFNs while total expert capacity equals 64 FFNs — a $64/2 = 32 \times$ gap. This is exactly the sparse-vs-dense lever: crank $N$ for capacity, hold $k$ for cost.

L4.2

Same layer ( $N = 64$ , $k = 2$ ), batch $= 4096$ tokens across the layer. You observe that in practice each expert receives up to $1.4 \times$ its fair share when the router is nearly balanced. Pick a capacity factor that produces zero drops in this near-balanced regime, and compute $C$ .

Recall Solution

With $k = 2$ , the total routing slots per batch $= k \times tokens = 2 \times 4096 = 8192$ . Fair share per expert $= 8192/64 = 128$ slots. Peak observed load $= 1.4 \times 128 = 179.2$ , round up to $180$ to cover it. Choose capacity factor $= 1.4$ so $C = 1.4 \times \frac{8192}{64} = 1.4 \times 128 = 179.2 \Rightarrow C = 180 (rounded up) .$ This meets the peak of $180$ tokens, giving zero drops in the near-balanced regime. Bigger factor = safer but more memory; this is the tightest choice that still avoids drops. Compare with Model Parallelism & Expert Parallelism where $C$ also fixes the fixed-size buffers each device allocates.

Level 5 — Mastery

L5.1

Prove that for fixed dispatch fractions $f_{i}$ with $\sum_{i} f_{i} = 1$ , the sum $\sum_{i} f_{i}^{2}$ is minimized exactly when all $f_{i} = 1/ N$ , and state its minimum value. (This is why the balanced case is the loss minimum in L3.2.)

Recall Solution

By the Cauchy–Schwarz inequality (or the QM–AM inequality) applied to the vector $f$ and the all-ones vector $1$ : $(\sum_{i = 1}^{N} f_{i} \cdot 1)^{2} \leq (\sum_{i} f_{i}^{2}) (\sum_{i} 1^{2}) = N \sum_{i} f_{i}^{2} .$ The left side is $(\sum f_{i})^{2} = 1^{2} = 1$ . Therefore $\sum_{i} f_{i}^{2} \geq \frac{1}{N} .$ Equality in Cauchy–Schwarz holds iff $f$ is proportional to $1$ , i.e. all $f_{i}$ equal; with $\sum f_{i} = 1$ that forces $f_{i} = 1/ N$ . Minimum value $= 1/ N$ . Hence $L_{aux} = α N \sum f_{i}^{2}$ bottoms out at $α N \cdot \frac{1}{N} = α$ under perfect balance — matching L3.2(a) where $α = 1$ gave $L = 1$ .

L5.2

A token's logits are $h = [c, c, c, c]$ (all equal) with $N = 4$ , $k = 2$ . What are the softmax probabilities? What are the Top-2 gates? What subtle issue does this degenerate case expose about Top- $k$ routing?

Recall Solution

Softmax of equal logits is uniform regardless of $c$ : $p_{i} = 1/4 = 0.25$ for all $i$ . Top-2 gates. Any two experts tie for "largest". After picking two and re-normalizing: each kept gate $= 0.25/ (0.25 + 0.25) = 0.5$ , so $g = 0.5, 0.5$ on the chosen pair. Subtle issue. With an exact tie, Top- $k$ selection is ambiguous — the choice of which two experts depends on tie-breaking (index order, floating-point noise). Routing becomes non-deterministic / discontinuous in the logits: an infinitesimal perturbation of one logit can flip which expert is chosen, giving a jump in $y$ . This is why the router is trained to produce peaked, well-separated logits, and why the gate values (not just the selection) receive gradient — smooth gates soften the otherwise hard, jumpy decision boundary.

L5.3

Under perfect load balance with $N$ experts and Top- $k$ routing, what fraction of the total expert parameters is touched on average by a single token, and what fraction of dispatch capacity is used? Give both as functions of $N, k$ , then evaluate for $N = 64, k = 2$ .

Recall Solution

Params touched per token. A token uses $k$ of the $N$ experts, so the fraction of expert parameters activated is $\frac{k}{N} = \frac{2}{64} = \frac{1}{32} = 0.03125 = 3.125%.$ Capacity used under perfect balance. Perfect balance means each expert receives exactly its fair share, so used capacity $= 100%$ (fraction $= 1$ ). No drops, no idle slots. Formally, load per expert $= \frac{k \cdot tokens}{N}$ , and with capacity factor $= 1$ that equals $C$ exactly, so utilization $= 1$ . Takeaway. Each token wakes only $3.125%$ of the expert weights (that's the conditional-computation saving), yet with balance the hardware is fully used — the ideal MoE regime. Sparse per token, dense per batch: that duality is the whole efficiency argument.

Recall Final self-check (cover and recite)

Compute scales with ::: $k$ (active experts), never $N$ . Re-normalize gates over ::: only the kept set $T$ , dividing by $\sum_{j \in T} p_{j}$ . Load-balancing loss form ::: $L_{aux} = α N \sum_{i} f_{i} P_{i}$ , minimized at uniform $1/ N$ . Capacity with Top- $k$ ::: driven by $k \times tokens$ routing slots, not tokens alone. Minimum of $\sum f_{i}^{2}$ ::: $1/ N$ , achieved when all $f_{i} = 1/ N$ .