Chapter 3.3 β€” Mixture of Experts

Contents

  1. From cheaper and longer attention to a bigger model at the same cost
  2. Conditional computation: the core idea
  3. The sparsely-gated MoE layer and top-k routing
  4. Load balancing: the central practical difficulty
  5. Switch Transformer: simplifying to top-1
  6. Mixtral: MoE as a mainstream production technique
  7. Interview angle
  8. Self-check questions
  9. Sources

1. From cheaper and longer attention to a bigger model at the same cost

The last two chapters were both about doing the same job more cheaply or over a longer sequence: Chapter 3.1 made attention itself less expensive, and Chapter 3.2 made a fixed model's positional understanding stretch further. This chapter changes the axis of the discussion entirely. Instead of asking "how do we do the same computation for less," it asks "how do we get a much larger, more capable model without paying a proportionally larger compute bill per token." That might sound like it's asking for something for nothing, and in the fully dense architecture that Part II described β€” where every token passes through every parameter in every feedforward layer β€” it would be. Mixture of Experts (MoE) is the architectural change that makes it possible, by breaking the assumption, quietly present in everything covered so far, that a token's forward pass has to touch every parameter the model owns.

The core move is conditional computation: rather than one feedforward network that every token is forced through, you build many feedforward networks β€” "experts" β€” and, per token, route that token to only a small subset of them. Total parameter count can then grow essentially without bound by adding more experts, while the compute cost of processing any single token stays roughly fixed, since a token only ever activates a small, fixed number of those experts. This decouples two quantities that were previously locked together in a dense Transformer: total model capacity (loosely, total parameter count) and per-token compute cost (loosely, FLOPs per forward pass). Everything in this chapter is really about how to make that decoupling work well in practice, because the idea is much easier to state than to execute β€” and the biggest practical obstacle, load balancing, will turn out to be almost as important to understand as the core routing idea itself.

2. Conditional computation: the core idea

To make conditional computation concrete, it helps to contrast it directly with the dense feedforward layer from Part II. In a standard Transformer block, the feedforward sublayer is a single two-layer MLP (typically an up-projection, a nonlinearity, and a down-projection) applied identically to every token's representation at that position in the network. If you wanted to increase that layer's capacity, your only lever was to make it wider or deeper β€” and every token would have to pay the full compute cost of that larger MLP, whether or not that particular token actually needed the extra capacity for that particular prediction. This is wasteful in a specific, identifiable way: not every token needs the same kind of processing. A token that's part of a rare technical term, a token continuing a simple grammatical pattern, and a token beginning a shift in topic plausibly benefit from qualitatively different transformations, and forcing all three through the identical dense computation is a blunt instrument.

MoE's answer is to replace that single dense feedforward network with many parallel feedforward networks β€” the experts, each with its own independent parameters β€” plus a small learned routing function, usually called the gate or router, that looks at each token's representation and decides which expert or experts should process it. Crucially, the token is only actually run through the small number of experts the router selects, not all of them; the un-selected experts contribute zero compute for that token. This is what makes MoE a genuine decoupling of capacity from cost: adding a hundred more experts to the layer adds a hundred times more total parameters sitting in the model, but if the router still only ever selects, say, two experts per token, the compute cost of processing any given token barely changes. The model becomes enormous in total size while remaining cheap to run per token β€” provided, and this caveat matters enormously for the rest of the chapter, that the routing decisions actually spread work sensibly across all those experts rather than collapsing onto just a few of them.

3. The sparsely-gated MoE layer and top-k routing

The modern formulation of this idea traces to Shazeer et al.'s sparsely-gated Mixture-of-Experts layer, which set the template that essentially all subsequent MoE work builds on. The layer consists of $n$ expert networks $E_1, \ldots, E_n$ and a trainable gating network $G$ that, given an input token representation $x$, produces a sparse vector of $n$ weights, one per expert, most of which are exactly zero. The layer's output is the weighted sum $y = \sum_{i=1}^{n} G(x)_i \, E_i(x)$, but because $G(x)$ is sparse β€” typically only its top $k$ entries are nonzero, for a small $k$ like 1 or 2 β€” the sum only actually needs to evaluate the $k$ experts whose gate weight is nonzero; the rest can be skipped entirely, since multiplying by a zero gate weight makes their contribution zero regardless of what $E_i(x)$ would have produced.

