Why Mixture-of-Experts?

Every dense transformer pays the same bill on every token. Whether the model is completing a Python function, translating Bengali, or finishing a Shakespeare line, the same parameters fire. As models grew from millions to billions to nearly a trillion parameters, that bill became unpayable. Mixture-of-Experts is the architectural answer to a brutally simple question: if most of the knowledge inside a giant model is irrelevant to any one token, why are we computing with all of it?

The bet of MoE. Grow the model's knowledge by 8× while keeping the compute per token roughly the same. You pay for capacity in memory, not in FLOPs.

The Dense Wall

Open up any decoder-only transformer block and roughly two-thirds of its parameters live inside one place: the feed-forward network (FFN). A standard FFN looks like $y = W_2 \, \mathrm{GELU}(W_1 x)$, where $W_1 \in \mathbb{R}^{d_{ff} \times d}$ and $W_2 \in \mathbb{R}^{d \times d_{ff}}$, with $d_{ff} \approx 4d$. The FFN holds most of the model's memorized knowledge, and it is also where most of the FLOPs go.

The trouble is what scaling actually buys you. To make a dense model 2× smarter, you roughly double the parameters, which doubles the FFN width, which doubles the FLOPs per token. Every doubling of capacity is a doubling of inference cost. By the time you reach a 70-billion-parameter dense model, every single token in every single forward pass touches every weight. Nothing is conditional. The model has no way to say this is a coding token, skip the Shakespeare weights.

The FFN is the elephant. It is also the layer where MoE intervenes — leaving attention untouched and replacing the single fat FFN with a population of smaller ones, only a few of which run per token.

The Specialist Intuition

Imagine a hospital where every patient is seen by every doctor — the cardiologist, the dermatologist, the orthopedic surgeon, the pediatrician — and each writes a small note that is then averaged. That is a dense FFN. It works, but it is absurdly wasteful: the dermatologist does not need to weigh in on a broken ankle.

A real hospital uses a triage desk. A nurse glances at the patient, picks the two most relevant specialists, and sends them in. The dermatologist still exists, still has years of training stored in her head, but she stays in her office unless you arrive with a rash. The hospital's total knowledge is huge; the active knowledge per patient is tiny.

Mixture-of-Experts is exactly that hospital. A small router examines each token and, based on what the token looks like in embedding space, picks the $k$ most relevant experts out of $E$. Only those $k$ experts run. The others sit idle, weights loaded but never multiplied. Common choices in real systems are $E \in \{8, 16, 64, 256\}$ with $k \in \{1, 2\}$.

Conditional Compute: The Math

A dense FFN computes a deterministic function of $x$: $\mathrm{FFN}(x) = W_2 \, \sigma(W_1 x)$.

An MoE FFN replaces that one function with a population of $E$ experts $\{E_1, E_2, \dots, E_E\}$, each one a small FFN with its own private $W_1^{(i)}, W_2^{(i)}$. A gate $g(x) \in \mathbb{R}^E$ produces a vector of weights — most of them zero — and the output is a weighted sum: $\mathrm{MoE}(x) = \sum_{i=1}^{E} g_i(x) \cdot E_i(x)$.

Here $g_i(x)$ is the gate value for expert $i$ on input $x$, and $E_i(x)$ is what expert $i$ would return on this input. If $g_i(x) = 0$, we never have to compute $E_i(x)$ — the term vanishes. The trick is to design $g$ so that most entries are exactly zero.

The router

The standard recipe is a single linear layer followed by top-$k$ selection and a softmax restricted to the survivors. Let $W_r \in \mathbb{R}^{E \times d}$ be the router weights. We compute logits, take the top $k$, and softmax only those:

$s = W_r x \in \mathbb{R}^{E}$, $\mathcal{T} = \mathrm{TopK}(s, k)$, $g_i(x) = \frac{e^{s_i}}{\sum_{j \in \mathcal{T}} e^{s_j}}$ if $i \in \mathcal{T}$, otherwise $g_i(x) = 0$.

Read the symbols carefully. $x$ is one token's hidden state. $W_r x$ is a cheap E-dimensional score vector — one score per expert. $\mathrm{TopK}$ picks the $k$ highest-scoring experts and discards the rest. The softmax over only those $k$ values ensures the gates form a proper convex combination ( $\sum_i g_i = 1$ ) without spending probability mass on experts we will not run.

