---
title: "Module 11: Mixture of Experts"
format:
html:
code-fold: false
toc: true
ipynb: default
jupyter: python3
---
{{< include ../_diagram-lib.qmd >}}
{{< include ../_components/step-control.qmd >}}
## Introduction
Every model we have built so far is **dense**: each token is pushed through
every single parameter. Double the parameters and you double the compute for
*every* token. That is the wall the frontier ran into — and **Mixture of
Experts (MoE)** is the most successful way around it.
An MoE replaces the one feed-forward network in a transformer block with *many*
parallel **expert** FFNs plus a small **router** that sends each token to only a
few of them. The model can then hold a huge number of parameters — all the
experts — while each token only pays for the handful it actually uses. Total
parameters and per-token compute are **decoupled**.
This is not a curiosity. **Mixtral 8×7B** (8 experts, top-2 routing),
**DeepSeek-MoE**, and several frontier flagships are MoE models: they serve the
quality of a very large model at the cost of a much smaller one.
### What You'll Learn
- Why **total parameters ≫ active parameters** is the entire point of MoE
- How a **top-k router** scores experts and dispatches each token
- How to build a sparse **`MoELayer`** from scratch (gather → run → scatter)
- Why routers **collapse** and how a **load-balancing loss** prevents it
- How **expert choice** and **auxiliary-loss-free balancing** (DeepSeek-V3) fix
load balance *without* the aux-loss gradient — one by construction, one with a
bias thermostat
- How **fine-grained** and **shared** experts (DeepSeek-MoE) sharpen
specialization at fixed compute — the granularity axis, not the routing axis
- How **Mixture-of-Depths** routes tokens *around the block entirely* — the same
top-k, but the choice is "the block" vs. "the residual", for a static compute cut
- The sparsity factor `E/k` and the memory-vs-compute trade-off it encodes
### Prerequisites
- [Module 06: Transformer](../m06_transformer/lesson.qmd) — the FFN that MoE replaces
- [Module 07: Training](../m07_training/lesson.qmd) — the loss the aux term is added to
- Comfort with `torch` tensor indexing (gather / scatter-add)
## Intuition: A Dense FFN Runs Everything, Every Time
Look back at the transformer block. Its feed-forward network is where most of
the parameters live — and it fires in full for every token, whether the token is
the word "the" or a rare technical term. Intuitively that is wasteful: different
tokens want different transformations, but a dense FFN forces them all through
the same one.
MoE turns that single FFN into a *committee* of experts and hires a **router** to
pick which experts see each token:
- **Experts** — several independent FFNs (say 8). Each can specialize.
- **Router** — a tiny linear layer that scores all experts for a token and keeps
the **top-k** (often just 1 or 2).
- **Combine** — the chosen experts' outputs are summed, weighted by the router.
Because only `k` of the `E` experts run per token, the compute is that of a
`k`-expert FFN, no matter how many experts the model *owns*. Add experts to grow
capacity; keep `k` fixed to keep the bill flat.
::: {.callout-note}
## Key Insight
MoE decouples **capacity** from **compute**. A dense model ties them together —
more parameters means more work per token. An MoE with `E` experts and top-`k`
routing has `E` experts' worth of parameters but does `k` experts' worth of work
per token. The ratio `E/k` is the **sparsity factor**: free capacity, paid for in
memory rather than FLOPs.
:::
## The Math: Routing and Load Balancing
**Routing.** For a token $x$, a linear gate produces one score per expert; a
softmax turns the scores into probabilities, and we keep the top $k$:
$$g(x) = \text{softmax}(x W_g) \in \mathbb{R}^E, \qquad \mathcal{T} = \text{top-}k(g(x))$$
The layer's output sums the chosen experts, weighted by their (renormalized)
gate values:
$$\text{MoE}(x) = \sum_{i \in \mathcal{T}} \tilde{g}_i(x)\, E_i(x), \qquad \tilde{g}_i = \frac{g_i}{\sum_{j \in \mathcal{T}} g_j}$$
where $E_i$ is the $i$-th expert FFN. Only the $k$ experts in $\mathcal{T}$ ever
run.
**Load balancing.** Left to itself the router **collapses** — it discovers a few
strong experts early and sends everything to them, starving the rest (a starved
expert gets no gradient and effectively dies). To counter this, we add an
auxiliary loss (Shazeer 2017; Switch Transformer, Fedus 2021). Over a batch of
$N$ tokens routed top-$k$:
$$\mathcal{L}_{\text{aux}} = E \cdot \sum_{i=1}^{E} f_i \, P_i$$
where $f_i$ is the fraction of dispatch slots sent to expert $i$ and $P_i$ is the
mean router probability for expert $i$. Both vectors sum to 1, so the loss is
**minimized at uniform routing** ($f_i = P_i = 1/E$), giving $\mathcal{L}_{\text{aux}} = 1$.
Any imbalance pushes it above 1. Training minimizes
$\mathcal{L}_{\text{task}} + \alpha\,\mathcal{L}_{\text{aux}}$ with a small $\alpha$,
so the router is nudged to keep every expert fed.
## Code: An Expert and a Router from Scratch
Everything lives in `moe.py`. First we load it and build the two pieces —
an `Expert` (an ordinary GELU FFN) and the top-k `MoERouter`.
```{python}
import importlib.util
import sys
from pathlib import Path
import torch
# Load moe.py (directory name starts with a digit, so use importlib)
spec = importlib.util.spec_from_file_location("moe", Path("moe.py").resolve())
moe = importlib.util.module_from_spec(spec)
sys.modules["moe"] = moe
spec.loader.exec_module(moe)
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"PyTorch {torch.__version__} on {device}")
router = moe.MoERouter(embed_dim=64, num_experts=8, top_k=2)
x = torch.randn(5, 64) # 5 tokens
out = router(x)
print(f"probs shape: {tuple(out.probs.shape)} (softmax over 8 experts)")
print(f"chosen experts: {out.topk_indices.tolist()}")
print(f"chosen weights: {out.topk_weights.round(decimals=3).tolist()}")
print(f"weights sum to 1? {out.topk_weights.sum(dim=-1).round(decimals=4).tolist()}")
```
Each token picked its 2 best experts, and the two kept weights were renormalized
to sum to 1 — the discarded probability mass is simply dropped.
## Code: The Sparse MoE Layer
The layer ties experts and router together. The trick is the **dispatch**: for
each expert, gather exactly the tokens routed to it, run the expert once on that
batch, then **scatter-add** the weighted results back. That is the same
gather → run → scatter pattern real MoE kernels use to avoid running every expert
on every token.
```{python}
moe_layer = moe.MoELayer(embed_dim=64, ff_dim=256, num_experts=8, top_k=2)
x = torch.randn(2, 16, 64) # (batch, seq, embed)
output, aux = moe_layer(x)
print(f"input shape: {tuple(x.shape)}")
print(f"output shape: {tuple(output.shape)} (same as input — a drop-in FFN)")
print(f"load-balancing aux loss: {aux.item():.4f} (1.0 = perfectly balanced)")
```
The forward pass returns **two** things: the output *and* the load-balancing
loss for this batch. During training you add `alpha * aux` to your task loss.
**Verifying the dispatch math.** With `top_k=1`, each token's output must be
*exactly* its one chosen expert applied to it (its weight renormalizes to 1). We
can check that directly:
```{python}
layer1 = moe.MoELayer(embed_dim=32, ff_dim=64, num_experts=4, top_k=1)
layer1.eval()
x = torch.randn(1, 6, 32)
out1, _ = layer1(x)
flat = x.reshape(-1, 32)
route = layer1.router(flat)
manual = torch.stack([
layer1.experts[route.topk_indices[t, 0]](flat[t:t+1])[0]
for t in range(flat.shape[0])
])
print("top-1 output matches the chosen expert exactly:",
torch.allclose(out1.reshape(-1, 32), manual, atol=1e-5))
```
## Total vs Active Parameters
Here is the payoff, made concrete. `count_parameters()` reports what the layer
*owns* versus what a single token *touches*:
```{python}
for E in (2, 8, 32):
c = moe.MoELayer(embed_dim=512, ff_dim=2048, num_experts=E, top_k=2).count_parameters()
print(f"E={E:2d}, k=2 → total {c['total']/1e6:6.1f}M | "
f"active {c['active']/1e6:5.1f}M | sparsity {c['sparsity_factor']:.1f}×")
```
Going from 2 to 32 experts multiplies total parameters ~16× while the **active**
parameters per token barely move. Drive the same trade-off yourself:
{{< include _moe-viz.qmd >}}
::: {.callout-tip}
## Try This
1. **Add experts, watch compute stay flat.** Slide `E` from 8 to 64 in the
sparsity meter. The *total* bar grows; the *active* bar barely moves. That gap
is capacity you get almost for free (in FLOPs — you still pay in memory).
2. **Raise top-k.** Push `k` from 1 to 4. Active compute climbs and the sparsity
factor `E/k` shrinks — top-k is the compute dial.
3. **Break the router.** In the routing simulator, drag the **skew** up. Tokens
pile onto a few experts, the others starve, and the load-balancing loss climbs
far above 1.0. This is the collapse the auxiliary loss exists to prevent.
4. **Balance it back.** Return skew to 0 and watch every bar settle near the
`1/E` line and the loss fall back to ≈ 1.0.
:::
## Expert Choice: Let Experts Pick the Tokens
Everything so far fights the same enemy: the router *wants* to collapse, and the
auxiliary loss is a leash that keeps it from doing so. Load balance is a
**soft** constraint we pay for with an extra loss term, a tuning knob ($\alpha$),
and — when an expert still overflows its capacity — **dropped tokens**.
**Expert Choice routing** (Zhou et al., 2022) removes the whole fight with one
idea: *flip the selection*. Instead of each **token** picking its top-$k$
**experts**, each **expert** picks its top-$k$ **tokens**.
Read that twice — it is the entire trick. If every expert takes exactly $k$
tokens, then every expert is busy with exactly $k$ tokens. Load balance is no
longer something we encourage; it is **true by construction**. No auxiliary loss.
No capacity overflow. No dropped tokens (under the budget).
What changes is *coverage*. In token choice, every token gets exactly $k$
experts. In expert choice, a token gets a **variable** number — a "hard" token
that scores high for many experts may be picked by several; an easy one may be
picked by none at all (its residual connection carries it through untouched).
Compute now flows to where the model thinks it is needed, instead of being spread
uniformly.
::: {.callout-note}
## Key Insight
Token choice and expert choice select from the **same** score matrix — they just
read it along different axes. Token choice takes the top-$k$ of each **row** (a
token's best experts); expert choice takes the top-$k$ of each **column** (an
expert's best tokens). Row-wise selection fixes *coverage* and leaves *balance*
to chance; column-wise selection fixes *balance* and leaves *coverage* to chance.
:::
### The Math: The Selection Matrix, Transposed
Start from the same router scores as before — a linear gate and a softmax give,
for a batch of $N$ tokens, a score for every expert:
$$S = \text{softmax}(X W_g) \in \mathbb{R}^{N \times E}$$
Token choice took the top-$k$ along the expert axis (each row). Expert choice
transposes and takes the top-$k$ along the **token** axis (each column of $S$, or
each row of $S^\top$):
$$G,\, I = \text{top-}k(S^\top, k), \qquad P = \text{onehot}(I)$$
- $I \in \mathbb{Z}^{E \times k}$ — `I[i, j]` is the $j$-th token expert $i$ chose.
- $G \in \mathbb{R}^{E \times k}$ — the matching gate weights (the scores
themselves; **not** renormalized — a token has no fixed budget to normalize
over).
- $P \in \{0,1\}^{E \times k \times N}$ — a one-hot **dispatch** tensor.
The capacity — how many tokens each expert keeps — is set by a **capacity
factor** $c$:
$$k = \Big\lceil \frac{N \cdot c}{E} \Big\rceil$$
At $c = 1$ there are $E \cdot k \approx N$ total slots, so on average each token
is covered once. The forward pass is then a gather → run → scatter, written as
two einsums over $P$:
$$X_{\text{in}} = P \cdot X, \qquad X_e^{(i)} = \text{Expert}_i(X_{\text{in}}^{(i)}), \qquad X_{\text{out}}[l] = \sum_{i,j} P[i,j,l]\, G[i,j]\, X_e^{(i)}[j]$$
Every expert runs once, on exactly its $k$ tokens. The scatter sums each token's
contributions from whichever experts chose it — zero, one, or many.
### Code: Expert Choice from Scratch
The new pieces live alongside the token-choice code in `moe.py`:
`ExpertChoiceRouter`, `ExpertChoiceMoELayer`, and the `expert_choice_route`
selection. First, the selection on a tiny hand-checkable example:
```{python}
import torch
# expert_choice_route(scores, capacity): each expert keeps its top-k tokens.
S = torch.tensor([[0.9, 0.1], # token 0 loves expert 0
[0.2, 0.8], # token 1 loves expert 1
[0.6, 0.4]]) # token 2 leans expert 0
G, I, P = moe.expert_choice_route(S, capacity=2)
print(f"expert 0 kept tokens {I[0].tolist()} with gates {G[0].tolist()}")
print(f"expert 1 kept tokens {I[1].tolist()} with gates {G[1].tolist()}")
print(f"P (dispatch) shape: {tuple(P.shape)} = (E, k, N)")
```
Now the full layer. Watch the two signatures of expert choice fall straight out
of the stats — **perfect balance** and **variable coverage**:
```{python}
ec = moe.ExpertChoiceMoELayer(embed_dim=64, ff_dim=256, num_experts=8,
capacity_factor=1.0)
x = torch.randn(2, 24, 64) # (batch, seq, embed) → N = 48 tokens
out, stats = ec(x)
tpe = stats.tokens_per_expert.int().tolist()
ept = stats.experts_per_token.int().tolist()
coverage = {c: ept.count(c) for c in range(max(ept) + 1)}
print(f"output shape: {tuple(out.shape)} (a drop-in FFN)")
print(f"capacity k: {stats.capacity}")
print(f"tokens per expert: {tpe}")
print(f" → identical for every expert: perfect balance, and no aux loss")
print(f"experts per token: {coverage} (count by coverage)")
print(f" → variable: {stats.uncovered_tokens} token(s) chosen by no expert")
```
There is no second return value for a load-balancing loss, because there is
nothing to balance — the layer returns routing **stats** instead. Compare that to
the token-choice `MoELayer`, which had to hand back `aux` on every forward.
### Two Routings, Side by Side
This is the whole story in one picture. Below is a **real** score matrix $S$ from
an `ExpertChoiceRouter`. Toggle the selection rule and watch the *same* scores
produce two completely different dispatches: token choice lights up each **row's**
best experts (even coverage, lumpy load); expert choice lights up each
**column's** best tokens (even load, lumpy coverage).
```{python}
#| output: false
#| echo: false
# Bridge a real routing example to the visualization.
torch.manual_seed(9)
N, E, top_k = 20, 5, 2
router = moe.ExpertChoiceRouter(embed_dim=48, num_experts=E)
tokens = torch.randn(N, 48)
scores = router(tokens, capacity=1).scores # (N, E) softmax scores
# Token choice: each token (row) keeps its top-k experts.
tc_idx = scores.topk(top_k, dim=-1).indices # (N, k)
tc_mask = torch.zeros(N, E)
tc_mask.scatter_(1, tc_idx, 1.0)
# Expert choice: each expert (column) keeps its top-k tokens (capacity_factor=1).
cap = moe.expert_choice_capacity(N, E, capacity_factor=1.0)
_, ec_idx, ec_perm = moe.expert_choice_route(scores, capacity=cap)
ec_mask = ec_perm.sum(dim=1).t() # (N, E), 1 where chosen
tc_load = tc_mask.sum(0).int().tolist() # tokens per expert (uneven)
ec_load = ec_mask.sum(0).int().tolist() # tokens per expert (all == cap)
ec_cover = ec_mask.sum(1).int().tolist() # experts per token (variable)
tc_cover = tc_mask.sum(1).int().tolist() # experts per token (all == top_k)
ojs_define(ecScores = scores.tolist())
ojs_define(ecTokenMask = tc_mask.tolist())
ojs_define(ecExpertMask = ec_mask.tolist())
ojs_define(ecTCLoad = tc_load, ecECLoad = ec_load)
ojs_define(ecTCCover = tc_cover, ecECCover = ec_cover)
ojs_define(ecCapacity = cap, ecTopK = top_k)
```
{{< include _expert-choice-viz.qmd >}}
::: {.callout-tip}
## Try This
1. **Flip the rule.** Switch the toggle from **Token choice** to **Expert
choice**. The highlighted cells jump from row-aligned to column-aligned, and
the load bars snap from ragged to a perfectly flat line at $k$.
2. **Read the columns.** Under expert choice, count the lit cells in any column —
it is always exactly the capacity $k$. That uniformity *is* the load balance,
with no loss term anywhere.
3. **Find the orphans.** In the coverage histogram, look at the `0` bucket: tokens
no expert wanted. Token choice never has these (everyone gets $k$); expert
choice trades that guarantee for its balance.
4. **Spot the hoarders.** Some tokens sit in the high-coverage buckets — chosen by
many experts at once. That is compute flowing to the tokens the model finds
hardest, which uniform token choice can't do.
:::
## Balancing Without a Loss: The Bias Thermostat
Expert choice bought perfect balance but spent something to get it: it needs the
**whole batch** present to let each expert pick its tokens. That is fine for
training, but **autoregressive decoding generates one token at a time** — there is
no batch to choose from (pitfall 6, below). So at inference we are pushed back to
**token choice**, and with it the auxiliary loss and all its baggage: a tuning
knob $\alpha$, and a balancing gradient that **fights** the language-modeling
gradient for control of the router's weights.
Is there a way to keep token choice — so the model still decodes one token at a
time — yet balance it **without** an auxiliary loss? **Loss-Free Balancing** (Wang
et al., 2024), the scheme **DeepSeek-V3** trains with, says yes, with one small
idea.
Give each expert a scalar **bias** $b_i$, and add it **only to the top-$k$
selection** — never to the gate value that weights the expert's output:
$$g_i = s_i \ \text{ if } (s_i + b_i) \in \text{top-}k(\{s_j + b_j\}_j), \quad \text{else } 0$$
Read that carefully: $s_i + b_i$ decides **which** experts win, but a winner's
combine weight is the **untouched** affinity $s_i$. Because $b_i$ never enters the
output, **no gradient ever flows through it**. It is not a learned parameter — it
is a **thermostat** the training loop nudges by hand: after each batch, push the
bias of any **overloaded** expert **down** (so it wins fewer tokens next time) and
any **starved** expert **up**.
::: {.callout-note}
## Key Insight
The auxiliary loss balances the router by **changing its gradient** — it competes
with the task for the gate weights. Loss-Free Balancing balances it by **changing
the selection threshold**, through a term the gradient never sees. One steers the
router with a force *inside* the loss; the other steers it with a knob *outside*
the loss entirely. That is why it is "loss-free": nothing about the balancing ever
touches the training objective.
:::
Here is the mechanism, one stage at a time. Notice the split: the bias flows into
the **Selection** stage, but the **Value** stage reads straight from the raw
affinity — the orange path skips the bias entirely.
```{ojs}
//| echo: false
viewof lfMechStep = stepControl({min: 0, max: 5, value: 0, label: "Stage"})
```
```{ojs}
//| echo: false
lfMechSteps = [
{title: "Affinity sᵢ", caption: "The softmax gate scores every expert for the token — the same sᵢ as token choice."},
{title: "Selection sᵢ + bᵢ", caption: "Add the per-expert bias — but only to decide the ranking, not the weight."},
{title: "Top-k pick", caption: "The biased scores choose which k experts run. A big −bᵢ can keep a popular expert out."},
{title: "Value = raw sᵢ", caption: "The winners' combine weights come from the UNbiased sᵢ. The bias is invisible here — no gradient touches it."},
{title: "Count load cᵢ", caption: "Tally how many tokens each expert received this batch."},
{title: "bᵢ ← bᵢ + u·sign(c̄ − cᵢ)", caption: "Overloaded (cᵢ > c̄) → bias down; starved → bias up. A gradient-free step, fed back to selection."}
]
```
```{ojs}
//| echo: false
lfMechDiagram = {
const theme = diagramTheme;
const width = 760, height = 320;
const svg = d3.create("svg")
.attr("viewBox", `0 0 ${width} ${height}`)
.attr("width", "100%").attr("height", height)
.style("max-width", `${width}px`)
.style("font-family", "var(--pg-mono)");
svg.append("rect").attr("width", width).attr("height", height)
.attr("fill", theme.bg).attr("rx", 12);
// Nodes: (x, y) is the TOP-LEFT corner; w, h the size. Top row = selection
// path; the value branch and the bias-update feedback drop below it.
const nodes = [
{id: 0, x: 30, y: 60, w: 120, h: 52, label: ["Affinity", "sᵢ"]},
{id: 1, x: 210, y: 60, w: 140, h: 52, label: ["Selection", "sᵢ + bᵢ"]},
{id: 2, x: 410, y: 60, w: 120, h: 52, label: ["Top-k", "pick"]},
{id: 4, x: 600, y: 130, w: 130, h: 52, label: ["Load cᵢ", "count"]},
{id: 3, x: 410, y: 200, w: 140, h: 52, label: ["Value", "= raw sᵢ"]},
{id: 5, x: 200, y: 200, w: 160, h: 52, label: ["bᵢ ← bᵢ +", "u·sign(c̄ − cᵢ)"]},
];
const byId = Object.fromEntries(nodes.map(n => [n.id, n]));
const cx = n => n.x + n.w / 2, cy = n => n.y + n.h / 2;
const active = lfMechStep;
const edges = [
{a: 0, b: 1, kind: "sel"}, // s → selection (bias added here)
{a: 1, b: 2, kind: "sel"}, // selection → top-k
{a: 2, b: 4, kind: "flow"}, // top-k → load
{a: 0, b: 3, kind: "val"}, // s → value (raw — bypasses the bias!)
{a: 3, b: 4, kind: "flow"}, // value → (combine) → load
{a: 4, b: 5, kind: "ctrl"}, // load → bias update
{a: 5, b: 1, kind: "ctrl"}, // bias update → selection (feedback)
];
svg.append("defs").append("marker").attr("id", "lfArrow")
.attr("viewBox", "0 -5 10 10").attr("refX", 9).attr("refY", 0)
.attr("markerWidth", 6).attr("markerHeight", 6).attr("orient", "auto")
.append("path").attr("d", "M0,-5L10,0L0,5").attr("fill", theme.edgeStroke);
function edgePath(e) {
const A = byId[e.a], B = byId[e.b];
if (e.kind === "val") {
// s drops down from its bottom, runs across, and enters the value box left.
return `M ${cx(A)} ${A.y + A.h} V ${cy(B)} H ${B.x}`;
}
if (e.a === 4 && e.b === 5) {
// load → bias: down from load's bottom, across to the update box right.
return `M ${cx(A)} ${A.y + A.h} V ${cy(B)} H ${B.x + B.w}`;
}
if (e.a === 5 && e.b === 1) {
// bias update → selection: up from the update box top into selection bottom.
return `M ${cx(A)} ${A.y} V ${B.y + B.h}`;
}
if (e.a === 3 && e.b === 4) {
// value → load: up-and-right into the load box left.
return `M ${A.x + A.w} ${cy(A)} H ${B.x - 20} V ${cy(B)} H ${B.x}`;
}
// straight, right side of A → left side of B.
return `M ${A.x + A.w} ${cy(A)} H ${B.x}`;
}
for (const e of edges) {
const on = (active >= 1 && active <= 2 && e.kind === "sel")
|| (active === 3 && e.kind === "val")
|| (active === 4 && e.a === 2 && e.b === 4)
|| (active === 5 && e.kind === "ctrl");
const color = e.kind === "val" ? theme.accent
: e.kind === "ctrl" ? theme.highlight : theme.edgeStroke;
svg.append("path").attr("d", edgePath(e)).attr("fill", "none")
.attr("stroke", color).attr("stroke-width", on ? 3 : 1.6)
.attr("opacity", on ? 1 : 0.4)
.attr("stroke-dasharray", e.kind === "ctrl" ? "6,4" : null)
.attr("marker-end", "url(#lfArrow)");
}
for (const n of nodes) {
const on = n.id === active;
// The value box is tinted with the accent when it lights (the "raw sᵢ" point).
const lit = on && n.id === 3 ? theme.accent : theme.highlight;
const g = svg.append("g").attr("transform", `translate(${n.x},${n.y})`);
g.append("rect").attr("width", n.w).attr("height", n.h).attr("rx", 8)
.attr("fill", on ? lit : theme.nodeFill)
.attr("stroke", on ? lit : theme.nodeStroke)
.attr("stroke-width", on ? 2.5 : 1.4)
.attr("filter", on ? `drop-shadow(0 0 8px ${lit})` : null);
g.append("text").attr("x", n.w / 2).attr("y", n.h / 2 - 7)
.attr("text-anchor", "middle").attr("dominant-baseline", "middle")
.attr("fill", on ? theme.bgOpaque : theme.nodeText)
.attr("font-size", "12px").attr("font-weight", "700").text(n.label[0]);
g.append("text").attr("x", n.w / 2).attr("y", n.h / 2 + 11)
.attr("text-anchor", "middle").attr("dominant-baseline", "middle")
.attr("fill", on ? theme.bgOpaque : theme.nodeText)
.attr("font-size", "11px").text(n.label[1]);
}
return svg.node();
}
```
```{ojs}
//| echo: false
html`<div style="
background: ${diagramTheme.highlightBg};
border-left: 3px solid ${diagramTheme.highlight};
border-radius: 0 6px 6px 0; padding: 10px 16px; margin-top: 10px;
font-family: var(--pg-mono); color: ${diagramTheme.nodeText};
">${lfMechSteps[lfMechStep].caption}</div>`
```
### The Math: A Bias on Selection, Not Value
The affinities are the same softmax as before, $s = \text{softmax}(x W_g)$. The
only change is the **selection set**: rank by the biased scores, but keep the
**raw** affinity as the weight.
$$\mathcal{T} = \text{top-}k(\{s_j + b_j\}_j), \qquad \text{MoE}(x) = \sum_{i \in \mathcal{T}} \tilde{s}_i\, E_i(x)$$
where $\tilde{s}_i = s_i / \sum_{j \in \mathcal{T}} s_j$ renormalizes the kept
**unbiased** affinities. After each batch of $N$ tokens, measure each expert's
**load** $c_i$ (how many tokens it received) and its **violation** from the mean
$\bar{c} = Nk/E$:
$$e_i = \bar{c} - c_i, \qquad b_i \leftarrow b_i + u\cdot\text{sign}(e_i)$$
$u$ is the **bias update rate** (DeepSeek uses $u = 0.001$). The update is pure
control — no loss, no backprop. `sign` means every expert's bias moves by exactly
$\pm u$ per step, so a perfectly balanced batch ($e_i = 0$) leaves the bias frozen.
Balance is scored by the **maximal violation**:
$$\text{MaxVio} = \frac{\max_i c_i - \bar{c}}{\bar{c}}$$
0 is perfect; 1 means the busiest expert carries twice its share.
::: {.callout-note}
## Key Insight
The bias is a **buffer, not a parameter**. In the code below it is a
`register_buffer`, so PyTorch's autograd never sees it — the "loss-free" property
is structural, not a convention you have to remember. The optimizer updates the
experts and the gate; the thermostat, separately, updates the bias.
:::
### Code: Loss-Free Balancing from Scratch
The new pieces live beside the other routings in `moe.py`. The heart is
`loss_free_select`: pick on `scores + bias`, but weight with the raw `scores`.
```{python}
import torch
# Three affinities; expert 2 is least attractive to this token.
S = torch.tensor([[0.5, 0.3, 0.2]])
# No bias → the top-2 are experts 0 and 1.
_, idx0 = moe.loss_free_select(S, torch.zeros(3), top_k=2)
# A big +bias on expert 2 pulls it INTO the top-2, displacing expert 1…
w2, idx2 = moe.loss_free_select(S, torch.tensor([0.0, 0.0, 1.0]), top_k=2)
print(f"no bias → experts {sorted(idx0[0].tolist())}")
print(f"bias E2 → experts {sorted(idx2[0].tolist())}")
print(f"…but E2's weight is its RAW 0.2, renormalized: {w2[0].tolist()}")
print(f" (0.2/0.7, 0.5/0.7) = {(0.2/0.7, 0.5/0.7)} — the bias never touched the value")
```
The thermostat is three tiny functions — count the load, measure the violation,
nudge the bias by its sign:
```{python}
idx = torch.tensor([[0, 1], [0, 2], [0, 0]]) # a batch that piles onto expert 0
counts = moe.expert_load(idx, num_experts=3)
print(f"load per expert: {counts.tolist()} (expert 0 overloaded)")
print(f"MaxVio: {moe.max_violation(counts):.2f}")
bias = torch.zeros(3)
bias = moe.bias_update(bias, counts, update_rate=0.1)
print(f"bias after 1 step: {[round(b, 2) for b in bias.tolist()]}"
f" (overloaded ↓, starved ↑)")
```
The full `LossFreeMoELayer` is still **token choice** — so it decodes one token at
a time — but returns routing **stats** instead of an auxiliary loss. A training
loop calls `update_bias` once per step, *after* the optimizer:
```{python}
layer = moe.LossFreeMoELayer(embed_dim=64, ff_dim=256, num_experts=8, top_k=2)
x = torch.randn(2, 16, 64)
out, info = layer(x)
print(f"output shape: {tuple(out.shape)} (a drop-in FFN — no aux loss returned)")
print(f"max violation: {info.max_violation:.2f}")
print(f"bias needs a gradient? {layer.router.bias.requires_grad} "
f"(it is a buffer — autograd never sees it)")
# One training step: backward + optimizer would go here, THEN the thermostat:
layer.update_bias(info) # gradient-free balancing step
```
### Watch the Thermostat Converge
The balancing mechanism is **separate** from gradient training — that is the whole
point — so we can watch it work on its own. Below we take a deliberately
**collapsed** router (a skewed affinity where two experts hoard every token) and
run only the bias control loop. No experts are trained; nothing but the bias
moves. Scrub the control step and watch the load bars flatten to the $\bar{c}$
line while `MaxVio` falls from 3.0 toward 0.
```{python}
#| output: false
#| echo: false
# Drive the control loop on a fixed, badly-skewed affinity and bridge the history.
torch.manual_seed(0)
E, k, N = 8, 2, 512
base = torch.tensor([3.0, 2.2, 0.4, 0.2, 0.0, -0.2, -0.4, -0.6]) # experts 0–1 hoard
scores = torch.softmax(base + 0.6 * torch.randn(N, E), dim=-1)
hist = moe.balance_experts(scores, top_k=k, update_rate=0.004, steps=80)
ojs_define(lfLoads = hist["loads"])
ojs_define(lfBiases = hist["biases"])
ojs_define(lfViolations = hist["violations"])
ojs_define(lfMean = hist["mean_load"])
ojs_define(lfSteps = hist["steps"])
```
{{< include _loss-free-viz.qmd >}}
::: {.callout-tip}
## Try This
1. **Watch the collapse heal.** Drag the control step from 0. At step 0 two
experts hold all 1024 dispatch slots (MaxVio 3.0); by the end every bar sits
near the $\bar{c}=128$ line — with **not one gradient step** taken.
2. **Read the bias.** The right panel is the thermostat's state: the two hoarders
are driven **negative** (harder to select), the starved experts **positive**.
The bias is doing the balancing that an aux-loss gradient would otherwise do.
3. **Find the wobble.** Late in the run `MaxVio` doesn't hit exactly 0 — it settles
into a small band. `sign`-based control always leaves a limit cycle; a smaller
$u$ tightens the band but converges slower (try it in `balance_experts`).
4. **Confirm the values never moved.** Nothing in this loop changed a single
combine weight — only *which* experts were selected. The gate values stayed the
raw affinities the whole time.
:::
## Shared and Fine-Grained Experts: A Better Expert, Not a Better Router
Every variant so far tuned *who picks whom* — token choice, expert choice, the
bias thermostat. **DeepSeek-MoE** (Dai et al., 2024) leaves the router alone and
changes *what an expert is*, with two moves that power the modern open-MoE
frontier (DeepSeek-V2/V3, and the recipe behind many 2024–2025 flagships):
- **Fine-grained segmentation** — cut each expert into `m` smaller ones (hidden
dim → `1/m`) and activate `m×` as many. Compute and parameter count are held
**exactly** constant, but the number of ways to pick the active set explodes.
- **Shared expert isolation** — reserve a few experts that *every* token uses,
unconditionally. They absorb the knowledge common to all tokens so the routed
experts stop re-learning it and are free to specialize.
::: {.callout-note}
## Key Insight
The routing tricks earlier in this module all fight the *load-balancing* problem.
DeepSeek-MoE fights a different one: **redundancy and coarseness**. A handful of
big experts each end up as generalists (re-encoding the same common knowledge);
many small experts, plus a shared generalist carved out explicitly, give sharper
specialists at the *same* cost.
:::
### Intuition: A Token Through the Layer
A DeepSeek-MoE layer runs two paths and sums them. The **shared** path is
ungated — every token flows through all `Kₛ` shared experts. The **routed** path
is the top-`kᵣ` gate you already built, over many small experts. Step through it:
```{ojs}
//| echo: false
viewof dsmFlowStep = stepControl({min: 0, max: 4, value: 0, label: "Stage"})
```
```{ojs}
//| echo: false
dsmFlowSteps = [
{title: "Token in", caption: "One token xₜ enters the FFN sublayer."},
{title: "Shared experts", caption: "It flows through ALL Kₛ shared experts — no gate, always on. They hold the common knowledge."},
{title: "Route", caption: "A gate scores the Nᵣ small routed experts; softmax → keep the top kᵣ (raw affinity, not renormalized)."},
{title: "Routed experts", caption: "Only those kᵣ routed experts run, each a fine-grained 1/m-width specialist."},
{title: "Sum", caption: "h = Σ shared + Σ gated routed. The block then adds the residual xₜ."},
]
```
```{ojs}
//| echo: false
dsmFlowDiagram = {
const theme = diagramTheme;
const width = 760, height = 340;
const svg = d3.create("svg")
.attr("viewBox", `0 0 ${width} ${height}`)
.attr("width", "100%").attr("height", height)
.style("max-width", `${width}px`)
.style("font-family", "var(--pg-mono)");
svg.append("rect").attr("width", width).attr("height", height)
.attr("fill", theme.bg).attr("rx", 12);
const step = dsmFlowStep;
const lit = (on) => on ? theme.highlight : theme.nodeStroke;
const litFill = (on) => on ? theme.highlightGlow : theme.nodeFill;
// token
const tokActive = step >= 0;
const tok = svg.append("g").attr("transform", "translate(60,150)");
tok.append("rect").attr("width", 70).attr("height", 40).attr("rx", 6)
.attr("x", -35).attr("y", -20)
.attr("fill", litFill(step === 0)).attr("stroke", lit(tokActive)).attr("stroke-width", 2);
tok.append("text").attr("text-anchor", "middle").attr("dy", 5)
.attr("fill", theme.nodeText).attr("font-family", "var(--pg-mono)").text("xₜ");
// shared experts (top row)
const sharedOn = step === 1 || step === 4;
const sx = 300, sy = 70;
svg.append("text").attr("x", sx).attr("y", sy - 44).attr("text-anchor", "middle")
.attr("fill", sharedOn ? theme.highlight : theme.nodeText).attr("font-size", "12px")
.attr("font-weight", 700).text("shared (always on)");
for (let i = 0; i < 2; i++) {
const g = svg.append("g").attr("transform", `translate(${sx - 60 + i * 120},${sy})`);
g.append("rect").attr("x", -42).attr("y", -20).attr("width", 84).attr("height", 40).attr("rx", 6)
.attr("fill", litFill(sharedOn)).attr("stroke", lit(sharedOn)).attr("stroke-width", 2);
g.append("text").attr("text-anchor", "middle").attr("dy", 5)
.attr("fill", theme.nodeText).attr("font-size", "12px").text(`FFNˢ${i + 1}`);
}
// router + routed experts (bottom row)
const routeOn = step === 2;
const routedOn = step === 3 || step === 4;
const gate = svg.append("g").attr("transform", "translate(210,250)");
gate.append("rect").attr("x", -34).attr("y", -22).attr("width", 68).attr("height", 44).attr("rx", 6)
.attr("fill", litFill(routeOn)).attr("stroke", lit(routeOn)).attr("stroke-width", 2);
gate.append("text").attr("text-anchor", "middle").attr("dy", 4)
.attr("fill", theme.nodeText).attr("font-size", "11px").text("gate");
const rx = 430, ry = 250, chosen = new Set([1, 3]);
svg.append("text").attr("x", rx + 30).attr("y", ry - 44).attr("text-anchor", "middle")
.attr("fill", routedOn ? theme.highlight : theme.nodeText).attr("font-size", "12px")
.attr("font-weight", 700).text("routed (top kᵣ of many small)");
for (let i = 0; i < 5; i++) {
const on = routedOn && chosen.has(i);
const g = svg.append("g").attr("transform", `translate(${rx - 20 + i * 55},${ry})`);
g.append("rect").attr("x", -22).attr("y", -18).attr("width", 44).attr("height", 36).attr("rx", 5)
.attr("fill", on ? theme.highlightGlow : (routeOn ? theme.nodeFill : theme.surface || theme.nodeFill))
.attr("stroke", on ? theme.highlight : theme.nodeStroke).attr("stroke-width", on ? 2.5 : 1)
.attr("opacity", routedOn && !on ? 0.4 : 1);
g.append("text").attr("text-anchor", "middle").attr("dy", 4)
.attr("fill", theme.nodeText).attr("font-size", "10px").text(`E${i}`);
}
// sum node
const sumOn = step === 4;
const sum = svg.append("g").attr("transform", "translate(700,150)");
sum.append("circle").attr("r", 24)
.attr("fill", litFill(sumOn)).attr("stroke", lit(sumOn)).attr("stroke-width", 2);
sum.append("text").attr("text-anchor", "middle").attr("dy", 7)
.attr("fill", theme.nodeText).attr("font-size", "20px").text("Σ");
// edges
const edge = (x1, y1, x2, y2, on) => svg.append("line")
.attr("x1", x1).attr("y1", y1).attr("x2", x2).attr("y2", y2)
.attr("stroke", on ? theme.highlight : theme.edgeStroke)
.attr("stroke-width", on ? 2.5 : 1.2).attr("opacity", on ? 1 : 0.5);
edge(95, 145, 200, 80, sharedOn); // token → shared
edge(95, 155, 176, 245, routeOn || routedOn); // token → gate
edge(244, 250, 408, 250, routedOn); // gate → routed
edge(360, 70, 676, 140, sumOn); // shared → sum
edge(560, 250, 676, 162, sumOn); // routed → sum
svg.append("text").attr("x", width / 2).attr("y", height - 12).attr("text-anchor", "middle")
.attr("fill", theme.nodeText).attr("font-size", "12.5px")
.text(`${dsmFlowSteps[step].title}: ${dsmFlowSteps[step].caption}`);
return svg.node();
}
```
### The Math: Shared Isolation and Fine-Grained Segmentation
With `Kₛ` shared experts and `Nᵣ` routed experts (of which the top `kᵣ` fire), the
layer output for token $x_t$ is
$$
\mathbf{h}_t = \underbrace{\sum_{i=1}^{K_s}\text{FFN}_i^{(s)}(x_t)}_{\text{shared, always on}}
+ \underbrace{\sum_{i=1}^{N_r} g_{i,t}\,\text{FFN}_i^{(r)}(x_t)}_{\text{routed, top-}k_r}
\; ,\qquad
g_{i,t} = \begin{cases} s_{i,t} & s_{i,t}\in\text{TopK}(\{s_{j,t}\}, k_r)\\ 0 & \text{otherwise}\end{cases}
$$
with $s_{i,t} = \text{softmax}_i(x_t^\top e_i)$ the routed affinities (the block
adds the residual $x_t$ afterward). One faithful detail worth flagging: the kept
routed experts are weighted by the **raw** affinity $s_{i,t}$ — DeepSeek-MoE does
**not** renormalize them to sum to 1, unlike Mixtral and our base `MoELayer`.
**Why segmentation is free.** Split each expert into `m` smaller ones and the
per-expert hidden width drops to $d_{ff}/m$, while the activated count rises to
$m\cdot k$. The FFN hidden units a token actually computes,
$$
\underbrace{(m\,k)}_{\text{active experts}}\times \underbrace{(d_{ff}/m)}_{\text{width each}} = k\,d_{ff},
$$
is **independent of $m$** — the same compute and parameters. What changes is the
number of ways to choose the active set: $\binom{16}{2}=120$ becomes
$\binom{64}{8}=4{,}426{,}165{,}368$ at $m=4$. More, smaller experts ⇒ finer
specialization at *identical* cost.
### Code: A DeepSeek-MoE Layer from Scratch
`DeepSeekMoELayer` (in `moe.py`) holds the two expert pools and sums their paths.
`from_base` builds a fine-grained, shared-isolated layer from a plain
`(N, k, d_ff)` baseline and a split factor `m`:
```{python}
ds = moe.DeepSeekMoELayer.from_base(
embed_dim=64, base_ff_dim=256, base_experts=16, base_active=2,
granularity=4, num_shared=2, # split ×4, carve out 2 shared experts
)
x = torch.randn(2, 16, 64)
out, info = ds(x)
print(f"shared experts (always on): {ds.num_shared}")
print(f"routed experts: {ds.num_routed}, activated per token: {ds.routed_active}")
print(f"active experts per token: {ds.count_parameters()['active_experts']} (= 2 shared + 6 routed)")
print(f"each expert width: {ds.expert_ff_dim} (¼ of the 256-wide baseline)")
print(f"output shape: {tuple(out.shape)}")
```
The shared path is **unconditional** — it is exactly the sum of the shared
experts, no matter what the router does. We can prove it by starving the routed
path (`routed_active=0`):
```{python}
shared_only = moe.DeepSeekMoELayer(64, 128, num_shared=2, num_routed=6, routed_active=0)
xf = x.reshape(-1, 64)
out0, _ = shared_only(x)
manual = sum(e(xf) for e in shared_only.shared_experts).reshape(x.shape)
print("routed off ⇒ output is exactly the shared experts:",
torch.equal(out0, manual))
```
### Explore: The Combinatorial Explosion and the Redundancy It Buys
The two widgets below drive the two ideas. **Left:** slide the segmentation factor
`m` and watch the routing combinations explode while the active compute (the grey
bar) never moves. **Right:** the shared-expert payoff — two layers with *identical*
parameter and compute budgets learn a `common(x) + specific(x)` target under fixed
routing; isolating one always-on expert reaches lower loss, because `common(x)` is
learned once instead of redundantly inside every routed expert.
```{python}
#| output: false
#| echo: false
# Fine-grained table: segment the top-2-of-16, d_ff=2048 baseline; compute stays fixed.
_ms = [1, 2, 4, 8, 16]
_cfgs = [moe.fine_grained_config(16, 2, _m, 2048) for _m in _ms]
ojs_define(
dsmM=_ms,
dsmExperts=[c["num_experts"] for c in _cfgs],
dsmActive=[c["active_experts"] for c in _cfgs],
dsmWidth=[c["expert_ff_dim"] for c in _cfgs],
dsmActiveWidth=[c["active_hidden_width"] for c in _cfgs],
dsmCombos=[float(c["combinations"]) for c in _cfgs],
)
# Shared-vs-plain training race at a matched budget.
_spec = moe.demonstrate_shared_specialization(verbose=False)
ojs_define(
dsmSteps=_spec["steps"],
dsmSharedLoss=_spec["shared_loss"],
dsmPlainLoss=_spec["plain_loss"],
dsmSharedFinal=_spec["shared_final"],
dsmPlainFinal=_spec["plain_final"],
)
```
{{< include _deepseek-moe-viz.qmd >}}
::: {.callout-tip}
## Try This
1. **Watch compute stay flat.** Slide `m` from 1 to 16. The combinations bar climbs
past $10^{30}$; the active-width bar does not budge. That is the whole trick:
more specialization, *zero* extra FLOPs.
2. **Find the shared-expert gap.** On the right, the DeepSeek curve settles below
the plain one. Both models have the *same* number of experts and activate the
same two per token — the only difference is that one expert is always on.
3. **Break it in code.** Call `demonstrate_shared_specialization(num_groups=8)` and
`(num_groups=2)`. More groups ⇒ more common knowledge to share ⇒ a wider gap.
4. **Confirm the raw gate.** Print `info.gates.gather(1, info.topk_indices).sum(-1)`
from the layer above — the kept routed weights do **not** sum to 1. DeepSeek-MoE
weights by the raw affinity; the shared path carries the rest of the signal.
:::
## Mixture of Depths: Route Around the Block, Not Just Between Experts
Every scheme so far routes **between experts** — token choice, expert choice,
loss-free balancing, fine-grained + shared. But look at what they share: *every
token still runs through some FFN in every layer*. They vary **who** and **what**;
they never ask **where** — whether a token needs this layer at all.
**Mixture-of-Depths** (Raposo et al., 2024) asks exactly that. A per-block router
gives each token a scalar weight; the block keeps the top-`k` tokens — a fixed
**capacity** `C` of the sequence — and the other `1 − C` skip straight down the
residual, doing *nothing* in this block. It is expert-choice routing where the two
"experts" are **the block's computation** and **the identity function**.
The paper's headline setting routes **every other block** at `C = 12.5%`: 87.5% of
tokens bypass each MoD block, the model matches a dense transformer's loss at a
fraction of the FLOPs, and it steps **up to 50% faster** when sampling.
::: {.callout-note}
## Key Insight
This is the third axis of the module. MoE decouples *parameters* from *compute* by
routing between experts (**who**/**what**); MoD decouples *depth* from *compute* by
routing around blocks (**where**). And because the capacity `k` is fixed in advance,
the compute is **static** — the tensor shapes never depend on the input, unlike
early-exit. You know the exact FLOP bill before the first token arrives.
:::
### Intuition: Not Every Token Needs Every Layer
Predicting the token after "the cat sat on the" barely needs the full stack — the
answer is nearly deterministic. A token that must resolve a long-range dependency
needs every ounce of depth. A dense model spends the *same* compute on both. MoD
gives each block a fixed budget and lets the router spend it on the tokens that
earn it — the rest ride the residual.
Below, a router has scored 16 tokens. Drag the **capacity** to set how many the
block keeps: the top-`k` by weight light up (they run the block); the rest bypass.
Notice the budget is whatever you set it to, *exactly* — that is the static graph.
```{python}
# | output: false
import importlib.util
import sys
from pathlib import Path
import torch
# Load mod.py the same way the lesson loads moe.py (directory starts with a digit).
spec = importlib.util.spec_from_file_location("mod", Path("mod.py").resolve())
mod = importlib.util.module_from_spec(spec)
sys.modules["mod"] = mod
spec.loader.exec_module(mod)
torch.manual_seed(3)
# One sequence of 16 tokens, scored by a scalar MoD router.
S = 16
mod_router = mod.MoDRouter(embed_dim=32)
x_demo = torch.randn(1, S, 32)
mod_weights = mod_router(x_demo, capacity_fraction=0.5).weights[0] # (S,)
ojs_define(modWeights=mod_weights.tolist())
ojs_define(modSeqLen=S)
```
{{< include _mod-routing-viz.qmd >}}
::: {.callout-tip}
## Try This
Set the capacity to **2 of 16** (`C = 12.5%`, the paper's value). Only the two
highest-weight tokens run the block; fourteen bypass. Now slide to **16** — every
token runs, and MoD becomes an ordinary block (a router-gated one). The slider is
the compute dial: FLOPs scale straight with it.
:::
### The Math: A Scalar Router and the Forward Equation
Each MoD block has its own scalar router — a single weight vector $w$:
$$
r_i = w^\top x_i \in \mathbb{R}
$$
one number per token. Collect them into $R = \{r_i\}$ and let $P_\beta(R)$ be the
cutoff that keeps the top $k = \lceil C \cdot S \rceil$ tokens (the $k$-th largest
weight). The block's output for token $i$ is:
$$
x_i^{\,l+1} =
\begin{cases}
\; r_i^{\,l}\, f\!\left(\tilde{X}^{\,l}\right)_i + x_i^{\,l}, & r_i^{\,l} > P_\beta(R^{\,l}) \quad\text{(selected)} \\[4pt]
\; x_i^{\,l}, & r_i^{\,l} \le P_\beta(R^{\,l}) \quad\text{(bypass)}
\end{cases}
$$
where $f$ is the block (self-attention + MLP) run on **only** the selected subset
$\tilde{X}$. Two things to notice. First, the residual $x_i$ is *always* kept — a
bypassed token is untouched, and a selected token adds the block on top. Second,
the router weight $r_i$ **multiplies** the block output. That is deliberate: the
hard top-$k$ choice is not differentiable, but $r_i$ is, so putting it on the
output path lets gradient descent shape *which* tokens the router prefers (a
selected token that helped the loss pulls its own weight up). This is exactly the
expert-choice gating trick from earlier in the module, applied to depth.
### Code: A Mixture-of-Depths Block from Scratch
`mod.py` builds it directly. `MoDRouter` is the scalar gate; `MoDBlock` wraps any
inner block, gathers the top-`k` tokens, runs the block on just those, and scatters
`x + r·f(·)` back — leaving the rest as pure residual.
```{python}
torch.manual_seed(0)
# Wrap a plain MLP block. In a real model this would be attention + MLP.
block = mod.MoDBlock(mod.MLPBlock(64), embed_dim=64, capacity_fraction=0.125)
x = torch.randn(2, 32, 64)
out, info = block(x)
print(f"input {tuple(x.shape)} → output {tuple(out.shape)}")
print(f"capacity C=12.5% → k={info.capacity} of 32 tokens run the block")
print(f"block FLOPs spent: {info.compute_fraction:.1%} of dense")
# The defining property: bypassed tokens are bit-for-bit the residual.
bypass = ~info.route.selected_mask
print(f"bypassed tokens unchanged: {torch.equal(out[bypass], x[bypass])}")
```
The router weight sits on the output, so a single `.backward()` trains it — even
though nobody ever differentiated the top-`k` selection itself:
```{python}
out.sum().backward()
g = block.router.gate.weight.grad
print(f"router receives gradient: {g is not None and not torch.all(g == 0)}")
```
### Compute: The Budget Is Static, and You Interleave It
`k = round(C · S)` depends only on the capacity and the length — never the input —
so the FLOP bill is fixed. The paper does **not** make every block a MoD block; it
interleaves, routing **every other** one, so half the layers stay dense:
```{python}
print("MoD schedule (6 layers, every other):",
mod.interleave_schedule(6, every=2))
flops = mod.mod_block_flops(2048, block_flops_per_token=1_000_000,
capacity_fraction=0.125)
print(f"per-block FLOPs: dense {flops['dense']:,.0f} → "
f"MoD {flops['mod']:,.0f} ({flops['ratio']:.1%})")
```
```{python}
# | output: false
# Bridge the FLOP ladder: dense, MoD-every-block, and MoD-interleaved
# (half the blocks dense), across a sweep of capacities.
caps = [0.125, 0.25, 0.5, 0.75, 1.0]
dense = 1.0
mod_every = [mod.mod_compute_fraction(2048, c) for c in caps]
mod_interleaved = [(dense + c) / 2 for c in mod_every] # half blocks stay dense
ojs_define(modCaps=[c * 100 for c in caps])
ojs_define(modDense=dense)
ojs_define(modEvery=mod_every)
ojs_define(modInterleaved=mod_interleaved)
```
{{< include _mod-flops-viz.qmd >}}
### The Catch: Top-k Is Non-Causal — and How to Decode Anyway
There is a subtlety hiding in "keep the top-`k` tokens". The top-`k` is over the
**whole sequence** — whether token $i$ makes the cut depends on the weights of
tokens $i+1, i+2, \dots$ that come *after* it. During training that is fine (the
sequence is all there). But autoregressive **decoding** produces one token at a
time: when token $i$ arrives you have no future tokens to compare it against, so
you cannot compute its top-`k` membership. The selection is non-causal.
Raposo et al. give two fixes; we build the cleaner one — a small **predictor** (a
second router) trained with a binary cross-entropy to guess, from a token's own
hidden state alone, whether it *would* be in the top-`k`. It never touches the
language-modeling loss, and it learns fast:
```{python}
# | output: false
history = mod.train_causal_predictor(
embed_dim=32, capacity_fraction=0.25, steps=300, verbose=False
)
print(f"predictor accuracy: {history['accuracy'][0]:.1%} → "
f"{history['final_accuracy']:.1%}")
ojs_define(modPredSteps=history["steps"])
ojs_define(modPredAcc=history["accuracy"])
ojs_define(modPredFinal=history["final_accuracy"])
```
{{< include _mod-predictor-viz.qmd >}}
A causal, per-token predictor recovers the non-causal group decision to well over
90% — the paper reports "upwards of 97%" — so a trained MoD model routes token by
token at decode time with no future in hand. The non-causality was never
fundamental; it was just a training-time convenience the predictor undoes.
::: {.callout-tip}
## Try This
1. **Dial the capacity to the compute.** In the FLOP ladder, drop `C` and watch
the MoD bar fall linearly; the interleaved bar falls half as fast because half
the blocks stay dense. `C = 12.5%` interleaved is the paper's ~56%-of-dense
per-block compute.
2. **Watch the predictor catch up.** The accuracy curve climbs from near the
50/50 base rate to the 90s in a few hundred steps — a causal stand-in learning
a non-causal rule. Call `mod.demonstrate_mod_predictor()` to see it print.
3. **Collapse it to a dense block.** Set `capacity_fraction=1.0` in `MoDBlock` and
confirm `info.route.selected_mask.all()` — no token bypasses, and the layer is
just a router-gated residual block.
:::
## Common Pitfalls
1. **Forgetting the auxiliary loss.** Without it the router collapses onto a few
experts within a few hundred steps, and the rest become dead weight. Always
add `alpha * aux_loss` (a common $\alpha$ is 0.01).
2. **Confusing total with active parameters.** An "8×7B" model does not run 56B
parameters per token — with top-2 it runs about 13B. Report both numbers; they
answer different questions (memory vs. FLOPs).
3. **Renormalizing (or not) the top-k weights.** After picking the top-k, their
gate values no longer sum to 1. Renormalize them (as here) so the combined
output is a proper weighted average and its scale is stable.
4. **Training instability from hard routing.** The top-k operation is
non-differentiable in the *choice* itself; gradients flow only through the
kept weights. This makes MoE routers more finicky to train than a dense FFN —
the load-balancing loss and careful initialization matter.
5. **Assuming MoE is always a win.** MoE trades FLOPs for memory and bandwidth.
If you are memory-bound (a single GPU, long context), holding many experts can
cost more than the compute it saves.
6. **Using expert choice at inference.** Expert choice needs the *whole batch* of
tokens present to pick each expert's top-$k$ — it is a batch-wise operation.
That is fine for training (and encoder/prefill), but autoregressive decoding
generates one token at a time, with no batch to choose from. Production MoEs
therefore tend to use expert choice (or its balancing ideas) during training
and fall back to token choice for causal decoding.
7. **Adding the bias to the gate value.** In loss-free balancing the bias $b_i$
goes into the top-$k$ *selection only*. Fold it into the combine weight and you
have quietly re-created an aux-loss-like distortion: the output now depends on
the balancing term, a gradient flows through it, and you have lost the whole
"loss-free" property. Weight with the raw affinity $s_i$; select with $s_i+b_i$.
8. **Updating the bias inside the forward/backward.** The thermostat step is
gradient-free and belongs *after* the optimizer step (`update_bias`), not inside
the loss. Run it under `no_grad` on a buffer, once per batch — a tiny `sign`
update, not another thing to backprop through.
9. **Reading fine-grained as "free memory".** Segmentation holds *compute* and
*parameters* constant — the FLOPs a token spends and the total weight count are
unchanged. It does **not** shrink the model: all `mN` experts still live in
memory, and more, smaller experts can cost more dispatch/gather overhead. The win
is specialization per FLOP, not a smaller footprint.
10. **Renormalizing the DeepSeek-MoE routed gate.** DeepSeek-MoE weights the kept
routed experts by the *raw* softmax affinity $s_{i,t}$, not a renormalized-to-1
weight (that is Mixtral / the base `MoELayer` above). Renormalizing here is a
real modelling change: the always-on shared path is meant to carry the "rest"
of the signal, so the routed weights are deliberately left summing to less
than 1. (DeepSeek-V3 later adds sigmoid gates + per-expert normalization — a
further variation, not this base form.)
11. **Confusing MoD with early-exit.** Both skip computation, but early-exit is
*dynamic* — a token stops when it is "done", so the amount of work depends on
the input and the batch has ragged shapes. MoD keeps a **fixed** budget `k`
every time; the graph is static and hardware-friendly. That is the whole
reason MoD trains and serves efficiently where early-exit is awkward.
12. **Forgetting the router weight on the output.** In `MoDBlock` the update is
`x + r·f(x)`, not `x + f(x)`. Drop the `r` and the top-`k` *choice* becomes the
only thing connecting the router to the loss — and that choice is
non-differentiable, so the router gets no gradient and never learns which
tokens to keep. The scalar gate on the output is what makes MoD trainable.
## Exercises
### Exercise 1: Sparsity of a real config
Mixtral 8×7B uses `E=8`, `top_k=2`. What is its sparsity factor? Roughly how many
parameters are active per token if the total is ~47B and the non-expert
(attention + embeddings) share is ~2B?
```{python}
# E/k sparsity factor; active ≈ non_expert + (k/E) * expert_params
total, non_expert = 47e9, 2e9
expert_params = total - non_expert
active = non_expert + (2 / 8) * expert_params
print(f"sparsity factor E/k = {8/2:.0f}×")
print(f"active params per token ≈ {active/1e9:.1f}B (of {total/1e9:.0f}B total)")
```
### Exercise 2: Watch a router collapse
Build a layer, feed it the *same* few tokens many times, and skew the gate by
hand to one expert. Print `expert_utilization` before and after — most of the
mass should move to a single expert.
### Exercise 3: The load-balancing floor
Show numerically that `load_balancing_loss` bottoms out at 1.0. Construct uniform
`probs` and round-robin `topk_indices` for `E` experts and confirm the loss is
≈ 1.0; then concentrate the routing and confirm it rises.
### Exercise 4: Expert choice needs no aux loss
Take the collapse scenario from Exercise 2 — scores that wildly favour one expert
— and route it two ways. With `MoERouter` + `load_balancing_loss`, confirm the
aux loss climbs far above 1.0. With `expert_choice_route` on the *same* scores,
confirm `tokens_per_expert` is still exactly `k` for every expert (perfect balance
with nothing to tune). Then print `experts_per_token` and find the cost: how many
tokens did no expert choose?
### Exercise 5: The thermostat's speed/precision trade
Run `balance_experts` on a fixed skewed score matrix at several `update_rate`
values (e.g. `0.001`, `0.004`, `0.02`). For each, plot or print `MaxVio` over the
control steps. Confirm the trade the pitfall named: a larger `u` reaches balance in
fewer steps but settles into a *wider* wobble band; a smaller `u` converges slowly
but ends up flatter. Which `u` gets `MaxVio` below `0.2` in the fewest steps?
```{python}
# Your implementation here — reuse the skewed `scores` from the section above.
```
### Exercise 6: Segmentation is compute-neutral
Use `moe.fine_grained_config(16, 2, m, 2048)` for `m ∈ {1, 2, 4, 8}`. For each,
print `num_experts`, `active_experts`, `expert_ff_dim`, `active_hidden_width`, and
`combinations`. Confirm two things: `active_hidden_width` is identical for every
`m` (compute is fixed), and `combinations` strictly increases. Which is the first
`m` to pass one billion routing combinations?
```{python}
# for m in (1, 2, 4, 8): print(moe.fine_grained_config(16, 2, m, 2048))
```
### Exercise 7: MoD's static budget
Build `mod.MoDBlock(mod.MLPBlock(32), 32, capacity_fraction=0.125)` and run *five*
different random `(1, 40, 32)` inputs through it. Print
`info.route.selected_mask.sum()` each time and confirm it is **identical** across
all five — the compute is fixed regardless of the tokens. Then compare with
`mod.mod_capacity(40, 0.125)`. Why is this the property that makes MoD
hardware-friendly where early-exit is not?
```{python}
# import mod; block = mod.MoDBlock(mod.MLPBlock(32), 32, capacity_fraction=0.125)
# for _ in range(5): _, info = block(torch.randn(1, 40, 32)); print(int(info.route.selected_mask.sum()))
```
## Summary
Key takeaways:
1. **MoE decouples capacity from compute.** Many expert FFNs, a router that keeps
only the top-`k` per token: total parameters ≫ active parameters.
2. **The router is a tiny top-k gate.** A single linear layer scores experts;
softmax + top-`k` selects them; the kept weights are renormalized to sum to 1.
3. **Dispatch is gather → run → scatter.** Each expert runs once, on exactly the
tokens routed to it — never all experts on all tokens.
4. **Routers collapse without help.** A load-balancing auxiliary loss
$E\sum_i f_i P_i$ (minimum 1.0 at uniform routing) keeps every expert fed.
5. **The sparsity factor is `E/k`.** It is free capacity paid for in memory, not
FLOPs — the trade-off behind Mixtral, DeepSeek-MoE, and modern frontier MoEs.
6. **Who picks whom is a design choice.** Token choice fixes coverage ($k$ experts
per token) and pays for balance with an auxiliary loss; expert choice
(`ExpertChoiceMoELayer`) fixes balance ($k$ tokens per expert, no aux loss) and
lets coverage float — even leaving some tokens untouched. Same score matrix,
read along opposite axes.
7. **You can balance token choice without a loss.** Loss-Free Balancing
(`LossFreeMoELayer`, DeepSeek-V3) adds a per-expert bias to the top-$k$
*selection only* — never to the gate value — and nudges it after each batch by
$b_i \mathbin{+}= u\cdot\text{sign}(\bar{c}-c_i)$. Because the bias is a buffer
the gradient never sees, load balance is maintained with **no auxiliary loss and
no gradient interference**, while routing stays token-choice and so still
decodes one token at a time — the property expert choice gave up.
8. **Granularity is a second axis (DeepSeek-MoE).** Beside *who picks whom*, you
can change *what an expert is*: **fine-grained** segmentation cuts each expert
into $m$ smaller ones and activates $m\times$ more — identical compute, but
$\binom{16}{2}=120 \to \binom{64}{8}\approx 4.4\text{B}$ routing combinations —
and **shared** experts, always on, absorb common knowledge so the routed experts
specialize. Both sharpen specialization at fixed cost, and the routed gate uses
the *raw* affinity (not renormalized).
9. **Depth is a third axis (Mixture-of-Depths).** The same top-k, but the choice is
*the block* vs. *the residual*: a scalar router keeps the top $k = C\cdot S$
tokens ($C=12.5\%$), and the rest bypass. The router weight multiplies the block
output so gradient reaches it; the budget is **static** (unlike early-exit); it
interleaves every other block; and a small **predictor** trained with BCE makes
the non-causal top-k decodable one token at a time.
## What's Next
We now have the modern FFN story: the dense gated FFN (SwiGLU, in m06) and the
sparse one (MoE, here) across three axes — **routing** (token choice, expert choice,
auxiliary-loss-free balancing), **granularity** (fine-grained + shared experts,
DeepSeek-MoE), and **depth** (Mixture-of-Depths, routing around blocks). The
frontier keeps composing these — DeepSeek-V3 stacks loss-free balancing *on top of*
fine-grained shared experts, and MoDE combines Mixture-of-Depths *with* experts.
Next come long-context attention and the alignment and reasoning stages that turn a
pretrained model into an assistant.
## Going Deeper
- Shazeer et al., *Outrageously Large Neural Networks: The Sparsely-Gated
Mixture-of-Experts Layer* (2017) — the modern MoE layer and its balancing loss.
<https://arxiv.org/abs/1701.06538>
- Fedus, Zoph & Shazeer, *Switch Transformers* (2021) — top-1 routing at scale;
the load-balancing loss used here. <https://arxiv.org/abs/2101.03961>
- Jiang et al., *Mixtral of Experts* (2024) — a strong open MoE (8 experts,
top-2), with the total-vs-active numbers this lesson quotes.
<https://arxiv.org/abs/2401.04088>
- Zhou et al., *Mixture-of-Experts with Expert Choice Routing* (2022) — the flip
that makes load balance exact by construction, no auxiliary loss.
<https://arxiv.org/abs/2202.09368>
- Dai et al., *DeepSeekMoE: Towards Ultimate Expert Specialization in
Mixture-of-Experts Language Models* (2024) — fine-grained segmentation
($\binom{16}{2}\to\binom{64}{8}$) + shared expert isolation, the section here.
<https://arxiv.org/abs/2401.06066>
- Wang et al., *Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts*
(2024) — the bias-on-selection thermostat this lesson builds
(`b_i += u·sign(c̄ − c_i)`, `u = 0.001`, the `MaxVio` metric).
<https://arxiv.org/abs/2408.15664>
- DeepSeek-AI, *DeepSeek-V3 Technical Report* (2024) — a 671B/37B-active open MoE
that trains with this loss-free scheme (sigmoid affinity, bias for routing only,
1 shared + 256 routed experts, top-8). <https://arxiv.org/abs/2412.19437>
- Raposo et al., *Mixture-of-Depths: Dynamically allocating compute in
transformer-based language models* (2024) — the depth axis: expert-choice top-k
over tokens ($C=12.5\%$, every other block), the scalar router on the output, and
the causal predictor for decoding. <https://arxiv.org/abs/2404.02258>