Mixture of Experts Architecture

The last two chapters focused on serving pressure: KV-cache memory, long contexts, batching, and GPU capacity. Mixture of Experts (MoE) attacks a different bottleneck inside the Transformer layer itself.

By the time you reach this article, you've seen how a Transformer layer works: self-attention mixes information across tokens, then a Feed-Forward Network (FFN) processes each token independently. That FFN is dense. Every single weight in it participates in every single forward pass, no matter what the input says. This is simple, but it's also expensive. If you want the model to know more facts, handle more syntax, or reason about more domains, the standard recipe is to make the FFN wider or deeper. Every extra neuron adds parameters and adds compute for every token.

Mixture of Experts (MoE) breaks that link. Instead of one giant FFN, an MoE layer keeps many smaller FFNs in parallel and uses a lightweight router to decide which ones to run for each token. This means the model can have hundreds of billions of total parameters while only spending compute on a small subset of them for any given input.

MoE is one widely used way to separate model capacity from per-token compute. Switch Transformer^[1] popularized sparse activation in trillion-parameter models, and open models such as Mixtral^[2], DeepSeek-V2^[3], and DeepSeek-V3^[4] show how MoE can scale to hundreds of billions of total parameters while keeping the active compute for each token much closer to a far smaller dense model.

From one FFN to many experts

In a standard Transformer layer, the FFN sublayer is a single dense network. In an MoE layer, that one FFN is replaced by $N$ parallel "expert" FFNs, plus a small router (also called a gating network) that decides which $k$ experts handle each token.^[5]

In many modern LLMs, the experts replace the FFN sublayer while attention stays dense and shared across all tokens.^[2][3] A common router design is a single linear projection from the token's hidden state down to $N$ scores.

A concrete routing walkthrough

Before we write a formula or code, let's trace one token through an MoE layer by hand. This is the shape of the computation that happens inside models like Mixtral on every single token.

Setup

• One token arrives with a hidden state $x$ of dimension 4.
• There are 4 experts, each a small FFN.
• The router uses top-$k = 2$ (it picks the 2 best experts).

Step 1: Router scores

The router multiplies the token by a learned weight matrix $W_g$ to get a raw score for each expert. Suppose the scores come out as:

Step 2: Softmax probabilities

Apply softmax so the scores become probabilities that sum to 1. The result is approximately:

Step 3: Top-k selection

Pick the top 2 experts. That's Expert 1 (0.52) and Expert 3 (0.26).

Step 4: Renormalize

The selected weights are renormalized so they sum to 1 again within the chosen subset:

• Expert 1: $0.52 / (0.52 + 0.26) \approx 0.67$
• Expert 3: $0.26 / (0.52 + 0.26) \approx 0.33$

Step 5: Expert computation

Only Experts 1 and 3 run. Each processes the same input token $x$ through its own FFN weights.

• $E_1(x)$ produces output vector $o_1$
• $E_3(x)$ produces output vector $o_3$

Step 6: Weighted sum

The final layer output blends the two expert outputs using the renormalized gate weights:

$$\text{Output} = 0.67 \cdot E_1(x) + 0.33 \cdot E_3(x)$$

That's it. The other two experts never process this token, so they contribute no expert-FFN FLOPs for this routing decision. Their weights still need to be stored or fetched by the serving system for other tokens that may select them.

The general MoE layer formula

We can write the same walkthrough as a compact equation. For an input token $x$, the MoE layer output is:

$$\text{MoE}(x) = \sum_{i \in \text{TopK}(g(x))} g_i(x) \cdot E_i(x)$$

Reading the formula: the router computes a gate distribution $g(x)$ over all $N$ experts. We keep only the configured top-$k$ entries, zero out the rest, and renormalize the remaining weights so they sum to 1. Mixtral uses top-2 out of 8, while other architectures choose different expert counts and routing policies.^[2][3] Each selected expert $E_i$ processes $x$ independently, then their outputs are blended using the gate weights $g_i$.

The architecture figure above is the visual form of this equation: the router scores every expert, keeps the top-$k$, runs only those FFNs, then blends the selected outputs.

Why sparse activation can pay off