The gating network itself is typically a simple learned linear (or lightly nonlinear) function of the token representation, followed by a softmax over all $n$ experts, from which only the top-$k$ scoring experts are kept and renormalized. This top-$k$ operation is what makes the whole scheme sparse rather than a soft weighted blend of every expert β€” without it, you'd be back to paying the compute cost of every expert on every token, since a genuine (not sparse) softmax weighting would still require every $E_i(x)$ to actually be computed before it could be weighted and summed. Shazeer et al.'s specific contribution was showing that this top-$k$ gating could be trained end-to-end with standard backpropagation despite the discrete, non-differentiable-seeming top-$k$ selection β€” handled by treating the gate values themselves as differentiable and adding noise to the routing logits during training, which both helps gradient flow through an otherwise discontinuous decision and helps distribute load across experts, a theme that turns out to need a great deal more attention than this initial mechanism alone provides.

4. Load balancing: the central practical difficulty

Getting a sparse gating network to converge on a well-functioning MoE layer turns out to be substantially harder than the mechanism in the previous section suggests, for a reason that's intuitive once stated: gradient descent is a rich-get-richer process, and if left alone, it will tend to make it worse. Suppose, early in training, one expert happens by chance to be slightly better than the others at handling some common category of token, purely due to random initialization. The router will notice this β€” its gradient signal will nudge it to route more of that category of token to that expert, since doing so slightly reduces the loss. But routing more tokens to that expert means that expert now receives more gradient updates than the others, on more data, and becomes even better at its job relative to the rest, which makes the router even more inclined to route to it, and so on. Left unchecked, this feedback loop tends to collapse onto a small subset of experts handling the overwhelming majority of tokens, while most of the other experts receive little to no training signal and end up contributing almost nothing to the model β€” which defeats the entire purpose of building a large bank of experts in the first place, since you're paying the memory cost of storing all of them while effectively only using a handful.

The standard fix is an auxiliary load-balancing loss, added to the main training objective specifically to counteract this collapse. The precise formulations vary across papers, but the underlying idea is consistent: construct a term that is minimized when tokens are routed roughly uniformly across all experts, and penalize the model when routing becomes too concentrated on a small subset. A common construction computes, for each expert, both the fraction of tokens actually routed to it and the average gating probability assigned to it across the batch, and then penalizes the product of these two quantities summed over experts β€” formally, $\mathcal{L}_{\text{aux}} = n\sum_{i=1}^{n} f_i \cdot P_i$, where $f_i$ is the fraction of tokens routed to expert $i$ and $P_i$ is its average gating probability β€” a formulation chosen specifically because it's differentiable everywhere (unlike the raw, hard, discrete routing decision) and because it's minimized precisely when load is spread evenly, giving the model a genuine gradient signal that actively counteracts the rich-get-richer dynamic rather than merely observing it after the fact. Getting the weight of this auxiliary loss right is itself a nontrivial hyperparameter choice: too small, and load imbalance creeps back in; too large, and the auxiliary objective starts to distort the primary language modeling objective the model is actually meant to optimize. Understanding that this tension exists, and that it's the central practical difficulty of training any MoE model rather than a minor implementation detail, is one of the more important things to internalize from this chapter.

Every routing scheme discussed so far is token-choice: each token independently picks its top-$k$ experts, and load balancing has to be imposed afterward via the auxiliary loss above, fighting the rich-get-richer dynamic rather than preventing it structurally. Zhou et al.'s expert-choice routing inverts the direction of selection to sidestep the problem at its source: rather than each token choosing experts, each expert is given a fixed capacity (a maximum number of token slots it can fill for the current batch) and independently picks its own top-scoring tokens from across the whole batch to fill those slots. Because every expert always fills the same fixed capacity by construction, load is balanced by definition, with no auxiliary loss needed at all to fight collapse. The trade is a different failure mode instead: a token could in principle be chosen by no expert at all (and get dropped, contributing nothing at that layer) or by several experts at once (processed redundantly), which is the opposite trade-off from token-choice routing's guarantee that every token gets processed by exactly $k$ experts, at the cost of needing an external mechanism to keep that guarantee from concentrating load.

5. Switch Transformer: simplifying to top-1

