Mixture of experts
Many parameters, little compute: how a router selects a small part of the model per token, and what that costs elsewhere.
The idea
Instead of one large department handling every request, there are many specialist units. A dispatcher decides per request which two units are responsible. The organisation as a whole knows a great deal, but each request occupies only a small part of it.
What that buys
| Dense model | Mixture of experts | |
|---|---|---|
| Total parameters | 70 bn | 140 bn |
| Computed per token | 70 bn | 17 bn |
| Memory need | 140 GB | 280 GB |
| Compute per token | high | low |
The router
Per token a small matrix computes a score for every expert. The best k are
selected and their outputs weighted and added. k = 1 or k = 2 is usual.
The real problem is even loading. Without countermeasures an initial preference reinforces itself: an expert chosen more often gets better, is therefore chosen more often still, and the rest atrophy.
- Add an auxiliary loss that penalises uneven load.
- Set a capacity limit per expert; surplus tokens are rerouted or dropped.
- Monitor the load distribution during training, not afterwards.
The formulas
The minimum of that term sits at an even distribution. Without it, routing collapses in practice within a few thousand steps.
The calculation
For a model with 64 experts, k = 2, feed-forward width f = 8h/3:
| Quantity | Value |
|---|---|
| Experts per layer | 64 |
| Computed per token | 2 |
| Share of active feed-forward parameters | 2/64 = 3.1 % |
| Compute per token versus dense | roughly like a model with 2/64 of the feed-forward |
In practice such a model reaches the quality of a dense model with about the
square root of the product of total and active parameters. At 140 billion total
and 17 billion active that corresponds roughly to a dense model of
√(140 × 17) ≈ 49 billion parameters, at the compute cost of a 17-billion one.
Operational consequences
- Memory need follows the total parameters. Unsuitable for edge devices and small installations.
- Distributed operation adds network traffic, because tokens are sent to the cards holding the relevant experts.
- Latency varies more than for dense models, because expert load varies per batch.
- Quantisation works well, because memory is the limiting factor. Such a model in int4 is often the only practical in-house form.
Related courses and sources
Switch Transformers
Only a fraction of the parameters compute per token. The paper that brought mixture of experts into wide use.
For understanding why large models do not compute every parameter per token.