Mixture of experts

Routing has to be learned, and kept fair

Nobody tells the router which expert to pick. It learns, along with the rest of the model. Each expert's output is scaled by its router probability. So when training adjusts the weights, the router learns which experts helped with which tokens.

Left alone, though, routing tends to collapse. Think of a supermarket with four checkout lanes. One cashier happens to be quick on the first morning, so more shoppers join that lane. With all that practice, the cashier gets even quicker, and the lane gets even longer. The other cashiers stand idle.

Experts do the same. An expert that does well early gets more tokens. It improves faster, so it attracts even more. The others sit unused.

To stop this, MoE training adds a load-balancing loss. It's an extra score added to the training loss, and it's smallest when tokens are spread evenly across the experts. Here is the version from Google's Switch Transformer. It uses three symbols. Here they are, one at a time:

• $E$ is the number of experts.
• $f_i$ is the fraction of tokens whose top choice is expert $i$.
• $P_i$ is the router's average probability for expert $i$, across all the tokens.

The $\sum$ sign (sigma) means "add up, for every expert $i$ from 1 to $E$".

In plain words: for each expert, multiply its share of the tokens by its average probability. Add those up, then multiply by the number of experts. The result is 1 when tokens are spread evenly. It rises toward $E$ when everything goes to one expert.

Try it with 4 experts. If every expert gets a quarter of the tokens and a quarter of the probability, each expert adds 0.25 × 0.25 = 0.0625. Four of those make 0.25, and 4 × 0.25 = 1. If every token goes to expert 0 with probability 1, only expert 0 adds anything: 1 × 1 = 1. Then 4 × 1 = 4, the worst score.

In Python, the sum is a loop over the experts:

E = 4
f = [0.25, 0.25, 0.25, 0.25]
P = [0.25, 0.25, 0.25, 0.25]
print(E * sum(fi * pi for fi, pi in zip(f, P)))   # 1.0

E = 2
f = [1.0, 0.0]
P = [0.9, 0.1]
print(E * sum(fi * pi for fi, pi in zip(f, P)))

def balance(f, P):
    E = len(f)
    return E * sum(fi * pi for fi, pi in zip(f, P))

print(balance([0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]))
print(balance([1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]))
print(balance([0.5, 0.5, 0.0, 0.0], [0.4, 0.4, 0.1, 0.1]))

Watch collapse happen

Here's a small experiment. The inputs come from four groups. Each group needs a different function to get the right answer. A 4-expert MoE learns them, sending each token to its single best expert (top-1, as in Switch Transformer). It trains twice: once without the balance loss, once with it.

In the code, balance_weight sets how much the balance loss counts. A weight of 0 switches it off. For each group, the code prints how many of its tokens went to each of the 4 experts.

import torch
import torch.nn as nn

class MoE(nn.Module):
    def __init__(self, d, n_experts, hidden):
        super().__init__()
        self.n = n_experts
        self.router = nn.Linear(d, n_experts, bias=False)
        self.experts = nn.ModuleList([nn.Sequential(nn.Linear(d, hidden), nn.GELU(), nn.Linear(hidden, d))
                                      for _ in range(n_experts)])

    def forward(self, x):                                   # top-1 routing, as in Switch Transformer
        probs = torch.softmax(self.router(x), dim=-1)
        weight, choice = probs.max(dim=-1)
        out = torch.zeros_like(x)
        for e, expert in enumerate(self.experts):
            chosen = choice == e
            if chosen.any():
                out[chosen] += weight[chosen].unsqueeze(-1) * expert(x[chosen])
        f = torch.bincount(choice, minlength=self.n).float() / len(x)
        balance = self.n * (f * probs.mean(dim=0)).sum()
        return out, choice, balance

def experiment(balance_weight):
    torch.manual_seed(0)
    centers = torch.randn(4, 8) * 3                         # four groups of inputs...
    maps = torch.randn(4, 8, 8)                             # ...each with its own target function
    group = torch.arange(400) % 4
    x = centers[group] + torch.randn(400, 8) * 0.5
    y = (x.unsqueeze(1) @ maps[group]).squeeze(1) / 3
    moe = MoE(8, n_experts=4, hidden=16)
    opt = torch.optim.Adam(moe.parameters(), lr=1e-2)
    for step in range(301):
        out, choice, balance = moe(x)
        error = ((out - y) ** 2).mean()
        loss = error + balance_weight * balance
        opt.zero_grad()
        loss.backward()
        opt.step()
    print(f"balance weight {balance_weight}: error {error.item():.3f}, balance {balance.item():.2f}")
    for g in range(4):
        print(f"  group {g} → tokens per expert {torch.bincount(choice[group == g], minlength=4).tolist()}")

experiment(0.0)
experiment(1.0)

Both runs learn to send each group to one expert. Routing by what the input looks like appears on its own, with no labels.

Without the balance loss, one expert ends up handling two groups while another gets nothing. The error is higher. With the balance loss, each group gets its own expert and the balance score is 1.00. The error is lower, because no expert is wasted.

In real models, balance matters for speed too. Experts usually live on different GPUs. An overloaded expert makes all the others wait.

What experts learn

The name suggests a "marketing expert" and a "coding expert". Studies of real MoE models find something less tidy. In Mixtral's paper, routing showed no clear split by subject such as math, biology or philosophy. It depended more on grammar and the type of token. The same expert was often picked for tokens that come one after another.

So experts are a cheap way to add capacity. They aren't a team of specialists you can call on by name.

The trade-offs

• Cheaper compute per token for a given total size. Generation is faster and cheaper. Training also gets more capacity for each unit of computing work (each training FLOP).
• More memory. Every expert must be loaded, because any token might pick any expert. Serving needs as much GPU memory as a dense model of the total size.
• Harder training and serving. You have to balance the load and pass data between GPUs that hold different experts. MoE models also have training instabilities that dense models don't have.

For marketers: when you compare models by price and speed, MoE explains why a model with a huge parameter count can still be fast and inexpensive per token. What you pay for at run time is mostly the active parameters.