Mixture of
Experts

Why would we ever want more parameters that we don't use?

Here is the central tension of modern deep learning: bigger models are smarter, but bigger models are also slower and more expensive. Every weight in a standard "dense" transformer participates in every computation, for every token, every time. We pay the full price of the model for every input we run through it — even when most of the model probably isn't relevant to the input at hand.

Think about how a hospital works. If every patient who walked in had to be examined by every doctor on staff — the cardiologist, the dermatologist, the pediatrician — the hospital would grind to a halt within an hour. Real hospitals use a triage nurse: a quick first pass that routes each patient to one or two specialists who actually know about their problem. The rest of the staff continues their work undisturbed.

A Mixture of Experts applies exactly that idea to a neural network. Instead of one giant feed-forward block that processes every token, we have many smaller specialist blocks (the "experts") and a tiny router network (the "triage nurse") that decides, on a per-token basis, which one or two experts get to weigh in.

For a concrete number to keep in your head: Mixtral 8×7B has roughly 47 billion total parameters but uses only about 12 billion of them for any given token. It has the knowledge capacity of a much larger model with the inference cost of a much smaller one. The tradeoff is that you still need enough memory to hold all 47B parameters — the experts are still sitting on the GPU, just sitting quietly.

Where exactly does MoE live inside a transformer?

Before we replace anything, let's be precise about what we're replacing. A standard transformer block has two sublayers:

1. Self-attention, where each token gathers information from the other tokens in the sequence.
2. A feed-forward network (FFN), where each token is transformed independently — typically a two-layer MLP that expands the hidden dimension by 4× and then projects back down.

The FFN is, perhaps surprisingly, where most of a transformer's parameters live. In a typical model, the FFN holds roughly two-thirds of the weights. It is also the part of the architecture that operates on each token in isolation — there is no information flowing between positions in the FFN, only attention does that.

   ┌─────────────────────────────────────────┐
   │           Transformer Block             │
   │                                         │
   │   ┌────────────────────────────┐        │
   │   │     Self-Attention         │        │
   │   └────────────────────────────┘        │
   │                ↓                        │
   │   ┌────────────────────────────┐        │
   │   │      FFN  ◄── we replace   │        │
   │   │           this with MoE    │        │
   │   └────────────────────────────┘        │
   └─────────────────────────────────────────┘

So the surgical operation we are about to perform is: swap the FFN for an MoE layer. Attention stays. Layer norms stay. Residual connections stay. Only the per-token feed-forward computation gets replaced by a router-and-experts ensemble.

A router and a panel of specialists

An MoE layer has two ingredients:

• Experts. A set of N small neural networks (in practice, just FFNs with the same shape). Call them $E_1, E_2, \ldots, E_N$.
• A gate network (also called the "router"). A tiny network — usually a single linear layer followed by a softmax — that takes a token's hidden state and outputs N weights, one per expert. The weights say "how much should this token use each expert?"

Here is the picture from the original Switch Transformer paper. Two tokens — "More" and "Parameters" — arrive at the layer. The router examines each one and dispatches them: "More" goes to expert FFN 2, "Parameters" goes to expert FFN 1. They each travel through their assigned expert and come out the other side as new hidden states.

That is genuinely the whole concept. The router is small and cheap; the experts are where the parameter count lives; and on any given forward pass, most experts are idle. The experts can be thought of as parallel candidate FFNs, and the router as the network's own learned scheduler.

Mathematically, if the router produces gating weights $G(x)_i$ for each expert $i$, the layer's output for input $x$ is:

This looks like a weighted ensemble — we run every expert and combine their outputs by the router's weights. If we did that, we'd save no compute at all. The trick comes next.

The whole point: most $G(x)_i$ are zero

Look at the formula again. If $G(x)_i = 0$ for some expert $i$, then that term contributes nothing to the sum, and we have no reason to actually run $E_i(x)$. We can simply skip the multiplication. This is called conditional computation: parts of the network become active or inactive depending on the input.

