Mixture of Experts (MoE)

Every scaling-law plot says the same thing: more parameters, lower loss. But in a dense transformer every single token pays for every single parameter. If you want a model that knows ten times more, you have to compute ten times more for the word "the". That is the waste MoE attacks.

Think about what a big MLP in a transformer block is doing. It is a lookup-and-transform machine: the token vector comes in, gets compared against thousands of learned directions, and gets pushed toward an output. Most of those directions are irrelevant to most tokens. A token in a Python snippet does not need the part of the network that has memorised French verb conjugations.

So here is the question this chapter answers: can we have a huge amount of knowledge stored in the weights, but only touch a small slice of it for each token? The answer is yes, and the trick is old (Jacobs et al. 1991), was revived for deep learning by Shazeer et al. (2017), and went mainstream with the Switch Transformer (Fedus et al. 2021), Mixtral, and DeepSeek-V3.

The problem: dense compute grows with knowledge

Let us put a number on the waste. In a dense transformer, the floating-point operations per token are roughly $2N$ where $N$ is the number of non-embedding parameters (one multiply and one add per weight). Double the parameters and you double the FLOPs for every token, whether that token is a rare chemistry term or a comma.

The MLP is where most of those parameters live. In a standard block with $d_{model}=4096$ and an MLP hidden size $d_{ff}=14336$ using SwiGLU (three weight matrices), the MLP alone has $3 \times 4096 \times 14336 \approx 176\text{M}$ parameters per layer, versus about $42\text{M}$ for attention with grouped-query heads. The MLP is 80% of the block.

So if we want to grow the model's capacity cheaply, the MLP is the obvious target. What if, instead of one big MLP that every token runs through, we had several MLPs and each token only ran through a couple of them?

The idea: E experts and a router

Replace the single MLP in a transformer block with $E$ separate MLPs called experts, each with the same shape as a normal MLP (or smaller). Add a tiny linear layer called the router (or gate) that looks at each token and decides which $k$ experts should process it. The token goes through only those $k$ experts and the results are combined with weights from the router.

Nothing else in the block changes. Attention, normalisation, and residual connections are exactly as before. An MoE model is just a transformer where some (or all) of the MLPs have been swapped for this "choose $k$ of $E$" layer.

Total versus active parameters

MoE forces us to separate two numbers that are the same in a dense model. Total parameters is how many weights exist and must be stored. Active parameters is how many weights are actually multiplied for a given token. Compute per token scales with active parameters; memory scales with total.

The arithmetic, per layer

Let a block have $A$ attention parameters and each expert MLP have $M$ parameters. With $E$ experts and $k$ active per token, plus $s$ "shared" experts that every token always uses (more on those later), the counts per layer are:

Multiply by the number of layers and add the embedding matrices, and you have the two headline numbers of any MoE model. Notice that attention is paid in full by every token; only the MLP part is sparse.

DeepSeek-V3 (DeepSeek-AI, 2024) pushes the ratio much further: it reports approximately 671B total parameters and 37B activated per token, with 256 small routed experts per layer, $k=8$, and one shared expert. That is an 18:1 ratio of stored to touched knowledge, compared with Mixtral's 3.6:1.

The presets are approximations: the calculator treats attention as a full $4d^2$ (real models use grouped-query or multi-head latent attention, which is smaller) and ignores details like DeepSeek-V3's first three dense layers. They land within a few percent of the reported 46.7B/12.9B and 671B/37B figures, which is the point: you can reconstruct the headline numbers of any MoE from six hyperparameters.

The router, slowly

The router is the only genuinely new piece, so let us build it from a single token. What we need is a function that turns a token vector $x$ into a small set of experts and a weight for each. The simplest thing that works is a linear layer followed by a softmax.

Step 1: logits and probabilities

First we score every expert with a dot product against a learned vector. The router weight matrix $W_g$ has one column per expert:

So $p$ is a probability distribution over experts for this token, exactly like the output of a tiny $E$-way classifier. The router has only $d \times E$ parameters, a rounding error next to the experts themselves.

Step 2: top-k selection and renormalisation

We do not want to run all $E$ experts, so we keep only the $k$ largest probabilities and rescale them so they still sum to one:

The renormalisation step matters more than it looks. Without it, a token whose top-2 probabilities were $0.3$ and $0.2$ would get an output scaled by $0.5$, while a confident token would get its output at full strength. Renormalising makes the MoE output the same scale as a dense MLP output regardless of router confidence, which keeps the residual stream well behaved. (The Switch Transformer with $k=1$ deliberately keeps the un-normalised $p_i$ so the router still receives gradient through the gate value; with $k \ge 2$ the renormalised weights already carry gradient.)