Think of a dense model like one fulfillment center where every active station processes every order. An MoE model keeps a larger menu of stations but sends each order through only a selected subset. It can spend less computation per order than a similarly sized dense menu, but the facilities still occupy space and coordination between them still costs time. The analogy maps to active FLOPs, total weight residency, and dispatch traffic respectively.

The main benefit is that MoE can scale total parameters faster than per-token compute. By adding more experts while keeping top-$k$ fixed, the model gets a larger pool of FFN parameters, but each token still runs only a small subset. The expert FLOPs stay roughly tied to the active experts, not the full expert pool. Routing, batching, and communication still add overhead, so "same active parameters" doesn't mean identical latency.

That doesn't mean serving cost scales only with active parameters. Compute per token is sparse, but the system still needs access to the full expert pool, so memory footprint and communication often scale with total parameters, not active parameters.

Only the FFN blocks are replicated across experts. Attention layers, embeddings, normalization layers, and the router are still shared. That's why Mixtral 8x7B exposes about 47B total parameters and about 13B active parameters per token, not the naive 56B and 14B you would get from multiplying everything by 8.^[2]

The key point is the split: active parameters mostly drive per-token expert FLOPs, while total parameters drive weight memory and checkpoint size. Network traffic depends on where selected experts live.

The router in code

Here's a basic PyTorch implementation of the standard top-k softmax router. It takes token embeddings, computes scores for each expert, and returns the indices and normalized weights for the top-$k$ choices.

Switch Transformer used top-1 routing, while Mixtral and DeepSeek-V2 use multiple selected routed experts per token.^[1][2][3] Selecting more experts changes quality, compute, capacity, and communication trade-offs; it is not automatically a better serving configuration.

Alternative: expert choice routing

While the standard Top-K router forces each token to choose its top $k$ experts, Expert Choice Routing^[6] flips the mechanism: each expert selects a fixed bucket of top-scoring tokens from the current batch.

This gives fixed expert-side load balancing because every expert processes the same number of tokens. However, it means some tokens might be processed by many experts while other tokens might be ignored entirely if no expert selects them.

By fixing the number of tokens each expert processes, expert choice routing avoids the same expert-side load-balancing problem. This can simplify the training objective and improve hardware utilization because every expert gets a predictable bucket. The risk moves to token coverage: a token can be selected by many experts, one expert, or none unless the design adds coverage rules. In practice, token-choice routing with balancing losses remains the simpler mental model for open-weight MoE LLMs, while expert-choice routing is useful to know when reading training-efficiency papers.

Why routers collapse without a balancing policy

During training, routers can drift toward expert collapse: many tokens route to a few "popular" experts while others are underutilized. Load statistics tell you whether this is occurring; they do not prove that a perfectly uniform router is best for model quality.

Imagine an e-commerce fulfillment center with four specialist teams: one for small electronics, one for apparel, one for bulky furniture, and one for perishable goods. On the first day, the electronics team happens to be slightly faster at handling common orders. The router notices this and sends more orders to the electronics team. That team gets more practice, so it improves. The router sends even more orders to them. Within a few weeks, the electronics team is handling 80% of all orders, while the furniture and perishables teams sit idle. The warehouse is still paying rent for all four teams, but it's effectively operating with one overworked team.

The router is trained jointly alongside the rest of the network.^[5] If a particular expert looks better for common tokens early in training, the router can send it more assignments. That expert then receives more gradient updates, which can reinforce the preference and leave other experts under-trained. The operational test is not a story about specialization: inspect assignment shares, overflow, task loss, and communication load together.

An auxiliary load-balancing loss

Switch Transformer uses a top-1 auxiliary loss that penalizes uneven routing.^[1] The following form is also a useful teaching calculation for top-$k$ routing when $f_i$ is normalized over routed assignments:

$$\mathcal{L}_{\text{balance}} = \alpha \cdot N \sum_{i=1}^N f_i \cdot p_i$$

Reading the formula: multiply each expert's routed-assignment fraction $f_i$ by its average routing probability $p_i$ and sum them up. If the router sends too many assignments to one expert (high $f_i$) and gives it high probability (high $p_i$), this product is large, so the loss penalizes that concentration. It is a balancing pressure, not proof that equal utilization maximizes task quality.

