Mixture of experts

The big idea

In a normal transformer, every token passes through every parameter. Make the model twice as big and every token costs twice as much compute. A mixture-of-expertsA layer with many parallel feed-forward networks (experts) and a router that sends each token to only a few of them, so total parameters can grow without growing the compute per token. (MoE) layer breaks that link. It holds many feed-forward networks, called expertsOne of the parallel feed-forward networks inside a mixture-of-experts layer. Each token is processed by only the few experts the router selects., and a small routerA small learned layer in a mixture-of-experts model that scores every expert for a token and picks the top few to process it. that sends each token to only one or two of them.

The result: a model can store far more parameters than it uses for any single token. More parameters means more capacity to memorise facts and patterns; using only a few per token keeps the cost of each token close to a much smaller dense model.

Think of a large hospital. It employs dozens of specialists, but a patient with a broken wrist sees a triage nurse and then one or two of them, not all fifty. The hospital's total knowledge is huge; the cost of each visit is small. The triage nurse is the router, and if the nurse sends everyone to the same doctor, the queue explodes while other doctors sit idle.

How it works

The router

The router is a single linear layer: it multiplies the token's vector by a matrix with one column per expert, giving a score per expert. Softmax turns the scores into probabilities, and the token goes to the top $k$ experts, usually 1 or 2. The chosen experts' probabilities are renormalised into gate weightsThe weights a mixture-of-experts router gives to the outputs of the experts it chose for a token, usually its softmax probabilities renormalised over the chosen experts., and the layer's output is the weighted sum of their outputs:

The router is trained along with everything else. Because the gate weights multiply the expert outputs, gradients tell the router which experts helped for which tokens.

Sparse activation: total vs active parameters

Each expert is an ordinary FFN. Only the experts a token is routed to run for that token, so the layer is sparsely activatedOnly a small subset of a model's parameters is used for each input, as in mixture-of-experts layers where each token runs through a few experts.. That is why MoE models quote two numbers:

• Total parameters: everything stored, all experts included. This sets memory needs.
• Active parameters: what one token actually uses, the shared parts (attention, embeddings) plus $k$ experts per layer. This sets compute per token.

With 8 experts and top-2 routing, each token uses 2/8 = 25 percent of the expert parameters. With 64 experts and top-2, about 3 percent. Attention layers are usually shared and dense, so the overall active fraction is somewhat higher than the expert fraction.

Load balancing

Left alone, routers collapse. Early in training a few experts happen to be slightly better, so they get more tokens, so they learn faster, so they get even more tokens. The rest starve. Two mechanisms fight this:

• A load-balancing lossAn extra training loss term in mixture-of-experts models that is smallest when tokens and router probability are spread evenly over experts, discouraging the router from overloading a few.. The Switch Transformer version is $\alpha \cdot N{\sum }_i{f}_i{P}_i$, where $f_i$ is the share of tokens sent to expert $i$, $P_i$ is the share of router probability given to it, and $N$ is the number of experts. It is smallest when both are spread evenly: $N\sum f_i{P}_i=1$ for a perfectly even router.
• An expert capacityThe maximum number of tokens an expert accepts per batch, usually the even share times a capacity factor. Tokens over the limit skip the expert.: each expert takes at most (tokens / experts) times a capacity factor per batch. Tokens beyond that overflow and skip the layer through the residual connection.

Try it

Try this:

1. At the defaults (8 experts, top-2, no skew) each expert gets between 118 and 139 of the 1,024 assignments, all under the capacity of 160, and the balancing loss is 1.001, essentially perfect.
2. Raise router skew to 1.5. The first expert now gets 366 assignments, 338 slots overflow and are dropped, and the loss jumps to 1.766. This is the collapse the balancing loss is trained to prevent.
3. Return skew to 0 and set capacity factor to 1.00. Even a well-balanced router overflows a little (18 dropped slots), because random routing is never perfectly even. That is why capacity factors above 1 are common.