Step 3: the weighted sum

Finally, the output is the gate-weighted sum of the chosen experts' outputs:

This is the whole layer. The gradient flows into the chosen experts through their outputs and into the router through the gates $g_i$. Experts that were not chosen get no gradient from this token, which is both the point (sparsity) and the source of every training headache in the rest of this chapter.

Two things to notice in the playground. First, as you sharpen the logits the gate weights go toward one-hot, and with $k=2$ the second expert's weight shrinks toward zero: the model can learn to ignore its second choice. Second, even with random logits the loads are rarely even, and a few rerolls will show you an expert that gets nothing. That is a preview of the central training problem.

In code

Here is a router in PyTorch that returns the gates, the chosen indices, and the full probabilities (needed for the balancing loss). This mirrors code/lumen/moe.py.

import torch, torch.nn as nn

class TopKRouter(nn.Module):
    def __init__(self, d_model, n_experts, k):
        super().__init__()
        self.gate = nn.Linear(d_model, n_experts, bias=False)
        self.k = k

    def forward(self, x):                            # x: [T, d]
        logits = self.gate(x).float()                # [T, E]  (fp32: see z-loss)
        probs = logits.softmax(dim=-1)
        top_p, top_i = probs.topk(self.k, dim=-1)    # [T, k] each
        gates = top_p / top_p.sum(dim=-1, keepdim=True)   # renormalise
        return gates, top_i, probs

And the layer itself. The loop over experts is the simple, readable version: gather the tokens that picked each expert, run them as one batch, scatter the weighted results back. Production kernels fuse this into grouped matrix multiplies, but the arithmetic is identical.

class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, n_experts, k):
        super().__init__()
        self.router = TopKRouter(d_model, n_experts, k)
        self.experts = nn.ModuleList([SwiGLU(d_model, d_ff) for _ in range(n_experts)])

    def forward(self, x):                            # x: [T, d]
        gates, idx, probs = self.router(x)           # [T,k], [T,k], [T,E]
        y = torch.zeros_like(x)
        for e, expert in enumerate(self.experts):
            tok, slot = (idx == e).nonzero(as_tuple=True)   # which tokens chose e
            if tok.numel() == 0:
                continue
            y[tok] += gates[tok, slot].unsqueeze(-1) * expert(x[tok])
        return y, probs, idx

Load balancing: why routers collapse

Here is the catch. Nothing in the router's objective says the experts should be used evenly. Suppose, by random initialisation, expert 3 is slightly better than the others at the start of training. Tokens routed to it get a lower loss, so the router is pushed to send more tokens to expert 3. Expert 3 then gets more gradient, gets better, and attracts still more traffic.

Meanwhile experts 1, 2, and 4 get few tokens, few gradients, and never catch up. The end state is a model that is paying for $E$ experts and using one. Shazeer et al. called this the "rich get richer" problem, and every MoE recipe has some mechanism to fight it.

The auxiliary load-balancing loss

The standard fix is to add a small extra loss that is minimised when the tokens are spread evenly. We need two quantities per expert over a batch of $T$ tokens. The first is the fraction of tokens actually dispatched to expert $i$:

The indicator counts a token if expert $i$ is among its top-$k$; dividing by $Tk$ makes the $f_i$ sum to 1. The second is the average router probability for expert $i$, before top-$k$:

Now the loss, from Fedus et al. (2021), is the dot product of these two vectors, scaled by $E$ so that its minimum is $1$:

Why the product of two things? $f_i$ is a hard count and has no gradient: you cannot differentiate through "which expert was picked". $P_i$ is a smooth function of the router weights and does have gradient. Multiplying them means the gradient on $P_i$ is proportional to $f_i$: the router is pushed to lower its probability for experts that are already overloaded, exactly in proportion to how overloaded they are. Under a mild assumption the sum $\sum_i f_i P_i$ is minimised when both are uniform at $1/E$, giving $E \cdot E \cdot (1/E)^2 = 1$.

The coefficient $\alpha$ is typically small, around $10^{-2}$ (Switch) down to $10^{-3}$ or less in later work. Too large and the router balances at the expense of routing quality; too small and you get collapse anyway. It is one of the more annoying knobs in MoE training, which is why the field has kept looking for alternatives.

Capacity factor and token dropping

Even with the loss, batches will not be perfectly even, and hardware wants fixed-size tensors. The Switch Transformer's answer is a hard cap: each expert gets a fixed capacity, the maximum number of tokens it will process per batch:

Here $Tk/E$ is the number of tokens each expert would receive if the load were perfectly even, and the capacity factor CF is a slack multiplier. With CF $= 1.0$, an expert that is even slightly over-subscribed will overflow; with CF $= 1.25$ it can absorb 25% extra.

What happens to a token that arrives at a full expert? It is dropped: the expert contributes nothing, and the token's output for this layer is just the residual passed through. The model still works (the residual stream carries the information forward), but that token got no MLP computation in this layer. Dropping a few percent of tokens is tolerable; dropping many degrades quality noticeably.

The trade-off is in the padding. Experts always compute on exactly $C$ slots (unused ones are zeros), so a large CF wastes compute on padding while a small CF drops tokens. Typical training values are $1.0$–$1.25$, with evaluation sometimes using $2.0$ to avoid drops. Many recent models (Mixtral, DeepSeek-V3) are trained dropless: they use variable-size grouped kernels and never drop a token, relying on balancing losses or biases alone.

Expert-choice routing

Token-choice routing (each token picks experts) is what we have described so far, and it is what makes balance a problem: tokens are free to all pick the same expert. Zhou et al. (2022) flipped the question: what if each expert picks its top-$C$ tokens instead?

Compute the same $T \times E$ score matrix, but now take the top-$C$ along the token axis for each expert column. Every expert receives exactly $C$ tokens by construction, so balance is perfect and no auxiliary loss is needed. The cost is that tokens are no longer guaranteed anything: a popular token may be picked by many experts while an unpopular one is picked by none (and passes through on the residual). Expert-choice also needs the whole batch of tokens before it can decide, which makes it awkward for autoregressive decoding, where tokens arrive one at a time. It is mostly used in encoders and in training.

Aux-loss-free balancing with bias terms

DeepSeek-V3 reports a third approach (Wang et al. 2024) that removes the balancing loss from the gradient entirely. Add a per-expert bias $b_i$ to the router score, but use it only for deciding the top-$k$, not for the gate value:

After each training step, look at the load: for every expert that received more than its fair share, decrease $b_i$ by a small step $\gamma$; for every under-loaded expert, increase $b_i$ by $\gamma$. The bias is a controller, not a learned parameter. It nudges over-popular experts out of the top-$k$ without distorting the gate weights or adding a gradient that fights the language-modelling loss. DeepSeek-V3 reports that this balances better than the auxiliary loss at the same quality; they still keep a very small sequence-level balancing loss as a safety net.

Fine-grained and shared experts

Mixtral-style MoE uses a handful of large experts, each the size of a dense MLP. DeepSeekMoE (Dai et al. 2024) asked whether that granularity is right, and argued for two changes that DeepSeek-V2 and V3 then adopted.

Fine-grained experts: more combinations for the same FLOPs

Split each expert into $m$ smaller experts (with $1/m$ the hidden width) and activate $mk$ of them instead of $k$. The active parameters and FLOPs are unchanged, but the number of possible expert combinations explodes. With 16 experts choose 2 there are $\binom{16}{2} = 120$ combinations; with 64 quarter-size experts choose 8, there are $\binom{64}{8} \approx 4.4 \times 10^9$. The idea is that finer experts can be composed more flexibly, so knowledge is stored in a more decomposed way with less redundancy across experts.

This is why DeepSeek-V3 has 256 experts per layer with an intermediate size of only 2048 (versus a 7168 model width), and activates 8 of them. Each expert is small, but a token gets eight of them.

Shared experts: always-on common knowledge

The second change: reserve $s$ experts that every token goes through, with no routing. The reasoning is that some computation is needed by all tokens (basic syntax, common transformations), and if there is no shared path, every routed expert has to learn it redundantly. Isolating that common knowledge in a shared expert frees the routed experts to specialise. DeepSeek-V3 uses one shared expert alongside its 8 routed ones per token, so $k + s = 9$ experts are active.

Training instabilities and the router z-loss

MoE models are known for being twitchier to train than dense ones. Part of it is the discrete routing: a tiny change in a logit can flip which expert a token goes to, which changes the output discontinuously. Part of it is numerical: the router logits are the input to a softmax that is then used to pick a hard argmax, and if those logits grow large, round-off in bf16 can change the argmax.

ST-MoE (Zoph et al. 2022) studied this carefully and proposed two fixes. First, compute the router in float32 even when the rest of the model is bf16 (that is the .float() in our router code). Second, add a router z-loss that penalises large logits:

The inner term is the log-partition function of the softmax, which grows when any logit grows. Squaring it and averaging gives a loss that gently pulls all router logits toward zero, keeping the softmax in its well-conditioned regime. A typical coefficient is $10^{-3}$. For our worked example logits $[1.0, 2.0, 0.5, -1.0]$, $\log \sum e^{h} = \log 12.12 = 2.49$ and the z-loss contribution is $2.49^2 = 6.22$; if the logits were ten times larger the contribution would be roughly a hundred times larger, which is the point.

Other reported stabilisers: initialising the router with a small scale, using a lower learning rate on router weights, and (in DeepSeek-V3) the bias-based balancing that avoids gradient fights. MoE models are also more sensitive to the fine-tuning recipe; ST-MoE reports that they overfit small fine-tuning sets more easily than dense models.

Systems: where the parameters live and how tokens travel

So far we have treated an MoE layer as a single-device object. In practice the whole reason to build one is that it is too big for a single device. Mixtral's 47B parameters are about 94 GB in bf16; DeepSeek-V3's 671B are over 1.3 TB. Since experts are independent by construction, the natural way to spread them out is expert parallelism (EP): put different experts on different GPUs.

Expert parallelism and all-to-all

With EP, each GPU holds the attention weights (replicated, or sharded by tensor parallelism) plus a subset of the experts. Tokens live on the GPU that ran their attention. The problem: a token on GPU 0 may be routed to an expert that lives on GPU 3.

So every MoE layer needs two communication rounds. In the dispatch all-to-all, every GPU sends each of its tokens to the GPU that owns its chosen expert. Experts then compute on whatever arrived. In the combine all-to-all, the outputs are sent back to the token's home GPU, where they are gate-weighted and added to the residual. This pattern was introduced for transformers by GShard (Lepikhin et al. 2020).

Why MoE is memory-heavy but compute-light

The stepper shows the shape of the cost. Compute per token is that of a $k$-expert model, but the parameters, the optimizer states, and the gradients for all $E$ experts must exist somewhere. In training, an Adam state alone is 8 bytes per parameter in fp32 (see the optimizers chapter), so DeepSeek-V3's 671B parameters imply several terabytes of optimizer state spread across the cluster, for a model whose forward pass costs about as much as a 37B dense model.

The all-to-alls are the other tax. They are latency-bound (many small messages) rather than bandwidth-bound, they happen twice per layer, and they cannot start until routing is done. Real systems overlap them with computation, group experts so that most traffic stays inside a node's fast interconnect, and (DeepSeek-V3 reports) limit each token to experts on at most a few nodes to bound cross-node traffic.

Inference implications

At serving time the memory point bites hardest. To generate one token, the model consults only $k$ experts per layer, but you do not know which ones in advance, so all $E$ must be resident (or fetchable). A 47B MoE needs the memory of a 47B dense model while doing the work of a 13B one. With batch size 1 this is a poor trade: you are paying for memory bandwidth to read $k$ experts' weights, and that bandwidth, not FLOPs, is the bottleneck of decoding anyway.

MoE wins at inference when the batch is large. With hundreds of sequences in flight, every expert gets some tokens, so all the weights you read are used, and the FLOPs saving translates into throughput. That is why MoE is attractive for high-volume API serving and less so for a single user on a laptop.

For the laptop case, expert offloading keeps the attention weights and a cache of recently used experts on the GPU and streams the rest from CPU memory or disk as needed. Because consecutive tokens tend to reuse experts (reported in the Mixtral analysis below), a small cache hits often enough to be practical, at some latency cost (see Eliseev & Mazur 2023).

Upcycling: dense to MoE without starting over

You do not have to train an MoE from scratch. Komatsuzaki et al. (2022) proposed sparse upcycling: take a trained dense checkpoint, copy its MLP $E$ times to make $E$ identical experts, add a freshly initialised router, and continue training. At the moment of conversion the model computes exactly what the dense model did (every expert is the same, so the weighted sum is the same), so nothing is lost. As training continues, the experts drift apart and the router learns to exploit the differences. The paper reports that this reaches a given quality with a fraction of the compute of training the MoE from scratch, and several later open models are reported to have been built this way. The catch is that upcycled experts start out perfectly correlated and may stay more similar to each other than experts trained from scratch.

What do the experts actually learn?

The word "expert" suggests that one expert handles code, another handles biology, and so on. The reported evidence says otherwise. Jiang et al. (2024) analysed Mixtral's routing across datasets (arXiv papers, GitHub code, PhilPapers, PubMed, StackExchange, Wikipedia) and found no obvious per-domain expert assignment: routing looked similar across most domains, with only code and mathematics showing a mildly different distribution.

