Mixture of Experts

Welcome to Part VII, the frontier. We begin with one of the most important architectural ideas behind today's largest models: the Mixture of Experts (MoE). Scaling laws (Chapter 16) told us that bigger models are better — but bigger models cost more to run, because every parameter is used for every token. MoE breaks this link: it lets a model have FAR more parameters while using only a FRACTION of them for any given token. More capacity, almost the same compute.

In a normal ('dense') model, every token passes through every parameter — all the weights are used for all the tokens. So doubling the parameters doubles the compute per token. This is the wall scaling runs into: more capacity means proportionally more cost, forever. MoE asks a radical question: what if each token only used the parameters it actually needs, rather than all of them?

This gives MoE its defining characteristic: a distinction between TOTAL parameters (all the experts combined — the model's full capacity) and ACTIVE parameters (the few used for a given token — what determines compute cost). An MoE model might have 8× the total parameters of a dense model but use only a fraction per token, getting much of the quality of a huge model at the cost of a small one.

Where exactly does MoE go in a Transformer? Recall from Chapter 13 that each Transformer block has an attention sub-layer and a feed-forward (FFN) sub-layer. MoE replaces the single FFN with MANY FFNs — the 'experts' — plus a small 'router' that decides which experts each token uses. Attention stays the same; only the FFN becomes a mixture of experts.

A sparse MoE layer has two parts. The EXPERTS are several independent feed-forward networks (say 8 of them), each identical in structure to the dense FFN it replaces. The ROUTER (or 'gate') is a small network that looks at each token and decides which experts should process it. For each token, only the chosen experts run — the rest are skipped entirely. That skipping is what makes it 'sparse'.

Typically, the FFN in EVERY Transformer block (or every other block) is replaced with an MoE layer, while attention remains dense and shared across all tokens. Since the FFN holds most of a Transformer's parameters, turning FFNs into MoE layers is where the huge parameter expansion comes from. The model becomes a stack of blocks, each with shared attention and a bank of routed experts.

The router is the brain of an MoE layer. For each token, it must decide which experts to use. The standard method is TOP-K ROUTING: the router scores all experts for the token, then sends the token to only the k highest-scoring ones (k is small — often 1 or 2). Let us see exactly how it works.

The router is just a small linear layer. For each token, it produces a score (a 'logit') for every expert. A softmax turns these into weights, and the top-k experts by weight are selected. The token is processed by those k experts, and their outputs are combined — weighted by the router's scores — into the layer's output. So the router both CHOOSES the experts and WEIGHTS their contributions.

Here is what routing looks like for a few tokens with 5 experts and top-2 routing — each token activates exactly 2 experts, and different tokens go to different experts:

Expert Route: Top-2 routing: each token to its 2 highest-scoring experts

Notice that each token (each row) lights up exactly 2 experts, and the assignment varies by token — the router learns to send different tokens to different specialists. Over a whole batch, all experts get used, but each individual token only pays for 2.

Let us make the compute savings concrete, because the numbers are what make MoE compelling. The key ratio is how many experts EXIST versus how many RUN per token. With 8 experts and top-2 routing, each token uses only 2 of 8 — the FFN does roughly a quarter of the work it would if all experts ran, yet the model holds 8 experts' worth of knowledge.

Mixtral 8x7B (Mistral, 2023) is the canonical open MoE. It has 8 experts per layer and uses top-2 routing. Its name suggests '8×7B = 56B' but the real numbers are subtler: because attention and embeddings are shared, the TOTAL parameter count is about 47B, while the ACTIVE parameters per token are only about 13B. So Mixtral runs at roughly the cost of a 13B dense model but has the capacity of a 47B one — and it performs accordingly, rivaling much larger dense models.

MoE has a notorious failure mode that does not exist in dense models, and understanding it is essential. Left to itself, the router tends to collapse — sending almost all tokens to just a few experts while the rest go unused. This 'expert collapse' wastes the model's capacity and is the central challenge of training MoE.

Expert collapse arises from a self-reinforcing feedback loop. Early in training, by chance, the router slightly favors a few experts. Those experts get more tokens, so they train faster and get better. Because they are better, the router favors them even more, so they get even more tokens. Meanwhile, the neglected experts get few tokens, barely train, stay bad, and get neglected further. The rich get richer; most experts wither. The model ends up using only a handful of its experts — throwing away the capacity MoE was supposed to provide.

A collapsed MoE looks like this — nearly all tokens funnel to two overloaded experts while the rest sit idle (grayed):

Expert Route: Expert collapse: most tokens funnel to a few experts

The standard cure is to ADD a penalty to the training loss that encourages BALANCED expert usage. This 'auxiliary load-balancing loss' is high when tokens are distributed unevenly across experts and low when they are spread out. By minimizing it alongside the main loss, training is pushed to use all experts roughly equally — breaking the rich-get-richer cycle. It is a gentle, constant pressure toward balance.

The intuition: the loss is minimized when every expert gets an equal share of tokens. By penalizing imbalance, it stops any expert from monopolizing the routing, forcing the router to spread tokens out and give every expert enough traffic to train. The coefficient (alpha) sets how strongly balance is enforced — too little and you risk collapse, too much and you override the router's useful specialization.

Even with load balancing, in any given batch some experts receive more tokens than others. For efficient hardware execution, each expert is usually given a fixed CAPACITY — a maximum number of tokens it can process per batch. This creates a new wrinkle: what happens when more tokens are routed to an expert than its capacity allows?

Hardware (GPUs/TPUs) is most efficient with FIXED, predictable tensor shapes. If experts could receive any number of tokens, the computation shapes would vary unpredictably, hurting efficiency. So each expert gets a fixed buffer sized for its expected share — the CAPACITY FACTOR controls how much slack above the average each expert gets. This keeps the computation regular and fast.

When more tokens are routed to an expert than its capacity, the overflow tokens are DROPPED — they skip that expert entirely (often passing through unchanged via the residual connection). This sounds alarming, but a modest drop rate is tolerable: the token still gets processed by the rest of the network and its OTHER chosen expert (with top-2 routing). Still, heavy dropping hurts quality, so the capacity factor and load balancing must keep drops low.

MoE's compute savings come with real systems costs. The experts must be STORED (lots of memory), and when experts are spread across multiple devices, tokens must be SENT to wherever their expert lives (lots of communication). These systems challenges are why MoE is harder to deploy than its compute numbers suggest, and connect directly to the distributed training of Chapter 18 and serving of Chapter 31.

The first cost is memory. Even though only k experts run per token, ALL the experts must be held in memory — the router might send the next token to any of them. So an MoE model's memory footprint is its TOTAL parameter count, not its active count. A 47B-total MoE needs ~47B params' worth of memory even though it computes like a 13B model. MoE trades the scarce resource (compute) for the more available one (memory) — but memory is still a real constraint.

When the experts don't all fit on one device, they are spread across many — 'expert parallelism' (a form of the model parallelism from Chapter 18). But then a token routed to an expert on a DIFFERENT device must be SENT there, processed, and the result sent back. This all-to-all communication — every device potentially sending tokens to every other — is a major bottleneck, and managing it efficiently is central to MoE systems.

The basic MoE recipe has been refined in many ways to improve quality, balance, and efficiency. Knowing the main variants helps you understand modern MoE models, which rarely use the vanilla recipe.

Instead of a few large experts, use MANY smaller ones (e.g. 64 experts, routing to 8). This gives finer-grained specialization — more distinct combinations of experts a token can use — and often better quality at the same active-parameter cost. DeepSeek's MoE models pioneered this fine-grained approach, finding that many small experts outperform few large ones.

A clever refinement (DeepSeek-MoE and others): designate one or more SHARED experts that ALWAYS process every token, alongside the routed experts. The shared experts learn common, general patterns that all tokens need, freeing the routed experts to specialize on the rest. This reduces redundancy (the routed experts don't each have to relearn common patterns) and improves both balance and quality.

Expert Route: Shared + routed experts: E0 always runs, plus top-2 routed

MoE models are trained much like dense models (Part IV), but the routing adds quirks and instabilities that require care. Understanding these helps explain why MoE, despite its appeal, took years to become reliable at scale.

MoE training is less stable than dense training. The discrete top-k routing decision is non-differentiable (you can't smoothly backprop through 'pick the top 2'), which is handled with approximations that add noise. Routing can oscillate, experts can collapse (Section 32.5), and the auxiliary loss must be tuned. MoE models are also more sensitive to hyperparameters and prone to loss spikes. Much of the engineering is about keeping training stable.

Beyond the load-balancing loss, MoE training often adds a 'router z-loss' that keeps the router's logits from growing too large (which would make routing overconfident and unstable). Together with careful initialization, gradient clipping (Chapter 15), and tuned auxiliary-loss weights, these stabilizers make large MoE training tractable. The recurring theme: the router needs babysitting that dense FFNs never required.

MoE models are notoriously harder to FINE-TUNE than dense models. With far more total parameters but limited fine-tuning data, MoE models overfit more easily, and the routing learned during pretraining can be disrupted by fine-tuning. Techniques like freezing the router, using higher auxiliary-loss weights during fine-tuning, or fine-tuning only some experts help, but MoE fine-tuning remains finickier than dense fine-tuning — a real practical consideration when choosing MoE.

MoE is not a research curiosity — it powers many of the most capable models in production. A tour of real MoE models grounds the concepts and shows how the variants combine in practice.

The Switch Transformer (Fedus et al., 2021) was a landmark that simplified MoE by routing each token to just ONE expert (top-1), maximizing sparsity. It showed MoE could scale to trillions of parameters and trained stably with the right load balancing. It established many of the techniques — the auxiliary loss, capacity factors — still used today, and demonstrated that simpler routing could work at massive scale.

Mixtral 8x7B (2023) brought MoE to the open community. With 8 experts and top-2 routing, it matched or beat much larger dense models while running at ~13B active cost. Its open release let everyone experiment with MoE, study routing, and build on it — making it the reference open MoE and the model most people first encounter when learning MoE hands-on.

While architectures of closed frontier models aren't officially disclosed, GPT-4 is widely reported to be a large MoE, and many other frontier models are believed to use experts. The reason is exactly this chapter's thesis: at the largest scales, MoE is how you get enormous capacity at a serving cost that remains feasible. MoE has become a standard tool in the frontier toolkit.

MoE is powerful but not always the right choice. Weighing it against a dense model is an important judgment, and the answer depends heavily on scale, hardware, and use case.

MoE gives more capacity per unit of compute, paid for with more memory, more systems complexity, more training instability, and harder fine-tuning. If you are pretraining or serving a very large general model on good hardware and compute-per-token is your constraint, MoE is often the right call. If you are working at smaller scale, on limited hardware, or need to fine-tune frequently, a dense model's simplicity and stability may win despite its lower capacity-per-FLOP.

Let us consolidate the chapter into one coherent picture of how a Mixture of Experts works and what it takes to make it succeed.

If you remember three things about MoE: First, MoE DECOUPLES capacity from compute — many total parameters, few active per token, by routing each token to a subset of experts. Second, ROUTING IS THE HEART AND THE HAZARD — a learned router enables specialization but risks expert collapse, fixed by load-balancing losses. Third, the COST IS RELOCATED, NOT REMOVED — MoE trades compute for memory and systems complexity, which is a great trade at the frontier and a questionable one at small scale.

Exercises 1–10 are pen-and-paper or derivations; 11–20 require code.

Further reading: “Adaptive Mixtures of Local Experts” (Jacobs et al., 1991) — the original idea. “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” (Shazeer et al., 2017) — MoE for deep learning. “Switch Transformers” (Fedus et al., 2021). “GLaM” (Du et al., 2021). “Mixtral of Experts” (Jiang et al., 2024). “DeepSeekMoE” (Dai et al., 2024) for fine-grained and shared experts. “MegaBlocks” (Gale et al., 2022) for dropless MoE. “Expert Choice Routing” (Zhou et al., 2022).

Next → Chapter 33: Long Context & Memory

MoE scaled a model's PARAMETERS efficiently. The next frontier is scaling how much a model can ATTEND to — its context length. Chapter 33 tackles long context and memory: why attention's quadratic cost makes long contexts expensive, the position-encoding tricks (RoPE scaling, ALiBi, YaRN) that let models extend far beyond their training length, efficient attention variants, and the external-memory approaches that let models recall information across vast or unbounded contexts. Having grown the model's capacity with experts, we now grow its window onto the world — and confront the 'lost in the middle' and quadratic-scaling limits that long context must overcome.