# Module 17: Mixture of Experts

> Mixture-of-experts layers: top-k routing, many small experts, shared experts, the load-balancing loss, expert capacity, and how routing fails.

- Part IV: The Modern Architecture · elective · Advanced · 3 hours
- By the end you can design an MoE layer and predict how routing fails.
- Builds on: [Module 15: The LLaMA Block](https://transformers-and-math.drmhse.com/md/15-the-llama-block.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/17-mixture-of-experts
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com

A dense model uses every weight for every token. A *mixture-of-experts* (MoE) model splits its feed-forward layers into many smaller networks, *experts*, and lets each token use only a few. It can hold far more parameters, and so more knowledge, at the per-token compute of a much smaller dense model. The catch is routing: left alone, the choice of experts tends to collapse onto a few favourites. Examples follow Alibaba's Qwen family of open models.

## 17.1 Routing

An MoE block replaces its feed-forward network with many **experts**, and a small **router** network picks a few for each token; only those run. Qwen's MoE models, such as Qwen2-57B-A14B, add two refinements: many small experts, and shared experts every token uses.

```mermaid
graph TD
    X["Token vector x"] --> Router["Router scores every expert"]
    X --> Shared["Shared expert: runs for every token"]
    Router --> Route1["Keep the top k, e.g. 8 of 64"]
    Route1 --> ExpA["Expert 4"]
    Route1 --> ExpB["Expert 19"]
    Route1 --> ExpC["Expert 52"]
    ExpA --> Combine((Weighted sum))
    ExpB --> Combine
    ExpC --> Combine
    Shared --> AddExpert((+))
    Combine --> AddExpert
    AddExpert --> Out["Layer output"]
```

> **In words.** The router scores every expert and keeps the best few, whose outputs are summed, each weighted by its score. A shared expert runs for every token and handles common material, so the routed experts can specialize.

## 17.2 Many Small Experts

Older MoE models used a few large experts, such as 8 with 2 active. Qwen splits the same capacity into many smaller ones, such as 64 with 8 active:

$$
\mathbf{y} = \operatorname{FFN}_{\text{shared}}(\mathbf{x}) + \sum_{i \in \operatorname{TopK}(\mathcal{G}(\mathbf{x}))} g_i(\mathbf{x}) \operatorname{FFN}_i(\mathbf{x})
$$

$\mathcal{G}(\mathbf{x})$ is the router's scores, $g_i(\mathbf{x})$ expert $i$'s weight, and $K$ the number of experts per token (unrelated to the key matrix). Mixtral applies the softmax to the $K$ best scores alone, so the weights sum to 1; Switch Transformer and Qwen2's MoE apply it over all experts and keep the $K$ largest, which sum to less than 1 unless renormalized.

Why smaller experts help: compute per token depends on how much capacity is active, not how it is cut up. Cut 16 experts (2 active) into 64 quarter-size ones (8 active) and active and total capacity are unchanged, but the possible active sets grow from $\binom{16}{2} = 120$ to $\binom{64}{8} \approx 4.43 \times 10^9$, giving each token a much more specific mix.

## 17.3 Shared Experts

In a plain MoE model every expert must learn general skills (grammar, punctuation, common words), wasting capacity on copies. A **shared expert** runs for every token without routing, so general knowledge lives there once and routed experts specialize, in code, chemistry or a language, say. The idea comes from DeepSeekMoE ("shared-expert isolation"). Qwen3's MoE models, such as Qwen3-30B-A3B (128 experts, 8 active), dropped it.

## 17.4 Router Collapse and the Load-Balancing Loss

An expert that gets slightly more tokens trains faster, so the router prefers it more, until a few experts get almost every token and the rest never train: *router collapse*. Training adds a loss term that is smallest when tokens are spread evenly:

$$
\mathcal{L}_{\text{aux}} = \alpha \cdot N_E \sum_{i=1}^{N_E} f_i \cdot P_i
$$

$f_i$ is the fraction of tokens sent to expert $i$, $P_i$ its average router probability, $N_E$ the number of experts, and $\alpha$ a small weight, often 0.01. The term is large when the same experts get many tokens and high probabilities; with one expert per token and a perfectly even split, $f_i = P_i = 1/N_E$ and it takes its smallest value, $\alpha$.

## 17.5 Expert Capacity and Dropped Tokens

Experts spread over GPUs (*expert parallelism*, section 11.7) exchange tokens in an all-to-all in every MoE layer, which needs fixed-size blocks. So each expert gets a fixed number of slots per batch, its **capacity**:

$$
C = \left\lceil \gamma \cdot \frac{K \cdot N_{\text{tokens}}}{N_E} \right\rceil
$$

The fraction is each expert's share under perfectly even routing, and the capacity factor $\gamma$ (typically 1.1 to 1.25) adds 10% to 25% spare room; the brackets round up. Tokens beyond $C$ are **dropped** by that expert, keeping only their other experts and the residual connection. A larger $\gamma$ drops fewer tokens but wastes memory on empty slots. And because the most overloaded expert sets the pace for the whole layer, an MoE model's speed depends on its balance.

> **Interactive figure** `moe-router`, on the web page: https://transformers-and-math.drmhse.com/modules/17-mixture-of-experts#figure-moe-router
>
> **Try:** Start the simulation with the load-balancing loss switched off. Before you press play, predict what the eight bars of expert usage will look like after a few hundred steps, and what happens to the dropped-token rate.
>
> **What happens:** A few experts take almost all the tokens and the rest go nearly silent: router collapse. With capacity limits, the overloaded experts overflow and the drop rate climbs. Switch the load-balancing loss on and run again: usage stays close to even and almost nothing is dropped.

> **In words.** A toy router assigns tokens to eight experts over many steps; experts with more tokens improve faster and the router prefers better experts, the feedback loop behind collapse. Bars show each expert's share and a line the dropped-token rate. Switches toggle the load-balancing loss and set the capacity factor.

> **Browser lab** `moe-balance`: an experiment that runs in the reader's browser on the web page: https://transformers-and-math.drmhse.com/modules/17-mixture-of-experts#lab-moe-balance

## Exercises

**Exercise (check).** What does a shared expert do that a routed expert does not?

*Solution:* It runs for every token with no routing, so general knowledge can live there once instead of being copied into every routed expert, leaving the routed experts free to specialize.

**Exercise (apply).** How many ways are there to choose 8 active experts out of 64? Give the answer in billions.

*Answer:* 4.426 (within 0.01)

*Solution:* $\binom{64}{8} = \frac{64!}{8!\,56!} = 4{,}426{,}165{,}368 \approx 4.43$ billion, against $\binom{16}{2} = 120$ for 2 of 16 larger experts with the same active capacity.

**Exercise (apply).** A batch has 4,096 tokens, each routed to $K = 8$ of $N_E = 64$ experts, with capacity factor $\gamma = 1.25$. What is each expert's capacity?

*Answer:* 640 (within 0)

*Solution:* $C = \lceil 1.25 \times 8 \times 4{,}096 / 64 \rceil = \lceil 1.25 \times 512 \rceil = 640$ token slots per expert.

**Exercise (apply).** With top-1 routing over $N_E = 8$ experts, perfectly even routing gives $f_i = P_i = 1/8$. What is $\mathcal{L}_{\text{aux}}$ with $\alpha = 0.01$? What if every token goes to expert 1 with probability 1?

*Answer:* 0.01 (within 0.0001)

*Solution:* Even: $\alpha N_E \sum_i f_i P_i = 0.01 \times 8 \times 8 \times \tfrac{1}{64} = 0.01$. Collapsed: $f_1 = P_1 = 1$ and the rest are 0, so the sum is 1 and the loss is $0.01 \times 8 \times 1 = 0.08$, eight times larger.

**Exercise (extend).** Explain the feedback loop behind router collapse, and why the load-balancing loss multiplies $f_i$ (a count, with no gradient) by $P_i$ (a probability, with a gradient).

*Solution:* An expert that receives slightly more tokens gets more gradient updates, becomes better, so the router learns to send it even more tokens: a rich-get-richer loop that ends with a few experts doing everything and the rest untrained. The fraction $f_i$ comes from a hard top-K choice, so it cannot be differentiated; the probability $P_i$ can. Multiplying them makes the loss's gradient push down the router probability of exactly the experts that are currently over-used, in proportion to how over-used they are.

### Core Academic References
- **Shazeer et al. (2017)**: *Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer*. ICLR 2017. [arXiv:1701.06538](https://arxiv.org/abs/1701.06538)
- **Fedus, Zoph & Shazeer (2022)**: *Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity*. JMLR 23(120): 1–39.
- **Lepikhin et al. (2021)**: *GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding*. ICLR 2021. [arXiv:2006.16668](https://arxiv.org/abs/2006.16668)
- **Dai et al. (2024)**: *DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models*. [arXiv:2401.06066](https://arxiv.org/abs/2401.06066)
- **Jiang et al. (2024)**: *Mixtral of Experts*. [arXiv:2401.04088](https://arxiv.org/abs/2401.04088)
- **Yang et al. (2024)**: *Qwen2 Technical Report*. [arXiv:2407.10671](https://arxiv.org/abs/2407.10671)
- **Yang et al. (2025)**: *Qwen3 Technical Report*. [arXiv:2505.09388](https://arxiv.org/abs/2505.09388)