where:

• $f_i = \frac{\text{assignments routed to expert } i}{\text{total routed assignments}}$; with top-$k$, the denominator is $T \times k$
• $p_i = \frac{1}{T}\sum_{t=1}^T \text{softmax}(W_g \cdot x_t)_i$ (average router probability)
• $\alpha$ is a configurable coefficient (Switch reports experiments including $10^{-2}$)
• $N$ is the number of experts

A balanced reference point is $f_i \approx p_i \approx 1/N$ (uniform distribution). A trained deployment still needs task-quality and systems measurements before you judge its observed specialization.

Here's a concrete top-1 numeric check. Suppose we have 4 experts and a mini-batch of 100 tokens:

If the load were perfectly uniform, each expert would get $f_i = p_i = 0.25$, and the sum would be $4 \times (0.25 \times 0.25) = 0.25$. With positive $\alpha$, this auxiliary term applies pressure against the imbalanced 0.400 case. Tune its weight with task-quality measurements because too much balancing pressure can conflict with the primary objective.

The auxiliary loss is a routing pressure, not the main learning objective. It says, "use the expert pool evenly enough," while the language-modeling loss still decides whether the model predicts the next token well.

Capacity factor and overflow policy

In a batched environment, many MoE implementations use fixed-size expert buffers so the layer can run predictable tensor operations. This size is determined by the capacity factor:

$$\text{Expert Capacity} = \frac{\text{Tokens per Batch} \times k}{N} \times \text{Capacity Factor}$$

This computes the average number of tokens each expert would receive under perfectly uniform routing, then multiplies by the capacity factor. A factor of $1.0$ means the buffer is exactly sized for uniform routing. Setting it to $1.2$ allows an expert to receive 20% more tokens than the average, providing a safety margin.

What happens if an expert is assigned more tokens than its capacity buffer allows depends on implementation:

1. Dropping or pass-through: Switch-style designs can skip overflowed expert computations and rely on the residual path for those token-layer assignments.^[1]
2. Rerouting: A system can send overflow to another eligible expert, changing load and communication.
3. Dropless execution: A system can execute all assignments with variable work, retaining quality paths while accepting less predictable workload.

This highlights the trade-off: a higher capacity factor reduces overflow risk but reserves more buffer space and may waste compute on empty slots. A lower factor is more compact but makes whichever overflow policy you selected more important.

How modern architectures scale the idea

The basic MoE pattern is consistent across modern models, but architects have refined it in interesting ways.

Mixtral 8x7B^[2]

MoE is applied only to the Feed-Forward Network (FFN) layers. Attention layers are shared across all tokens, which is the standard approach. It uses a SwiGLU (Swish Gated Linear Unit) activation function, a variant of the Gated Linear Unit (GLU) that's been shown to improve Transformer performance.

DeepSeek-V2^[3]

DeepSeek-V2 uses several extensions to coarse-grained MoE:

• Fine-grained experts: 160 smaller routed experts instead of a handful of coarse experts, which gives the router many more combinations to compose per token.^[3]
• Shared experts: 2 experts are always active for all tokens. The design aims to capture common computation and reduce redundancy among routed experts; every token gets 2 shared expert outputs in addition to the 6 routed ones.^[3]
• Device-limited routing: The selected routed experts for each token are constrained to a small number of devices, reducing cross-device communication overhead.^[3]
• Multi-level balancing losses: DeepSeek-V2 adds separate expert-level, device-level, and communication-balance losses to keep routing efficient at cluster scale.^[3]
• Multi-head Latent Attention (MLA): Compresses the KV cache into a low-rank latent vector, which the paper reports as a large KV-memory and throughput win relative to DeepSeek 67B.^[3]

The following diagram contrasts Mixtral's coarse-grained top-2 routing with DeepSeek-V2's fine-grained routed-plus-shared design. The DeepSeek side is schematic, not a full 162-expert grid.

DeepSeek-V3^[4]

DeepSeek-V3 retains the DeepSeekMoE pattern at a larger reported scale.

