The Mixture-of-Experts Layer

"They built me a room with eight specialists and a receptionist who reads two words of your problem before deciding which two of us you get. I am the specialist nobody's two words ever point to."

An Expert Nobody Routed To

In Chapter 15 every worker held a full copy of the same dense model and the only thing that differed across machines was the data. A Mixture-of-Experts model breaks that symmetry: the model itself contains parts that most tokens never touch. To see how, we zoom all the way in to a single transformer block and change exactly one sublayer. Recall that a transformer block alternates an attention sublayer, which mixes information across positions, with a position-wise feed-forward network (FFN) that transforms each token independently. The FFN is where Mixture-of-Experts does its work, because it is both the heaviest part of the block and the part that already treats every token on its own.

1. From One Feed-Forward Network to Many Experts Beginner

The dense feed-forward sublayer applies the same two-layer network to every token vector $x \in \mathbb{R}^{d}$, expanding to a hidden width $h$ and projecting back: $\mathrm{FFN}(x) = \sigma(x W_1) W_2$, with $W_1 \in \mathbb{R}^{d \times h}$ and $W_2 \in \mathbb{R}^{h \times d}$. Because $h$ is typically four times $d$, this one sublayer holds roughly two thirds of the parameters in a transformer block, and every token pays for all of them on every forward pass. That is the cost a Mixture-of-Experts layer attacks.

The construction is direct. Instead of one FFN, install $E$ separate feed-forward networks, called experts, each with its own weights $W_1^{(e)}, W_2^{(e)}$ and identical shape. Add a small gating network $g$ that reads the token and decides which experts should handle it. The layer's output for a token is then a weighted sum over the experts the gate selected, not over all of them. With $E$ experts the layer stores $E$ times as many feed-forward parameters as a dense block, yet if the gate picks only $k$ of them per token, the compute per token rises by a factor of $k$, not $E$. Choosing $k = 2$ while $E = 64$ buys thirty-two times the feed-forward capacity for twice the feed-forward compute.

2. The Gate: Scoring and Top-k Selection Intermediate

The gating network is deliberately tiny, usually a single linear layer. For a token $x$ it produces one score per expert, $s = x W_g$ with $W_g \in \mathbb{R}^{d \times E}$, giving a vector of $E$ logits. The gate then keeps the $k$ largest of those logits, the token's top-$k$ experts, and discards the rest. Crucially, the softmax that turns scores into combination weights is taken over only the selected experts, not all $E$, so the weights of the chosen experts sum to one and the unchosen experts contribute nothing. Writing $\mathcal{T}(x)$ for the set of top-$k$ expert indices, the gate weight for a selected expert $e$ is

and $G_e(x) = 0$ for every expert outside $\mathcal{T}(x)$. The layer output is the gate-weighted combination of the selected experts' feed-forward outputs,

This is the gated combination at the heart of the layer, and it is worth reading carefully. The sum runs over only $k$ experts, so only $k$ feed-forward networks are ever evaluated for this token. The weights $G_e(x)$ are produced by the gate from the same token, so the routing is learned end to end: the gradient flows back through $G_e(x)$ into $W_g$, teaching the gate which experts serve which tokens. We treat the design of that learning signal, including how to stop the gate from collapsing onto a few favorite experts, as its own subject in Section 17.3. Here the gate is just the mechanism that selects and weights.

3. Why the Feed-Forward Sublayer Is the Place to Sparsify Intermediate

One might ask why the experts replace the feed-forward sublayer rather than the attention sublayer. Three properties of the FFN make it the natural target. First, it carries the parameters: with hidden width $h \approx 4d$, the two feed-forward matrices dominate the block's weight count, so multiplying that sublayer by $E$ is where extra capacity actually lands. Second, it is already position-wise. The FFN transforms each token independently of the others, so routing each token to its own experts changes nothing about the layer's structure; attention, by contrast, mixes tokens together, and splitting it per token would break the very communication it exists to provide. Third, the routing decision is cheap relative to the FFN it gates. A single linear gate over $E$ experts costs $d \cdot E$ multiplies per token, negligible beside the $\mathcal{O}(d h)$ of even one expert, so selection adds almost nothing to the bill it controls.

The result is a layer that is a drop-in replacement: the surrounding attention, residual connections, and normalization are untouched, and only the FFN becomes a routed set of FFNs. A common pattern places an MoE layer in every other transformer block, leaving the rest dense, so a model is sparse where capacity pays off and dense where it does not. This is the same partition-the-heavy-part instinct that drove sharding in Chapter 16; there we split one big matrix across devices, here we split capacity across many specialized matrices and route to a few.

4. Shared and Routed Experts Intermediate

Pure top-$k$ routing has a subtle inefficiency: knowledge that every token needs, basic grammar, common formatting, must be relearned and duplicated inside many experts, because no single expert is guaranteed to be on the route for any given token. Several recent designs fix this by splitting the experts into two kinds. A small number of shared experts are always active, run for every token regardless of the gate, and absorb the common knowledge. The remaining routed experts are selected top-$k$ as before and specialize. The layer output adds the shared contribution to the routed combination,

where $\mathcal{S}$ is the fixed set of shared experts and $\mathcal{T}(x)$ the token's routed top-$k$. The DeepSeek family popularized this shared-plus-routed structure, isolating common patterns in the always-on experts so the routed experts are free to specialize harder. The cost is modest, one or two extra FFN evaluations per token for the shared experts, and in exchange the routed experts stop wasting capacity on redundancy. We will see in Section 17.3 that this also stabilizes the gate, since the shared path keeps gradients flowing even to tokens whose routing is briefly poor.

5. The Forward Pass, Token by Token Intermediate

