Mixture-of-Experts: Sparse Computation

Chapter 1: The Sparsely-Gated Layer

Chapter 0 sketched the shape of the fix. Now build the actual component. It has exactly two parts: a set of expert networks and one gating network that decides, per input, which experts get to run.

Formally: n expert networks E₁,…,E_n, each a simple feed-forward network with its own separate parameters (same architecture, different weights — think of them as n copies of the same blueprint, independently trained). A gating network G(x) reads the input x and outputs an n-dimensional vector, one weight per expert. The layer's output is:

Read this literally: run every expert, multiply each one's output by its gate weight, add them up. If that were the whole story there would be no savings — you would still evaluate all n experts every time, just to throw most of the results away after multiplying by near-zero weights. The savings live in one sentence from the paper, easy to skim past: “Wherever G(x)_i = 0, we need not compute E_i(x).”

The layer's structure, before the gate is filled in

Write down the module before worrying about how G decides anything — the structure alone already tells you where the parameter count lives:

python
import torch.nn as nn

class MoELayer(nn.Module):
    def __init__(self, d_model, d_hidden, n_experts):
        super().__init__()
        # n independent feed-forward networks, identical shape, separate weights
        self.experts = nn.ModuleList([
            nn.Sequential(nn.Linear(d_model, d_hidden), nn.ReLU(), nn.Linear(d_hidden, d_model))
            for _ in range(n_experts)
        ])
        self.gate = NoisyTopKGate(d_model, n_experts)   # built in Chapter 2

    def forward(self, x):                 # x: (d_model,) — one token's hidden state
        gates = self.gate(x)                # (n_experts,), mostly exact zeros
        y = 0
        for i, g in enumerate(gates):
            if g == 0: continue          # the actual skip — this is where compute is saved
            y = y + g * self.experts[i](x)
        return y

That loop processes one token at a time, which is the clearest way to see the skip — but it is not how a real system runs it. A training batch holds thousands of tokens at once, and different tokens in the same batch route to different experts. The shape that matters in practice is not “one token, some experts skipped,” it is “a batch of tokens, partitioned across experts, each expert seeing only its own slice”:

python
def forward_batched(self, x):                  # x: (B, d_model) — B tokens at once
    gate_w, gate_idx = self.gate(x)             # both (B, k) — top-k weights and expert ids per token
    y = torch.zeros_like(x)                       # (B, d_model), filled in per expert below
    for e in range(self.n_experts):
        # which rows of the batch picked expert e, and at which of their k slots
        rows, slots = (gate_idx == e).nonzero(as_tuple=True)
        if rows.numel() == 0: continue       # expert e got nobody this batch — skip entirely
        w = gate_w[rows, slots].unsqueeze(-1)     # (n_e, 1) — this expert's weight for its own rows
        y[rows] += w * self.experts[e](x[rows])  # expert e runs ONCE, on its own n_e-row sub-batch
    return y

Two things change once you think in batches. First, the skip from the single-token version (if g == 0: continue) becomes a gather: instead of deciding per token whether to call an expert, you first group all the tokens that chose expert e together into one sub-batch of shape (n_e, d_model), then call that expert exactly once on the whole group. This is strictly better than looping per token — one matmul over n_e rows is far cheaper on a GPU than n_e separate matmuls over one row each. Second, n_e, the number of rows expert e actually receives, is now a random variable that depends on what the whole batch happened to contain — and how large or small n_e tends to be is exactly the subject of Chapter 4's shrinking-batch problem and Chapter 5's load-balancing losses. The single-token version hid that variable; the batched version is where it becomes visible and where it starts to matter.

Notice what is not parameterized here: nothing about which expert specializes in what. Every one of the n_experts modules is initialized the same way, from the same distribution, with the same shape. Specialization — the “plays a critical role” expert you will meet in Chapter 6 — is not designed in; it emerges purely from which examples each expert's weights happen to get gradient updates from, which is itself a consequence of the gate's early, mostly-random choices reinforcing themselves into consistent patterns. Keeping every expert structurally identical is itself a design decision, and a deliberate simplification: nothing stops an expert from having a different hidden size or depth than its siblings, but identical shapes keep the systems side of Chapter 4 (placing experts on devices, batching them uniformly) enormously simpler, at the cost of assuming every “kind” of specialization needs roughly the same amount of capacity to represent. The paper does not explore heterogeneous experts; it is a design space this architecture leaves open.

Where this layer actually sits

The paper's language models are two stacked LSTM layers with one MoE layer sandwiched between them. Crucially, the MoE is called once per position in the text — convolutionally, the same way a convolution slides a fixed operation across every location. This is a detail worth sitting with: the gate does not pick one combination of experts for an entire document. It re-decides at every single token, so a sentence's opening word and its closing word can be routed to completely different experts.

Push that fact to its logical conclusion with a small combinatorial estimate. A single token only ever touches k experts — that is the whole compute-saving point of Chapter 0. But a single sentence of T words, each independently choosing k experts, can touch as many as T×k distinct experts across the whole sentence (fewer if some words happen to choose overlapping experts, but no more). For a twenty-word sentence with k = 4, that ceiling is 20×4 = 80 distinct experts consulted somewhere over the course of one sentence, out of whatever n the model was built with. Zoom out to an entire training batch of thousands of sentences, and it becomes clear why, in aggregate, a well-trained gate ends up exercising every one of even a very large n over the course of training — even though any single token's own forward pass stays exactly as cheap as Chapter 0 promised.

Concept → realization: trace one expert's shape

Abstractions are cheap; a real expert has a real parameter count. In the paper's 1-billion-word language model, the embedding layer, each LSTM layer, and the MoE layer's input/output dimension are all d = 512. Each expert is a feed-forward network with one ReLU hidden layer of size 1024:

That is the unit of capacity this whole architecture trades in: roughly a million parameters, bought and paid for independently for every one of the n experts. Stack 256 of them and the MoE layer alone holds ≈256 million parameters; stack 4,096 and it holds ≈4 billion — while each token still only pays for the k that were actually selected. Chapter 6 pushes this same arithmetic all the way to 131,072 experts and 137 billion parameters.

The same arithmetic, with different inputs, sizes the experts used later in this lesson for machine translation (Chapter 6). There, the paper widens the hidden layer to h = 2,048 while keeping the same d = 512 input/output width:

Doubling the hidden width did not double-then-double-again the parameter count — it exactly doubled it, because both matrices scale linearly in h. Hold onto that fact loosely for now; Chapter 4 explains why the translation experts specifically need this wider, more expensive expert, and it has nothing to do with wanting a better learner and everything to do with a bandwidth constraint that Chapter 0's obstacle table already flagged.

When n gets too large to gate directly: hierarchical MoE

A single flat gate choosing among tens of thousands of experts has to output a score for every expert before it can even find the top few — that scoring step alone stops being free at large n. The paper's fix is a two-level hierarchical MoE: a primary gate picks a sparse combination of groups, and each group is itself a small ordinary MoE with its own secondary gate.

Read it as “first pick which group, then pick which expert inside that group, then multiply the two picks together as the final weight.” With a groups of b experts each, the total expert count is a×b, but the primary gate only ever has to score a things and each secondary gate only ever scores b things — a square-root-shaped reduction in gating cost instead of a linear one. Chapter 6 uses this exact equation with the paper's own real branching factors, reproduced in full here so the pattern is visible across every scale the paper actually trained:

Two things worth noticing in this table before Chapter 6 puts real perplexity numbers next to it. First, these particular branching factors — 32, 32, 64, 128, 256, 256 — are the paper's own choices for its largest-scale runs (Appendix D), and the paper does not spell out a rule for how it picked them; a = 32 is simply where the smallest of these configurations starts. It is worth being precise about this because a different, smaller family of models in the same paper (Appendix C.1, the 8-million-operations-per-timestep 1-billion-word models with 256, 1,024, and 4,096 total experts) uses a different first-level branching factor entirely — a = 16 — and there the paper is explicit about why: “the first level branching factor was 16, corresponding to the number of GPUs in our cluster.” That GPU-count justification is real, but it belongs to the 16-way split of the smaller models, not to the 32-way split of the table above; the two configurations were tuned on different clusters for different experiments, and the paper does not repeat that same reasoning for the larger runs. Second, the ratio a÷b needs to stay close to 1 for the “square-root-shaped reduction” claim above to actually hold — a lopsided split (say a = 4, b = 32,768) would leave one of the two gates doing almost all the scoring work anyway, defeating the point of splitting in the first place.

What hierarchy costs, when you don't strictly need it yet

Section B introduced hierarchical gating as the fix for when n gets too large to score directly. It is worth asking the opposite question too: at an expert count small enough that a flat gate is still perfectly affordable, does switching to a hierarchical gate anyway cost anything? The paper answers this directly, because it trained both versions at the same total expert count. A flat MoE-256 model (one gate scoring all 256 experts) reaches 35.7 test perplexity. A hierarchical MoE-256-h model, built as 16 groups of 16 experts using the exact y_H machinery from this chapter, reaches 36.0 — very slightly worse, at the same total expert count and the same overall compute budget.

That small gap is the honest price of the coarser, two-stage decision: a hierarchical gate first commits to a group, then can only ever choose among the experts inside that one group, whereas a flat gate is free to compare every expert against every other expert directly, with no group boundary to accidentally trap the best choice on the wrong side of. Read together with Section B's own square-root-shaped scoring-cost argument, the design tradeoff becomes precise rather than vague: hierarchy is not a free win, it is a small, measured quality cost paid in exchange for a large, necessary reduction in gating cost — worth paying once n is too large to gate any other way, not worth paying before then. Chapter 6 revisits this exact 256-expert comparison alongside the paper's full scaling results.

Chapter 2: Noisy Top-K Gating

Chapter 1 promised that G(x) is exactly zero for all but k experts. Now build a G that actually keeps that promise, starting from the simplest gate that does not.

The starting point: ordinary softmax gating

The simplest possible gate multiplies the input by a learned weight matrix and takes a softmax:

This is dense. Softmax outputs are strictly positive everywhere — no entry is ever exactly zero, only smaller or larger. Every expert would need to run on every input, forever, no matter how small its weight. This gate is a fine mixture but a useless conditional-computation layer. Two things need to be bolted onto it: a mechanism that produces true, hard zeros, and (as Chapter 5 will need) a source of randomness for load balancing.

Building Noisy Top-K Gating, piece by piece

The paper's actual gate has three moving parts, and each earns its place:

Part 1 — the clean logit, (x·W_g)_i: exactly the ordinary softmax gate's score for expert i, a learned linear read of the input. This is the signal — “how good a fit is expert i for this input, as far as the model currently believes.”

Part 2 — the noise magnitude, Softplus((x·W_noise)_i): a second, independently learned linear read of the input, passed through Softplus (a smooth, always-positive relative of ReLU: Softplus(z) = ln(1+e^z)). This is not a fixed hyperparameter — the network learns, per expert and per input, how much randomness to inject into that expert's logit. Softplus rather than ReLU matters here for a specific reason: ReLU is exactly zero for any negative input, which would let the model learn to switch noise off completely for some expert and defeat the load-balancing purpose noise exists for in the first place (Chapter 5). Softplus is always strictly positive, so noise never fully vanishes.

Part 3 — the actual jitter, StandardNormal(): a fresh Gaussian sample per expert per forward pass, scaled by that learned magnitude. This is the randomness that Chapter 5's load-balancing loss will lean on.

Part 4 — KeepTopK: sort H(x), keep the k largest entries as-is, and set every other entry to −∞. Not a small number — literal negative infinity. That is the detail that makes sparsity exact rather than approximate: Softmax(−∞) evaluates to exactly 0, every time, with no approximation error to worry about. If you have seen masked attention, this is the identical trick — mask with −∞ before the softmax, never after.

“No approximation error” is worth confirming rather than accepting on faith. Softmax turns a masked logit m into e^m in its numerator, so the question is really: how fast does e^m shrink as m gets more negative, and does it ever actually reach zero before m reaches negative infinity?

