24  Mixture of Experts

TipTL;DR

Scaling laws say more parameters help — but dense models pay for every parameter on every token. Mixture of Experts breaks that link: replace the feed-forward layer with \(E\) expert FFNs and a router that sends each token to only its top-\(k\) experts. The model has billions of parameters but activates only a small fraction per token: \[y = \sum_{i \in \text{top-}k} g_i(x)\, E_i(x).\] Mixtral 8×7B and DeepSeek-V3 show this delivers the quality of a much larger dense model at a fraction of the inference compute.

Depends on: Scaling Laws & Emergence

24.1 Why this matters

The scaling laws create a dilemma: quality tracks parameter count, but a dense model’s compute also tracks parameter count, so each step toward a better model gets proportionally more expensive to train and serve. MoE decouples the two — parameter count (model capacity) from active compute per token (cost). You can grow capacity ~10× while keeping per-token FLOPs roughly flat.

The FFN is the natural target: it holds most of a Transformer’s parameters (The Transformer) and runs independently per token, so swapping it for a sparse bank of experts is a clean substitution.

24.2 The mechanism

24.2.1 Sparse experts and a router

Replace one FFN with \(E\) expert FFNs \(E_1, \dots, E_E\) plus a lightweight router (gating network) \(g\). For each token \(x\), the router produces scores, keeps the top-\(k\) experts, softmaxes over those, and combines:

\[ g(x) = \operatorname{softmax}\big(\operatorname{top\text{-}k}(x W_g)\big), \qquad y = \sum_{i \in \text{top-}k} g_i(x)\, E_i(x). \]

  • Shazeer et al. (Shazeer et al., 2017) introduced the sparsely-gated MoE layer (top-\(k\), \(k>1\)) for LSTMs and early Transformers.
  • Switch Transformer (Fedus et al., 2022) simplified to top-1 routing — each token goes to exactly one expert — which maximises sparsity and simplifies the implementation.
  • Mixtral 8×7B (Jiang et al., 2024) brought MoE to a strong open decoder-only LLM: 8 experts, top-2 routing, ~47B total parameters but only ~13B active per token.
  • DeepSeek-V2/V3 (DeepSeek-AI, 2024) pushed fine-grained experts plus shared experts and large \(E\), reaching frontier quality at low active cost.
token x router top-k Expert 1 ✓ Expert 2 Expert 3 ✓ Expert E weighted sum → y Only the top-k experts (green) run for this token; the rest (grey) are skipped. Capacity ≫ active compute.

24.2.2 Load balancing: the central difficulty

Left alone, the router collapses: a few popular experts get all the tokens while the rest starve and never train. MoE training therefore adds an auxiliary load-balancing loss that penalises uneven expert usage, encouraging the router to spread tokens evenly across experts.

WarningPitfall: MoE saves compute, not memory

A sparse model activates few parameters per token, but all experts must be resident in memory (any token might route to any of them). Mixtral 8×7B needs the VRAM of a ~47B model while computing like a ~13B one. MoE trades memory (cheap, parallelisable across devices) for compute and quality — which is exactly the right trade at serving scale, but it is not a way to fit a big model on a small GPU.

TipGoing deeper

24.2.3 Minimal sketch

Top-2 routing over a bank of toy experts, with per-token gate weights.

import numpy as np

rng = np.random.default_rng(0)
d, E, k = 4, 4, 2
experts = [rng.normal(scale=0.3, size=(d, d)) for _ in range(E)]
W_g = rng.normal(size=(d, E))

def softmax(z):
    e = np.exp(z - z.max()); return e / e.sum()

def moe(x):
    scores = x @ W_g
    top = np.argsort(scores)[-k:]            # indices of top-k experts
    g = softmax(scores[top])                 # gate over chosen experts only
    return sum(g_i * (x @ experts[i]) for g_i, i in zip(g, top)), top

x = rng.normal(size=d)
y, chosen = moe(x)
print("experts used:", sorted(chosen), "  out of", E)   # 2 of 4 active
print("output dim:", y.shape)
experts used: [np.int64(0), np.int64(1)]   out of 4
output dim: (4,)

24.3 Application & impact

Idea Introduced by Lives on as
Sparsely-gated experts Shazeer 2017 The MoE layer
Top-1 routing Switch Maximal sparsity, simpler kernels
Open MoE LLM Mixtral Quality of a big model, cost of a small one
Shared + fine-grained experts DeepSeek-V2/V3 Frontier-quality efficient LLMs
  • MoE is how frontier models scale economically — capacity grows without per-token FLOPs growing, the dominant lever once dense scaling got too expensive to serve.
  • Routing is the new research surface — load balancing, expert specialisation, and routing stability are active problems unique to sparse models.
NoteKey takeaway

MoE decouples capacity from compute: many expert FFNs, but a router activates only top-\(k\) per token, so a model can be huge yet cheap to run. The price is memory (all experts resident) and a new failure mode (router collapse, fought with load-balancing losses). It is the efficiency story that keeps the scaling laws affordable.

→ Next: Beyond Quadratic Attention