Mixture of experts

Problem

Bigger transformers are better, but in a dense model every parameter takes part in every token's forward pass, so compute per token grows in step with parameter count. The feed-forward blocks hold about two thirds of a dense transformer's non-embedding parameters, and each token uses all of them.

Mixture of experts breaks that link. Keep E copies of the feed-forward block, the experts, and let a learned router pick k of them per token, with k usually 1 or 2. Parameters grow with E; compute per token grows with k. The cost moves elsewhere: memory to hold every expert, communication to send tokens to wherever their experts live, and a routing problem that has to be trained to stay balanced.

Intuition

Think of a help desk with many specialists and a receptionist. Each question goes to the one or two specialists the receptionist thinks fit best, and the answer is a blend of theirs, weighted by the receptionist's confidence. The desk as a whole knows a great deal; each question only costs two specialists' time. Two things can go wrong: the receptionist can send everyone to the same specialist, who then has a queue while the others sit idle, and a specialist with a full queue has to turn people away. The load-balancing loss and the capacity factor below deal with exactly those two problems.

Four experts hold four feed-forward blocks' worth of parameters; this token's forward pass runs two of them.

Mechanics

Replacing the feed-forward block

In a transformer block, attention mixes information across positions and the feed-forward network (FFN) then transforms each position on its own. Because the FFN is applied token by token, nothing stops different tokens from using different FFNs. A mixture-of-experts layer keeps the attention sublayer and swaps the FFN for E expert FFNs E₁ … E_E (each the same shape as the original) plus a router matrix Wᵣ with E rows:

h = Wᵣ x,   p = softmax(h),   y = Σ_{i ∈ TopK(h)} gᵢ(x) · Eᵢ(x)

The router is one linear layer, so it costs E · d multiply-adds per token, tiny next to an expert. The gates gᵢ multiply each expert's output, which is how the router gets gradient: the choice of experts is discrete, but the gate values are differentiable, and making a gate larger for an expert whose output lowered the loss is what teaches the router. Not every block has to be sparse: GShard replaces every other FFN, Mixtral replaces all of them.

Top-k gating

The router keeps the k largest logits. Two gate conventions are common:

• Switch (k = 1). Each token goes to argmax h, and the gate is that expert's full-softmax probability pᵢ, which is below 1. Switch Transformers argues that one expert per token is enough and halves the routing computation and communication compared with k = 2.
• Renormalized top-k (Mixtral, k = 2). g = softmax over the k selected logits only, so the gates sum to 1. This equals taking the full softmax and dividing the chosen probabilities by their sum, because the shared normalizer cancels.

The original sparsely-gated layer (Shazeer et al.) added tunable Gaussian noise to the logits before the top-k, which spreads early routing decisions; GShard used top-2 with the second expert chosen at random in proportion to its weight. Later models mostly use the plain deterministic top-k above.

The load-balancing loss

Left alone, routing collapses: whichever experts start slightly better get more tokens, train faster and get more tokens still, while the rest never learn. The Switch Transformer auxiliary loss pushes back. Over a batch of T tokens, let fᵢ be the fraction of tokens whose top choice is expert i, and Pᵢ the mean router probability for expert i:

L_aux = α · E · Σ_i fᵢ · Pᵢ,   α = 0.01 in the paper

fᵢ is a count, so it has no gradient; Pᵢ does, and the gradient of fᵢ · Pᵢ with respect to Pᵢ is fᵢ, so the loss pushes probability away from the experts that are already busiest. When the routing is perfectly uniform, every fᵢ is 1/E and the sum is 1/E whatever the Pᵢ, so E · Σ fᵢ Pᵢ = 1; when every token goes to one expert with probability 1, it is E. The factor E keeps the value comparable as the number of experts changes. ST-MoE adds a router z-loss, the mean over tokens of (log Σⱼ exp hⱼ)², with a small coefficient (10⁻³ in the paper). It keeps the router's logits from growing large, which the authors found made training more stable, especially in low precision.

Capacity and dropped tokens

On accelerators, each expert processes a fixed-size buffer so that shapes are static and the work can be split evenly across devices. Switch Transformers sets

expert capacity = (tokens per batch / number of experts) × capacity factor

A factor of 1 leaves exactly enough slots for perfect balance; values above 1 leave slack for uneven routing, at the price of padding and memory. When more tokens choose an expert than it has slots, the overflow tokens are dropped: they skip the expert computation, and the residual connection carries their input forward unchanged. Which tokens are dropped depends on the order slots are filled; the simplest rule, used below, fills them in position order, so tokens late in the batch are the ones lost. Dropless implementations (MegaBlocks, and the loop-over-experts style in the implementation section) instead give each expert a variable-size batch and drop nothing, at the cost of irregular shapes.

Total versus active parameters

A sparse model is described by two numbers. Total parameters count every expert, and decide the memory needed to hold the model. Active parameters count what one token's forward pass touches: the shared parts (embeddings, attention, norms, routers) plus k experts per MoE layer, and decide the compute per token. Mixtral 8x7B, worked out below, has about 46.7B total and 12.9B active parameters, not 56B, because only the FFNs are replicated. DeepSeek-V3 takes the ratio further, with 671B total and 37B active, using many small experts plus shared experts every token visits.