The compute equation

Suppose a dense FFN has roughly $2 d \cdot d_{ff}$ parameters and the same order of FLOPs per token. An MoE layer with $E$ experts of the same size stores $E \cdot 2 d \cdot d_{ff}$ parameters but only runs $k$ of them per token, so FLOPs per token are $k \cdot 2 d \cdot d_{ff}$. The ratio is $\text{FLOPs}_{\text{MoE}} / \text{FLOPs}_{\text{dense}} = k / E$. With $E = 8, k = 2$ you keep 25% of the cost; with $E = 64, k = 2$ you keep just 3%.

Manual Numerical Walkthrough

Let us route one toy token through a 4-expert MoE with $k = 2$. We will compute every number by hand so the mechanism is fully visible.

Visualizing Sparse Routing

The interactive diagram below is the picture you should hold in your head every time you see the letters "MoE". Toggle between dense and MoE; slide $k$ up and down; switch which token is being routed. Watch which experts light up, which sit idle, and what happens to the compute bar.

Three observations are worth pausing on. First, when you switch tokens, the lit experts move. The router has learned (in our toy example by construction; in real models, through training) that coding tokens prefer some experts and poetry tokens prefer others. Second, the green compute bar tracks $k / E$ exactly — sliding $k$ from 8 to 1 drops cost from 100% to 12.5%. Third, the parameter count (purple) is independent of $k$. You always pay full memory; you only sometimes pay full compute.

Plain Python: An MoE Layer From Scratch

Before the PyTorch version, let us write the whole thing in NumPy so there is nowhere for the mechanism to hide. The code below is small enough to fit in your head and complete enough to actually run.

(4, 4)

If logits = [0.7, -0.2, 1.3, 0.4], topk_idx = [2, 0] (the 2 largest).

1import numpy as np
2
3num_experts, d_model, d_ff, k = 4, 4, 8, 2
4
5rng = np.random.default_rng(0)
6W1 = rng.standard_normal((num_experts, d_ff, d_model)) * 0.1
7W2 = rng.standard_normal((num_experts, d_model, d_ff)) * 0.1
8W_router = rng.standard_normal((num_experts, d_model)) * 0.1
9
10def moe_forward(x):
11    logits = W_router @ x                     # (E,)
12    topk_idx = np.argsort(-logits)[:k]        # indices of best k experts
13    topk_logits = logits[topk_idx]
14
15    z = topk_logits - topk_logits.max()
16    gates = np.exp(z) / np.exp(z).sum()       # (k,) sums to 1
17
18    y = np.zeros(d_model)
19    for g, i in zip(gates, topk_idx):
20        h = np.maximum(0, W1[i] @ x)          # expert i: ReLU(W1 x)
21        y += g * (W2[i] @ h)                  # weighted by gate
22    return y
23
24x = np.array([0.5, -0.2, 0.9, 0.1])
25print(moe_forward(x))

Every interesting thing about MoE is on the screen above. The router $W_r x$ is a single matvec. The top-k selection is one argsort. The softmax runs over $k$ numbers. The expensive part — actually evaluating an expert FFN — happens only inside the loop over $k$ survivors, not inside a loop over all $E$.

Sanity check. Set $k = E$ in the snippet above. The MoE layer becomes equivalent to a softmax-weighted average of all experts — exactly what you would get from a fully-dense ensemble. The whole compute saving comes from making $k \ll E$.

PyTorch: From Toy to Production

