GPT-4, Mixtral, and DeepSeek all use Mixture of Experts. Explain how MoE works mechanically — what it replaces in the transformer, how routing works, why it enables larger models without proportional compute cost, and what the practical challenges are in training and serving.

Question

Accepted Answer

What MoE replaces

In a standard transformer, every token passes through the same FFN layer at each depth. In an MoE transformer, each FFN layer is replaced with N expert FFN networks. Only top-K experts (typically K=2) process each token. Attention layers remain dense and shared.

Dense FFN:   output = FFN(x)

MoE layer:   scores = softmax(x · W_gate)         # router: shape [d_model → N]
             top2 = argtop(scores, k=2)
             output = Σ scores[i] · Expert_i(x)   for i in top2

Why it scales (conditional computation)

Parameters scale with N experts but FLOPs per token stay constant — only K/N experts fire.

Mixtral 8x7B: 8 experts, each ~7B parameters, top-2 routing. Total parameters: ~47B. Activated parameters per token: ~13B — similar FLOPs to a 13B dense model, but representational capacity closer to a 47B model. You pay compute for a small model, get quality from a large one.

Routing and expert collapse

The router is a learned linear projection: scores = softmax(W_gate · x). Without constraints, training collapses: the router favors a few experts early, those experts receive more gradient, improve, attract more tokens — a feedback loop. Remaining experts starve.

Fix: auxiliary load-balancing loss

L_aux = α · Σ_i f_i · P_i

- f_i = fraction of tokens routed to expert i in this batch
- P_i = mean routing probability for expert i
- α ≈ 0.01 — small enough not to override the main loss

Penalizes uneven load, pushing the router toward uniform distribution. Without this, you often end up with 1-2 hot experts and the rest unused.

Expert specialization

Without explicit design, trained MoE models develop emergent specialization. Analysis of Mixtral shows experts that specialize by domain (code, math, general language), syntactic function, and reasoning step. The router learned to specialize implicitly from gradient signal alone.

Expert parallelism in multi-GPU serving

Experts are sharded across GPUs — each GPU holds a subset. The forward pass requires:
1. Router scores computed per token
2. All-to-all dispatch: tokens sent to the GPU holding their assigned expert
3. Expert computation on each GPU
4. All-to-all gather: results returned to originating GPU

All-to-all scales O(N_experts × batch) and becomes the bottleneck on slow interconnects. MoE models are more sensitive to GPU topology than dense models — NVLink vs PCIe matters significantly.

Token-choice vs expert-choice routing

Token-choice (default): each token selects top-K experts. Risk: some experts overloaded, some idle within a batch.

Expert-choice: each expert selects top-K tokens. Guarantees perfect load balance, but tokens can be dropped if no expert selects them — problematic for rare but important tokens.

Tradeoffs summary

| | Dense | MoE |
|---|---|---|
| FLOPs/token | All params active | K/N params active |
| GPU memory | Proportional to params | Same — all experts loaded |
| Communication | Minimal | All-to-all per MoE layer |
| Training stability | Stable | Needs load-balancing loss |
| Quality per FLOP | Baseline | Higher at scale |
| Serving latency variance | Low | Expert dispatch adds jitter |