# Mixture of experts: big models, small bills · AI in Bed

Advanced level · 8 min · https://www.aiinbed.com/learn/advanced/mixture-of-experts

*How some of the largest models only use a small slice of themselves for each token, and what that saves and costs.*

In a normal transformer, every token flows through every weight in every layer. A **mixture of experts**, or MoE, breaks that rule. Some layers are split into many smaller networks, and each token only visits a few of them.

### Experts and the router

An MoE layer replaces one feed-forward block with several copies, called **experts**. A tiny network called the **router** looks at each token and gives every expert a score.

The token is sent to the top-k experts, often just one or two. Their outputs are blended, weighted by the router's scores, and the rest of the experts sit idle for that token.

**A router picking top-2 experts**

```python
import math, random

random.seed(7)
num_experts, top_k = 8, 2

def softmax(xs):
    exps = [math.exp(x) for x in xs]
    total = sum(exps)
    return [e / total for e in exps]

for token in ["The", "cat", "sat", "on"]:
    scores = [random.gauss(0, 1) for _ in range(num_experts)]  # router scores
    probs = softmax(scores)
    best = sorted(range(num_experts), key=lambda i: probs[i], reverse=True)[:top_k]
    share = sum(probs[i] for i in best)
    picks = ", ".join(f"expert {i} ({probs[i] / share:.0%})" for i in best)
    print(f"{token:>4} -> {picks}")

print(f"Each token used {top_k} of {num_experts} experts.")
```

Real routers learn their scores during training. Here they are random, just to show the softmax and top-k steps.

### Total versus active parameters

**Total parameters** count every weight in the model. **Active parameters** count only the weights a single token actually uses.

An MoE model can have a very large total but a much smaller active count. That is the trick: lots of stored knowledge, but the compute cost per token of a far smaller model.

**Quick check:** An MoE model has 8 experts per layer and uses the top 2 for each token. What mostly sets its compute cost per token?

A. All 8 experts, because they are all stored
B. The 2 chosen experts plus the shared parts of the model
C. Only the router
D. The length of the training data

Answer: B. Compute follows the active parameters. Idle experts take up memory but do no work for that token.

### Keeping the load balanced

Left alone, a router may fall in love with a few favourite experts and starve the rest. Those unused experts never improve, and the busy ones become a bottleneck.

Training adds a **load-balancing loss**: a small penalty that nudges the router to spread tokens fairly. Systems also cap how many tokens each expert can take at once.

> Experts rarely line up with neat human topics like "maths" or "French". Studies of trained routers often find they specialise in quieter patterns, such as punctuation or particular kinds of word.

- Memory: every expert must still be loaded, so MoE models need lots of memory even though each token uses little.
- Routing instability: training can wobble as the router's choices shift, which takes careful tuning.
- Communication: experts are often spread across many chips, so tokens have to travel between them.
- Fine-tuning can be trickier, as small datasets may overfit some experts.

**Quick check:** What is the main job of a load-balancing loss in an MoE model?

A. To make the model's answers shorter
B. To encourage the router to spread tokens across experts instead of overusing a few
C. To reduce the total number of parameters
D. To remove the need for a router

Answer: B. Without it, a few experts can grab most tokens, leaving others undertrained and creating bottlenecks.

**Key takeaways**

- A router sends each token to only a few experts per MoE layer.
- Compute follows active parameters; memory follows total parameters.
- Load balancing and stable routing are the hard parts.

**Go deeper (free sources)**

- [Switch Transformers: Scaling to Trillion Parameter Models](https://arxiv.org/abs/2101.03961) by Fedus, Zoph and Shazeer
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://arxiv.org/abs/1701.06538) by Shazeer et al.
- [Mixture of Experts Explained](https://huggingface.co/blog/moe) by Hugging Face
- [Mixtral of Experts](https://arxiv.org/abs/2401.04088) by Mistral AI