Mixture-of-Experts — making a model wider without using all of it for every token

Over the last few lessons, we have focused on how to distribute dense Transformer computation across many GPUs.

Today we change the model itself.

A normal Transformer MLP applies the same feed-forward network to every token. A Mixture-of-Experts, or MoE, layer instead contains many different MLPs and uses a small router to choose only a few of them for each token.

The core idea is:

This is called conditional computation.

It lets a model increase parameter capacity much faster than it increases the amount of computation required for one token.

Start with a dense Transformer MLP¶

A simplified Transformer block contains:

attention
   ↓
MLP
   ↓
next layer

The MLP might be:

Every token passes through the same matrices \(W_1\) and \(W_2\).

If the model processes:

"The cat sat on the mat"

then every token uses the same MLP parameters:

"The" → MLP
"cat" → MLP
"sat" → MLP
"on"  → MLP
"the" → MLP
"mat" → MLP

A larger dense MLP gives the model more capacity, but it also increases computation for every token.

Roughly:

MoE tries to loosen that relationship.

Replace one MLP with many experts¶

Suppose instead of one MLP we create eight:

Each expert is itself an ordinary feed-forward network.

But each token visits only two experts.

Conceptually:

                   ┌→ Expert 1
                   ├→ Expert 2
token → router ────┼→ Expert 3
                   ├→ ...
                   └→ Expert 8

router chooses only 2

For token representation \(x\), a router computes scores:

After a softmax:

we select the top \(k\) experts. For top-2 routing:

If the router assigns:

Expert 1: 0.04
Expert 2: 0.51
Expert 3: 0.02
Expert 4: 0.31
Expert 5: 0.03
Expert 6: 0.01
Expert 7: 0.05
Expert 8: 0.03

the token goes to Experts 2 and 4.

The output can be a weighted combination:

usually with routing weights renormalized or otherwise handled according to the architecture.

The important part is that Experts 1, 3, 5, 6, 7, and 8 do no MLP computation for this token.

Why this can create huge models efficiently¶

Imagine each expert contains:

With 64 experts, the expert bank contains:

But if each token activates only two experts, it uses only:

for that layer's computation.

So there are two different parameter counts worth distinguishing:

and:

This is why an MoE model can have a very large headline parameter count without requiring dense-model-equivalent FLOPs for every generated token.

A rough mental model is:

dense model:
100 workers are hired
100 workers work on every job

MoE model:
1,000 specialists are hired
20 appropriate specialists work on each job

The organization has much more stored expertise, but each job does not pay the computational cost of all 1,000 specialists.

Tokens can choose different experts¶

The routing happens per token, not necessarily per sequence.

Suppose the input is:

The derivative of x squared is 2x.

A hypothetical router might produce:

"The"        → experts 3, 7
"derivative" → experts 2, 5
"of"         → experts 3, 7
"x"          → experts 2, 5
"squared"    → experts 2, 5
"is"         → experts 3, 7
"2x"         → experts 2, 5

It is tempting to interpret Expert 2 as a "math expert" and Expert 7 as a "grammar expert."

Sometimes experts do develop recognizable specialization.

But we should be careful.

The router is not normally given human labels such as:

math
French
programming
history

It learns routing patterns jointly with the model.

An expert may specialize in combinations of syntax, token frequency, semantic features, position, or patterns that are not easily named.

So "a panel of human-like specialists" is a useful intuition, but not a literal description.

The distributed-systems problem appears immediately¶

Suppose we have eight GPUs and one expert on each GPU:

GPU 0 → Expert 0
GPU 1 → Expert 1
GPU 2 → Expert 2
GPU 3 → Expert 3
GPU 4 → Expert 4
GPU 5 → Expert 5
GPU 6 → Expert 6
GPU 7 → Expert 7

Now GPU 0 initially holds some tokens.

The router decides:

token A → Expert 6
token B → Expert 2
token C → Expert 6
token D → Expert 1

Those tokens must physically move to the GPUs holding their selected experts.

So an MoE layer creates a communication pattern like:

tokens on GPUs
     │
     ▼
   router
     │
     ▼
all-to-all token exchange
     │
     ▼
experts compute
     │
     ▼
all-to-all return exchange
     │
     ▼
tokens continue

This is called expert parallelism.

The parameters are distributed by expert, and the tokens travel to the parameters they need.

Notice how different this is from ordinary tensor parallelism.

In tensor parallelism, several GPUs cooperate on the same token's matrix multiplication.

In expert parallelism, different tokens may be sent to entirely different GPUs.

Why all-to-all communication is challenging¶

Suppose every GPU begins with a batch of tokens.

After routing, any token may need any expert.

The communication pattern can therefore look like:

GPU 0 → GPU 1, GPU 4, GPU 7
GPU 1 → GPU 0, GPU 3, GPU 6
GPU 2 → GPU 1, GPU 5, GPU 7
...

This is much less regular than a simple pipeline boundary.