A small MoE layer in NumPy, with the same balancing loss:

# A tiny mixture-of-experts layer: top-2 of 8 experts per token.
import numpy as np

rng = np.random.default_rng(0)
d, hidden, E, k, n = 16, 64, 8, 2, 1000
W_router = rng.normal(size=(d, E)) * 0.5
experts = [(rng.normal(size=(d, hidden)) * 0.1,
            rng.normal(size=(hidden, d)) * 0.1)
           for _ in range(E)]
x = rng.normal(size=(n, d))                       # 1,000 token vectors

def moe(x, W_router):
    logits = x @ W_router
    probs = np.exp(logits - logits.max(1, keepdims=True))
    probs /= probs.sum(1, keepdims=True)
    top = np.argsort(-logits, axis=1)[:, :k]      # chosen experts per token
    out = np.zeros_like(x)
    for t in range(len(x)):
        gate = probs[t, top[t]] / probs[t, top[t]].sum()  # gate weights
        for g, e in zip(gate, top[t]):
            W1, W2 = experts[e]
            out[t] += g * (np.maximum(x[t] @ W1, 0) @ W2)
    f = np.bincount(top.ravel(), minlength=E) / top.size
    P = probs.mean(0)
    return out, np.bincount(top.ravel(), minlength=E), E * np.sum(f * P)

per_expert = 2 * d * hidden
print(f"expert params: total {E * per_expert:,}, "
      f"active per token {k * per_expert:,}")

out, load, aux = moe(x, W_router)
print("tokens per expert:", load.tolist(), f"balance loss {aux:.3f}")

skewed = W_router.copy(); skewed[:, 0] += 1.0  # favour expert 0
x_shift = x + 0.8  # shared direction
_, load, aux = moe(x_shift, skewed)
print("skewed router:    ", load.tolist(), f"balance loss {aux:.3f}")

expert params: total 16,384, active per token 4,096
tokens per expert: [257, 224, 209, 302, 269, 264, 237, 238] balance loss 1.016
skewed router:     [994, 60, 18, 29, 3, 72, 36, 788] balance loss 3.916

The layer stores 16,384 expert parameters but each token uses 4,096. With a neutral router the 2,000 assignments spread from 209 to 302 per expert (loss 1.016). With a router biased toward expert 0 and inputs that share a direction, expert 0 takes 994 and expert 4 only 3: the loss of 3.916 is the signal training would push back against.

Going deeper

What experts specialise in

It is tempting to imagine one expert for biology and another for law. In practice, routing tends to pick up on lower-level patterns such as token types, punctuation, or syntax as much as on topics, and specialisation is fuzzy. Treat experts as a way to add capacity, not as a set of human-readable departments.

Why MoE is harder to serve

All experts must be in memory, even though each token uses a few, so an MoE model needs memory for its total parameters and compute for its active ones. Experts are often spread across several GPUs (expert parallelism), which means tokens travel between devices twice per MoE layer. At small batch sizes, different tokens hit different experts, so the GPU reads many experts' weights for little work, and decode stays memory-bound ( Context windows and the KV cache ). MoE shines at large batch sizes where every expert gets plenty of tokens.

Training instabilities

MoE models have been harder to train than dense ones: router logits can blow up, and lower precision makes it worse. Common fixes include keeping the router in higher precision, adding a small penalty on large router logits, and careful initialisation. The Switch Transformer paper reported that its simplifications made bfloat16 training work.

Trade-offs

Common mistakes

• Comparing an MoE model's total parameters with a dense model's parameters. Compare active parameters for compute and total for memory.
• Assuming MoE reduces memory. It increases it: every expert must be loaded.
• Forgetting the balancing loss (or setting it to zero) and wondering why a few experts do all the work.
• Expecting experts to map to clean human topics.
• Benchmarking MoE throughput at batch size 1 and concluding it is slow, or only at huge batch sizes and concluding it is always fast.

In the wild

Check yourself

Recap