The PyTorch version moves from one-token-at-a-time to a real batch. The pattern we want is: route every token, then for each expert gather the tokens that picked it, run that expert once on a packed mini-batch, and scatter the result back. This keeps GPU matmuls large and idle FLOPs near zero.

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class MoEFFN(nn.Module):
6    def __init__(self, d_model: int, d_ff: int, num_experts: int, k: int):
7        super().__init__()
8        self.num_experts = num_experts
9        self.k = k
10        self.experts = nn.ModuleList([
11            nn.Sequential(
12                nn.Linear(d_model, d_ff),
13                nn.GELU(),
14                nn.Linear(d_ff, d_model),
15            )
16            for _ in range(num_experts)
17        ])
18        self.router = nn.Linear(d_model, num_experts, bias=False)
19
20    def forward(self, x: torch.Tensor) -> torch.Tensor:
21        B, T, D = x.shape
22        x_flat = x.reshape(B * T, D)                          # (N, D)
23
24        logits = self.router(x_flat)                          # (N, E)
25        topk_w, topk_idx = logits.topk(self.k, dim=-1)        # (N, k), (N, k)
26        topk_w = F.softmax(topk_w, dim=-1)                    # (N, k)
27
28        y = torch.zeros_like(x_flat)
29        for i in range(self.num_experts):
30            mask = topk_idx == i                              # (N, k)
31            if not mask.any():
32                continue
33            tok_rows, slot = mask.nonzero(as_tuple=True)
34            gate = topk_w[tok_rows, slot].unsqueeze(-1)       # (M, 1)
35            y[tok_rows] += gate * self.experts[i](x_flat[tok_rows])
36        return y.reshape(B, T, D)

Two things are subtly different from the NumPy version, and both matter for real training:

1. We loop over experts, not over tokens. If we looped per token we would issue thousands of tiny matmuls — terrible for GPU utilization. By gathering the tokens that picked each expert and running one large $(M, D) \to (M, D)$ matmul, we keep the GPU saturated.
2. topk is non-differentiable in its indices but differentiable in its values. The gradient flows back through the gate weights $g_i(x)$, which teaches the router to put more probability mass on experts that produced good outputs and less on those that did not. The discrete "which experts" decision is updated indirectly, as a side effect of the values flowing back.

What Changes at Massive Scale

The toy example has 4 experts and runs on a laptop. A frontier MoE language model has hundreds. Three numbers tell the story:

DeepSeek-V3 is the cleanest illustration. The model carries 671 billion parameters of knowledge but spends only 37 billion parameters of compute per token — a roughly 18× decoupling. That same decoupling applies during training: gradients only flow into the experts that were activated for each token, so the gradient compute is also $k / E$ of a same-size dense model.

The bottleneck shifts. With dense models, the limit was compute. With MoE, the limit becomes memory bandwidth and cross-GPU communication. Every parameter still has to live somewhere on the cluster; routing means tokens have to travel to whichever GPU owns their chosen experts. We will spend several sections of this chapter on the machinery that makes that travel cheap.

The scaling-law angle

Empirically, language-model loss scales as a power law in compute and parameters (Kaplan et al., 2020; Hoffmann et al., 2022). Dense models tie those two axes together. MoE lets you push the parameter axis while leaving compute roughly fixed — and the experimental evidence is that this still buys you a scaling-law gain in loss, just along a different curve. You do not get something for nothing (the memory and communication costs are real), but you do get to choose where to spend.

The Engineering Reality of MoE

The math above is clean. The systems are not. If you write the naive MoE in production, three things will go catastrophically wrong, and each of them is the subject of an entire later section in this chapter:

1. Load imbalance. The router's natural tendency, untreated, is to send most tokens to a small subset of popular experts. The unpopular experts never train; the popular ones bottleneck the GPU. This is routing collapse, and it is the central failure mode of MoE training (covered in Chapter 6).
2. Expert parallelism and all-to-all communication. If the model has 256 experts and your cluster has 256 GPUs, every token might need to travel to a different GPU. That cross-device shuffle is an all-to-all collective — by far the most expensive primitive in distributed training (covered in section 5 of this chapter).
3. Inference latency. Conditional compute is great for training throughput but adds an extra hop at inference. KV-cache savings still apply, but you now pay a routing decision per token, plus the all-to-all if experts are sharded. Serving stacks like vLLM and SGLang have specialized MoE schedulers; we will look at them in the inference chapter.

Despite these costs, the trade is overwhelmingly worth it. Once the routing, balance, and parallelism problems are solved well, MoE models reliably beat dense models of the same compute budget — and that is precisely why DeepSeek-V3, Mixtral, Qwen-MoE, and most of the frontier open-weight models of 2024-2025 are all sparse mixtures, not dense stacks.

For now, the one sentence to carry forward is this: MoE is the architectural move that broke the link between parameter count and compute cost, and every model with hundreds of billions of parameters running on anything less than a planet-scale datacenter is built on top of it.