Mixture of Experts (MoE)

Many experts, a few at a time

The feed-forward layer holds most of a Transformer's parameters. A mixture-of-experts (MoE) layer keeps several copies, each with its own weights, called experts, and adds a router: a small linear layer that scores each expert for the current token.

Mixtral's router computes $G(x) = \text{Softmax}(\text{TopK}(x \cdot W_g))$: it keeps the top-K expert scores, sets the others to minus infinity, and takes a softmax, so only K experts receive non-zero weight Established. The layer's output is

$$y = \sum_{i \in \text{top-}K} G(x)_i \, E_i(x).$$

Tiny example. Eight experts, top-2. A token's router scores are $[1.2, -0.3, 2.0, 0.1, 0.4, -1.0, 1.9, 0.0]$. The top two are expert 3 (2.0) and expert 7 (1.9). Softmax over just those: $e^{2.0}/(e^{2.0}+e^{1.9}) \approx 0.52$ and $0.48$. The output is $0.52\,E_3(x) + 0.48\,E_7(x)$; the other six experts are never run for this token. The next token, or the same token in the next layer, can choose different ones.

Total versus active parameters

Shazeer and colleagues reported over 1,000× improvements in model capacity with only minor losses in computational efficiency, using MoE layers of up to 137 billion parameters between LSTM layers Established. Mixtral 8x7B has 8 experts per layer and routes each token to 2; each token has access to 47B parameters but uses only 13B during inference, and Mixtral matched or outperformed Llama 2 70B and GPT-3.5 on the benchmarks reported Established.

The catch is memory. Mixtral's authors note that the memory cost of serving it is proportional to its 47B total parameters, still smaller than Llama 2 70B, and that routing adds overhead and extra memory loads Established. Every expert must sit in GPU memory even though each token uses two.

DeepSeek-V3 has 671B total parameters with 37B active per token; each MoE layer has 1 shared expert and 256 routed experts, of which 8 are activated per token Established.

Keeping the experts busy

Left alone, a router tends to favour a few experts, which then get most tokens while others idle and learn little. The Switch Transformer routes each token to a single expert and adds an auxiliary load-balancing loss, minimized when tokens are spread uniformly across experts, with a coefficient of 0.01 Established. Each expert also has a capacity, a maximum number of tokens per batch. In Switch, tokens beyond an expert's capacity are dropped from that layer: their computation is skipped and the representation passes on through the residual connection Established. With these changes Switch models pretrained up to 7× faster than T5-Base with the same compute Established. DeepSeek-V3 instead balances load without an auxiliary loss Established.

Token:catsatthe

1. Score every expert2. Keep the top 23. Softmax over those two4. Run only those experts

1. Expert 11.2idle
2. Expert 2-0.3idle
3. Expert 30.52runs
4. Expert 40.1idle
5. Expert 50.4idle
6. Expert 6-1.0idle
7. Expert 70.48runs
8. Expert 80.0idle

Output = 0.52 × expert 3 + 0.48 × expert 7. Six experts do no work for this token, but all eight stay in memory. Try another token: it can choose different experts.

Mixtral 8x7B, per token

stored (memory)47B

used (compute)13B

Router scores are made up for illustration; the top-2 selection and softmax are computed. Mixtral’s parameter counts are from its paper.

The hardware reality

Experts are usually spread across GPUs (expert parallelism), so each MoE layer sends tokens to the GPUs that hold their experts and gathers the results back. That all-to-all communication, the uneven load and the memory for idle experts are the costs of the saved arithmetic. MoE pays off when compute per token is the bottleneck and there is memory and network bandwidth to spare.