Shazeer et al.'s original formulation typically routed each token to more than one expert (top-$k$ for $k > 1$), on the reasoning that blending a couple of experts' outputs would give a smoother, more robust signal than committing entirely to a single expert's judgment. Fedus, Zoph, and Shazeer's Switch Transformer pushed in the opposite direction, simplifying routing all the way down to top-1: every token is sent to exactly one expert, full stop. This is a genuinely bold simplification, since it removes any blending of expert outputs and puts the router's decision entirely on the hook for correctness, but the Switch Transformer authors found it to be not just adequate but preferable at the scale they were targeting, for reasons that are worth understanding rather than just memorizing.

Top-1 routing roughly halves (or more) the compute and communication cost of a top-2 scheme, since only one expert's forward pass has to be computed and, in a distributed setting, only one expert's worth of data needs to be sent across the network per token rather than two. That savings compounds enormously at the scale Switch Transformer targeted: the paper demonstrated models with over a trillion total parameters, distributing experts across many accelerators, and at that scale, halving the per-token routing overhead is not a marginal optimization β€” it's frequently the difference between a training run being practically feasible or not. The paper also introduced refinements to the load-balancing objective and to training stability (careful initialization and precision choices to prevent the notoriously fragile large-scale sparse-routing training process from diverging), and demonstrated that despite the aggressive simplification to a single expert per token, the resulting sparse models could match or exceed the quality of much more expensive dense models at equivalent training compute β€” a striking, concrete demonstration that the capacity-versus-compute decoupling promised in section 2 is real and pays off in practice, not just in theory.

6. Mixtral: MoE as a mainstream production technique

Everything up to this point in the chapter could, fairly, be read as "an interesting and validated research direction." Mixtral closes that gap by being a widely deployed, openly released production model built on exactly this template, which is worth treating as a distinct and important data point: MoE is not a research curiosity that never left the lab, it's a technique used in models people actually run in production today. Mixtral uses eight experts per MoE layer and routes each token to its top-2 experts β€” a more conservative choice of expert count than Switch Transformer's trillion-parameter regime, and a return to top-2 rather than top-1 routing, reflecting a different point on the same design space rather than a departure from it. The result is a model whose total parameter count is much larger than the compute cost per token would suggest β€” the standard way to describe this is that Mixtral has a much larger total ("sparse") parameter count than its active ("dense-equivalent") parameter count per forward pass, and it is precisely that gap between total and active parameters that is the entire point of the architecture.

Mixtral is worth studying not just as an example of the routing mechanism working, but as an illustration of the practical, systems-level considerations that come with deploying MoE models, which don't show up at all in a purely mathematical description of the gating layer. Because different tokens in a batch route to different experts, efficient serving requires either replicating all experts on every accelerator (simple but memory-expensive, since you're storing every expert everywhere even though any given token only uses a couple of them) or distributing experts across accelerators and routing tokens' activations to wherever their chosen expert lives (memory-efficient but requiring careful, latency-sensitive communication between devices for every token, every layer). This tension between memory efficiency and communication cost is a genuinely new serving-time problem that dense models never had to contend with, and it's a natural bridge into the next chapter's subject: whatever the accounting benefits of MoE look like on paper, actually serving any of these models β€” MoE or dense β€” to real users at scale surfaces its own distinct set of inference-time engineering problems around memory management, batching, and latency, which is exactly where the book turns next.

DeepSeekMoE, described by Dai et al. and used throughout the DeepSeek model line, pushes the granularity of this design further in two directions at once, and both are worth knowing since they've since been adopted well beyond DeepSeek's own models. First, fine-grained expert segmentation: instead of a handful of large experts, split the same total FFN parameter budget into many smaller experts and activate more of them per token, which gives the router a finer-grained set of combinations to compose β€” the same total capacity, but assembled from a larger palette of narrower specializations rather than a small number of broad ones. Second, shared experts: a small number of experts are exempted from routing entirely and process every token unconditionally, alongside whichever routed experts a token's gate selects. The reasoning is that some computation β€” general, broadly useful transformations every token needs regardless of its specific content β€” is exactly the kind of thing that shouldn't have to be re-learned redundantly by several different routed experts each independently discovering it; a shared expert learns that common component once, freeing the routed experts to specialize more cleanly on what's actually token-specific.