The standard recipe is top-k routing: the router computes a score for each expert, then we keep only the top $k$ scores (say $k=2$) and zero out the rest. We renormalize and run only those $k$ experts. With $N=8$ experts and $k=2$, we activate 25% of our experts per token — but during training those experts also get specialized, because the router learns to send each token to experts that handle it well.

This is why sparsity matters so much:

• Total parameters: $N \times$ (size of one expert) — large.
• Active parameters per token: $k \times$ (size of one expert) — much smaller.
• Compute (FLOPs) per token: scales with $k$, not $N$. Doubling $N$ doesn't double the compute.

You get the modeling capacity of a huge network at the inference cost of a small one. The catch — and there are several — is that all $N$ experts must still be loaded into memory, because any token might need any expert. Sparse compute, dense memory.

Where the idea came from

The MoE idea is older than the transformer by nearly three decades. It first appeared in a 1991 paper by Jacobs, Jordan, Nowlan, and Hinton — Adaptive Mixtures of Local Experts. The original setup was much like what we just described, except in the language of small neural networks doing supervised classification: a gating network, several expert networks, and a learned routing rule. Each expert specializes in a region of input space; the gate decides which region the input is in.

Two threads of research kept the idea alive in the 2010s:

• Experts as components, not whole models. Eigen, Ranzato, and Sutskever (2013) showed that MoE could be a layer inside a deep network, not the whole architecture — a precondition for everything that followed.
• Conditional computation. Bengio and collaborators argued that activating only parts of a network on a per-input basis was the natural way to scale models without scaling compute.

The breakthrough for modern NLP came in 2017 with Shazeer et al.'s "Outrageously Large Neural Networks." They wedged an MoE layer between LSTM layers and trained a 137-billion-parameter translation model, which was enormous for the time. This is the figure they used:

They demonstrated for the first time that sparsely-gated MoEs could be trained at large scale on a real NLP task, and they introduced two techniques we still use today: noisy top-k routing, and an auxiliary load-balancing loss. We will get to both shortly.

How does the router actually decide?

The simplest possible router is a single matrix. Take the token's hidden state $x$, multiply by a learned weight matrix $W_g$, apply softmax:

This gives a probability distribution over experts. We could just sample from it, or take the argmax, or take the top-k. In practice we want the routing to be a learned, gradient-friendly choice — we want to backpropagate through it — so we typically combine top-k with renormalization.

Noisy top-k gating

Shazeer's recipe adds two refinements. First, inject Gaussian noise into the gating logits before selecting top-k. Second, scale that noise by another learned linear projection of $x$, so the model can decide how confident vs. exploratory it wants to be per token. The whole thing looks like this:

Setting all non-top-k entries to $-\infty$ before the softmax guarantees they come out as exactly zero — clean sparsity, no rounding issues.

Why $k > 1$, historically?

Early MoEs always used $k \geq 2$. The reasoning was subtle but important: if you only ever pick one expert, the gradient signal flowing back through the gate is weak and unstable. By picking two and combining their outputs with weights, you get a real signal about how much each expert contributed to a good prediction. As we'll see, the Switch Transformer eventually challenged this and got away with $k=1$, but only with careful tricks.

The router has a "rich get richer" problem

Here is the failure mode that kept MoE research stuck for years. Imagine you start training with random weights. By pure chance, the router slightly prefers expert 3 for a few inputs. Expert 3 thus gets more training data, and improves faster than the others. The router notices expert 3 is doing well, and routes even more tokens to it. The other experts starve, never improving, never getting traffic. You end up with one bloated expert and seven dead ones — effectively a dense FFN with a confusing wrapper.

Two ideas, used in combination, fix this:

1. Auxiliary load-balancing loss. Add an extra term to the training objective that penalizes uneven token-to-expert distributions. The model's main task and the auxiliary load-balance task are optimized together. Get the routing balanced or pay a price.
2. Expert capacity. Set a hard cap on how many tokens any single expert is allowed to process per batch. If expert 3 fills up, additional tokens that wanted expert 3 are overflowed — passed forward through the residual connection unchanged, or routed to a backup expert, or in some implementations dropped entirely.

