Module 11: Mixture of Experts

Introduction

Every model we have built so far is dense: each token is pushed through every single parameter. Double the parameters and you double the compute for every token. That is the wall the frontier ran into — and Mixture of Experts (MoE) is the most successful way around it.

An MoE replaces the one feed-forward network in a transformer block with many parallel expert FFNs plus a small router that sends each token to only a few of them. The model can then hold a huge number of parameters — all the experts — while each token only pays for the handful it actually uses. Total parameters and per-token compute are decoupled.

This is not a curiosity. Mixtral 8×7B (8 experts, top-2 routing), DeepSeek-MoE, and several frontier flagships are MoE models: they serve the quality of a very large model at the cost of a much smaller one.

What You’ll Learn

  • Why total parameters ≫ active parameters is the entire point of MoE
  • How a top-k router scores experts and dispatches each token
  • How to build a sparse MoELayer from scratch (gather → run → scatter)
  • Why routers collapse and how a load-balancing loss prevents it
  • The sparsity factor E/k and the memory-vs-compute trade-off it encodes

Prerequisites

Intuition: A Dense FFN Runs Everything, Every Time

Look back at the transformer block. Its feed-forward network is where most of the parameters live — and it fires in full for every token, whether the token is the word “the” or a rare technical term. Intuitively that is wasteful: different tokens want different transformations, but a dense FFN forces them all through the same one.

MoE turns that single FFN into a committee of experts and hires a router to pick which experts see each token:

  • Experts — several independent FFNs (say 8). Each can specialize.
  • Router — a tiny linear layer that scores all experts for a token and keeps the top-k (often just 1 or 2).
  • Combine — the chosen experts’ outputs are summed, weighted by the router.

Because only k of the E experts run per token, the compute is that of a k-expert FFN, no matter how many experts the model owns. Add experts to grow capacity; keep k fixed to keep the bill flat.

NoteKey Insight

MoE decouples capacity from compute. A dense model ties them together — more parameters means more work per token. An MoE with E experts and top-k routing has E experts’ worth of parameters but does k experts’ worth of work per token. The ratio E/k is the sparsity factor: free capacity, paid for in memory rather than FLOPs.

The Math: Routing and Load Balancing

Routing. For a token \(x\), a linear gate produces one score per expert; a softmax turns the scores into probabilities, and we keep the top \(k\):

\[g(x) = \text{softmax}(x W_g) \in \mathbb{R}^E, \qquad \mathcal{T} = \text{top-}k(g(x))\]

The layer’s output sums the chosen experts, weighted by their (renormalized) gate values:

\[\text{MoE}(x) = \sum_{i \in \mathcal{T}} \tilde{g}_i(x)\, E_i(x), \qquad \tilde{g}_i = \frac{g_i}{\sum_{j \in \mathcal{T}} g_j}\]

where \(E_i\) is the \(i\)-th expert FFN. Only the \(k\) experts in \(\mathcal{T}\) ever run.

Load balancing. Left to itself the router collapses — it discovers a few strong experts early and sends everything to them, starving the rest (a starved expert gets no gradient and effectively dies). To counter this, we add an auxiliary loss (Shazeer 2017; Switch Transformer, Fedus 2021). Over a batch of \(N\) tokens routed top-\(k\):

\[\mathcal{L}_{\text{aux}} = E \cdot \sum_{i=1}^{E} f_i \, P_i\]

where \(f_i\) is the fraction of dispatch slots sent to expert \(i\) and \(P_i\) is the mean router probability for expert \(i\). Both vectors sum to 1, so the loss is minimized at uniform routing (\(f_i = P_i = 1/E\)), giving \(\mathcal{L}_{\text{aux}} = 1\). Any imbalance pushes it above 1. Training minimizes \(\mathcal{L}_{\text{task}} + \alpha\,\mathcal{L}_{\text{aux}}\) with a small \(\alpha\), so the router is nudged to keep every expert fed.

Code: An Expert and a Router from Scratch

Everything lives in moe.py. First we load it and build the two pieces — an Expert (an ordinary GELU FFN) and the top-k MoERouter.

import importlib.util
import sys
from pathlib import Path

import torch