A separate, practical technique worth naming because it changes how an MoE model actually gets trained rather than how it's routed: sparse upcycling, described by Komatsuzaki et al., initializes a new MoE model from an already-pretrained dense checkpoint rather than starting from scratch, by copying the dense model's single FFN into every expert slot as an identical starting point and then continuing training, letting the router and the now-duplicated experts differentiate from that shared initialization rather than learning routing and per-expert specialization simultaneously from random weights. This reuses the (often substantial) compute already sunk into the dense checkpoint instead of discarding it, and it's a big part of why turning an existing, well-validated dense model into an MoE has become a common shortcut rather than every new MoE being trained from a random initialization.

7. Interview angle

"Explain the core idea of Mixture of Experts in one or two sentences, precisely." A strong answer states the decoupling directly: MoE replaces a single dense feedforward layer, which every token passes through in full, with many independent expert feedforward networks plus a learned router that selects a small subset (top-$k$) of experts per token, so that total model capacity (total parameters across all experts) can scale up independently of per-token compute cost (which stays fixed at roughly $k$ experts' worth of computation regardless of how many total experts exist).

"Why does naive MoE training tend to collapse onto using only a few experts, and how is this fixed?" A strong answer describes the rich-get-richer feedback loop precisely: a marginally better expert (even by chance) gets routed more tokens by the gradient signal, receives more training as a result, becomes genuinely better, and gets routed even more tokens, spiraling toward using only a handful of experts while the rest go untrained. The fix is an auxiliary load-balancing loss added to the training objective, typically built from each expert's fraction of routed tokens times its average gate probability, summed over experts and minimized when routing is roughly uniform, giving the router a genuine counteracting gradient signal rather than relying on hope.

"Why did Switch Transformer move to top-1 routing when the original sparsely-gated MoE work used top-k for k greater than 1?" A strong answer explains the tradeoff rather than treating top-1 as simply "better": top-1 routing roughly halves the per-token compute and cross-device communication cost relative to top-2, which becomes a decisive advantage at the scale Switch Transformer targeted (over a trillion parameters distributed across many accelerators), even though it removes the smoothing effect of blending two experts' outputs. It's a scale-dependent engineering tradeoff, not a universal improvement, which is why Mixtral later chose top-2 again at a different point in the design space.

"What's the difference between a model's 'total' and 'active' parameter count in an MoE model, and why does that distinction matter for someone estimating serving cost?" A strong answer explains that total parameters count every expert's weights across the whole model (which determines memory footprint β€” you generally need enough memory or distributed storage to hold every expert, since you don't know in advance which tokens will route where), while active parameters count only the experts actually invoked per token (which determines the FLOPs, and therefore roughly the latency, of a single forward pass). An MoE model can have a memory footprint like a much larger dense model while running with the per-token compute cost of a much smaller one, and conflating the two numbers leads to badly wrong intuitions about either memory requirements or inference latency.

"What new serving-time problem does MoE introduce that a dense model doesn't have?" A strong answer names the routing-and-communication problem directly: because different tokens in the same batch can route to different experts, efficient serving needs either full expert replication on every device (wasting memory) or distributing experts across devices and shipping token activations to wherever their selected expert lives, which introduces per-token, per-layer communication overhead that a dense model, where every token takes the identical computational path, never has to deal with.

8. Self-check questions

  1. In your own words, what specific assumption present in every dense Transformer feedforward layer does Mixture of Experts abandon, and what does abandoning it buy you?
  2. Write out the sparsely-gated MoE layer's output as a sum over experts, and explain precisely why the un-selected experts contribute zero compute rather than merely zero weight.
  3. Walk through the feedback loop that causes naive MoE training to collapse onto a small number of experts, starting from a small, purely random initial advantage for one expert.
  4. What quantity does a typical auxiliary load-balancing loss penalize, and why is it constructed to be differentiable rather than operating directly on the hard top-k routing decision?
  5. Why does Switch Transformer's move to top-1 routing matter more at trillion-parameter scale than it would at a smaller scale?
  6. Explain the difference between "total parameters" and "active parameters" in an MoE model like Mixtral, and why both numbers are needed to reason about its resource requirements.
  7. Why does distributing experts across multiple accelerators, as a real MoE serving system must eventually do, introduce a communication cost that has no equivalent in a dense model's serving path?