Each MoE layer has 1 shared expert and 256 routed experts, with 8 routed experts activated per token.^[4]

The paper also highlights a major training change: instead of leaning on auxiliary balancing losses, DeepSeek-V3 uses an auxiliary-loss-free load-balancing strategy.^[4] That's important because it shows how modern MoE work is not only about adding more experts. Router stability, communication efficiency, and keeping the main training objective clean matter just as much.

The hardware reality of serving an MoE

Implementing MoE introduces complexities not present in dense models. The compute per token is sparse, but the system still faces several hard serving constraints.

Communication overhead

In distributed training and inference, experts often reside on different devices to meet memory requirements. Routing a token to an expert on another GPU requires transferring the token's hidden state over a high-bandwidth interconnect such as NVLink within a node or InfiniBand across nodes.

This communication can quickly become a bottleneck if not managed carefully. The process usually has an all-to-all dispatch phase that scatters tokens to the GPUs hosting their selected experts, followed by an all-to-all combine phase that returns the expert outputs to the original token order. To maintain high throughput, system designers try to overlap those transfers with useful compute so the network isn't stalling the whole layer.

Memory requirements

Unlike dense models where there's only one FFN per layer, MoE models carry many expert FFNs. Fast GPU deployments generally keep those expert weights resident somewhere in the serving pool so routing does not stall on weight fetches.

This creates the core MoE paradox: low active compute per token, but very high total weight residency. A model like Mixtral 8x7B only activates about 13B parameters per token, yet the serving system still needs the full ~47B parameter checkpoint available across the cluster.^[2] Tensor Parallelism (TP) splits individual layer computations across multiple GPUs, allowing larger shared layers to fit into memory, but it doesn't solve the unique structure of MoE. Expert Parallelism (EP) is needed to place different full experts on different GPUs while keeping every expert reachable when a token needs it.

Training instability

MoE training can be more sensitive than training a comparable dense Transformer. The top-$k$ routing step introduces discontinuities at expert-selection boundaries, and small routing biases can snowball into expert collapse. Gradients still flow through the selected experts and their gate weights, but the router is noisier to optimize than a dense FFN.

To combat this, implementations may add stabilizers. Router z-loss^[8] penalizes unnecessarily large routing logits; the ST-MoE paper reports improved training stability from this term. Router jitter can add exploration pressure early in training. The relevant evidence is training loss, routing distribution, overflow, and downstream quality rather than assuming one stabilizer is mandatory for every MoE.

Expert parallelism layout

Large distributed MoE deployments commonly use Expert Parallelism (EP). While Tensor Parallelism (TP) splits individual matrix multiplications, EP maps entire experts to specific GPUs. The layout below illustrates how a 96-expert model might be distributed across three GPUs, with each device holding a distinct subset of experts while sharing the attention parameters:

Each GPU stores a subset of experts. When a token routes to an expert on another GPU, an all-to-all communication sends the token there and retrieves the result.

Expert prefetching

For CPU-offloaded serving (for example, running a quantized MoE checkpoint on consumer hardware), one possible design is:

1. Keep attention weights and router weights in GPU memory (always needed).
2. Keep expert weights on CPU or SSD.
3. After a layer's router decides which experts are needed, transfer those expert weights before its expert computation.
4. Optionally use prediction or caching to stage likely later-layer experts, accepting misses or unused transfers.

This can make large MoE checkpoints fit where full GPU residency cannot, especially when combined with weight quantization, but it trades memory for transfer-dependent latency. You cannot know an exact next-layer route until that layer receives its hidden state, so unconditional "prefetch the next experts" is not free.

Batch routing efficiency

In batch serving, different tokens in a batch may route to different experts. The key optimization is batched expert execution:

1. Router assigns all tokens in the batch to their top-k experts.
2. Group tokens by expert assignment.
3. Execute each expert as a batched matrix multiply over its assigned tokens.
4. Scatter results back to original token positions.

This turns $N$ small matrix multiplies into a smaller number of large, efficient ones, which is critical for achieving high GPU utilization with MoE models.