Why is the capacity cap necessary in addition to the loss? Because GPU tensors have static shapes. We need to allocate, in advance, a fixed amount of buffer space per expert. If we don't cap, we don't know how much memory to set aside. The capacity factor formula will appear shortly when we get to Switch Transformers.

GShard, then Switch simplifies

The first transformer-based MoEs at really large scale were Google's GShard (2020) and Switch Transformer (2021). They each made a different design choice that's worth understanding.

GShard — every other FFN, top-2 routing

GShard scaled transformers past 600 billion parameters. The recipe: replace every other FFN layer with an MoE layer (so attention layers and half the FFN layers stay shared and dense), and use top-2 routing. They also added two refinements: random routing (the second expert is sampled with probability proportional to its weight, rather than always the deterministic runner-up), and expert capacity with overflow.

Switch Transformer — top-1 is enough

The Switch Transformer paper made a heretical move: drop $k$ all the way to 1. Each token goes to exactly one expert. This had been considered too aggressive — the gradient signal would be too weak — but Switch made it work, and the simplification paid for itself:

• The router is half as expensive (only one expert is selected).
• Each expert's batch is at least twice as large (no token is split between two experts), which is good for hardware utilization.
• Communication cost between machines drops, since each token's hidden state only needs to travel to one expert.

Quality, remarkably, was preserved. The headline result: a 4× pretraining speedup over T5-XXL at the same model size budget.

The capacity factor formula

Switch made the capacity formula explicit and tunable:

If routing were perfectly uniform, every expert would receive exactly tokens / N tokens per batch. The capacity factor (typically 1.0 to 1.25) is a buffer on top of that ideal — slack that absorbs natural unevenness in routing. Higher capacity factor means fewer dropped tokens but more memory and more inter-device communication. Lower capacity factor means tighter packing but more overflow. A common sweet spot for Switch is 1.25 in training.

Selective precision

One more Switch trick worth knowing: train the experts in bfloat16 (half precision) but keep the router in full precision. The router has a softmax with exponentials, and exponentials magnify roundoff error. Precision-mixing here was a stability win at no quality cost:

What do experts actually learn to be good at?

Here is the question every student asks first: if I have 8 experts, do I get a "verb expert," a "noun expert," a "punctuation expert"? Is there a "French expert" in a multilingual model?

The answers are surprisingly mixed.

The ST-MoE paper carefully analyzed which tokens get routed to which experts in a trained model. Their finding: encoder experts do specialize, but along surprising axes — punctuation, conjunctions, proper nouns, numbers, visual descriptors — i.e. shallow lexical and syntactic categories rather than deep semantic ones. Decoder experts specialize less; routing in the decoder looks more uniform.

For multilingual models the story is even more counter-intuitive: there is no clean per-language expert. You might have hoped that "expert 3 handles French and expert 7 handles Mandarin," and that would be tidy and human-interpretable. It doesn't happen. The auxiliary load-balancing loss actively prevents it — load balance forces every expert to see tokens from every language, so each expert ends up multilingual by construction.

Inside a real MoE: instrumenting Qwen1.5-MoE

Theory tells us the router learns to specialize, that experts diverge, and that load balancing is a constant battle. But none of that is quite as convincing as opening up a working production model and watching the routing happen on real prompts. Let's do exactly that.

Our subject is Qwen1.5-MoE-A2.7B-Chat, an open-weight production MoE released by Alibaba. It's small enough to fit on a free Colab T4 GPU, but architecturally serious. The headline specs:

• 14.3 billion total parameters, of which only ~2.7 billion are active per token. The "A" in "A2.7B" stands for active — Qwen telegraphs the sparse-vs-dense distinction right in the model name.
• 24 transformer layers, every single one of them MoE. No interleaving with dense FFNs (unlike GShard). Routing happens at every depth.
• 60 experts per layer, top-4 routing. So 4/60 ≈ 6.7% of the experts in any layer fire for any given token.
• A shared expert alongside the routed pool, which runs for every token regardless of routing — a hybrid trick borrowed from DeepSeek-MoE that captures common knowledge cheaply while leaving the specialists free to specialize. This is why active params are ~2.7B rather than the strict (4/60) × 14.3B you'd compute from sparsity alone.