Collectively, systems often implement this with an all-to-all exchange.

Each GPU sends different token buffers to several other GPUs and receives routed tokens in return.

That means the speed of an MoE layer depends not only on expert FLOPs but also on:

routing distribution
network bandwidth
message sizes
expert placement
token packing
load balance

Once again, model architecture and cluster topology become inseparable.

The router can create traffic jams¶

Suppose there are eight experts and 8,000 tokens.

Perfect balance would send roughly:

tokens to each expert.

But imagine the router instead chooses:

Expert 0: 4,700 tokens
Expert 1:   800
Expert 2:   700
Expert 3:   600
Expert 4:   400
Expert 5:   300
Expert 6:   300
Expert 7:   200

Expert 0 becomes a bottleneck.

Even if the other GPUs finish quickly, the layer cannot complete until the overloaded expert processes its tokens.

This is analogous to the pipeline lesson:

But here the imbalance is dynamic because the router decides where work goes.

Load-balancing losses encourage better routing¶

MoE training therefore usually includes mechanisms that encourage tokens to spread across experts.

Conceptually, we want expert usage closer to:

for \(N\) experts rather than collapsing onto a few favorites.

An auxiliary load-balancing objective may penalize routing distributions where some experts receive far more traffic than others.

The total training objective can look schematically like:

The language-model loss says:

predict tokens well.

The balancing term says something like:

while doing that, do not route almost everything to the same experts.

The coefficient \(\lambda\) controls the tradeoff.

This is interesting because the loss function is now partly optimizing systems efficiency, not merely predictive quality.

Experts also have capacity limits¶

Real systems often allocate a finite token capacity to each expert.

Suppose an expert is provisioned to process at most:

tokens in a microbatch.

If routing sends 1,700 tokens to it, the system needs a policy.

Depending on the architecture, it might:

drop excess assignments
route them elsewhere
increase capacity
pad buffers and accept inefficiency

A useful quantity is the capacity factor, which provisions some headroom above perfectly balanced routing.

More headroom reduces overflow risk but wastes memory and compute when buffers are underfilled.

Again we find a systems tradeoff:

MoE changes the meaning of parameter efficiency¶

For a dense model, parameter count and per-token computation are strongly coupled.

For MoE, they separate.

Suppose two models both perform roughly similar active computation per token:

Dense model:
20B total parameters
20B active parameters

MoE model:
200B total parameters
20B active parameters

The MoE model has access to a much larger parameter pool while spending something closer to 20B-parameter compute on each token.

This does not mean it behaves exactly like a dense 200B model for the cost of a dense 20B model.

Routing overhead, communication, training dynamics, memory, and expert utilization all matter.

But conditional computation gives model designers a new scaling axis:

can grow faster than:

That is the architectural attraction of MoE.

Inference has a different bottleneck¶

During inference, expert sparsity saves arithmetic because only selected experts run.

But all expert parameters may still need to live somewhere in the serving system.

So an MoE model can be:

For example, a server may need enough aggregate GPU memory to hold hundreds of billions of expert parameters even though each token activates only a fraction of them.

And decode introduces another complication.

At each generated token:

compute router
→ route token
→ transfer to selected experts
→ run experts
→ return result
→ continue layer

Because autoregressive decode is sequential, routing and communication latency can appear on the per-token critical path.

So sparse FLOPs do not automatically imply proportionally faster inference.

This mirrors a lesson we have seen repeatedly:

FLOP count alone is not a latency model.

A concrete mental model: a hospital with specialists¶

Imagine a hospital with 64 specialists.

A dense architecture would be like requiring every patient to visit all 64 doctors:

patient
→ cardiologist
→ neurologist
→ dermatologist
→ orthopedist
→ ...

That is enormously expensive.

An MoE hospital first sends the patient to triage:

patient
   ↓
triage
   ↓
choose 2 specialists

One patient might go to:

cardiology + radiology

while another goes to:

orthopedics + neurology

The hospital can employ many specialists without making every patient consult all of them.

But now the operational problems become obvious.

If triage sends half the hospital to cardiology:

cardiology: enormous queue
other departments: idle

throughput collapses.

And if the specialists are located in different buildings, moving patients between them takes time.

The model therefore needs both:

good specialist selection
and
good traffic engineering

That is MoE in one picture.

The deeper lesson¶

Dense Transformers make a simple assumption:

Mixture-of-Experts breaks that assumption.

Instead, it asks:

That transforms scaling.

We can add parameters without forcing every token to use all of them.

But sparsity does not eliminate cost. It changes where the cost appears.

Dense models concentrate cost in:

matrix multiplication

while MoE models introduce additional cost in:

routing
communication
load balancing
expert memory
irregular execution

So today's main principle is:

And this connects directly to the distributed-systems lessons we have just covered. Once tokens dynamically choose experts, the network connecting GPUs is no longer merely transporting tensors between fixed model partitions. It is carrying a changing stream of token traffic whose destination is decided by the model itself.