Every finite mask value, no matter how negative, produces some strictly positive e^m — it just gets small fast enough that a computer's floating-point representation rounds it to zero somewhere around m ≈ −700. That rounding is a hardware accident, not a mathematical guarantee; a higher-precision computer would keep a nonzero (if minuscule) value for longer. Only the literal −∞ in the KeepTopK step guarantees the zero by construction, independent of any floating-point implementation detail. This is the precise sense in which Chapter 1's promise — “not small, not near-zero, exactly 0.0” — is not a rhetorical flourish but a specific, checkable claim about how the mask value is chosen.

Checking the “Softplus never fully vanishes” claim by hand

Part 2's argument for Softplus over ReLU is worth confirming with actual arithmetic rather than taking on faith. Softplus is defined as Softplus(z) = ln(1 + e^z). Plug in a few values of the raw noise-logit z and compare against ReLU(z) = max(0, z) side by side:

Read the bottom row carefully: even when the noise-magnitude weights push z as low as −10, ReLU would give exactly zero noise, while Softplus still gives a tiny but strictly positive number. That gap between “exactly zero” and “a number so small it is nearly zero” is precisely the distinction Part 2 warned about. With ReLU, gradient descent could find a value of z low enough that noise becomes exactly zero and stays there — a dead, permanently silent noise channel. With Softplus, no matter how negative z gets pushed, there is always some noise left, so the load-balancing pressure in Chapter 5 always has a channel to act through. The gap shrinks toward zero, but by construction it can never reach it.

A hand-worked example: routing one input through four experts

Small enough to check every digit by hand. Four experts, k = 2. For clarity, assume the noise magnitude has already been computed by the Softplus term and happens to equal 1.0 for every expert on this particular input — that keeps the arithmetic in one variable instead of two. The clean logits (already computed as x·W_g) and one drawn noise sample:

Notice already how much the noise moved things: expert 2's clean logit (1.0) beats expert 3's (0.5), but after noise, expert 3 (0.6) beats expert 2 (0.8)… wait — check again: 0.8 > 0.6, so expert 2 still wins that particular pair. But expert 1's clean lead over expert 2 (2.0 vs 1.0) widened to 2.3 vs 0.8. The two largest entries of H(x) are expert 1 (2.3) and expert 2 (0.8) — keep those two, mask experts 3 and 4 to −∞.

Now softmax over only the two kept values:

Only experts 1 and 2 run. Experts 3 and 4 contribute the number zero to y, cost the number zero FLOPs, and the model never touches their weights on this token. That four-entry vector, mostly zero, is the entire mechanism Chapter 1 promised.

Rerun the exact same H(x) = [2.3, 0.8, 0.6, −2.0], but with k = 1 instead of 2. KeepTopK now keeps only expert 1 (2.3) and masks experts 2, 3, and 4 all to −∞. Softmax over a single surviving value is, by the definition of softmax, always exactly 1 — e^2.3÷e^2.3 = 1, regardless of whether that 2.3 had been a 0.001 or a 1,000,000. So G(x) = [1, 0, 0, 0]: expert 1 runs at full weight, and the softmax step contributed nothing but a constant. Hold onto this exact number — a flat, contentless 1.0 — because Chapter 3 builds its entire argument for why k must exceed 1 directly from what a constant means for a gradient.

The same input, a different noise draw

“A fresh Gaussian sample per forward pass” from Part 3 is easy to read past. Rerun the identical input — same clean logits, same W_g, same everything except the noise — with a different draw and watch the winner change:

This time the two largest entries of H(x) are expert 3 (1.5) and expert 1 (1.4) — expert 3, which had the second-lowest clean logit of the four, wins a spot in the top-2 this round, bumping expert 2 out entirely. Softmax over the new survivors: e^1.5≈4.482, e^1.4≈4.055, sum≈8.537, giving G(x)₃≈0.525 and G(x)₁≈0.475 — almost an even split, and expert 3 is the one running, not expert 2.