What they did find was structure at the syntactic level. Tokens like self in Python, or indentation whitespace, were consistently sent to the same experts, and consecutive tokens were routed to the same expert far more often than chance would predict. Other analyses of open MoE models (for example OLMoE, Muennighoff et al. 2024) report similar patterns: specialisation exists, but along token-level and positional lines more than topical ones, and it varies by layer.

Putting it together: an MoE transformer block

Everything in this chapter fits into one block that drops in for a dense one. The pieces: a router, $E$ experts, an optional shared expert, and the two extra losses returned to the training loop.

class MoEBlock(nn.Module):
    def __init__(self, d_model, d_ff, n_experts, k, n_shared=0):
        super().__init__()
        self.ln1, self.ln2 = RMSNorm(d_model), RMSNorm(d_model)
        self.attn = CausalSelfAttention(d_model)
        self.moe = MoELayer(d_model, d_ff, n_experts, k)
        self.shared = nn.ModuleList([SwiGLU(d_model, d_ff) for _ in range(n_shared)])

    def forward(self, x):                                  # x: [B, T, d]
        x = x + self.attn(self.ln1(x))
        h = self.ln2(x)
        flat = h.reshape(-1, h.size(-1))                   # [B*T, d]
        y, probs, idx = self.moe(flat)
        for s in self.shared:
            y = y + s(flat)                                # shared experts: every token
        x = x + y.view_as(x)
        aux = load_balance_loss(probs, idx, self.moe.router.gate.out_features, self.moe.router.k)
        return x, aux

def load_balance_loss(probs, idx, n_experts, k):
    T = probs.shape[0]
    f = torch.zeros(n_experts, device=probs.device)
    f.index_add_(0, idx.flatten(), torch.ones(T * k, device=probs.device))
    f = f / (T * k)                                        # fraction dispatched
    P = probs.mean(dim=0)                                  # mean router prob
    return n_experts * (f * P).sum()                       # == 1.0 when balanced

In the training loop, the total loss is $\mathcal{L}_{\text{LM}} + \alpha \sum_{\ell} \mathcal{L}_{\text{aux}}^{(\ell)} + \beta \sum_{\ell} \mathcal{L}_{z}^{(\ell)}$ summed over MoE layers. code/lumen/moe.py includes the z-loss, a capacity-limited variant that drops tokens, and a test that checks the layer reduces to a dense MLP when $E = k = 1$.

When MoE is and is not the right choice

MoE is the right tool when you are compute-limited but not memory-limited: a large training cluster, or a serving fleet with big batches. The reported gains are substantial: Switch reports reaching the same loss as a dense T5 several times faster in wall-clock; Mixtral matches or beats Llama 2 70B on most benchmarks with 5× fewer active parameters; DeepSeek-V3 reports training at a total cost well below comparable dense models.

MoE is the wrong tool when memory is the constraint (a single consumer GPU), when your batches are small, or when you want the simplest possible training run. It adds a router to tune, two extra losses, a communication pattern that stresses the interconnect, and a fine-tuning process that is reported to be more fragile. It also complicates every downstream technique: quantisation must handle rarely-used experts, LoRA has to decide which experts to adapt, and distillation from an MoE teacher is common precisely because deploying the MoE is inconvenient.

Practice

Further reading

• Shazeer et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. The paper that made MoE practical for deep learning; the noisy top-$k$ gate and load losses start here.
• Fedus, Zoph & Shazeer (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. $k=1$ routing, the capacity factor, the auxiliary loss in the form used here.
• Lepikhin et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. Expert parallelism and all-to-all for transformers.
• Zoph et al. (2022). ST-MoE: Designing Stable and Transferable Sparse Expert Models. Router z-loss and a careful study of instabilities and fine-tuning.
• Zhou et al. (2022). Mixture-of-Experts with Expert Choice Routing. Let experts pick tokens for guaranteed balance.
• Jiang et al. (2024). Mixtral of Experts. The open 8×7B model and the routing analysis discussed above.
• Dai et al. (2024). DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. Fine-grained and shared experts.
• Wang et al. (2024). Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts. The bias-controller approach used in DeepSeek-V3.
• DeepSeek-AI (2024). DeepSeek-V3 Technical Report. 671B/37B, 256 fine-grained experts, node-limited routing, and the systems work that makes it train.
• Komatsuzaki et al. (2022). Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints. Turning a trained dense model into an MoE.
• Muennighoff et al. (2024). OLMoE: Open Mixture-of-Experts Language Models. A fully open MoE with extensive routing and specialisation analysis.
• Eliseev & Mazur (2023). Fast Inference of Mixture-of-Experts Language Models with Offloading. Expert caching and offloading for consumer hardware.