The experiment

We register PyTorch forward hooks on every Qwen2MoeTopKRouter module — the small (60, 2048) linear layer at the heart of each MoE block — and intercept its outputs as the model runs. For each token, at each layer, we record the raw logits, the post-softmax probabilities, and which 4 experts were chosen.

Then we run four prompts, each from a different domain:

   math       →  "What is 2 + 2?"
   code       →  "def fibonacci(n):"
   creative   →  "Once upon a time in a dark forest,"
   factual    →  "The capital of France is"

Each prompt is a few tokens long. With four prompts and 24 layers each making top-4 selections, we end up with a few hundred routing events to inspect — small but enough to see real patterns.

Finding 1 — Routing diverges by domain at layer zero

You might expect early layers to do generic, domain-agnostic processing — recognize that there are tokens here, do shallow lexical work — and only later layers to specialize. The data says no. The four prompts choose different top-4 experts at layer 0 already. By the time a token has been embedded and run through one self-attention sublayer, the router has enough signal to disagree based on input domain. There is no "we haven't decided yet what kind of input this is" warm-up phase.

Finding 2 — Utilization on a small probe is visibly skewed

Across all four prompts × 24 layers, the routers made 384 expert selections. If routing were perfectly uniform, every one of the 60 experts would receive 384/60 ≈ 6.4 selections. In practice:

• 54 of 60 experts received at least one selection.
• Six experts received zero traffic from these prompts.
• The most-chosen expert (ID 52) was selected 22 times — roughly 3.4× the uniform ideal.
• The 5th-most-chosen got 13 selections — so the gap from #1 to #5 is 22 vs. 13.

This isn't catastrophic collapse — most experts are participating — but it is exactly the lopsided utilization the auxiliary load-balance loss is designed to flatten during training. At training scale, with millions of tokens flowing through, the distribution does flatten substantially. On a four-prompt probe like this, the unevenness is unmissable.

Finding 3 — The router's softmax is almost flat. Top-k is what makes routing decisive.

This is the most counter-intuitive observation in the whole notebook, and a great teaching moment. The router emits 60 logits, applies softmax, and the entropy of the resulting probability distribution is ~99.9% of the theoretical maximum of $\log_2(60) \approx 5.91$ bits. The probability assigned to the "best" expert is barely larger than the probability assigned to the "worst." The softmax is nearly uniform.

So how does routing manage to be decisive at all? Because the router's actual output isn't a sample from the softmax — it's an argmax over the top 4. A nearly-flat softmax can still have a definite top-4. Tiny logit differences, of order milli-units in probability, are amplified by the discrete top-k operation into hard yes/no expert assignments.

Finding 4 — Experts are genuinely distinct, not redundant copies

An MoE could fail in a different way than collapse: all 60 experts could end up learning nearly the same transformation. Routing would still be balanced, but it wouldn't matter who got picked because every expert would do the same thing. To check, the notebook computes pairwise cosine similarity between expert weight matrices (the actual MLP parameters of each expert) within a single layer.

The result is a strongly diagonal-dominant matrix: each expert is similar to itself (cosine 1.0) and only weakly similar to its peers. The off-diagonal cells are low. The 60 experts have learned 60 genuinely different transformations. That's what makes the routing analysis meaningful — if the experts' weights were near-identical, the heatmaps and utilization plots would be empty noise.

Finding 5 — Math and code share their experts; creative goes elsewhere

For each prompt, average the router's probability vectors across all tokens and all 24 layers, then pick out the top-3 most-favored experts. The result is striking:

Math and code share the same top-3 experts. Both are structured, symbolic, low-ambiguity inputs, and the model has effectively lumped them together as "formal language." Creative narrative routes through an entirely different set of three. Connect this to Finding 1 — the router really does carve up the input space into domain-shaped regions, just not always along the boundaries you might predict.

Finding 6 — Routing similarity has a characteristic depth profile

If you instead compute the cosine similarity between two prompts' routing distributions layer by layer (rather than averaged across the whole model), you typically see a pattern like this:

This is the same pattern that probing studies of dense models have shown for years: middle layers are where domain-specific representation lives; early and late layers are more universal. The MoE router, given a free choice of which experts to consult per layer, exhibits the same depth structure.

Finding 7 — Token path diversity tracks token type

For every token in the input, count how many distinct experts handled it across the 24 layers. With top-4 routing the theoretical maximum is min(4 × 24, 60) = 60. In the notebook's runs, content-rich tokens (numbers, code keywords, vivid nouns) tend to see more diverse expert paths; structural tokens (punctuation, articles, whitespace) often touch the same handful of experts at every layer. The router treats "the" the same way at layer 5, layer 15, and layer 23 — but it consults different specialists for "fibonacci" at different depths.

Combined with Finding 4 (experts are genuinely distinct), this means high path diversity is meaningful: a content-rich token is being passed between specialists doing genuinely different transformations, not shuffled among redundant copies.

That's a lot of empirical findings packed into one experiment. The larger lesson for a student of MoE: this architecture is far less mysterious than its reputation. Standard PyTorch hooks expose every routing decision. You can verify domain-dependent routing, expert distinctness, near-uniform softmax, depth-dependent specialization, and per-token path diversity in a single Colab session. If you remember nothing else from this section, remember that a router is interrogable — and that interrogating it is the fastest way to build intuition for what an MoE actually does.

Why MoEs are tricky to fine-tune (and what helps)

Pretraining MoEs is a triumph. Fine-tuning them, for years, has been a headache. The core issue is that sparse models overfit much more easily than their dense equivalents on small downstream datasets. Look at this learning curve from ST-MoE:

Several practical findings have accumulated about how to fine-tune MoEs well:

• Different regularization for the sparse layers. Use one dropout rate for the dense parts of the model and a higher one inside the experts. The expert layers have lots of capacity and need more taming.
• Small batch sizes, larger learning rates. Counter to dense-model intuition. The figure below summarizes ST-MoE's recommendation.
• You can freeze the MoE layers entirely. Surprisingly: if you fine-tune everything except the experts, you lose almost no quality, and you save a lot of memory and time. The reverse — freezing everything except the experts — performs much worse.

The instruction-tuning surprise

The most encouraging recent result: MoEs benefit more from instruction tuning than dense models do. The 2023 paper "MoEs Meet Instruction Tuning" trained both a dense T5 and an MoE T5 in three regimes — single-task fine-tune, multi-task instruction-tune, and instruction-tune-then-fine-tune. The MoE lagged the dense model under plain single-task fine-tuning, but pulled clearly ahead once instruction tuning was involved.

This makes intuitive sense if you stop and think: many small tasks ↔ many specialists is a natural fit. A model with eight experts has a much richer capacity to allocate "this kind of input goes to that kind of expert" when there are many task types in the training mix. The dense model has no such structure to exploit.

Making MoE go brrr on real hardware

An MoE's compute-vs-memory tradeoff only pays off if you can actually run experts in parallel across GPUs without the GPUs sitting idle waiting for tokens to arrive. This is a systems problem as much as a modeling one.

Expert parallelism

The standard parallelism strategies in deep learning are:

Data parallelism

Same weights replicated on every GPU; each GPU sees a different slice of the batch.

Model parallelism

Different chunks of the model on different GPUs; the same data flows through all of them.

Expert parallelism (the new one)

Different experts live on different GPUs. When a token needs expert i, its hidden state is sent (over the network) to whichever GPU holds expert i, processed there, and sent back.

Expert parallelism introduces a new bottleneck: the all-to-all communication step. Every GPU must send its tokens to the GPUs holding the right experts, then receive the results back. If your network is slow, this dominates everything else. If your routing is unbalanced, some GPUs sit idle while others are swamped. This is why load balance is not just a quality concern — it's a throughput concern.

Capacity factor — a knob you can dial

Recall that capacity factor controls how much overflow buffer each expert gets. In training, ~1.25 is typical. At inference time, you can change the capacity factor — turn it down to save compute when serving, or up if you have spare capacity and want fewer dropped tokens. It's a runtime knob, not a fixed architectural choice.

