Mixture of Experts (MoE): Sparse Scaling

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. MoE Architecture
5. Gating & Routing Mechanisms
6. Load Balancing Loss
7. Switch Transformer
8. Mixtral & Modern MoE LLMs
9. Expert Specialization & Analysis
10. Training Challenges
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Derive the MoE forward pass with top-k routing.
2. Explain load balancing loss and why it's necessary.
3. Compare Switch Transformer, GShard, and Mixtral architectures.
4. Analyze expert specialization patterns.
5. Identify and mitigate routing collapse and training instability.

Notation

• $E$ — number of experts
• $K$ — number of active experts per token (top-K)
• $G(\mathbf{x})$ — gating function (router)
• $f_i(\mathbf{x})$ — expert $i$'s output

Core Intuition

Standard transformers use ALL parameters for EVERY token — wasteful since different tokens need different processing. MoE activates only a SUBSET of parameters per token (e.g., 2 out of 8 experts). This gives the CAPACITY of a much larger model (total parameters) with the COMPUTE of a smaller one (active parameters). Mixtral 8x7B has 46.7B total parameters but only uses 12.9B per token — matching 70B dense model quality at 7B inference cost.

MoE Architecture

Replace dense FFN with sparse MoE layer:

where $e_1, \ldots, e_K$ are the top-K experts selected by the router.

Typical design:

• Attention layers: SHARED (all tokens use same attention).
• FFN layers: REPLACED by MoE (each token routed to K experts).
• Every other layer or every layer can be MoE.

Parameters: $E$ copies of the FFN = $E \times$ FFN parameters. But only $K$ copies activated per token.

Compute: $K/E$ fraction of FFN compute per token (e.g., 2/8 = 25%).

Gating & Routing Mechanisms

Top-K softmax router:

Select top-K experts; set others to zero; renormalize gate values.

Token-choice routing (standard): Each token picks its top-K experts.

• Simple, differentiable (through selected experts).
• Risk: popular experts get overloaded; unpopular experts starve.

Expert-choice routing (Zhou et al., 2022): Each expert picks its top-C tokens.

• Guaranteed load balance (each expert processes exactly C tokens).
• Risk: some tokens may not be selected by any expert.

Capacity factor: Maximum tokens per expert = $C \times N/E$ where $N$ is batch tokens and $C \geq 1$ is capacity factor. Excess tokens are dropped (routed to residual).

Load Balancing Loss

Problem: Without intervention, routing collapses — a few experts handle most tokens while others are unused ("dead experts").

Auxiliary load balancing loss:

where:

• $f_i = \frac{\text{tokens routed to expert } i}{\text{total tokens}}$ (fraction of tokens).
• $p_i = \frac{1}{T}\sum_t G(\mathbf{x}_t)_i$ (average router probability for expert $i$).
• $\alpha \approx 0.01$ (small weight to not dominate main loss).

Goal: Minimize $\mathcal{L}_{\text{bal}}$ encourages uniform routing (equal load across experts). If all experts get equal tokens: $f_i = p_i = 1/E$ and $\mathcal{L}_{\text{bal}} = \alpha$.

Switch Transformer

Fedus et al. (2022, Google): Top-1 routing (K=1) with simplifications:

• Only ONE expert per token (simplest routing).
• Hard routing (no weighted combination — just one expert's output).
• Large capacity factor (1.25-2.0) to handle load imbalance.
• Apply MoE to every other FFN layer.

Result: 7x speedup in pre-training compared to T5 at same quality. 1.6T parameter model trained with effective compute of a 200B dense model.

Key insight: Top-1 routing works as well as top-2 with half the compute.

Mixtral & Modern MoE LLMs

Mixtral 8x7B (Mistral, 2024):

• 8 experts per MoE layer, top-2 routing.
• 32 transformer layers, every layer is MoE.
• Total: 46.7B parameters. Active: 12.9B per token.
• Matches LLaMA-2 70B quality at 3x lower inference cost.

Design choices:

• Shared attention, expert FFN.
• Sliding window attention (4096 tokens).
• Byte-level BPE tokenizer.

DeepSeek MoE:

• Fine-grained experts (more, smaller experts = 64).
• Shared experts + routed experts (some experts always active).
• Better load balance and expert utilization.

Grok, DBRX, Arctic: All modern frontier models use MoE for cost-efficient scaling.

Expert Specialization & Analysis

Do experts specialize? Partially yes:

Observed specialization patterns:

• Syntax experts (handle function words, punctuation).
• Domain experts (activate for code vs natural language).
• Positional experts (specialize in early vs late positions).
• Language experts (in multilingual models: one expert per language cluster).

But: Specialization is noisy and overlapping. No single expert is solely responsible for any capability. Removing any single expert degrades ALL capabilities slightly.

Analysis tools:

• Expert activation frequency by domain.
• Token type distribution per expert.
• Expert similarity (cosine between expert weights).

Training Challenges

1. Routing collapse: All tokens go to 1-2 experts; others die.

• Fix: Load balancing loss + expert dropout + noise in routing.

2. Training instability: MoE models are less stable than dense models.

• Fix: Lower learning rate, stronger gradient clipping, router z-loss.

3. Expert imbalance during inference: Some experts are overloaded.

• Fix: Capacity factor, token dropping, expert parallelism.

4. Memory: Total parameters are $E \times$ larger. A 8x7B MoE uses 46.7B parameter memory (not 7B).

• Implication: Need expert parallelism (distribute experts across GPUs).

5. Communication overhead: In distributed training, tokens must be sent to the GPU hosting their selected expert (all-to-all communication).

Common Pitfalls

Pitfall 1. Assuming MoE inference cost equals active parameters. While compute is $K/E$ of dense, MEMORY still holds ALL parameters. An 8x7B MoE needs 47B parameters in memory (not 13B).

Pitfall 2. Setting load balancing weight $\alpha$ too high. If $\alpha$ is too large, the model prioritizes balance over quality — routing becomes uniform but unintelligent.

Pitfall 3. Using MoE at small scale. MoE benefits emerge at scale (above 1B parameters). Small MoE models often underperform equivalent dense models due to routing overhead and training instability.

Summary

• MoE: Activate $K$ of $E$ experts per token; capacity of large model, compute of small.
• Router: Learned linear → softmax → top-K selection.
• Load balancing: Auxiliary loss prevents routing collapse.
• Switch Transformer: Top-1, simplest MoE; 7x training speedup.
• Mixtral 8x7B: Top-2 routing; matches 70B dense at 3x cheaper inference.
• Challenges: Routing collapse, instability, memory (total params), communication.

Exercises

Exercise 1. For Mixtral 8x7B: compute the FLOPs per token vs a dense 47B model. What's the effective compute reduction?

Exercise 2. Derive the gradient of the load balancing loss with respect to router weights.

Exercise 3. Design an expert-choice routing scheme for a batch of 2048 tokens with 8 experts and capacity factor 1.5.

Exercise 4. Compare memory requirements: 70B dense model vs 8x7B MoE (both in BF16). How much GPU memory does each need?

Exercise 5. Propose a method to detect and recover dead experts during training (experts that receive less than 1% of tokens).