mixture of experts routing
**Mixture of Experts (MoE) Routing and Load Balancing** addresses the **critical challenge of efficiently distributing tokens across expert networks in sparse MoE architectures** — where a gating network must learn to route each input token to the most appropriate subset of experts while maintaining balanced utilization, avoiding expert collapse, and minimizing communication overhead in distributed training.
**MoE Architecture**
```
Input token x
↓
Gating Network: g(x) = Softmax(W_g · x) → [score_1, ..., score_E]
↓
Top-K selection (typically K=1 or K=2 of E experts)
↓
Output = Σ(g_i(x) · Expert_i(x)) for selected experts
```
In practice, MoE replaces the MLP in every N-th transformer layer (e.g., every other layer in Mixtral, every layer in Switch Transformer), keeping attention layers dense.
**Routing Strategies**
| Strategy | Description | Used By |
|----------|------------|--------|
| Top-K | Select K experts with highest gate score | Mixtral (K=2), GShard |
| Switch | Top-1 routing (simplest, most efficient) | Switch Transformer |
| Expert Choice | Each expert selects its top-K tokens (inverted) | Expert Choice (Google) |
| Soft MoE | Weighted average of all experts (fully differentiable) | Soft MoE (Google) |
| Hash routing | Deterministic routing via hash function | Hash Layer |
**The Load Balancing Problem**
Without intervention, routing collapses: a few experts receive most tokens while others are underutilized. This happens because:
- Popular experts get more gradient updates → become even better → attract more tokens (rich-get-richer)
- Underutilized experts stagnate or become effectively dead
- In distributed training, imbalanced routing causes GPU idle time (bottleneck = most loaded expert)
**Auxiliary Load Balancing Loss**
```python
# Switch Transformer auxiliary loss
# f_i = fraction of tokens routed to expert i
# P_i = mean gate probability for expert i
# Ideal: f_i = P_i = 1/E for all experts (uniform)
aux_loss = alpha * E * sum(f_i * P_i for i in range(E))
# This loss is minimized when routing is perfectly uniform
# alpha typically 0.01 to 0.1 (balances vs. task loss)
```
**Expert Capacity and Token Dropping**
To bound computation, each expert has a fixed **capacity factor** C:
```
Expert buffer size = C × (total_tokens / num_experts)
C = 1.0: exact uniform capacity
C = 1.25: 25% overflow buffer
C > 1: more flexibility but more computation
Tokens exceeding capacity → dropped (pass through residual only)
```
**Expert Parallelism**
Distributing experts across GPUs:
```
Data Parallel: each GPU has all experts, different data
Expert Parallel: each GPU hosts a subset of experts
→ All-to-All communication: tokens routed to correct GPU
→ GPU 0: Experts 0-3, GPU 1: Experts 4-7, ...
→ Forward: AllToAll(tokens→experts) → compute → AllToAll(results→tokens)
```
The All-to-All communication is the primary overhead — proportional to (tokens × hidden_dim) across GPUs. Modern systems combine expert parallelism with data and tensor parallelism (e.g., DeepSeek-V2 uses EP=8 × DP=many).
**MoE routing and load balancing are the engineering linchpin of sparse model architectures** — the gating mechanism must simultaneously learn task-relevant specialization AND maintain computational efficiency, making routing strategy design one of the most impactful decisions in scaling language models beyond dense transformer limits.