Module 17 · Part IV · The Modern Architecture

Mixture of Experts

Routing and Top-K Selection, Many Small Experts, Shared Experts, the Load-Balancing Loss, and Expert Capacity

By the end you can design an MoE layer and predict how routing fails.

Builds onModule 15·
Check yourself ▸

Before you start, can you already…

  • ☐turn the 2017 block into a 2024 block and justify each change. (Module 15)

If any is shaky, the linked module teaches it.

Open in Colab
3 Hours·Elective·Advanced·0/5

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 RoutingRun this section

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.

Architecture Flowchart
Diagram: Token vector x; Router scores every expert; Shared expert: runs for every token; Keep the top k, e.g. 8 of 64; Expert 4; Expert 19; Expert 52; Weighted sum; +; 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 ExpertsRun this section

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:

y=FFN⁡shared(x)+∑i∈TopK⁡(G(x))gi(x)FFN⁡i(x)\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})

G(x)\mathcal{G}(\mathbf{x}) is the router's scores, gi(x)g_i(\mathbf{x}) expert ii's weight, and KK the number of experts per token (unrelated to the key matrix). Mixtral applies the softmax to the KK best scores alone, so the weights sum to 1; Switch Transformer and Qwen2's MoE apply it over all experts and keep the KK 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 (162)=120\binom{16}{2} = 120 to , 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 LossRun this section

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:

Laux=α⋅NE∑i=1NEfi⋅Pi\mathcal{L}_{\text{aux}} = \alpha \cdot N_E \sum_{i=1}^{N_E} f_i \cdot P_i

fif_i is the fraction of tokens sent to expert ii, PiP_i its average router probability, NEN_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, fi=Pi=1/NEf_i = P_i = 1/N_E and it takes its smallest value, α\alpha.

17.5 Expert Capacity and Dropped TokensRun this section

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=⌈γ⋅K⋅NtokensNE⌉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 CC 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.

Figure 17.5
Router collapse, and the loss that prevents it

A toy simulation of a top-1 router over eight experts. Bars are each expert’s share of the tokens in the latest batch; the line under them is the share of tokens dropped for exceeding capacity.

Toy dynamics, not a trained model: experts that receive more tokens improve faster, and the router learns to prefer better experts. Even routing is 12.5% per expert.
Predict, then 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.

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

With and without the load-balancing loss, five seeds each

The router simulation of the figure, run for 400 steps with capacity factor 1.25, five seeds per setting. Section 17.4 predicts that without the auxiliary loss usage collapses onto a few experts and many tokens overflow capacity, and that with it usage stays near the even 12.5% each.

A toy: the experts here are a single 'skill' number each, not networks. It shows the feedback loop and the effect of the loss on it; it does not measure anything about real MoE models.

Exercises

Check

What does a shared expert do that a routed expert does not?

Apply

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

Apply

A batch has 4,096 tokens, each routed to K=8K = 8 of NE=64N_E = 64 experts, with capacity factor γ=1.25\gamma = 1.25. What is each expert's capacity?

Apply

With top-1 routing over NE=8N_E = 8 experts, perfectly even routing gives fi=Pi=1/8f_i = P_i = 1/8. What is Laux\mathcal{L}_{\text{aux}} with α=0.01\alpha = 0.01? What if every token goes to expert 1 with probability 1?

Extend

Explain the feedback loop behind router collapse, and why the load-balancing loss multiplies fif_i (a count, with no gradient) by PiP_i (a probability, with a gradient).

Core Academic References

  • Shazeer et al. (2017): Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017. arXiv: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
  • Dai et al. (2024): DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. arXiv:2401.06066
  • Jiang et al. (2024): Mixtral of Experts. arXiv:2401.04088
  • Yang et al. (2024): Qwen2 Technical Report. arXiv:2407.10671
  • Yang et al. (2025): Qwen3 Technical Report. arXiv:2505.09388