Mixture of Experts

Decoupling parameters from compute

In a dense transformer every token passes through every weight, so cost per token grows with parameter count. Scaling laws (7.9 Scaling laws and compute-optimal training) say more parameters mean lower loss — but we pay for all of them on every token. A mixture of experts (MoE) breaks the link: replace the FFN of each block (7.4 Inside the block: residuals, LayerNorm, FFN and embeddings) with $E$ expert FFNs and a small router that sends each token to only $k$ of them. The model has the capacity of $E$ FFNs but the per-token compute of $k$. Shazeer et al. (2017) made this work at scale inside LSTMs; Switch Transformer (2021) brought it to transformers with trillion-parameter models; Mixtral 8×7B (2023) made it popular in open weights; DeepSeek-V3 (2024) showed a 671B-parameter MoE could match dense frontier models while activating only 37B parameters per token.

The router

For token vector $\mathbf{x}$ the router computes gate probabilities with a linear layer $\mathbf{W}_r\in\mathbb{R}^{E\times d}$:

Switch uses $k=1$ (one expert per token — the simplest routing and communication), Mixtral $k=2$ of $E=8$, DeepSeek-V3 $k=8$ of 256 small "fine-grained" experts plus one always-on shared expert for common knowledge. The $\operatorname{TopK}$ selection is not differentiable, but the gate values $g_i$ that scale the chosen experts' outputs are, and that is how the router learns which experts suit which tokens.

Load balancing

Left alone, routers collapse: a few experts receive most tokens, train faster, and get chosen even more. Switch adds an auxiliary loss that pushes the token distribution toward uniform. With $f_i$ the fraction of tokens dispatched to expert $i$ and $P_i$ the mean router probability of expert $i$ over the batch,

which is minimized when both are uniform ($f_i = P_i = 1/E$ gives $\mathcal{L}_{\text{aux}}=\alpha$), with $\alpha\approx 0.01$ so it does not dominate the language-modelling loss. Hardware adds a capacity factor: each expert's buffer holds a fixed number of tokens per batch, and overflow tokens are dropped (they pass through the residual connection only). DeepSeek-V3 replaced the auxiliary loss with per-expert bias terms adjusted online to balance load without distorting the gradient.

Math check. If one expert takes all tokens, $f_1 = 1$ and $P_1\to 1$, so $\mathcal{L}_{\text{aux}}\to \alpha E$ — $E$ times larger than at balance. The loss grows with imbalance exactly as intended.

Experts live on different GPUs

With hundreds of experts, no single device holds them all. Expert parallelism places experts on different GPUs and routes tokens between them with all-to-all communication — the dominant cost of MoE training and the reason routing is kept simple. Inference needs all experts in memory even though each token uses few: Mixtral 8×7B has about 47B parameters (≈94 GB in fp16) but only 12.9B active per token, so it needs big-GPU memory yet runs at the speed of a 13B model. Memory, not compute, is the MoE tax.

Intuition. A dense model is one generalist doing every job. An MoE is a large staff of specialists plus a receptionist (the router) who sends each token to the two most relevant people. You pay everyone's salary (memory) but each token takes only two people's time (compute).

Parameter counting with experts

Per block, attention stays $4d^2$ (7.2 Scaled dot-product and multi-head attention, fully derived); the FFN becomes $E$ FFNs. Mixtral: $d=4096$, $d_{ff}=14{,}336$ with SwiGLU → one expert has $3\times 4096\times 14{,}336 = 176$M parameters; eight experts are 1.41B per block, plus 67M for attention; 32 blocks give about 47B total. Active per token: attention plus two experts ≈ 0.42B per block → about 13B. "Active parameters" predicts latency and compute; "total parameters" predicts memory and, roughly, knowledge capacity. MoE models follow their own scaling laws: better loss per training FLOP than dense models, worse loss per parameter.

A tiny MoE layer

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

class MoE(nn.Module):
    def __init__(self, d, d_ff, n_experts=8, k=2):
        super().__init__()
        self.k = k
        self.router = nn.Linear(d, n_experts, bias=False)
        self.experts = nn.ModuleList([nn.Sequential(nn.Linear(d, d_ff), nn.GELU(), nn.Linear(d_ff, d))
                                      for _ in range(n_experts)])

    def forward(self, x):                                   # x: (tokens, d)
        gates = F.softmax(self.router(x), dim=-1)           # (tokens, E)
        topv, topi = gates.topk(self.k, dim=-1)             # (tokens, k)
        topv = topv / topv.sum(-1, keepdim=True)            # renormalize over the chosen experts
        y = torch.zeros_like(x)
        for e, expert in enumerate(self.experts):           # production code groups tokens per expert
            rows, slot = (topi == e).nonzero(as_tuple=True)
            if rows.numel():
                y[rows] += topv[rows, slot, None] * expert(x[rows])
        return y, gates

The loop over experts is the "dispatch": real systems sort tokens by expert, run one batched matrix multiplication per expert, and scatter the results back. Add the load-balancing loss from gates and a count of tokens per expert.

Trade-offs and 2026 practice

MoE wins when you are compute-bound in training and memory-rich in serving — the situation of frontier labs, which is why most frontier models and many open models are MoEs. Costs: training instability from routers, sensitivity during fine-tuning (small datasets can unbalance experts), communication overhead, and the memory footprint for inference. Upcycling (initializing experts from a dense checkpoint) lowers the entry cost; expert-choice routing (experts pick tokens, guaranteeing balance) is an alternative to token-choice; shared plus fine-grained experts are the current design consensus.

Where this shows up later

• MoE is the main reason "parameters" and "cost" diverged in the model landscape — 8.1 The LLM lifecycle and model landscape; it interacts with scaling laws (7.9 Scaling laws and compute-optimal training) and serving memory (7.10 Efficient inference: KV cache, GQA, FlashAttention, speculative decoding, quantization).
• Fine-tuning MoE models with LoRA has its own quirks (adapters on experts vs attention) — 8.2 Fine-tuning and PEFT.
• Routing a token to specialists is the same idea as the routing workflow pattern for agents — 10.5 Agentic workflow patterns.