Megablocks: when the static-shape assumption hurts

Standard MoE implementations use batched matrix multiplication, which assumes every expert processes the same number of tokens (so the batched tensor has a uniform shape). In reality, expert assignment is uneven — even with load balancing, some experts get more tokens. The mismatch is usually papered over by either dropping overflow tokens or padding with zeros. Both are wasteful.

Megablocks (2022) reformulates the whole MoE layer as a single block-sparse matmul, with custom GPU kernels that handle the irregular block shapes. The result: no token dropping, no wasted padding, and substantial speedups over the naive batched-matmul approach.

Distillation, for serving

Even with all of these tricks, an MoE is memory-hungry to serve. A common workaround at deployment time is to distill the trained MoE back into a dense model with fewer parameters but similar quality. The Switch Transformer paper showed you can keep 30–40% of the sparsity benefit this way. You get a small dense model that is much easier to serve, while the MoE was just a temporary scaffolding for getting there cheaply.

When should you actually reach for an MoE?

Crucially: do not directly compare parameter counts between dense and sparse models. A dense 13B parameter model and a sparse "8x7B" MoE are not comparable in any single number. The MoE has ~47B total parameters, ~12B active per token, and learns differently from either a 13B or a 47B dense baseline. Pick a metric — quality at fixed pretraining compute, latency at fixed batch size, memory at fixed quality — and compare on that metric.

Watch some tokens get routed

Type a sentence below. Each word becomes a token; we'll simulate a top-k router with a tunable number of experts and show you which expert(s) each token gets dispatched to. (The routing here is a fixed deterministic hash of each word, so the same word always lands at the same expert — much like a reasonably trained router would do — but a real one would also be using the surrounding context.)

Notice: short, common words ("the", "of", "a") tend to land at consistent experts — a router would learn to specialize an expert on function words. Open-class words spread across experts. Try changing $N$ from 8 to 4: each expert now handles a wider variety of tokens. Try $k=1$ vs $k=2$: at $k=2$, each token's compute doubles, but you get richer mixing.

Takeaways & exercises

The six things to remember

1. An MoE replaces the FFN sublayer of a transformer block with a router and a panel of expert FFNs. Attention and the rest are unchanged.
2. Sparsity decouples capacity from compute. Total parameters can be huge; per-token compute scales with $k$, not $N$. Memory still scales with $N$.
3. Load balancing is not optional. Without an auxiliary loss and a capacity cap, the router collapses to a few favored experts and the architecture loses its purpose.
4. Pretraining is the easy part; fine-tuning is the hard part. Sparse models overfit easily on small datasets; instruction tuning works exceptionally well; many practical tricks (freezing, dropout asymmetry, smaller batches) exist.
5. MoE is a systems-level architecture. Its real-world efficiency depends on expert parallelism, all-to-all communication, capacity factor tuning, and (often) custom kernels like Megablocks.
6. You can watch it work. A few PyTorch hooks on the router expose every routing decision. Domain-dependent routing, expert distinctness, near-uniform softmax, and depth-dependent specialization are all directly observable on a real model in a single Colab session.

Exercises for the reader

Further reading, in order

• Jacobs, Jordan, Nowlan, Hinton (1991) — Adaptive Mixtures of Local Experts. The original.
• Shazeer et al. (2017) — Outrageously Large Neural Networks. The first big NLP MoE.
• Lepikhin et al. (2020) — GShard. First production-scale transformer MoE.
• Fedus, Zoph, Shazeer (2021) — Switch Transformer. The simplification to top-1.
• Zoph et al. (2022) — ST-MoE. The stability and analysis paper.
• Gale et al. (2022) — Megablocks. The block-sparse-matmul rewrite.
• Shen et al. (2023) — MoEs Meet Instruction Tuning. The fine-tuning rehabilitation.
• P. Kulkarni — live-moe-demo notebook. The hands-on instrumentation of Qwen1.5-MoE used in §10. Run it yourself in Colab on a T4.