9. Sources

  • Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017. arXiv:1701.06538. https://arxiv.org/abs/1701.06538
  • Fedus, W., Zoph, B., & Shazeer, N. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR 2022. arXiv:2101.03961. https://arxiv.org/abs/2101.03961
  • Jiang, A. Q., et al. (2024, Mistral AI). Mixtral of Experts. arXiv:2401.04088. https://arxiv.org/abs/2401.04088
  • Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A. M., Le, Q. V., Laudon, J., et al. (2022). Mixture-of-Experts with Expert Choice Routing. NeurIPS 2022. arXiv:2202.09368. https://arxiv.org/abs/2202.09368
  • Dai, D., Deng, C., Zhao, C., Xu, R. X., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., et al. (2024). DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. ACL 2024. arXiv:2401.06066. https://arxiv.org/abs/2401.06066
  • Komatsuzaki, A., Puigcerver, J., Lee-Thorp, J., Ruiz, C. R., Mustafa, B., Ainslie, J., Tay, Y., Dehghani, M., & Houlsby, N. (2022). Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints. ICLR 2023. arXiv:2212.05055. https://arxiv.org/abs/2212.05055

Self-check quiz

Test your understanding of this chapter with a short quiz β€” a mix of concept questions and quick calculations.

What core assumption does Mixture of Experts abandon, relative to a dense Transformer feedforward layer?

Explanation: MoE replaces one dense feedforward network with many independent experts plus a router that activates only a small subset per token, decoupling total capacity from per-token compute cost.

In a sparsely-gated MoE layer with top-k routing, why do the un-selected experts contribute zero compute rather than merely zero weight?

Explanation: Since G(x) is sparse, only the top-k nonzero-weight experts need to be evaluated at all β€” this is exactly what makes MoE compute-cheap rather than just a soft blend of every expert's output.

Why does naive MoE training tend to collapse onto using only a few experts?

Explanation: Left unchecked, this feedback loop concentrates almost all tokens on a handful of experts, leaving the rest undertrained β€” exactly what the auxiliary load-balancing loss is designed to counteract.

What is the key difference between a model's "total" and "active" parameter count in an MoE model?

Explanation: An MoE model can have the memory footprint of a much larger dense model while running at the per-token compute cost of a much smaller one β€” that gap is the entire point of the architecture.

Mixtral's MoE layers use 8 experts per layer with top-2 routing. What fraction of the total experts does any single token actually activate?

Explanation: 2 out of 8 experts = 2/8 = 0.25, i.e. 25%.

An MoE layer has 8 experts, each with 0.2 billion parameters. What is the total parameter count across all experts (ignoring the router), in billions?

Explanation: 8 Γ— 0.2 = 1.6 billion total parameters.

Using the same layer (8 experts, 0.2 billion parameters each) with top-2 routing, how many billion parameters are actively used for any single token's feedforward computation?

Explanation: 2 selected experts Γ— 0.2 billion parameters each = 0.4 billion active parameters, regardless of the other 6 experts sitting unused for that token.

A top-2 MoE layer costs 4.0 GFLOPs per token. Switch Transformer's top-1 routing roughly halves the compute cost relative to top-2. Approximately how many GFLOPs per token would an equivalent top-1 layer cost?

Explanation: Halving 4.0 GFLOPs gives approximately 2.0 GFLOPs β€” only one expert's forward pass needs to be computed instead of two.

In the common auxiliary load-balancing loss formulation $\mathcal{L}_{\text{aux}} = n\sum_{i=1}^{n} f_i \cdot P_i$, why is $P_i$ (the average gating probability) used instead of relying only on $f_i$ (the fraction of tokens actually routed to expert $i$)?

Explanation: The hard top-k routing decision behind $f_i$ isn't differentiable, so the loss multiplies it by the router's soft gating probability $P_i$, giving a term that's minimized exactly when load is spread evenly and that can actually be backpropagated through.

An MoE layer has $n = 4$ experts. Expert 2 receives fraction $f_2 = 0.4$ of tokens and has average gating probability $P_2 = 0.5$. Using $\mathcal{L}_{\text{aux}} = n\sum_i f_i \cdot P_i$, what is expert 2's individual contribution $n \cdot f_2 \cdot P_2$ to the auxiliary loss?

Explanation: $n \cdot f_2 \cdot P_2 = 4 \times 0.4 \times 0.5 = 0.8$ β€” a disproportionately high $f_i \cdot P_i$ product like this is exactly what the auxiliary loss penalizes to counteract routing collapse.