# Load moe.py (directory name starts with a digit, so use importlib)
spec = importlib.util.spec_from_file_location("moe", Path("moe.py").resolve())
moe = importlib.util.module_from_spec(spec)
sys.modules["moe"] = moe
spec.loader.exec_module(moe)

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"PyTorch {torch.__version__} on {device}")

router = moe.MoERouter(embed_dim=64, num_experts=8, top_k=2)
x = torch.randn(5, 64)                      # 5 tokens
out = router(x)
print(f"probs shape:        {tuple(out.probs.shape)}   (softmax over 8 experts)")
print(f"chosen experts:     {out.topk_indices.tolist()}")
print(f"chosen weights:     {out.topk_weights.round(decimals=3).tolist()}")
print(f"weights sum to 1?   {out.topk_weights.sum(dim=-1).round(decimals=4).tolist()}")
PyTorch 2.12.1+cu130 on cpu
probs shape:        (5, 8)   (softmax over 8 experts)
chosen experts:     [[5, 2], [4, 0], [4, 7], [5, 2], [7, 4]]
chosen weights:     [[0.503000020980835, 0.4970000088214874], [0.5040000081062317, 0.4959999918937683], [0.5149999856948853, 0.48500001430511475], [0.5600000023841858, 0.4399999976158142], [0.5109999775886536, 0.48899999260902405]]
weights sum to 1?   [1.0, 1.0, 1.0, 1.0, 1.0]

Each token picked its 2 best experts, and the two kept weights were renormalized to sum to 1 — the discarded probability mass is simply dropped.

Code: The Sparse MoE Layer

The layer ties experts and router together. The trick is the dispatch: for each expert, gather exactly the tokens routed to it, run the expert once on that batch, then scatter-add the weighted results back. That is the same gather → run → scatter pattern real MoE kernels use to avoid running every expert on every token.

moe_layer = moe.MoELayer(embed_dim=64, ff_dim=256, num_experts=8, top_k=2)
x = torch.randn(2, 16, 64)                  # (batch, seq, embed)
output, aux = moe_layer(x)

print(f"input  shape: {tuple(x.shape)}")
print(f"output shape: {tuple(output.shape)}   (same as input — a drop-in FFN)")
print(f"load-balancing aux loss: {aux.item():.4f}   (1.0 = perfectly balanced)")
input  shape: (2, 16, 64)
output shape: (2, 16, 64)   (same as input — a drop-in FFN)
load-balancing aux loss: 1.0143   (1.0 = perfectly balanced)

The forward pass returns two things: the output and the load-balancing loss for this batch. During training you add alpha * aux to your task loss.

Verifying the dispatch math. With top_k=1, each token’s output must be exactly its one chosen expert applied to it (its weight renormalizes to 1). We can check that directly:

layer1 = moe.MoELayer(embed_dim=32, ff_dim=64, num_experts=4, top_k=1)
layer1.eval()
x = torch.randn(1, 6, 32)
out1, _ = layer1(x)

flat = x.reshape(-1, 32)
route = layer1.router(flat)
manual = torch.stack([
    layer1.experts[route.topk_indices[t, 0]](flat[t:t+1])[0]
    for t in range(flat.shape[0])
])
print("top-1 output matches the chosen expert exactly:",
      torch.allclose(out1.reshape(-1, 32), manual, atol=1e-5))
top-1 output matches the chosen expert exactly: True

Total vs Active Parameters

Here is the payoff, made concrete. count_parameters() reports what the layer owns versus what a single token touches:

for E in (2, 8, 32):
    c = moe.MoELayer(embed_dim=512, ff_dim=2048, num_experts=E, top_k=2).count_parameters()
    print(f"E={E:2d}, k=2 → total {c['total']/1e6:6.1f}M | "
          f"active {c['active']/1e6:5.1f}M | sparsity {c['sparsity_factor']:.1f}×")
E= 2, k=2 → total    4.2M | active   4.2M | sparsity 1.0×
E= 8, k=2 → total   16.8M | active   4.2M | sparsity 4.0×
E=32, k=2 → total   67.2M | active   4.2M | sparsity 15.9×

Going from 2 to 32 experts multiplies total parameters ~16× while the active parameters per token barely move. Drive the same trade-off yourself:

TipTry This
  1. Add experts, watch compute stay flat. Slide E from 8 to 64 in the sparsity meter. The total bar grows; the active bar barely moves. That gap is capacity you get almost for free (in FLOPs — you still pay in memory).
  2. Raise top-k. Push k from 1 to 4. Active compute climbs and the sparsity factor E/k shrinks — top-k is the compute dial.
  3. Break the router. In the routing simulator, drag the skew up. Tokens pile onto a few experts, the others starve, and the load-balancing loss climbs far above 1.0. This is the collapse the auxiliary loss exists to prevent.
  4. Balance it back. Return skew to 0 and watch every bar settle near the 1/E line and the loss fall back to ≈ 1.0.

Common Pitfalls

  1. Forgetting the auxiliary loss. Without it the router collapses onto a few experts within a few hundred steps, and the rest become dead weight. Always add alpha * aux_loss (a common \(\alpha\) is 0.01).

  2. Confusing total with active parameters. An “8×7B” model does not run 56B parameters per token — with top-2 it runs about 13B. Report both numbers; they answer different questions (memory vs. FLOPs).

  3. Renormalizing (or not) the top-k weights. After picking the top-k, their gate values no longer sum to 1. Renormalize them (as here) so the combined output is a proper weighted average and its scale is stable.

  4. Training instability from hard routing. The top-k operation is non-differentiable in the choice itself; gradients flow only through the kept weights. This makes MoE routers more finicky to train than a dense FFN — the load-balancing loss and careful initialization matter.

  5. Assuming MoE is always a win. MoE trades FLOPs for memory and bandwidth. If you are memory-bound (a single GPU, long context), holding many experts can cost more than the compute it saves.

Exercises

Exercise 1: Sparsity of a real config

Mixtral 8×7B uses E=8, top_k=2. What is its sparsity factor? Roughly how many parameters are active per token if the total is ~47B and the non-expert (attention + embeddings) share is ~2B?

# E/k sparsity factor; active ≈ non_expert + (k/E) * expert_params
total, non_expert = 47e9, 2e9
expert_params = total - non_expert
active = non_expert + (2 / 8) * expert_params
print(f"sparsity factor E/k = {8/2:.0f}×")
print(f"active params per token ≈ {active/1e9:.1f}B  (of {total/1e9:.0f}B total)")
sparsity factor E/k = 4×
active params per token ≈ 13.2B  (of 47B total)

Exercise 2: Watch a router collapse

Build a layer, feed it the same few tokens many times, and skew the gate by hand to one expert. Print expert_utilization before and after — most of the mass should move to a single expert.

Exercise 3: The load-balancing floor

Show numerically that load_balancing_loss bottoms out at 1.0. Construct uniform probs and round-robin topk_indices for E experts and confirm the loss is ≈ 1.0; then concentrate the routing and confirm it rises.

Summary

Key takeaways:

  1. MoE decouples capacity from compute. Many expert FFNs, a router that keeps only the top-k per token: total parameters ≫ active parameters.

  2. The router is a tiny top-k gate. A single linear layer scores experts; softmax + top-k selects them; the kept weights are renormalized to sum to 1.

  3. Dispatch is gather → run → scatter. Each expert runs once, on exactly the tokens routed to it — never all experts on all tokens.

  4. Routers collapse without help. A load-balancing auxiliary loss \(E\sum_i f_i P_i\) (minimum 1.0 at uniform routing) keeps every expert fed.

  5. The sparsity factor is E/k. It is free capacity paid for in memory, not FLOPs — the trade-off behind Mixtral, DeepSeek-MoE, and modern frontier MoEs.

What’s Next

We now have both halves of the modern FFN story: the dense gated FFN (SwiGLU, in m06) and the sparse one (MoE, here). The frontier keeps building on these — MoE routing variants (expert choice, shared experts), long-context attention, and the alignment and reasoning stages that turn a pretrained model into an assistant.

Going Deeper

  • Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (2017) — the modern MoE layer and its balancing loss. https://arxiv.org/abs/1701.06538
  • Fedus, Zoph & Shazeer, Switch Transformers (2021) — top-1 routing at scale; the load-balancing loss used here. https://arxiv.org/abs/2101.03961
  • Jiang et al., Mixtral of Experts (2024) — a strong open MoE (8 experts, top-2), with the total-vs-active numbers this lesson quotes. https://arxiv.org/abs/2401.04088