Active parameters set the FLOPs but not the cost of serving. At small batch sizes decoding is limited by reading weights from memory, and different tokens in the batch touch different experts, so a batch can read most of the total. Across many devices, expert parallelism places different experts on different devices and exchanges tokens with two all-to-all communications per layer, one to dispatch tokens and one to bring the outputs back.

Worked example

Red squares are dropped at capacity factor 1.0: they pass through the residual connection without any expert. Doubling the factor saves two of them and leaves experts 1 to 3 three-quarters empty.

Compute follows the lower bar, memory the upper one.

Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class SwiGLU(nn.Module):
    def __init__(self, d, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d, d_ff, bias=False)
        self.w3 = nn.Linear(d, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class MoE(nn.Module):
    def __init__(self, d, d_ff, num_experts, k=2, capacity=None):
        super().__init__()
        self.router = nn.Linear(d, num_experts, bias=False)
        self.experts = nn.ModuleList(SwiGLU(d, d_ff) for _ in range(num_experts))
        self.k, self.capacity = k, capacity

    def forward(self, x):                                  # x: (T, d), batch and sequence flattened
        logits = self.router(x)                            # (T, E)
        probs = logits.softmax(-1)
        top_logits, top_idx = logits.topk(self.k, dim=-1)  # (T, k), largest first
        gates = top_logits.softmax(-1)                     # renormalized over the k chosen
        out = torch.zeros_like(x)
        for e, expert in enumerate(self.experts):
            token, slot = (top_idx == e).nonzero(as_tuple=True)   # in token order
            if self.capacity is not None:
                token, slot = token[:self.capacity], slot[:self.capacity]  # the rest are dropped
            if token.numel():
                out.index_add_(0, token, gates[token, slot, None] * expert(x[token]))
        E = len(self.experts)
        f = F.one_hot(top_idx[:, 0], E).float().mean(0)   # share of tokens whose first choice is e
        aux = E * (f * probs.mean(0)).sum()                # 1 when balanced, E when collapsed
        z = torch.logsumexp(logits, -1).pow(2).mean()
        return out, aux, z

# In the block: h = h + moe_out; in the loss: loss = task_loss + 0.01 * aux + 0.001 * z

The loop over experts gathers each expert's tokens, runs the expert once on that sub-batch and scatters the gated results back with index_add_, so a token with two experts gets the sum of both. With capacity=None nothing is dropped. Production kernels replace the Python loop with a sort of tokens by expert and a grouped matrix multiply, and across devices with an all-to-all exchange, but the arithmetic is the same. Two details matter in practice: compute the router in float32 even under mixed precision, since small logit errors flip discrete choices, and count the auxiliary losses per MoE layer, summing them across layers.

Tradeoffs

• Quality per FLOP improves, quality per byte does not. An MoE model usually beats a dense model with the same active parameters, and usually loses to a dense model with the same total parameters. Choose it when compute, not memory, is the limit.
• Serving is harder. All experts must be resident, batches spread across experts unevenly, and at low batch sizes memory bandwidth, which scales with total parameters read, dominates. Quantization and expert offloading help.
• Training is less stable. Routing can collapse or oscillate; the balancing and z losses, float32 routers and careful initialization are standard. Fine-tuning a sparse model has been reported to overfit more readily than its dense counterpart.
• Experts do not specialize the way the name suggests. Analyses of Mixtral found little specialization by topic and more by syntax and token type, with consecutive tokens often going to the same expert.

Connections

The transformer page builds the block whose feed-forward sublayer the experts replace, and decoder-only language models are where MoE is now most used. Expert parallelism is one more axis alongside the data, tensor and pipeline parallelism of distributed training, and the memory-bound serving problem is the one on latency and throughput. The idea of combining specialists by a learned weighting goes back to ensembles like random forests, with the difference that a router picks a few members per input instead of averaging all of them.

References

• Jacobs, Jordan, Nowlan and Hinton. Adaptive Mixtures of Local Experts. Neural Computation 1991. doi:10.1162/neco.1991.3.1.79
• Shazeer et al. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017. arXiv:1701.06538
• Lepikhin et al. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. 2020. arXiv:2006.16668
• Fedus, Zoph and Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR 2022. arXiv:2101.03961
• Zoph et al. ST-MoE: Designing Stable and Transferable Sparse Expert Models. 2022. arXiv:2202.08906
• Gale et al. MegaBlocks: Efficient Sparse Training with Mixture-of-Experts. MLSys 2023. arXiv:2211.15841
• Jiang et al. Mixtral of Experts. 2024. arXiv:2401.04088
• Dai et al. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. 2024. arXiv:2401.06066
• DeepSeek-AI. DeepSeek-V3 Technical Report. 2024. arXiv:2412.19437
• PyTorch v2.15.0: nn.ModuleList, nn.Linear, torch.topk, index_add, F.softmax.