This table is a hypothesis to benchmark, not a universal crossover. Exact outcomes depend on kernels, expert placement, interconnect, and request mix. MoE adds router computation, expert selection, token grouping, and sometimes all-to-all communication; batching can amortize that overhead, but only measurements on the target system establish whether MoE wins latency or throughput.

Common pitfalls

"Active FLOPs are small, so deployment memory will be small too"

Symptom: FLOP estimates look cheap, but deployment still runs out of VRAM. Cause: Team budgeted only for active parameters per token and forgot that the full expert pool must still be resident somewhere in the serving system. Fix: Separate active expert FLOPs from total weight residency. Size memory for the full expert pool, not only the routed path.

"Router imbalance will fix itself later"

Symptom: One or two experts receive most tokens by early training. Cause: Router found a slightly better path, that expert got more gradients, and collapse reinforced itself. Fix: Track routing histograms from the start and add balancing pressure such as auxiliary loss, z-loss, jitter, or capacity constraints before unused experts go stale.

"Sparse compute guarantees lower latency"

Symptom: MoE benchmark wins disappear at batch size 1. Cause: Sparse expert FLOPs are not the whole latency story. Router work, grouping, dispatch, and all-to-all traffic can dominate single-request latency. Fix: Evaluate MoE on the real serving regime. Use dense models for low-latency single requests when routing overhead outweighs sparse compute gains.

"Expert placement is a later infra detail"

Symptom: Throughput drops sharply after experts are split across GPUs. Cause: Expert placement created communication hotspots. Tokens now bounce across devices faster than grouped GEMMs can amortize the transfers. Fix: Treat routing and placement as one systems problem. Use expert parallelism, device-aware routing, and batching to keep all-to-all traffic under control.

"Balanced histograms mean routing is healthy"

Symptom: Quality drops even though router histograms look balanced. Cause: Experts may be hitting capacity limits, and the configured overflow policy is dropping, rerouting, or delaying assignments. Fix: Inspect overflow counts and policy directly. Evaluate capacity factor, routing balance, burstiness, and throughput before assuming the issue is elsewhere.

"Expert-choice routing covers every token automatically"

Symptom: Expert-choice routing leaves some tokens underprocessed. Cause: Experts select their favorite tokens first, so token coverage is no longer guaranteed automatically. Fix: Add explicit coverage rules or use token-choice top-k routing when every token must reliably receive expert computation.

"Experts are neat human-labeled topic buckets"

Symptom: Teams describe experts as clean topic buckets like "math expert" or "biology expert." Cause: Human-friendly labels are being projected onto behavior that often follows token patterns, formatting, or structural clusters instead. Fix: Validate specialization with routing traces, activation analysis, and targeted prompts that test a claimed specialty before telling a story about it.

Mastery check

Key concepts

• Sparse versus dense FFN computation
• Router logits, top-k selection, and renormalization
• Expert collapse and balancing pressure
• Capacity factor and dropped-token overflow
• Total weight residency versus active FLOPs
• Expert parallelism, dispatch, and all-to-all cost

Evaluation rubric

• Choose dense vs MoE for a serving regime, and justify that choice with latency, throughput, memory residency, and communication cost.
• Trace one token by hand: compute router scores, pick top-$k$, renormalize, then blend expert outputs.
• Diagnose collapse from routing histograms and say which mitigation you would try first.
• Compute capacity risk from batch tokens, top-$k$, number of experts, and capacity factor, then explain what overflow does to quality.
• Explain why tensor parallelism helps shared layers but expert parallelism and placement strategy are needed for routed FFNs.
• Defend when shared experts, device-limited routing, or finer-grained expert pools buy enough systems value to justify extra complexity.

Follow-up questions

Your team wants a 48-expert MoE because active FLOPs look like a 13B dense model. Why can memory still look like a much larger checkpoint?

Active FLOPs follow only the experts selected for each token, but the full expert pool still has to live somewhere in the serving system. That means VRAM or cluster-wide weight residency is driven by total parameters, not only active parameters. MoE separates compute cost from capacity, not memory from capacity.

One expert takes 70% of tokens by step 500. Why is that failure self-reinforcing, and what do you inspect first?