It is worth tracing the whole forward pass once at the level of a single token, because every later complication is a distributed version of these four steps. A token vector arrives. The gate scores all $E$ experts and selects the top $k$ (gate). Those $k$ expert FFNs are evaluated on the token, and no others (select and dispatch). Their outputs are scaled by the gate weights and summed (combine). The result replaces what the dense FFN would have produced, and the block continues as usual. The code below builds exactly this on a small batch with pure NumPy, four experts and top-2 routing, so the routing and the combine are fully visible rather than hidden behind a framework.

import numpy as np

rng = np.random.default_rng(7)
T, d, h, E, k = 6, 8, 16, 4, 2        # tokens, model dim, hidden dim, experts, top-k

# One small FFN expert: x -> relu(x W1) W2, each expert has its OWN weights.
W1 = rng.standard_normal((E, d, h)) * 0.3   # per-expert first layer
W2 = rng.standard_normal((E, h, d)) * 0.3   # per-expert second layer

def expert_forward(e, x):                    # FFN of expert e on row-batch x
    return np.maximum(x @ W1[e], 0.0) @ W2[e]

# The gate: a single linear scorer shared by all tokens, one logit per expert.
Wg = rng.standard_normal((d, E)) * 0.5
X = rng.standard_normal((T, d))              # a batch of T token vectors

logits = X @ Wg                              # (T, E) raw expert scores per token
# Top-k selection: keep the k highest logits per token, softmax over JUST those.
topk_idx = np.argsort(-logits, axis=1)[:, :k]          # (T, k) chosen expert ids
topk_logits = np.take_along_axis(logits, topk_idx, 1)  # (T, k) their scores
g = np.exp(topk_logits - topk_logits.max(1, keepdims=True))
gate_w = g / g.sum(1, keepdims=True)                   # (T, k) softmax gate weights

# Forward pass: each token is dispatched ONLY to its k experts, then combined.
Y = np.zeros((T, d))
for t in range(T):
    for j in range(k):
        e = topk_idx[t, j]
        Y[t] += gate_w[t, j] * expert_forward(e, X[t:t+1])[0]

# A dense baseline would run ALL E experts for every token. Count the saving.
sparse_calls = T * k
dense_calls = T * E

print("tokens T, experts E, top-k :", T, E, k)
print("expert-FFN evals (sparse)  :", sparse_calls)
print("expert-FFN evals (dense)   :", dense_calls,
      f"  -> {sparse_calls/dense_calls:.0%} of dense work")
print()
print("per-token routing (expert: gate-weight):")
for t in range(T):
    parts = ", ".join(f"E{topk_idx[t,j]}:{gate_w[t,j]:.2f}" for j in range(k))
    print(f"  token {t}  ->  {parts}")
print()
print("output norms per token     :", np.round(np.linalg.norm(Y, axis=1), 3))

tokens T, experts E, top-k : 6 4 2
expert-FFN evals (sparse)  : 12
expert-FFN evals (dense)   : 24   -> 50% of dense work

per-token routing (expert: gate-weight):
  token 0  ->  E2:0.70, E1:0.30
  token 1  ->  E1:0.85, E3:0.15
  token 2  ->  E3:0.76, E2:0.24
  token 3  ->  E3:0.79, E1:0.21
  token 4  ->  E3:0.93, E0:0.07
  token 5  ->  E2:0.54, E0:0.46

output norms per token     : [1.197 1.095 2.692 1.186 1.313 2.454]

The output makes the layer's character concrete. Token 0 leans on expert 2 with weight 0.70 and expert 1 with 0.30; token 4 sends almost all of its weight, 0.93, to expert 3. No token runs all four experts, the evaluation count is exactly $T \times k = 12$ rather than $T \times E = 24$, and each token's output is a convex combination of just its two specialists. Notice already a hint of the load problem this chapter must solve: expert 3 appears on four of the six tokens while some experts appear rarely, an imbalance that becomes a real cost once experts sit on separate machines and a popular expert's machine is swamped. That is the thread Section 17.6 picks up with load-balancing losses.

import torch, torch.nn as nn, torch.nn.functional as F

class MoE(nn.Module):
    def __init__(self, d, h, E, k):
        super().__init__()
        self.k = k
        self.gate = nn.Linear(d, E, bias=False)                 # the router
        self.experts = nn.ModuleList(                           # E independent FFNs
            nn.Sequential(nn.Linear(d, h), nn.ReLU(), nn.Linear(h, d))
            for _ in range(E))

    def forward(self, x):                                       # x: (tokens, d)
        w, idx = self.gate(x).topk(self.k, dim=-1)              # top-k scores + ids
        w = F.softmax(w, dim=-1)                                # gate weights over k
        y = torch.zeros_like(x)
        for j in range(self.k):                                 # one slot at a time
            for e in range(len(self.experts)):
                m = idx[:, j] == e                              # tokens routed to e
                if m.any():
                    y[m] += w[m, j:j+1] * self.experts[e](x[m])
        return y

6. What the Layer Sets Up for the Rest of the Chapter Beginner

We now have the anatomy: an MoE layer is $E$ expert FFNs, a linear gate, top-$k$ selection, and a gate-weighted combine, optionally with a few always-on shared experts. Two questions follow immediately, and each owns a later section. First, how should the gate learn to route, and how do we keep it from sending nearly every token to the same few experts? That is routing and gating, the subject of Section 17.3. Second, what happens when the experts are too numerous to fit on one machine and must be spread across the cluster? Then a token's trip to its experts crosses the network, the combine becomes a collective, and load imbalance becomes a hot-machine problem. That is expert parallelism, developed in Sections 17.4 and 17.5, where the all-to-all foreshadowed above becomes the engine of token routing. The local layer in this section is the unit that those sections replicate, shard, and connect.