Nothing about W_g changed between these two passes — the same input, evaluated twice, routed to a genuinely different pair of experts purely because of which random numbers landed where. This is precisely the mechanism Chapter 5 leans on: an expert that is currently behind on the clean logit alone (expert 3, at 0.5, trailing expert 2's 1.0) is not permanently locked out. Often enough noise favors it, it wins the slot anyway, receives a real gradient update from that token, and its own clean logit has a chance to improve. Turn the noise off and this second table simply could not have happened — expert 3 would trail expert 2 on every single pass, forever, with no mechanism to ever change that ranking.

It is worth asking, briefly, why the surviving weights get softmax-normalized at all rather than something simpler — say, dividing each kept logit by the sum of all kept logits, an ordinary linear renormalization. Two problems rule that out immediately. First, logits can be negative (expert 4's clean logit above is −0.5), and dividing by a sum that could itself be negative or near zero produces nonsense weights, possibly even negative ones, which have no sane interpretation as “how much of this expert's output to use.” Softmax sidesteps this entirely because e^anything is always strictly positive, no matter how negative the input. Second, softmax is not just a way to force positivity, it specifically amplifies the gap between large and small logits — recall 2.3 vs 0.8 becoming 0.818 vs 0.182, a wider split than the raw logits' own 2.3-to-0.8 ratio would suggest. A gate that is more confident (bigger logit gap) ends up more decisively weighted toward its favorite, which is exactly the behavior Chapter 3 needs for its gradient argument to have any teeth.

What noise buys you beyond load balancing

Rerun the identical example with the noise scale turned down to zero — H(x) then equals the clean logits exactly: [2.0, 1.0, 0.5, −0.5]. Top-2 is still experts 1 and 2, and the softmax weights change only slightly (e^2.0≈7.389, e^1.0≈2.718, giving G(x)₁≈0.731, G(x)₂≈0.269). Nothing seems to break.

The danger is not in any single forward pass; it is over the course of training. Early on, W_g is close to its random initialization, so the clean logits for any given input are close to arbitrary — but whichever ordering they happen to start in is the ordering KeepTopK will keep selecting, every single time that input (or one like it) appears, with zero noise to ever let a different expert get a turn. An expert that starts slightly ahead by chance gets picked, gets the only gradient updates, and stays ahead — the exact self-reinforcing collapse Chapter 5 names and measures. Noise gives every expert a nonzero chance of being sampled into the top-k even while it is behind on the clean logit alone, which is what lets a temporarily-unlucky expert's turn eventually come around and prove itself. Chapter 5's load-balancing losses are the trained, deliberate version of this pressure; the noise term is what makes that pressure possible to apply at all, by keeping the selection genuinely stochastic instead of deterministically locked in from initialization.

An equivalent way to write the same gate, in the paper's own appendix

The −∞-then-softmax version above is the clearest way to understand the gate, but it is not the only way to write it. The paper's Appendix F gives an algebraically equivalent reformulation, built around an explicit binary mask rather than an infinite penalty. Starting from the plain softmax gate G_σ(x) = Softmax(x·W_g) from the top of this chapter, multiply it component-wise by a mask M and renormalize:

where the ordinary top-k mask is 1 for the k largest entries and 0 everywhere else. Multiply-then-renormalize and mask-to-−∞-then-softmax land on the same numbers — both zero out the same n−k entries and redistribute the remaining probability mass over the survivors in the same proportions. The reason to know this second form at all is that it generalizes in a direction the −∞ version cannot: because M is just some function that decides which entries are 1, you can swap in a different mask without touching anything else about the gate. Chapter 4 uses exactly this hook — the paper's translation experiments, for infrastructure reasons specific to that setup, swap the ordinary per-example top-k mask for a mask that guarantees every expert receives the exact same batch size instead.

Chapter 3: Training the Gate

A gate that hard-selects k out of n experts sounds like it should be untrainable — “which experts get picked” is a discrete decision, and discrete decisions famously have zero gradient almost everywhere. Yet the paper trains the whole system, gate included, with plain back-propagation. Here is why that works.

What is actually discrete, and what is not

Separate two questions that look like one. Question A: which k experts get nonzero weight? That is discrete — a hard cutoff produced by KeepTopK. Question B: how much weight does each of those k chosen experts get, relative to each other? That is continuous — it is a softmax over real numbers. Question A has no useful gradient. Question B has a perfectly ordinary one, because softmax is a smooth function of its inputs.

The paper states this precisely: “If we choose k>1, the gate values for the top k experts have nonzero derivatives with respect to the weights of the gating network.” Read the condition carefully — k>1 is doing real work in that sentence.

The path gradients actually take

Trace it forward, then backward. Forward: x feeds both the experts and the gate; the gate's softmax weights scale each kept expert's output; those get summed into y; y feeds the rest of the network and eventually a loss. Backward: the loss gradient reaches y, splits into a gradient on each G(x)_i and a gradient on each E_i(x) (for the kept experts only — masked experts received no forward computation, so they receive no backward computation either, and no update). The gradient on G(x)_i flows on through the softmax and the KeepTopK's surviving entries, into W_g and W_noise, and further still into whatever produced x.

That last clause matters: the gate's gradient does not stop at the gating weights. It continues into x itself — meaning the representation that feeds the gate (an LSTM hidden state, in the original paper) is shaped in part by “does this help the gate route well,” not only by the primary next-word-prediction loss. Routing pressure and language-modeling pressure co-train the same representation.

Trace where that x-directed gradient actually goes. x is read by three separate places in this layer: the kept experts (through E_i(x)), the clean-logit read (through x·W_g), and the noise-magnitude read (through x·W_noise). Ordinary backpropagation sums contributions from every path that touches a variable, so x's total gradient is the sum of all three: however the kept experts wanted x to change to make their own outputs more useful, plus however the gate wanted x to change to make its routing decision more confident, plus however the noise pathway wanted x to change. Whatever produced x — an LSTM layer, in the original paper — receives that combined signal, with no way to separate “compute this well” pressure from “route this well” pressure after the fact. That entanglement is not a flaw; it is the mechanism by which routing pressure and language-modeling pressure end up co-training the same representation, exactly as the paragraph above claimed, now traced down to which specific paths the gradient actually sums over.

Concept → realization: the chain rule, symbol by symbol

“Gradients flow backward through the softmax into W_g” is easy to say and easy to leave vague. Write out the actual chain, one derivative at a time, for a single kept expert i. Start from the loss L at the very end of the network and work backward to W_g, the thing actually being updated:

Each factor is something you have already derived or can read straight off an earlier equation in this lesson. ∂L/∂y is whatever the rest of the network hands back — it is not this layer's business, only the layer's input from upstream. ∂y/∂G(x)_i is, from Chapter 1's y = ∑G(x)_iE_i(x), simply E_i(x) — the kept expert's own output. ∂G(x)_i/∂H(x)_j is the ordinary softmax Jacobian, the thing computed by hand two paragraphs below. And ∂H(x)_j/∂W_g is, from Chapter 2's H(x) = (x·W_g) + noise, simply x itself, since the noise term does not depend on W_g at all. Multiply the four factors together and you have the update W_g actually receives — nothing here is a special case invented for Mixture-of-Experts; it is the same chain rule that trains every other linear layer in the network, applied to a gate instead of a hidden layer.

Put a number through it. Say the loss gradient arriving at y happens to be ∂L/∂y = 0.1 (a single scalar, for a one-dimensional toy y), and reuse Chapter 2's kept expert 1 with gate weight G(x)₁ = 0.818 and, say, E₁(x) = 2.0 for this token. ∂y/∂G(x)₁ = E₁(x) = 2.0, so the gradient reaching G(x)₁ itself is 0.1×2.0 = 0.2. That 0.2 then meets the softmax Jacobian derived two sections below (p₁(1−p₁) ≈ 0.149 for the logit-gap direction), giving a gradient on H(x)₁ of roughly 0.2×0.149 ≈ 0.030. Finally, since H(x)₁ is a linear read of x through W_g, that 0.030 becomes the actual update applied to the relevant column of W_g, scaled by x itself. Every one of these numbers is small and easy to lose track of symbolically; multiplied out concretely like this, the whole chain is nothing more exotic than three ordinary multiplications in sequence.

python
# forward pass, shapes annotated (batch dim omitted for clarity)
clean_logits  = x @ Wg                                  # (n,)
noise_scale   = softplus(x @ Wnoise)                     # (n,)  always > 0
noisy_logits  = clean_logits + torch.randn(n) * noise_scale   # (n,)

topk_vals, topk_idx = noisy_logits.topk(k)               # k values, k indices
masked = torch.full((n,), -float('inf'))
masked[topk_idx] = topk_vals
gates = torch.softmax(masked, dim=-1)               # (n,), exactly n-k zeros

y = 0
for i in topk_idx:                                       # only k experts actually run
    y = y + gates[i] * expert[i](x)

Every line here is doing something you can now name: the two independent linear reads, the Softplus that protects against a silent noise shutoff, the −∞ mask that guarantees exact zeros, the softmax that gives the surviving experts a differentiable relative weighting, and the loop that only touches the k selected experts — the physical realization of Chapter 1's promise.

Both kept experts learn, and by how much depends on the gap between them

Reuse Chapter 2's worked numbers: G(x) = [0.818, 0.182, 0, 0]. Softmax has a standard derivative — for a two-element softmax with weights p₁ and p₂ = 1−p₁, the gradient of p₁ with respect to the gap between the two logits is p₁(1−p₁). Plugging in 0.818: 0.818×0.182 ≈ 0.149. Compare that to what the gradient would be if the gate were nearly certain, say p₁ = 0.99: 0.99×0.01 = 0.0099, fifteen times smaller. The lesson generalizes: a gate that is already confident receives a weaker training signal on its relative weighting than a gate that is still genuinely deciding between two nearly-tied experts. Both expert 1 and expert 2 receive real gradient on their own weights either way — the sensitivity that shrinks near certainty belongs specifically to the gating weights that decide the split between them, not to the experts' own parameters.

That p₁(1−p₁) term was only the diagonal of the full softmax Jacobian — the part answering “how does p₁ change if I nudge expert 1's own logit?” The full two-expert Jacobian has an off-diagonal term too, answering the complementary question “how does p₁ change if I nudge expert 2's logit instead?” For a two-way softmax that term is −p₁p₂:

Same magnitude, opposite sign. Read that as: pushing expert 1's logit up by some amount and pushing expert 2's logit down by that same amount have exactly the same effect on p₁, because with only two experts kept, more trust in one is definitionally less trust in the other — p₁ + p₂ is pinned at 1 no matter what W_g does. This is why the earlier framing called it a “split”: the gradient genuinely measures relative confidence between the two survivors, not some independent absolute score for each.

Read the last column as one consistent story rather than seven disconnected numbers: a single incoming gradient of 0.1 splits unevenly across the kept experts in proportion to their gate weights (0.0818 to expert 1's own parameters, a smaller 0.0182 to expert 2's), reaches the gating weights themselves scaled down further by the softmax Jacobian (≈0.030, an order of magnitude smaller again), and never reaches the masked experts or their gating columns at all. Every number in this row traces back to the same single scalar arriving at y, just redistributed by how much each parameter actually influenced this particular token's output.

One more variant worth running by hand: keep the same H(x) = [2.3, 0.8, 0.6, −2.0] but set k = 3 instead of 2. Now experts 1, 2, and 3 all survive (2.3, 0.8, 0.6), only expert 4 is masked:

Compare this to k = 2's G(x) = [0.818, 0.182, 0, 0]. Adding a third survivor did not just insert a new 0.130 out of nowhere — it pulled weight away from experts 1 and 2 as well (0.818 → 0.711, 0.182 → 0.159), because softmax always renormalizes the whole surviving set to sum to 1. That is the general shape of what changing k does to every kept expert's gradient, not just whether a given expert is kept at all: a larger k spreads the same total probability mass thinner across more experts, softening the gradient each individual kept expert receives, while a smaller k concentrates it more sharply on fewer winners. Chapter 4 and Chapter 6 both use k = 2 or k = 4 in their real configurations — now you can see, numerically, part of what that choice trades off.

And the two masked experts, 3 and 4, are not merely down-weighted — they are structurally absent from the backward pass entirely. Look again at the loop in the code above: for i in topk_idx means experts outside the top-k are never called, so there is no computation graph connecting them to the loss at all on this particular token. Zero forward compute, zero backward compute, zero optimizer update, on this example. They may well be updated heavily on the next token, if that token's gate happens to prefer them — which is exactly the per-token re-routing Chapter 1 described.

It is worth being precise about why REINFORCE-style training is noisier, in general terms, since the comparison otherwise stays vague. A REINFORCE-family estimator for a discrete choice looks, in its general textbook shape, something like reward × ∂(log probability of the choice actually made) ÷ ∂(parameters) — it only ever tells the parameters “make the choice you happened to sample more or less likely,” scaled by how good that one sampled choice turned out to be. Every other choice you didn't sample this step contributes nothing to the gradient at all, so the estimate's quality depends entirely on how lucky or representative that one sample was — hence the high variance, and hence why it needs a carefully chosen baseline (a running average reward, typically) just to be usable in practice. Chapter 3's softmax-over-survivors gradient, by contrast, gets to see the exact relative weighting of every kept expert on every single step, no sampling and no baseline required — it is not an estimate of a gradient, it is the gradient.

How to tell, from the outside, whether the gate is actually learning

Every implementation detail in this chapter has an observable symptom if it is wrong, and it is worth knowing what to look for before Chapter 5 introduces two more moving parts. If W_g receives gradient but never changes its ranking of experts for a given input across training, suspect a k = 1 configuration — recall this chapter's central result: softmax over one surviving value is a constant, so W_g's gradient with respect to relative trust is genuinely zero, not just small. If routing looks like it is learning early in training but then freezes into the same k experts for nearly every input, suspect the noise term has decayed to nearly nothing (Chapter 2's Softplus floor exists specifically to prevent this) or that Chapter 5's balancing losses are weighted too weakly to counteract the self-reinforcing collapse Chapter 5 names directly. And if gradients on an expert's own weights look healthy but the gate never seems to prefer it over its neighbors, check that the expert is actually reachable within the top-k at all — an expert whose clean logit sits so far below its peers that noise can never lift it into contention will show textbook “dead unit” behavior indistinguishable, from the loss curve alone, from a bug.

One further symptom, and its likely cause: if training loss looks healthy overall but a small handful of experts' individual parameter norms grow far faster than the rest, suspect the self-reinforcing dynamic Chapter 5 names directly — those experts are winning the top-k cut disproportionately often, receiving disproportionately many gradient updates, and growing accordingly. This particular symptom will not show up in this chapter's own machinery at all; it is Chapter 5's balancing losses, not anything about the gate's gradient path, that exist specifically to prevent it. Knowing which chapter's fix addresses which symptom is itself a useful debugging skill: a k = 1 misconfiguration, a decayed noise term, and a missing balancing loss all eventually show up as “the gate seems stuck,” but they are three different bugs with three different fixes, living in three different chapters of this lesson.

Chapter 4: Batches & Bandwidth

The math works. Now make it run fast — this chapter is Section 3 of the paper, the systems engineering that turns a correct idea into a trainable one on real GPU clusters.

Every fix in this chapter answers one of the two hardware obstacles Chapter 0's table named up front: GPUs favoring dense arithmetic over branching, and the shrinking-batch problem that conditional computation creates by construction. Nothing here changes what the gate computes — Chapters 2 and 3's math is untouched. What changes is how that math gets scheduled across a real cluster of real, imperfect devices.

The shrinking-batch problem

Large batches matter on GPUs because they amortize the fixed overhead of loading parameters and applying updates — a batch of 1 wastes almost all its time on overhead, a batch of 1,024 barely notices it. Here is the problem conditional computation creates: if a global batch of b examples gets routed across n experts with k active per example, each individual expert only receives, on average,

Worked example, with round numbers you can check by hand: b = 2,048 examples in the batch, n = 512 experts, k = 4 active per example.

Sixteen. Out of a global batch of over two thousand, a single expert sees, on average, sixteen examples — a batch size small enough to waste most of a GPU's throughput on overhead rather than useful arithmetic. This is the shrinking-batch problem, and it gets worse, not better, as you add more experts to chase more capacity.

Fix 1: mix data parallelism with model parallelism

The paper's solution reframes the cluster itself. In ordinary data parallelism, d devices each hold a full copy of the model and process independent batches asynchronously. Here, instead, the standard layers and the gate are replicated data-parallel across the d devices as usual — but each expert is kept as a single shared copy, model-parallel-style, living on one device. Crucially, the d devices process their local batches synchronously, so that the relevant examples from all d local batches can be pooled together before being routed to each expert.

Continue the same worked example, now with d = 16 devices, each still contributing a local batch of 2,048:

Sixteen became two hundred fifty-six — a factor-of-16 improvement, exactly matching d, exactly as the formula predicts. Add devices to chase more experts, and the per-expert batch size holds steady rather than shrinking further. This is the mechanism that makes the 131,072-expert model in Chapter 6 trainable at all.

Fix 2: exploit convolutionality

Because the MoE is applied identically at every position in a sequence (Chapter 1's convolutional placement), there is a second, free multiplier available: wait for an entire unrolled sequence to finish the previous layer, then call the MoE once on all those timesteps stacked together as one combined batch, instead of calling it once per timestep. A sequence of length T turns one call into an effective ×T on the batch size handed to the MoE layer.

This trick has a real limit worth naming honestly: it depends on the MoE's input at timestep t not depending on the MoE's own output at timestep t−1. If you ever wanted a recurrent MoE — the gate itself feeding back through time — this convolutional stacking breaks, because you cannot batch timesteps you have not computed yet. The same shrinking-batch problem resurfaces decades later in a different guise: a modern MoE language model generating one token at a time at inference has exactly the tiny-per-expert-batch problem this section solves for training, which is why production serving stacks for MoE models lean hard on large continuous batching across many concurrent users.

The paper does not stop at naming this limitation — it points to a specific way around it. Replacing an LSTM's own weight matrices with a MoE would let a network route differently at every timestep of a recurrence, not just at every position of a convolution, but batching such a model is hard for exactly the reason above: timestep t + 1 cannot be batched with timestep t because it has not been computed yet. The paper's suggested fix borrows a technique from Gruslys et al. 2016 for training recurrent networks with a much smaller memory footprint, by recomputing forward activations during the backward pass instead of storing every one of them for every timestep of the unrolled sequence. Trading stored activations for recomputed ones frees up exactly the memory that a larger batch would otherwise need — the identical trade this chapter reaches for again, on a different axis, two sections from now.

Network bandwidth: the ratio that decides efficiency

Because each expert lives on one device (model-parallel), most communication in this system is shipping inputs and outputs across the network to wherever an expert sits. GPU compute can outrun network bandwidth by roughly a thousand to one, so an expert is only worth calling remotely if its compute-per-byte-moved clears that bar. For a one-hidden-layer expert with input/output size s and hidden size h, the two matrix multiplies cost work proportional to s·h, while the bytes that must cross the network are proportional to s (input in, output out). The ratio:

The hidden-layer size h is, almost embarrassingly directly, the knob that buys bandwidth-efficiency. This is exactly why the paper's experts use hidden layers with thousands of ReLU units — not because bigger experts are inherently better learners, but because a bigger hidden layer makes each expert do proportionally more computation for every byte it costs to ship across the cluster.

Put the ratio through the actual expert size from Chapter 1: input/output width s = 512, hidden width h = 1,024. Compute for the two matmuls is proportional to s·h twice = 2×512×1,024 = 1,048,576. Bytes moved across the network are proportional to shipping the input in and the output back out: 2×s = 1,024 numbers. The ratio:

Compare that 1,024 against the paper's own stated hardware ratio — GPU compute can outrun bandwidth by “thousands to one.” A 512-in/512-out, 1,024-hidden expert clears a 1,000:1 bar only narrowly. This is precisely why the machine-translation experts in Chapter 6 use a larger hidden layer (h = 2,048, doubling the expert to ≈2 million parameters) than the language-modeling experts do: translation's encoder-decoder architecture places more network hops between MoE layers, and a fatter hidden layer is the direct, mechanical answer to a worse compute-to-bandwidth ratio on that particular architecture.

The same k·b·d÷n batch-pooling formula from earlier in this chapter shows up again, with different real numbers, in the paper's multilingual translation model: n = 512 experts, k = 2 active, each with a widened hidden layer of size 8,192 specifically to keep the bandwidth ratio healthy at that larger expert count — which doubles the model's per-token computational budget from 85 million to 102 million operations per timestep. Bigger n needs the same batch-pooling engineering to keep each expert fed, and it often needs a bigger h too, to keep each expert worth shipping across the network in the first place.

A stricter alternative: forcing exactly equal batches

Chapter 2 showed an equivalent way to write the gate using an explicit mask M, multiplied component-wise into the plain softmax gate and renormalized. The paper's Appendix F explains why that reformulation exists at all: at the time some of the machine-translation experiments were run, a peculiarity of the paper's own infrastructure made training run faster if every expert received exactly the same number of examples per batch, not merely close to the same number. The ordinary top-k mask from Chapter 2 does not guarantee that — it guarantees each example goes to k experts, but says nothing about how many examples land on any given expert.

The fix flips which axis gets sorted. Instead of, per example, keeping the top k expert-scores (sorting along the expert axis), the batchwise mask sorts along the batch axis, separately for each expert: for expert i, look at its score across every example in the batch, and keep only the top m of them, where

with |X| the batch size. Every expert ends up with exactly m examples, no more, no less — by construction, not on average. Compare this to the very first formula in this chapter, k·b÷n: m is that same expected-value formula, just enforced as a hard constraint instead of an expectation. This is the harder cousin of the shrinking-batch problem's soft, statistical fix: instead of pooling batches across devices and hoping the random routing spreads out roughly evenly, force it to spread out exactly evenly, every single batch. The paper is candid that this comes at a cost worth naming: a batch-dependent mask like this one needs care during training, since (as also observed independently by Ioffe & Szegedy 2015, of batch-normalization fame) a function whose output for one example depends on which other examples happen to be in the same batch behaves differently at training time than at inference time, when a “batch” of one example has no other examples to compare against.

Worked example: pooling a real translation-scale batch

Run the batch-pooling formula from earlier in this chapter through the translation model's own reported numbers instead of the toy ones. The paper states its translation training setup directly: “Training was done synchronously on a cluster of up to 64 GPUs… Each training batch consisted of a set of sentence pairs containing roughly 16,000 words per GPU.” Take the single-language-pair encoder MoE (n = 512 experts, k = 2 active) and plug in b ≈ 16,000 words per GPU, d = 64 GPUs:

A single GPU's local batch alone would have handed each expert only 2×16,000÷512 = 62.5 words — the same shrinking-batch story as the language-modeling example earlier in this chapter, just with different round numbers. Pooling across all 64 GPUs working in lockstep multiplies that by 64, landing at 4,000 words per expert per batch: large enough to make full use of each device's arithmetic throughput, on a model that would otherwise have been starved by exactly the problem this chapter opened with. Notice, too, that the units here are words, not whole sentences — the MoE is applied per word-piece position (Chapter 1's convolutional placement), so “examples” in this formula means individual token positions, exactly as it did for the language-modeling case.

Two more memory tricks, for when the parameter count itself is the bottleneck

At the largest scales in Chapter 6 (thousands of experts on a single device), raw memory — not bandwidth, not batch size — becomes the binding constraint, and the paper reaches for two further optimizations to fit up to a billion parameters per GPU.

Recompute instead of store. A normal backward pass needs each layer's forward activations kept around in memory, to compute local gradients. The paper skips storing the experts' hidden-layer activations and instead recomputes them during the backward pass — trading a second forward pass through each expert for the memory that pass's activations would otherwise have occupied. This only pays off because the experts are cheap per-call (Chapter 1's ≈1M-parameter budget) and memory, not compute, is what is scarce at this scale.

Notice this is the identical trade the paper reached for two sections earlier, for a different reason. The recurrent-MoE discussion above cited Gruslys et al. 2016 for exactly this technique — recompute forward activations during the backward pass instead of storing them all — there in service of unlocking a bigger batch size, here in service of fitting more parameters per device. Same trade, memory for compute, applied to two different bottlenecks in two different parts of the same paper. Once you have the trade in hand, it is worth actively looking for — anywhere activations are being kept around only to save a second forward pass, that memory is a candidate for this exact swap.

Shrink the optimizer, not just the model. Adam keeps a running first and second moment estimate for every parameter, which roughly triples the memory footprint of the parameters themselves. The paper cuts this down two ways: dropping the first-moment estimator entirely (setting its decay β₁ = 0), and replacing the full second-moment matrix for each expert's weight matrix with a factored approximation — storing only a row-wise average vector and a column-wise average vector, and reconstructing the full matrix on the fly as their outer product divided by their mean. Two vectors standing in for an entire matrix of per-parameter statistics is a substantial memory win precisely where it is needed most: multiplied across thousands of experts, per-parameter optimizer state is exactly the kind of cost that scales with n, not with k, and so gets no help at all from the sparsity Chapter 1 bought everywhere else.

Put a real expert's shape through the factored-vs-full comparison. Chapter 1's language-modeling expert has two weight matrices, 512×1024 and 1024×512, together holding ≈1,048,576 parameters (Chapter 1's worked calculation). A full, unfactored second-moment estimator would need one number per parameter — the same ≈1,048,576 count, per expert, on top of the parameters themselves. The factored version instead stores one row-wise average and one column-wise average per matrix: for the first matrix, that is 512 + 1024 = 1,536 numbers; for the second matrix, another 1024 + 512 = 1,536; total 3,072 numbers standing in for what would otherwise be roughly a million.

Multiply that saving across a 65,536-expert model (Chapter 6's headline configuration) and the factored approximation is the difference between an optimizer state that fits on the same hardware as the parameters it tracks, and one that does not fit at all. This is not a free lunch — reconstructing the full matrix on the fly as an outer product of two averages assumes each parameter's second moment is well-approximated by the product of its row's and column's typical scale, which is an approximation, not an identity. But at the scale Chapter 6 operates at, a 341× memory reduction bought for the price of that approximation is a trade the paper is clearly willing to make.

Hold that separation in mind heading into Chapter 5, because it cuts the other way there: the next chapter's fix is an algorithm change, two new loss terms added directly to training, not a systems trick. Chapters 2–3 built correct math; this chapter made that math affordable to run; Chapter 5 makes it affordable to run well, without collapsing onto a handful of favorite experts. Three different kinds of problem, three different kinds of fix, in three consecutive chapters.

Chapter 5: Load Balancing

Leave the gate to train on its own and something predictable and bad happens: “the gating network tends to converge to a state where it always produces large weights for the same few experts.” This is not a rare failure mode — it is self-reinforcing. An expert that gets picked more often receives more gradient updates, gets better faster at whatever it was already doing, and so gets picked even more. Left unchecked, a thousand-expert model can collapse into a two-expert model wearing a thousand-expert costume — all that Chapter 0 capacity, sitting on the shelf, untouched.

Unlike every fix in Chapter 4, this one is not a systems trick — it is a change to the training objective itself, sitting directly alongside the ordinary next-token loss the rest of the network already optimizes, and it is the last piece this lesson needs before Chapter 6 can put real numbers on the table. The fix is two auxiliary losses, added to the ordinary training loss, that push in the opposite direction. They target two different notions of “balance,” and the paper is explicit that one alone is not enough — a real, concrete case: “one expert may receive a few examples with large weights, and another may receive many examples with small weights.” Same total weight, wildly different example count, and on real distributed hardware it is the example count that determines memory pressure and step time per device.

Loss 1 — Importance: balance the summed weight

Define the importance of expert i over a batch X as the batch-wise sum of its gate values:

CV is the coefficient of variation — standard deviation divided by mean — of the importance vector across all n experts. It is scale-invariant (it does not care whether every expert's importance is large or small, only whether they are spread out relative to each other), it is zero exactly when every expert has identical importance, and squaring it produces a smooth, always-nonnegative penalty that grows quickly as imbalance grows. Minimizing this loss pushes every expert toward equal total gate-weight across a batch.

CV is easy to name and easy to leave unexamined. Compute it by hand on two tiny 3-expert toy importance vectors (illustrative numbers, chosen only to make the mechanics concrete) to see exactly what “spread relative to the mean” means as arithmetic. First, a balanced case, Importance = [10, 10, 10]:

Every expert identical, CV lands exactly at zero — the penalty vanishes, exactly as the definition promised. Now an unbalanced case, Importance = [25, 10, 1]:

Compare that 0.68 against Chapter 5's own real numbers a few paragraphs below: the paper's completely unbalanced 256-expert run measured CV(Importance) ≈ 3.04, an even more extreme spread across many more experts than this 3-expert toy. Both numbers are answering the identical question — how spread out is this vector, relative to its own average — just at different scales. Once CV² stops being a name and becomes an arithmetic recipe you can run yourself, the rest of this chapter's tables stop being numbers to trust and start being numbers you could, in principle, reproduce.

Making the paper's warning concrete with two tiny toy batches

The paper's own sentence — “one expert may receive a few examples with large weights, and another may receive many examples with small weights” — is easiest to trust once you have built the numbers yourself. Construct two illustrative two-expert batches (small, invented numbers, purely to make the distinction concrete, not paper-reported data) that have identical importance but very different load.

Both batches give expert A and expert B the exact same Importance — 0.90 apiece — so L_importance alone sees nothing wrong with either one. But Batch 2's expert B did ten times more work than Batch 1's: it was invoked for nine separate forward passes instead of one, occupying nine slots of memory and compute instead of one, on real distributed hardware. Importance is blind to this distinction by design — it only ever sums weight, never counts calls — which is exactly the gap Load(X)_i and its P(x,i) machinery below exist to close.

Loss 2 — Load: balance the example count, via a derivation you can check by hand

Importance can be balanced while example counts stay lopsided, as the quote above shows. What is really wanted is a loss on how many examples each expert receives — but that count is a discrete, non-differentiable quantity you cannot backpropagate through directly. The paper's fix: replace the count with a probability that is smooth in the gating weights.

Define P(x,i) as the probability that expert i's gate value would still be nonzero if you redrew only expert i's noise sample, holding every other expert's already-sampled noise fixed. Expert i survives the top-k cut exactly when its noisy logit beats the k-th largest logit among all the other experts — call that threshold kth_excluding(H(x),k,i). Since the only randomness left is expert i's own Gaussian noise, this probability has a closed form:

where Φ is the standard normal CDF — the probability that a standard Gaussian sample falls below a given value. Now the same worked example from Chapter 2, continued: reuse H(x) = [2.3, 0.8, 0.6, −2.0], noise scale 1 for every expert, k = 2.

Step 1 — find kth_excluding for each expert. For expert i, drop it from H(x), sort the rest, take the 2nd-largest (since k = 2):

Step 2 — convert each margin to a probability via Φ. Using the clean logits from Chapter 2 (2.0, 1.0, 0.5, −0.5) and noise scale 1, the argument to Φ is simply (clean logit − kth_excluding):

Step 3 — sum across the batch, then penalize imbalance exactly like Importance did:

Same coefficient-of-variation recipe as Importance, but now measured on expected example count rather than summed weight — a genuinely different quantity that can diverge from Importance exactly the way the paper's quote warned.

What happens if you skip these losses — the paper's own numbers

The paper ran the identical 256-expert model with different loss weightings and measured what happened after 10 training epochs:

With both losses off, the single most-loaded expert receives 17.8× the average load — a near-total collapse onto a handful of favored experts, and perplexity is visibly worse (39.8 vs ≈35.6–35.7 with either loss on). Turning on either loss alone already fixes most of the damage; using both together is only marginally better than one, but keeps both notions of balance in check simultaneously. This table is the direct, measured cost of skipping Chapter 5.

The last two rows answer a question the first four leave open: how sensitive is this to the exact weight chosen? Dial the weights down to 0.01 — a tenth of the 0.1/0.1 row — and imbalance creeps back (CV(Importance) climbs to 0.48, more than the 0.06 seen at 0.2/0.0), even though perplexity barely moves (35.7). Dial them up to 1.0 — ten times the 0.1/0.1 row — and balance only gets tighter still (max÷mean of 1.07, the best in the table) with, again, no perplexity cost. Read across the whole table and a pattern emerges that is genuinely reassuring for anyone tuning this in practice: too little weight on these losses visibly hurts (both the 0.0/0.0 row's perplexity and its CV), but there is no matching penalty for too much — perplexity holds essentially flat from 0.1 all the way to 1.0, a hundred-fold range. The failure mode this loss guards against is asymmetric: erring on the side of too much balancing pressure costs essentially nothing, while erring on the side of too little costs real perplexity.

Balancing at every level, once the gate is hierarchical

Chapter 1 introduced the hierarchical y_H equation for when n is too large for one flat gate. Importance and Load need a matching hierarchical version, because a primary gate and each secondary gate are both capable of collapsing independently. The paper's Appendix B extends both definitions directly. Hierarchical importance for expert (i,j) — group i, expert j within that group — multiplies the two gates' contributions before summing over the batch:

Read this as “how much total weight expert j inside group i received, accounting for both how often group i itself was picked and how often expert j was picked within it.” The hierarchical load is a little more subtle, because a secondary gate's balance only makes sense measured over the examples that actually reached that group in the first place:

where X⁽ⁱ⁾ is the subset of the batch routed to group i, and Load_i is the ordinary, Chapter-5-style load function computed by group i's own secondary gate, but only over that subset. Dividing by |X⁽ⁱ⁾| converts “expected count within the group” back into the same units as the rest of the batch, so that a group receiving few examples does not automatically register as more or less balanced than a group receiving many. The takeaway is simple even though the formula looks busier than Chapter 5's flat version: nothing new is invented here, the same CV² machinery just gets applied twice, once to the primary gate's choice of group and once to each secondary gate's choice of expert within its group — a tree of balancing pressures matching the tree of gating decisions from Chapter 1.

Look closely at that division by |X⁽ⁱ⁾| in Load_H and one sharp edge case stands out: what happens if group i receives zero examples in a given batch? |X⁽ⁱ⁾| would be zero, and dividing by zero is undefined. This is not a corner case the formula's authors overlooked — it is exactly the failure mode the next section's fix exists to make vanishingly unlikely in the first place. If every group starts with equal odds of being selected (the all-zero initialization below), no group's population collapses to literally zero on the very first batch, and the ongoing CV² pressure keeps every group's population away from zero for the rest of training. The hierarchical balancing math and the initialization trick are not two independent fixes bolted together by coincidence; the second one is partly what keeps the first one's formula from ever having to divide by zero in practice.

A subtlety the losses alone cannot fix: getting started

CV² penalties are a soft constraint — gradient descent nudges toward balance over many steps, it does not enforce balance instantly. That creates a chicken-and-egg problem at the very first training step: if W_g starts at a typical random initialization, some experts will, purely by chance, start with noticeably higher initial logits than others, and on a large cluster that imbalance can be severe enough to blow through per-device memory budgets before the soft losses have had time to pull things back into line. The paper's fix, described in Appendix A, is almost too simple to expect: initialize both W_g and W_noise to all zeros. With every weight at zero, every expert's clean logit and noise scale start identical for every input — H(x) starts as pure symmetric noise, so the top-k selection starts uniform across experts by construction, and the soft balancing losses only ever have to maintain a balance that already existed at step zero, rather than having to claw one back from an already-lopsided start.

Three fixes, three different timescales, is the shape of this whole chapter in miniature: the noise term from Chapter 2 keeps any single forward pass from being deterministically locked in; the CV² losses in this chapter apply gentle, continuous pressure across every training step; and the all-zero initialization handles the one moment — step zero — where neither of the other two has had time to act yet. Remove any one of the three and a different failure mode opens up, at a different point in training.

Chapter 7 turns this exact three-part story into something you can watch happen rather than only read about — a live histogram of experts either collapsing or holding steady, with the same self-reinforcing dynamic this chapter derived on paper now animated on a canvas, step by step.

Chapter 6: 137 Billion Parameters

Every earlier chapter was mechanism. This one is the payoff: does decoupling capacity from compute actually buy better models, and how far does it scale before something breaks?

Two datasets, two tasks, one architecture unchanged throughout: the 1-billion-word and 100-billion-word language modeling corpora first, then machine translation. Watching the same gate, the same noisy top-k mechanism, and the same two balancing losses succeed on genuinely different data and a genuinely different task is the strongest version of “this is not a language-modeling trick” this lesson can offer — not an assertion, but a pattern repeated across independent experiments.

First, a controlled experiment: fixed compute, varied capacity

On the 1-Billion-Word Language Modeling Benchmark, the paper holds computation nearly fixed at ≈8 million operations per timestep and varies only the number of experts — 4, 32, 256 (flat), and 256, 1024, 4096 (hierarchical), always with k = 4 active. This isolates the one variable Chapter 0 cared about: capacity, with compute held constant. The 4-expert model (no real sparsity — all 4 always run) performs about the same as compute-matched dense baselines, which is itself an important negative result: MoE is not free magic from nothing, the gain has to come specifically from the sparse capacity growth. The 4,096-expert model, at the same 8M-ops/timestep budget, achieves 24% lower test perplexity.

“Compute-matched dense baseline” deserves to be more than a phrase. The paper trained four separate baselines specifically to control for every alternative explanation of why MoE might be winning, holding total compute fixed at the same 8M-ops/timestep budget as every MoE variant above:

MoE-1-Wide and MoE-1-Deep are the control this experiment needed: both replace the MoE layer with a single dense “expert” sized (wide, or deep) to use exactly the same compute budget as the sparse models, ruling out the boring explanation that MoE wins simply because “more parameters arranged into more matrix multiplies, however you shape them, helps.” They do not win — 46.1 and 45.7, essentially tied with the plain 4xLSTM-512 baseline at 46.0, all clustered close together. What actually moves the needle, sharply, is more experts specifically: MoE-32 at 39.7 is a full 5–6 perplexity points better than every single-shape control, at the identical compute budget. The gain is not from bigger matrices in any particular shape; it is from the conditional-computation structure Chapter 1 built, letting many small specialists exist where one large generalist used to sit.

Worked example: beating the best published model on 6% of its compute

The paper's own Table 1 makes the capacity/compute decoupling concrete enough to check by hand. The best previously published dense model (a 2-layer LSTM, 151 million parameters) scored 34.7 test perplexity after 10 epochs, at a cost of 151 million ops/timestep. The paper's Low-Budget MoE model scored better — 34.1 — while using only 8.9 million ops/timestep, despite holding 4,303 million total parameters:

Read those two lines together: 28.5× more parameters, using roughly 6% of the compute, and it still wins on quality. That is Chapter 0's promise, verified with the paper's own reported numbers, not a hypothetical.

Does more compute still help, once you already have a large MoE?

Chapter 0 decoupled capacity from compute; it did not claim compute stops mattering. The paper checks this directly by training two further models — MoE-34M and MoE-143M, named for their ops/timestep budgets — that hold MoE capacity roughly fixed at ≈4 billion parameters while giving the surrounding LSTM layers substantially more computation to work with (MoE-143M's LSTM layers have 4,096 units, versus 512 in the low-budget model). The result: 31.3 and 28.0 test perplexity respectively — both improvements over the 34.1 low-budget number, and the paper's own comparison against the previously-best published model is direct: “our model has a lower test perplexity by 18%.” Capacity and compute are decoupled, not substitutes for each other — a model benefits from more of both, and the sparse-gating trick is what lets you afford to buy each one independently instead of always paying for both at once.

A negative result worth not skipping past: MoE is not automatically efficient

Chapter 0 already flagged the 4-expert model as a negative result on quality — no real sparsity, no real gain. The paper's full results table (Appendix C's Table 7) reports a second negative result, this time on raw hardware efficiency, easy to miss if you only look at perplexity. The MoE-4 model, run on the same infrastructure as every other row, achieves only 0.52 TFLOPS/GPU — worse than the plain 4-LSTM dense baseline's 1.07 TFLOPS/GPU at the identical 8.4M-ops/timestep budget. A model with a gate and conditional routing bolted on, running slower per FLOP than the dense model it is supposed to beat, is exactly what Chapter 0's obstacle table warned about: conditional computation does not pay for itself automatically, it has to be engineered to.

The flip side of that same table is the best evidence Chapter 4's bandwidth argument was right. MoE-143M — the highest-compute, widest-expert model in the language-modeling suite — reaches 1.56 TFLOPS/GPU, higher than any dense baseline in the table, including the 1.07–1.29 TFLOPS/GPU range the paper reports for its non-MoE models. Bigger matrices per expert do not just buy a better compute-to-bandwidth ratio in theory (Chapter 4's h = compute÷bytes derivation); here, measured on real hardware, they buy a model that is not merely competitive with a dense network's efficiency but exceeds it. Read the two ends of this table together and the lesson is precise, not vague: a small, narrow MoE can be less efficient than doing nothing sparse at all; a large, wide one, engineered the way Chapter 4 describes, can be more efficient than a dense network of any size the paper tested.

Scaling up: the 100-Billion-Word Google News Corpus

With more training data available, the paper pushed capacity much further — up to 131,072 experts, using hierarchical MoE (Chapter 1's y_H equation) with branching factors chosen per configuration:

The 65,536-expert row is the paper's headline result: a 68-billion-parameter MoE layer, hierarchically organized as 256 primary groups of 256 experts each (256×256 = 65,536, using the exact y_H equation from Chapter 1), achieving 39% lower perplexity than a compute-matched dense baseline after training on the full 100 billion words — while running at a respectable 0.72 TFLOPS/GPU.

“Respectable” is worth anchoring to a hard ceiling rather than taking on faith: the paper reports its dense baselines running at 1.07–1.29 TFLOPS/GPU on the same Tesla K40 hardware, against a theoretical maximum of 4.29 TFLOPS/GPU that NVIDIA claimed for that chip. A model that is 99.994% sparse per forward pass still delivering 0.72 of that same GPU's throughput — roughly a sixth of the theoretical ceiling, and comparable to more than half of what the fully-dense baselines themselves achieved — is the entire payoff of Chapter 4's batching and bandwidth engineering made visible in one number. Without the data/model-parallel batch pooling, the convolutional batching trick, and the wider hidden layers, this same architecture would be memory-bound and network-bound long before it got anywhere near that figure.

Worked example: how sparse is 99.994%, exactly?

The paper describes the 65,536-expert model as “99.994% layer sparsity.” Derive that number instead of trusting it. With k = 4 active experts out of n = 65,536:

Matches the paper exactly. Notice, too, the cautionary tail of the table: the 131,072-expert model, with more total parameters, scores slightly worse (29.2 vs 28.9) than the 65,536-expert model, and its computational efficiency collapses to 0.30 TFLOPS/GPU. The paper's own hedge: “possibly a result of too much sparsity.” Chapter 4's batch-pooling trick keeps per-expert batches healthy up to a point; push n far enough and even pooling across every device in the cluster cannot keep each of a hundred-plus thousand experts fed with enough examples to train well. Capacity does not help without bound.

Beyond language modeling: translation, and what experts actually learn

The same architecture, inserted into a machine-translation encoder-decoder (2,048 experts, ≈2M parameters each — a larger hidden layer than the language-modeling experts, per Chapter 4's bandwidth argument), beat Google's strong GNMT baseline on WMT’14 English→French by +1.34 BLEU (40.56 vs 39.22) and on English→German by +1.12 BLEU (26.03 vs 24.91) — with no reinforcement-learning refinement, the technique GNMT itself used to get its best numbers. A single multilingual MoE model, trained on twelve language pairs at once, beat a same-capacity multilingual dense baseline by 19% lower dev perplexity ((4.14−3.35)÷4.14 = 0.1908) and outperformed it on BLEU for 11 of the 12 language pairs, by as much as +5.84 points on Korean→English. This architecture is not a language-modeling trick; it generalizes.

The comparison worth sitting with is not just “MoE beat the other multilingual model” — it is that a single multilingual MoE model beat eight of the twelve separately-trained, single-language-pair GNMT models, each of which had an entire dedicated model (and dedicated training run) for just that one language pair. One shared model with a sparse gate outperformed most of a dozen specialized ones, at a fraction of the aggregate training cost. The paper is candid about the one pair where multilingual MoE lost ground badly — English→Korean — and names a specific, plausible cause rather than shrugging: for rarer language pairs, a small number of real training examples were “highly oversampled” to balance the combined dataset, and the paper attributes English→Korean's weak result to severe overtraining on that oversampled, thin slice of data. That is a data problem, not an architecture problem — a useful reminder that not every gap in a results table traces back to the mechanism this lesson is teaching.

One small, honest caveat belongs here, not because it changes the results but because it is the kind of detail this lesson's own standards insist on surfacing rather than smoothing over: the paper's Appendix G notes that, “for performance reasons,” its translation models use a slightly different attention function than GNMT's own published one — a modified form of the same attend-from-source-to-target computation, chosen for engineering convenience rather than a claimed accuracy improvement. It is a small architectural difference between the MoE model and the GNMT baseline it is compared against, on top of the MoE layer itself, and the paper is upfront about it rather than hiding it in a footnote-free comparison.

Worked comparison: what it cost, in hardware, to beat GNMT

BLEU points alone do not say anything about cost. The paper's own Table 2 reports the training setup for both sides of the English→French comparison side by side, and it is worth reading as a hardware story, not just a quality one. The MoE model (longer-trained variant) reaches 40.56 BLEU at 85 million ops/timestep, holding 8.7 billion total parameters, trained for six days on a cluster of 64 Tesla K40 GPUs. Google's own GNMT baseline reaches 39.22 BLEU at 214 million ops/timestep — more than double the MoE model's per-token compute — holding a comparatively tiny 278 million total parameters, trained for the same six days but on 96 Tesla K80 GPUs, a newer and individually more powerful chip than the K40s the MoE model trained on.

Read those two setups side by side and the MoE model wins on every axis reported except raw parameter economy: a higher BLEU score, less than half the per-token compute, and fewer (and older) GPUs, in the same wall-clock training time. The one caveat worth being honest about, in the same spirit as this lesson's earlier correction about GPU-count claims: K40 and K80 are different chip generations with different per-chip throughput, so “64 K40s beat 96 K80s” is not a perfectly apples-to-apples hardware comparison — the paper does not equalize for that difference, and neither will this lesson pretend to. What the comparison does support cleanly, without needing that caveat resolved, is the ops/timestep and parameter-count story: 85 million versus 214 million operations per token, decoupled from 8.7 billion versus 278 million total parameters, is Chapter 0's capacity-compute decoupling showing up again, this time on a completely different task and against a different research group's own strongest published system.

And the specialization Chapter 1 promised is real, not folklore. From the paper's own inspection of the translation model's 2,048 experts, sorting each expert's inputs by how strongly it was selected:

Three different experts, three different organizing principles, none of them designed in. Expert 752 has converged on an idiom — a fixed grammatical slot (“plays a ___ role”) regardless of which adjective fills it. Expert 2004 has converged on a semantic field — speed and rapidity — scattered across completely different grammatical constructions (“rapidly,” “swiftly,” “drastically,” “fastest”). Expert 381 sits somewhere between the two, clustering around a topic (research and innovation) rather than either a fixed syntactic slot or a single part of speech. The gate was never given labels for “idiom,” “speed-word,” or “research-topic” — it only ever received a training signal that said, indirectly, “did routing this token to you help translate the sentence correctly.” Three different useful groupings of language fell out of that one signal, applied identically to 2,048 initially-identical experts.

No one told the first of these experts to specialize in the “plays a [important-role adjective]” idiom. It emerged purely from the interaction of routing, the load-balancing pressure from Chapter 5 keeping every expert alive and competing for examples, and ordinary gradient descent. Hold onto this concrete fact — it is exactly what makes Chapter 8's discovery possible seven years later.

Step back and every number in this chapter is answering one of two questions Chapter 0 opened with: does the capacity-compute decoupling actually buy quality (yes, repeatedly, on two unrelated tasks and three unrelated benchmarks), and does it scale without limit (no — the 131,072-expert row's own regression is the honest ceiling this chapter found). Both answers matter equally to trusting the architecture: a decoupling that only worked at one scale, or that never showed its own limits, would be far less convincing than one that shows both its wins and exactly where it starts to break.

Chapter 7: Router Showcase

Chapter 5 derived the losses and showed you a table of what happens with and without them. Now watch the mechanism run. This is a simplified simulation of the qualitative dynamics described in Section 4 — not a literal re-run of the paper's own numbers — built so you can watch the self-reinforcing collapse happen, and watch balancing stop it, in real time.

Each “step” below routes a fresh batch of synthetic tokens through n experts. Every expert carries a strength that starts equal. With balancing off, an expert's strength grows whenever it gets over-selected and shrinks when under-selected — the self-reinforcing loop from Chapter 5, made visible. With balancing on, strength gets pulled back toward the average every step, mimicking what the importance/load losses do to the real gating weights. Watch the load histogram at the bottom: with balancing off it should collapse onto one or two tall bars; with it on it should stay roughly flat, the way Table 6's “max÷mean” column did in Chapter 5.

How to read this simulator, mechanically

“Strength” is a deliberate simplification standing in for something Chapter 2 and Chapter 5 built out of several separate pieces — W_g, W_noise, and the resulting logits H(x) that KeepTopK sorts. Collapsing all of that into one scalar per expert throws away the token-by-token noise draws and the exact CV² loss shape, but it keeps the one property that actually drives the qualitative behavior this chapter is here to show: whatever quantity determines an expert's selection odds this round should go up when that expert wins more often than its share, and down when it wins less. With balancing off, each step multiplies an expert's own log-strength by a term proportional to how far its last round's load sat above or below the mean — the simulator's stand-in for “more gradient updates make an already-favored expert more favored.” With balancing on, every expert's strength is pulled halfway back toward a neutral baseline after each step, a crude but faithful analogue of what a CV² penalty does over many real training steps: it does not reset an expert's weights, it steadily discourages the imbalance those weights would otherwise drift toward.

Spelled out as the actual update rule the code runs, once per synthetic batch of 40 tokens: each token has a soft affinity toward a randomly-placed “topic” on a ring of experts (closer experts score higher, farther ones score lower, plus a little Gaussian jitter — a rough stand-in for Chapter 2's clean-logit-plus-noise shape). With balancing off, that affinity score gets an extra term added on top: log(strength) × 1.1, so an expert already running hot from previous rounds gets a further boost this round, on top of whatever the token's own affinity says — the self-reinforcement, made literal in one line of arithmetic. After the batch, each expert's strength updates by how far its load sat above or below the mean: strength × (1 + 0.18×(load÷mean − 1)), clamped so it never drops below a small floor. With balancing on, that whole update is replaced by pulling strength halfway toward 1 every round, regardless of how the last batch went — strength × 0.5 + 0.5, a hard-coded restoring force with no dependence on this round's load at all.

Why a ring of experts, rather than experts with no particular relationship to each other? Purely to make the “topic” visual and legible on a small canvas — placing experts around a ring and letting each synthetic token's affinity fall off with ring-distance from a random topic point gives every step a recognizable, roughly-bell-shaped pattern of who-wants-what, instead of pure unstructured noise. Nothing about Chapter 2's real gate has any notion of experts being “close to” or “far from” each other — the n experts in the real architecture are an unordered set, and W_g is free to route any input to any expert regardless of index. Treat the ring purely as a rendering convenience for this widget, not as a claim about how real MoE experts relate to one another spatially.

One more bridge worth drawing explicitly: Chapter 5's real w_importance and w_load are continuous dials, not an on/off switch — Table 6 showed perplexity and balance both responding smoothly as those weights moved from 0 to 0.01 to 0.1 to 1.0. This showcase's balancing toggle collapses that whole continuous dial down to a binary choice, on purpose, so the qualitative difference between “present” and “absent” balancing pressure is impossible to miss on a small canvas. If you wanted to extend this simulator yourself, the natural next step would be replacing the toggle with a slider that blends between the two update rules — interpolating strength's pull-toward-1 by some weight w, the same way L_importance and L_load are themselves scaled by w_importance and w_load in the real loss. The qualitative shape of the result would not change; what would change is exactly how much pressure it takes before the collapse stops happening, which is precisely the question Chapter 5's Table 6 answered with real numbers.

Reading the histogram precisely: what the colors and the dashed line mean

The bars at the bottom track cumulative load — the running total of tokens each expert has received across every step so far, not just the most recent one — because a single step's load is noisy enough to obscure the trend this chapter exists to show. Watching the running total makes the direction of drift unmistakable: a genuinely balanced system's bars grow together, roughly in lockstep, step after step; a collapsing one shows a small number of bars pulling steadily ahead of the rest, the gap widening every step rather than merely fluctuating around it.

Color is doing real diagnostic work, not decoration. An expert whose cumulative load has climbed past 2.2 times the mean turns red — a rough visual analogue of the Table 6 danger zone, since the paper's own worst case measured a max÷mean ratio of 17.80, far past any bar this simulator will show turning merely warm-colored. An expert whose load has fallen below 0.4 times the mean fades to a dim, muted color — the visual signature of an expert being systematically starved, the flip side of another expert's over-selection. Everything in between stays teal: neither hoarding tokens nor being frozen out. The dashed horizontal line marks the mean load itself, the reference every bar is implicitly being measured against; in a balanced run, every bar's top edge should sit close to that line, and in a collapsing run you can watch a handful of bars climb steadily above it while the dashed line itself barely moves, since it is only ever the average of everything happening below it.

Two experiments to run, with real numbers to check your run against

This showcase has no quiz — the simulation is the test — but “play with the sliders” is more useful with a specific target to compare against. Chapter 5's Table 6 gives two real anchor points: a completely unbalanced 256-expert model measured a max÷mean load ratio of 17.80; the same model with either balancing loss switched on measured a ratio between 1.07 and 1.47, depending on which loss and how strongly weighted.

Push a third variant yourself: raise n toward 16 with balancing off and watch how much faster the collapse happens with more experts competing for the same fixed pool of synthetic tokens — a hands-on feel for why Chapter 5 called the imbalance problem self-reinforcing rather than merely occasional, and why a thousand-expert model left untended is a much worse bet than a small one.

A fourth variant, worth trying specifically because it connects back to a real number from Chapter 6: set n down to 4, the same expert count as the paper's own MoE-4 model. Recall Chapter 6's verdict on that configuration — with only 4 experts and k = 4 active, every expert runs on every token, so there is no real sparsity and, mechanically, nothing for collapse to even act on. This simulator's slider bottoms out at n=4 for a related reason: below a handful of experts, the qualitative story this showcase is built to demonstrate (a few experts hoarding load while others starve) stops being a meaningful distinction from “every expert gets roughly a quarter of the traffic by default.” The interesting dynamics this chapter cares about live specifically in the regime Chapter 5's real experiments used — n large enough that a gate genuinely could specialize unevenly, which is also exactly the regime where balancing losses earn their keep.

That, ultimately, is the one-sentence summary this whole showcase exists to leave you with, independent of which exact experiment you ran: balance is not the default outcome of training a gate, it is a property that has to be actively fought for, on every single step, for as long as training continues. Turn the fight off and the system does not stay balanced by inertia — it drifts, and the drift compounds. Chapters 2 and 5 are the two places in this lesson where that fight is actually implemented; this chapter is only here to make sure you have felt why the fight is necessary before reading the machinery that wages it.

If you came to this chapter after reading Chapters 0–6 straight through, notice how differently “17.80” lands now compared to when Chapter 5 first printed it in a table. A number in a table is a fact to accept; a number you have just watched a histogram climb toward, one step at a time, is closer to an experience. That gap — between reading a result and watching it happen — is the entire reason this lesson budgets a showcase chapter for every major mechanism at all.

Before moving on, it's worth being honest one more time about scale: this simulator runs at n up to 16, while the paper's own real experiments went as far as n = 131,072. Nothing about the qualitative dynamics changes at that larger scale — self-reinforcement is self-reinforcement whether it is happening among 16 competitors or 131,072 of them — but the speed of collapse and the severity of the imbalance both get worse as n grows, which is exactly why Chapter 5's real losses, not this toy demonstration, are what actually has to hold at the scale Chapter 6 operates at.

That is also the honest limit of what any small, interactive widget can teach about a system meant to run at 100,000-expert scale: intuition transfers, exact numbers do not. Trust this chapter for the shape of the argument — self-reinforcement is real, and a fix aimed directly at it works — and trust Chapter 5's own measured table for the magnitudes.

With that, this lesson's core mechanism is complete: Chapters 1–3 built it, Chapters 4–5 made it affordable and stable, Chapter 6 proved it at scale with real numbers, and this chapter let you watch its most fragile piece — balance — hold or fail in real time. Chapters 8 and 9 pick up an entirely different thread from here, one the original 2017 authors never anticipated.

Take one thing from this chapter specifically into that jump: the balancing pressure you just watched, in the “on” runs, keeping every expert genuinely competing rather than idle. Chapter 8 leans on that exact property — a router that stays balanced is a router whose choices still depend on what it is looking at, which turns out to be the one fact that makes everything in the next two chapters possible.

Nobody who built this loss in 2017 was thinking about embeddings; they were thinking about not wasting a GPU cluster. The property that ends up mattering most, seven years later, is a side effect of an engineering fix for an entirely unrelated problem — worth remembering the next time a training instability's fix seems like it only matters for the metric it was built to fix, since you rarely know in advance which side effect will turn out to be the interesting one.

What this simulation deliberately leaves out

Being honest about a simplification's edges is part of using it well. Three things this simulator does not model, on purpose, so that watching it does not accidentally teach the wrong lesson about the real mechanism:

None of these simplifications change the qualitative shape of the story — collapse without correction, stability with it — which is the one thing this showcase exists to make you feel rather than just read. For the exact mechanism behind each simplification, Chapters 2 and 5 are where the real machinery, and the real numbers, live.

A few patterns worth recognizing when you see them

Run the simulator enough times with balancing off and a small set of recurring shapes tends to show up, each telling you something slightly different about the underlying dynamics:

None of these three outcomes is more “correct” than another — they are all instances of the same underlying rule (strength compounds strength) playing out on slightly different random draws of where the synthetic topics happened to land. That variability is itself worth noticing: Chapter 5's real losses do not promise to prevent any one particular pattern of collapse, they promise to prevent all of them at once, by attacking the compounding mechanism directly rather than any one symptom of it.

Notice, too, what balancing being switched on does not do: it does not force every expert to receive the exact same load on every single step, and you should not expect to see perfectly identical bars. Chapter 5's real CV² losses are a soft, statistical pressure toward balance across a batch, not a hard per-step guarantee — the strictly-equal-batches alternative from Chapter 4's Appendix F detour is the only mechanism in this whole lesson that forces exact equality, and it comes with its own tradeoffs discussed there. What you should expect from balancing on, here and in the real training run alike, is the running totals staying close to each other and never drifting apart, not perfect lockstep on every individual step. That distinction — bounded, self-correcting fluctuation, versus runaway divergence — is the entire qualitative difference this showcase exists to make visible.

Chapter 8: MoE as Free Embedding

Jump seven years. Chapters 0 through 6 spent their entire budget on one question: given a fixed compute budget, how do you buy more capacity? Chapter 8 asks a completely different kind of question about the exact same mechanism — not “how do we make this cheaper,” but “what else is this secretly good for.” Every chapter so far treated the gate as a means — a way to decide which experts should do the actual work. In 2024, Li & Zhou (University of Maryland) asked a different question: what if the gate's choices themselves, thrown away after routing in every system built so far, are a usable signal in their own right?

The 2024 paper opens by naming exactly the paper Chapters 0–6 of this lesson just spent six chapters on, by citation: it observes that “MoE models have been predominantly used in multitask learning and efficient large-scale training scenarios (Shazeer et al. 2017)” and that “their potential for generating instance-level embedding has been underexplored.” That is this lesson's own arc, stated by the second paper's own authors as their motivation: everyone who read Shazeer et al. for seven years read it as a compute-efficiency paper, because that is what it is. Li & Zhou are the first to ask what else the same mechanism might be quietly good for. Their own abstract states the finding plainly: “the expert routers in MoE LLMs can serve as an off-the-shelf embedding model with promising performance on a diverse class of embedding-focused tasks, without requiring any finetuning.” Every word in that sentence earns its place — “off-the-shelf” because nothing about the base model changes, “diverse class of tasks” because Chapter 9's six-task-family results back it up, and “without requiring any finetuning” because that is the entire distinction this chapter draws against SimCSE's whole lineage.

The problem: decoder-only LLMs make mediocre embedding models by default

An embedding model turns a piece of text into one fixed vector, useful for search, clustering, or classification — sentences that mean similar things should land near each other in that vector space. Modern LLMs are trained for one job: predict the next token. The standard way to repurpose one as an embedder, with no extra training, is to grab the hidden state (HS) of the last token at the last layer.

That hidden state was optimized to predict what comes next, not to summarize everything already said. The paper's own example: two inputs with subtly different meanings can end up needing the same predicted next token or label, so the hidden state that drives that prediction can end up nearly identical for both — even though a human would say the two inputs mean different things. The information distinguishing them was never forced to survive into that one vector.

This is not a new problem, and it did not go unaddressed before MoEE. An older line of work builds embedding models by training specifically for it, going back to SkipThought (Kiros et al. 2015), which trained a model to predict a sentence's neighboring sentences and used the resulting representation as its embedding: the embedding was a designed target, not a byproduct. SimCSE and similar contrastive methods (alongside related contrastive frameworks like InfoNCE and MoCo, and LLM-specific descendants such as Sentence-T5) took a different, more modern route to the same designed-target philosophy — pull two augmented views of the same sentence together in vector space and push unrelated sentences apart, purpose-built and often outperforming naive HS extraction by a wide margin. But that route, in every one of these forms, costs a dedicated training run, a specific objective, and (for the contrastive methods) pairs of augmented data — exactly the kind of “per-capability engineering,” in Chapter 0's terms, that a foundation model was supposed to let you skip. MoEE's contrast with that whole lineage is precise: it is training-free, not training-cheap. Nothing about the base MoE model changes; the signal was already sitting in every forward pass, waiting to be read instead of discarded.

Written formally, for an input sequence x = [x₁,…,x_T] with last-layer hidden states H^(L)∈ℝ^T×d, the two standard training-free ways to extract e_HS are:

One common trick sharpens the last-token version further: PromptEOL (Jiang et al. 2023) wraps the input in an instruction before extracting HS, so the “next token to predict” is forced to be a single summarizing word rather than whatever token happens to follow in a raw continuation. The exact template:

That forces the model's next-token machinery to actually compress the input into something summary-like before the hidden state is read off — and it measurably helps, as Chapter 9's results table shows. But it is still working entirely within the next-token-prediction paradigm, coaxing a better summary out of the same mechanism. RW does not need coaxing, because summarizing the input well was never optional for it — a router that routes badly makes worse predictions, full stop, so the training signal shaping RW's usefulness was already there from the start.

The discovery: a router's job is already “summarize this input, then act on the summary”

Recall Chapters 1–2: at every layer, the gate g^(l)(H^(l)) looks at the current representation and decides which experts are the right fit for this specific input, at this specific depth of the network. That decision is necessarily a compressed judgment about what kind of input this is — not because anyone designed it to be, but because routing well requires making exactly that judgment. MoEE's core move: stop discarding that judgment after using it to route, and use it directly as an embedding coordinate.

Write out what g^(l) actually is, one layer at a time, before concatenating anything. At layer l, with N^(l) experts and per-expert logits z^(l)(H^(l)), the routing weight for expert i is an ordinary softmax:

This is precisely Chapter 2's ordinary softmax gate G_σ, the dense one from the very start of that chapter, before KeepTopK ever masked anything to −∞. That is worth pausing on: MoEE is not introducing new machinery at all, at the per-layer level. It is taking a quantity every MoE layer already computes on the way to Chapter 2's Part 4, before the masking step throws most of it away, and simply not throwing it away. Concatenate one such vector per layer, in order, and the full embedding follows:

One important departure from Chapter 2: for embedding purposes, MoEE uses the full dense softmax distribution over all N^(l) experts at each layer — not the sparse, top-k-masked gate used for actual routing. That is a deliberate choice. Routing needs hard zeros to save compute (Chapter 1's whole point); an embedding wants the opposite — every expert's degree of relevance, including the close second-choices that top-k discards, because those runner-up scores carry information too.

Concept → realization: trace the actual dimensionality

Three real MoE models, three real embedding sizes, computed by concatenating one softmax vector per layer:

Every one of these vectors is a genuinely free byproduct: it is already sitting in memory after an ordinary forward pass through a model you already have, extracted for exactly the same computational cost as generating one token. No fine-tuning run, no labelled pairs, no contrastive objective — the title of the paper is not an exaggeration.

Scale the free part of “free byproduct” against the model it comes from, since the two numbers are easy to leave disconnected. DeepSeekMoE-16B carries its parameter count in its own name — roughly 16 billion parameters — and produces, at zero extra cost, an embedding of 1,792 numbers. That is a compression ratio of roughly 16,000,000,000 ÷ 1,792 ≈ 8.9 million to one: a model with nine billion-ish parameters' worth of learned structure, reduced to under two thousand numbers that still, according to Chapter 9's results, capture enough of what the model “thinks” about an input to beat purpose-trained embedding models on several tasks. That ratio is not a coincidence of this one model either — it is exactly what you would expect from Chapter 0's whole architecture: a gate only ever has to be as wide as the expert count, never the expert capacity, so the embedding built by concatenating gates inherits that same narrow width no matter how many billions of parameters sit behind it.

Walk OLMoE-1B-7B's 1,024 dimensions through, one construction step at a time, since it is the smallest of the three and easiest to hold in your head fully. A sentence goes in; the model runs its ordinary forward pass, one layer at a time, exactly as it would to generate a next token. At layer 1, the gate produces a 64-entry softmax vector — one weight per expert at that layer, summing to 1. That vector is set aside, untouched, rather than immediately masked down to a sparse top-k for routing. The same thing happens at layer 2, then layer 3, all the way through layer 16, each one contributing its own 64-entry vector. After the forward pass finishes, the 16 saved vectors are concatenated end to end, in layer order: 16 × 64 = 1,024 numbers total, one single vector representing the entire input. Nothing about this walkthrough required a second forward pass, a second model, or a single additional weight — every number in that 1,024-dimensional vector was already computed on the way to the model's ordinary next-token prediction; MoEE's contribution is entirely in the decision to keep it.

One more empirical detail worth flagging before Chapter 9's numbers: the paper found that using only the last token's routing weights beats averaging routing weights across every token in the input. This mirrors the same “use the last token” convention that hidden-state embedding already follows, but for a different underlying reason — it is examined directly in Chapter 9's ablation. The paper's own summary of that finding, stated plainly: focusing on the last token, for either signal, “consistently delivers the best performance,” because the last token captures the most critical semantic information while pooling across tokens or layers introduces noise instead of adding signal. Averaging is not free extra information here — it is a way of diluting a sharp signal with a lot of less-relevant ones.

It is worth being specific about what “MTEB” means here too, rather than leaving it as an unglossed acronym. The Massive Text Embedding Benchmark is a standardized suite spanning several distinct kinds of embedding task, each scored with its own appropriate metric specifically so that no single task's quirks can make an embedding method look better or worse than it really is across the board:

Chapter 9's headline numbers average across six of these task families (all but Retrieval, which the paper reports separately); treat that average the way you would any benchmark average — a useful single number, but one worth checking task-by-task before trusting it blindly, which Chapter 9 does.

Why this actually encodes semantics, not noise

This is where Chapter 5 pays for itself in an entirely different currency. Recall the collapse scenario: with the balancing losses off, the router converges to always favoring the same few experts, regardless of input — Table 6's CV(Importance) ≈ 3, max÷mean load ≈ 17.8. If a router had actually collapsed that badly, its routing-weight vector would be nearly the same for every input — close to a constant, and a constant vector is a useless embedding, since a useless embedding is precisely one that fails to vary with the thing it is supposed to represent.

It is exactly the training-time pressure that Chapter 5 derived — forcing every expert to stay competitive, forcing routing decisions to actually depend on input content rather than settling on a favorite — that guarantees the router's output varies meaningfully from sentence to sentence at inference time. The 2017 fix for a training pathology turns out, unintentionally, to be the exact property that makes the 2024 discovery possible: a router that never collapsed is a router whose choices can tell two different inputs apart.

Push this one step further and a small paradox resolves itself. Chapter 1's promise was that only k of n experts run per token — the whole point of the architecture is that most of the gate's output is exactly zero. Chapter 8's embedding, by contrast, wants the opposite: every expert's degree of relevance, dense, not sparse. These are not in tension, because they are never the same tensor. The sparse, top-k-masked G(x) from Chapters 1–3 is what actually gets multiplied against expert outputs and determines compute cost. The dense g^(l)(H^(l)) that MoEE reads is the softmax computed one step before KeepTopK ever touches it — the same forward pass, a different, earlier tap point. Reading it costs nothing extra precisely because it was already sitting there, computed and then normally discarded, before the sparsity Chapter 1 promised ever gets applied.

Put another way: Chapter 1 built sparsity as a subtraction from a dense computation — start with the full softmax over all n experts, then subtract away everything outside the top-k. MoEE simply reads the value that exists in the instant right before that subtraction happens. Nothing about this requires modifying the model, retraining it, or even understanding it any differently than before — it requires knowing precisely where, in the sequence of operations Chapter 2 already defined, the useful intermediate value lives.

That is, in miniature, the whole discovery this chapter is built around: not a new mechanism, but a new place to look inside an old one. Chapter 9 puts that new signal side by side with the old one, hidden state, and asks how much they actually agree with each other, and what happens when you use both at once.

Chapter 9: Connections

Chapter 8 established that routing weights (RW) and hidden states (HS) are two different signals sitting in the same forward pass. Now: how different, exactly, and what happens when you combine them?

Measuring the overlap

The paper clusters the same set of inputs twice — once by RW, once by HS — and compares the two clusterings directly:

Read these together, not in isolation, and it helps to know what each metric is actually asking. Mutual information (in both its adjusted and normalized forms here) measures how much knowing one clustering tells you about the other — a score of 1.0 would mean the two clusterings are, structurally, the same partition of the data; a score of 0.0 would mean knowing one tells you nothing at all about the other. 0.29 mutual information says the two clusterings share some structure but are mostly independent. Jaccard similarity asks a stricter, more literal question — of the actual sentence-pairs grouped together by one method, what fraction are also grouped together by the other — and it is the more damning number here: 0.06, close to zero, meaning RW and HS are, for the most part, grouping the same sentences by genuinely different criteria. A companion Spearman-correlation analysis across nine different prompt templates finds the correlation between RW-based similarity and HS-based similarity is the lowest of any comparison in the study — 0.51, versus 0.63 for RW-vs-RW under different prompts and 0.52 for HS-vs-HS. RW is also markedly more stable across different prompt wordings than HS is, which matters in practice: a supposedly-fixed embedding that swings with how you phrase the instruction is a liability.

The nine prompts behind that stability check are not exotic — they are the same kind of one-line instruction template PromptEOL uses, aimed at nine different aspects of the sentence, so that the stability comparison is not accidentally measuring sensitivity to just one style of question. All nine, verbatim from the paper's own table:

Read across the list and notice what varies: sentiment, tone, intent, complexity, subjectivity, register, grammar — nine genuinely different lenses on the same sentence, not nine rewordings of the same question. That variety is what makes the stability result meaningful: RW's lower variance across this list is not “RW happens to like one particular phrasing,” it is RW staying comparatively steady across a real spread of different framings.

Swap only the instruction and re-embed the same sentence, and HS moves noticeably more than RW does — a concrete way to picture what “more stable” meant in the paragraph above. RW is being computed from the router's actual specialization structure, which does not reorganize itself just because the wrapper text asked a differently-worded question about the same underlying sentence; HS, still chasing whatever word comes next, is more easily redirected by a differently-worded question.

Where each one wins, in the paper's own examples

A qualitative case study bears this out with real sentence pairs the paper reports. HS correctly identifies pairs where the surface form changes but the structure barely does:

RW, by contrast, wins on pairs where the words change substantially but the underlying idea does not — paraphrase, not restatement:

Different failure modes, different strengths — textbook complementarity, not redundancy.

These are not the only two examples either — the paper's own case-study tables list five sentence pairs where HS wins and RW misses, and five more where RW wins and HS misses; the pair quoted above from each side is just the first row of its table. The paper's own summary of the pattern across all ten pairs matches what those two examples already suggested: HS embeddings excel at “capturing formal linguistic consistency, particularly when sentence structure undergoes only superficial changes,” while RW “performs better when handling paraphrasing, synonym use, and nuanced stylistic shifts” — cases where wording changes substantially but meaning does not.

Notice the shape of both examples: HS wins when the surface form (digits, punctuation, exact wording) changes while the underlying structure of the sentence barely does; RW wins when the surface form changes substantially while the underlying idea stays fixed. That is not a coincidence of these two particular examples, it is the direct consequence of what each signal was built to track — HS chases whatever token comes next, which is sensitive to exact phrasing; RW tracks which kind of input this is, at every layer, which is comparatively indifferent to phrasing as long as the meaning underneath stays the same.

Combining them: concatenate, or blend the scores?

The paper tries two strategies. Concatenation: e_final = [e_HS ; e_RW], simple, but it forces two embeddings with genuinely different native geometries into one shared vector space before any similarity metric ever looks at them — HS lives in an activation space shaped by next-token prediction, RW lives in something closer to a stacked probability-simplex, and nothing guarantees cosine similarity treats those two halves fairly once merged.

Weighted sum of similarities sidesteps that alignment problem entirely by never merging the vectors at all:

Each similarity score is computed inside the space where it already means something, and the combination step only has to blend two scalars — a far easier problem than aligning two incompatible vector geometries. Empirically, this wins.

The results, without any fine-tuning

Averaged over six MTEB task families (classification, clustering, pair classification, re-ranking, semantic similarity, summarization). Averages can hide as much as they reveal, so break DeepSeekMoE-16B's 35.36/35.91/40.03/43.30 row down into its six underlying task scores before trusting it:

The average's story holds up under this closer look, but with real texture: MoEE (sum) is the best method on five of the six task families, sometimes by a wide margin (Pair Classification: 52.12 vs HS's 44.34; STS: 48.34 vs HS's 34.54). Classification is the one exception the paper is upfront about — HS's 44.79 nearly matches MoEE (sum)'s 48.74, a much narrower gap than elsewhere, and the paper's own explanation is precise rather than hand-wavy: the final layer's hidden state is, by construction, aligned with whatever output-specific feature the model's next-token head reads off it, which happens to be exactly what a classification label needs. RW's whole advantage comes from capturing something HS does not capture well — so on the one task where HS's native strength is already well-matched to the job, adding RW has less room to help.

For DeepSeekMoE-16B, MoEE (sum) is a 22.45% relative gain over HS alone ((43.30−35.36)÷35.36). Add the PromptEOL prompting trick (a one-word-summary instruction template) on top and the gap widens further — DeepSeekMoE's improvement grows to 25.96%, and OLMoE's MoEE (sum) reaches an overall average of 55.16, ahead of even the fully-supervised SimCSE-BERT baseline's 53.53 — a byproduct of routing, with zero embedding-specific training, beating a model built and trained specifically to produce embeddings.

One more grounded, slightly surprising number from the paper's ablation: routing weights from the last token alone (STS average 61.18) already beat hidden states from the last token of the last layer (60.40) — before any combination at all. Averaging routing weights across every token in the sequence, meanwhile, hurts badly (46.03) — the same “the last position concentrates the useful signal” pattern that governs HS extraction turns out to hold for RW too, for reasons the paper leaves as future work.

Run the full 2×2 the paper reports and the picture sharpens further. HS has a fallback RW does not: pool across every layer instead of every token, keeping only the last token's representation at each depth and averaging those. That variant, “HS — last token, all layers,” scores 55.03 — worse than plain last-token-last-layer HS (60.40), suggesting the earlier layers' representations are still too tied to low-level, non-semantic features to help a similarity task once mixed in. Pool across both tokens and layers together and HS recovers slightly, to 56.34 — still short of RW's 61.18. No matter which axis you pool across, this specific ablation keeps landing on the same conclusion: for both signals, concentrate on the single last token, and reach for MoEE's combination rather than any form of averaging to do better than either alone.

One row of that same ablation table is worth calling out on its own, because it is the most dramatic number in it: “HS — all tokens, last layer” — ordinary mean pooling, the same recipe this lesson's own e_HS formula offered as an alternative to last-token extraction — scores only 32.78. That is not a small regression, it is the worst score anywhere in the table, worse than every other HS variant and worse than RW averaged across tokens (46.03). Mean-pooling a decoder-only model's hidden states at only the final layer combines two weaknesses at once: it averages away whatever sharp, last-token signal the model does carry, and it does so using only the shallowest useful semantic layer available. Put together with the earlier rows, the ablation's real message is not “pooling is mildly worse,” it is “pooling can be catastrophically worse, and which pooling axis you get wrong matters.”

The best row in the same table, for comparison, is MoEE (best) at 71.75 — the paper's own best-configuration combination of RW and HS on this specific STS12–16 slice, roughly 10 points above RW alone (61.18) and nearly 20 points above HS alone at its best (60.40). The gap between the worst row (32.78) and the best row (71.75) in one small table, built from the same underlying model with no retraining anywhere, is a concrete measure of how much extraction-method choice alone can matter — a bigger swing, here, than the gap you would expect from switching to a meaningfully different model entirely.

From a 2017 layer to today's frontier models

The exact noisy top-k gate from Chapter 2 and the exact importance/load losses from Chapter 5 are still the starting point for essentially every production Mixture-of-Experts language model running today — the specific numbers of experts, active counts, and additional tricks like shared experts and fine-grained splitting have moved on, and those specifics belong to their own lessons.

The 2017 paper's own closing line is worth reading now that Chapters 0–6 have earned it: “This work is the first to demonstrate major wins from conditional computation in deep networks… While we focused on text, conditional computation may help in other domains as well, provided sufficiently large training sets.” Be precise about what that sentence is and is not predicting: it is a bet on conditional computation spreading to other data modalities — images, audio, anything with enough training data — not a prediction that the gate itself would one day be repurposed for an entirely different job within the same modality. Chapter 8's discovery is a different, and in some ways stranger, kind of generalization than the one the original authors were imagining: not the mechanism moving to a new domain, but staying in exactly the same domain and turning out to already be computing something nobody had asked it for. Continue with Mixture of Experts for a from-scratch build of a modern MoE layer, or CS336 Lecture 4 for how these ideas scale into a 671B-parameter, 37B-active model. For a paper-grade read on that exact model, see the DeepSeek-V3 veanor, or the follow-on architecture in DeepSeek-V4. If Chapter 8–9 left you wanting the embeddings side of the story from the ground up, start with Vector Embeddings.