It is self-reinforcing because the favored expert gets more tokens, more gradients, and therefore improves faster, which makes the router favor it even more. First inspect routing histograms, average router probabilities, dropped-token rates, and balancing-loss terms. Those signals show whether collapse is still reversible or already starving the rest of the expert pool.

Why might a dense model beat MoE on latency for batch size 1 while an MoE candidate wins throughput at batch size 64?

At batch size 1, MoE pays router, grouping, and often all-to-all overhead without enough routed tokens to amortize it. Dense models avoid that coordination overhead and can win on raw latency. At larger batch sizes, grouped expert execution can make sparse compute competitive on throughput. Benchmark both claims on the actual hardware.

A serving stack already uses tensor parallelism for attention. Why is that not enough for MoE layers?

Tensor parallelism shards individual matrix multiplications, but MoE also needs whole-expert placement, token dispatch, and result recombination. That is why expert parallelism exists: the problem is not only splitting math, but moving tokens to whichever GPU holds their selected experts and then restoring original token order.

A batch sends 4096 tokens through top-2 routing with 32 experts and capacity factor 1.0. What is the nominal per-expert buffer, and what failure appears if one expert gets 400 assignments?

The nominal per-expert buffer is (4096 * 2) / 32 = 256 token assignments. If one expert gets 400 assignments, then 144 assignments overflow that buffer. A Switch-style dropping policy skips the overflowed expert computation and may hurt quality; rerouting or dropless designs pay different compute and communication costs. Inspect the configured policy.

Your interconnect is weak, and all-to-all traffic is already the bottleneck. What do you change before adding even more experts?

Start by constraining and measuring communication, not by assuming more experts solve it. Device-limited routing, expert placement that keeps common routes local, shared experts for common computation, and larger grouped batches are candidate mitigations. Adding routed experts changes capacity and traffic; benchmark that change against the existing bottleneck.^[3][4]

Try it yourself: manual router exercise

Here's a self-contained pencil-and-paper exercise that checks whether you can trace the exact computation an MoE router performs. Think of it as routing a single customer order through the right fulfillment specialists. Do this without looking back at the worked example above.

Setup

• One token arrives with hidden state $x = [0.5, -0.2, 0.1]$.
• The router weight matrix is $W_g = \begin{bmatrix} 0.4 & -0.1 & 0.2 \\ -0.2 & 0.3 & 0.1 \\ 0.1 & 0.1 & -0.3 \end{bmatrix}$.
• There are 3 experts. Top-$k = 2$.
• Suppose Expert 1 outputs $o_1 = [1.0, 0.5]$, Expert 2 outputs $o_2 = [0.2, 1.2]$, and Expert 3 outputs $o_3 = [0.8, -0.1]$.

Questions

1. Compute the raw router scores: $s = W_g \cdot x$. (Hint: this is a 3-vector.)
2. Apply softmax to get probabilities $p_1, p_2, p_3$.
3. Which two experts are selected?
4. Renormalize the selected probabilities so they sum to 1.
5. Compute the final layer output as a weighted sum of the selected expert outputs.

Solution sketch

1. $s = [0.4(0.5) + (-0.1)(-0.2) + 0.2(0.1), \dots] = [0.24, -0.15, 0.00]$
2. Softmax gives roughly $p \approx [0.406, 0.275, 0.319]$
3. Selected: Expert 1 and Expert 3 (top 2 probabilities).
4. Renormalized: $g_1 = 0.406 / (0.406 + 0.319) \approx 0.560$, $g_3 = 0.319 / (0.406 + 0.319) \approx 0.440$
5. Output $\approx 0.560 \cdot [1.0, 0.5] + 0.440 \cdot [0.8, -0.1] \approx [0.912, 0.236]$

If your numbers are within 0.02 of these, you understand the routing math well enough to read any MoE paper.

What this unlocks next

You now understand the central trick behind large sparse MoE models: replace selected dense FFNs with routed expert sets and activate only configured paths per token. You can trace a token through the router by hand, explain why total parameters do not equal compute cost, and evaluate balance, overflow policy, memory residency, and communication before claiming a serving win.