mixture of experts training
**Mixture of Experts (MoE) Training** is the **specialized training methodology for sparse conditional computation models where only a subset of parameters (experts) are activated per input** — requiring careful handling of expert load balancing, routing stability, communication patterns across devices, and auxiliary losses to prevent expert collapse, with techniques like expert parallelism, top-k gating, and capacity factors enabling models like Mixtral 8x7B, GPT-4 (rumored MoE), and Switch Transformer to achieve dense-model quality at a fraction of the per-token compute cost.
**MoE Architecture**
```
Standard Transformer FFN:
x → [FFN: 4096 → 16384 → 4096] → y
Every token uses ALL parameters
MoE Layer (8 experts, top-2 routing):
x → [Router/Gate network] → selects Expert 3 and Expert 7
x → [Expert 3: 4096 → 16384 → 4096] × w_3
+ [Expert 7: 4096 → 16384 → 4096] × w_7 → y
Each token uses only 2 of 8 experts (25% of FFN params)
```
**Key Training Challenges**
| Challenge | Problem | Solution |
|-----------|---------|----------|
| Expert collapse | All tokens route to 1-2 experts | Auxiliary load balancing loss |
| Load imbalance | Some experts get 10× more tokens | Capacity factor + dropping |
| Communication | Experts on different GPUs → all-to-all | Expert parallelism |
| Training instability | Router gradients are noisy | Straight-through estimators, jitter |
| Expert specialization | Experts learn redundant features | Diversity regularization |
**Load Balancing Loss**
```python
# Auxiliary loss to encourage balanced expert usage
def load_balance_loss(router_probs, expert_indices, num_experts):
# f_i = fraction of tokens routed to expert i
# p_i = average router probability for expert i
f = torch.zeros(num_experts)
p = torch.zeros(num_experts)
for i in range(num_experts):
mask = (expert_indices == i).float()
f[i] = mask.mean()
p[i] = router_probs[:, i].mean()
# Loss encourages uniform f_i (each expert gets equal tokens)
return num_experts * (f * p).sum()
```
**Expert Parallelism**
```
8 GPUs, 8 experts, 4-way data parallel:
GPU 0: Expert 0,1 | Tokens from all GPUs routed to Exp 0,1
GPU 1: Expert 2,3 | Tokens from all GPUs routed to Exp 2,3
GPU 2: Expert 4,5 | Tokens from all GPUs routed to Exp 4,5
GPU 3: Expert 6,7 | Tokens from all GPUs routed to Exp 6,7
GPU 4-7: Duplicate of GPU 0-3 (data parallel)
all-to-all communication: Each GPU sends tokens to correct expert GPU
```
**MoE Model Comparison**
| Model | Experts | Active | Total Params | Active Params | Quality |
|-------|---------|--------|-------------|--------------|--------|
| Switch Transformer | 128 | 1 | 1.6T | 12.5B | T5-XXL level |
| GShard | 2048 | 2 | 600B | 2.4B | Strong MT |
| Mixtral 8x7B | 8 | 2 | 47B | 13B | ≈ Llama-2-70B |
| Mixtral 8x22B | 8 | 2 | 176B | 44B | ≈ GPT-4 class |
| DBRX | 16 | 4 | 132B | 36B | Strong |
| DeepSeek-V2 | 160 | 6 | 236B | 21B | Excellent |
**Capacity Factor and Token Dropping**
- Capacity factor C: Maximum tokens per expert = C × (total_tokens / num_experts).
- C = 1.0: Perfect balance, may drop tokens if routing is uneven.
- C = 1.25: 25% buffer for imbalance (common choice).
- Dropped tokens: Skip the MoE layer, use residual connection only.
- Training: Some dropping is acceptable. Inference: Never drop (use auxiliary buffer).
**Training Tips**
- Router z-loss: Penalize large logits to stabilize gating → prevents routing oscillation.
- Expert jitter: Add small noise to router inputs during training → prevents collapse.
- Gradient scaling: Scale expert gradients by 1/num_selected_experts.
- Initialization: Initialize router weights small → initially uniform routing → gradual specialization.
MoE training is **the methodology that enables trillion-parameter models with affordable compute** — by activating only a fraction of parameters per token and carefully managing expert load balancing, routing stability, and communication across devices, MoE architectures achieve the quality of dense models 5-10× larger while requiring only the inference compute of much smaller models, making them the dominant architecture choice for frontier language models.