Mixture of Experts in LLMs

The Scaling Problem MoE Solves

Larger language models can learn more patterns because they have more parameters, but a standard dense Transformer uses the same full set of layer weights for every token. Adding parameters therefore increases both model capacity and the computation required to process each token. At large scale, this coupling becomes expensive during training and inference.

Mixture of Experts, or MoE, addresses this problem with conditional computation. Instead of one feed-forward network, an MoE layer contains several independent feed-forward networks called experts. A learned router examines each token representation and selects only a small number of those experts to process it.

This separates two quantities that are identical in a dense layer: total parameters, which determine how much expert capacity the model stores, and active parameters, which determine how much of that expert capacity one token uses. The model can increase the first quantity without increasing the second at the same rate.

1. Start With a Dense Feed-Forward Network

A Transformer block has two main sublayers. Self-attention moves information between token positions. A feed-forward network, or FFN, then transforms each token position independently. In a dense Transformer, every token passes through the same FFN weights.

Many current LLMs use a gated variant such as SwiGLU, which has three large matrices instead of two. These FFN matrices account for a substantial share of a Transformer's parameters and arithmetic. That makes the FFN a natural place to add conditional computation.

2. Where MoE Fits Inside a Transformer

A decoder-only Transformer is built by repeating the same basic unit many times. That unit is called a Transformer block or Transformer layer. Each block receives one vector for every token, updates those vectors with attention and an FFN, and passes the result to the next block.

The data flow through one block

In the common pre-normalization design, this flow can be written compactly. Let $x$ be the input to the block, $h$ the representation after attention, and $y$ the final block output:

What changes in an MoE block?

The attention sublayer, both residual connections, and normalization layers remain in place. The dense FFN in step 5 is replaced by an MoE sublayer containing a router and several expert FFNs:

Model designers decide which blocks use this replacement. A model may use MoE in every block, alternate dense and MoE blocks, or begin with several dense blocks before switching to MoE blocks. Regardless of the schedule, an expert is only an FFN inside one block. It is not a complete Transformer block and does not contain its own attention layer.

Sparse MoE therefore does not make attention sparse. Long-context attention, the KV cache, and attention communication remain separate costs. MoE changes which token-wise FFN parameters are active inside the blocks that use it.

3. Router, Experts, and Combiner

A basic sparse MoE layer has three conceptual pieces. Suppose a batch contains $T$ token states, each with width $d$, and the layer has $E$ experts.

4. The Routing Math

For one token state $x\in\mathbb{R}^{d}$, a linear router produces one logit per expert:

The model keeps only the $k$ experts with the largest scores. Let $S(x)$ be that selected set. A common implementation renormalizes the selected probabilities and forms a weighted sum:

With top-1 routing, one expert produces the output. With top-2 routing, two experts run and their results are blended. Some architectures use sigmoid routing rather than a single softmax across all experts, and details such as score normalization, expert grouping, routing bias, and shared experts vary by model.

Interactive: Top-k Expert Routing

Select different tokens, switch between top-1 and top-2 routing, and reveal the raw router logits. Notice that the router makes a new decision for every token at every MoE layer.

5. One MoE Forward Pass, End to End

An MoE forward pass has one mathematical path and one systems path. Mathematically, the router selects expert functions and combines their outputs. Operationally, the runtime reorganizes token states into efficient expert batches, moves them when experts are sharded, and restores the original token order before the residual addition.

The visualization uses four experts and top-2 routing, so each token creates two expert assignments. Capacity checks happen before dispatch. Inter-device exchange is required only when the selected experts are stored on different devices.

6. What an Expert Actually Computes

Each expert is typically an ordinary FFN with its own weights. If Expert 3 and Expert 7 receive the same vector, they can produce different outputs because training has changed their matrices in different ways.

The two-dimensional transforms above are a visual analogy. A real expert operates in thousands of dimensions, uses nonlinear activations, and may contain a gated FFN such as SwiGLU. The important point is that the router chooses a parameterized transformation, not a stored answer or a database entry.

7. Tokens Route Independently

Routing happens for token representations, not whole prompts. The tokens in one sentence can visit different experts, and the same token string can route differently when its context changes its hidden state. The decision is also repeated at every MoE layer, so a token can take a different route deeper in the model.

8. Total Parameters vs Active Parameters

The most useful MoE distinction is between parameters stored by the model and parameters used for one token. Suppose every MoE layer has $E$ experts, each expert has $P_{\text{expert}}$ parameters, and the router selects $k$ of them.

If a layer has eight equally sized experts and uses top-2 routing, a token activates two of the eight expert FFNs. That is one quarter of the expert bank, plus the shared attention, router, normalization, embedding, and output parameters elsewhere in the model.

MoE only has a fair compute comparison when expert width, top-k, layer placement, and shared dense work are specified. Two same-sized active experts can cost more arithmetic than one dense FFN. Model designers often adjust expert width or the number of MoE layers to meet a target compute budget.

9. How an MoE Model Learns

The language-model objective stays familiar. The model predicts the next token, computes cross-entropy, and backpropagates the loss. MoE changes the path that each token takes through selected FFN parameters.

This creates a feedback loop. Early routing decisions determine which experts receive examples. Those experts improve on their assigned traffic, which can make the router prefer them again. Useful specialization can emerge, but uncontrolled feedback can also cause expert collapse.

10. Expert Collapse and Load Imbalance

If the router sends most tokens to the same expert, several problems arrive together. The popular expert becomes a compute bottleneck, its capacity fills, other experts learn slowly, and the model wastes most of its stored parameters. This failure mode is often called expert collapse.

Perfectly uniform routing is not always the goal. Real data is not uniform, and some experts may deserve more traffic. The engineering goal is to avoid severe hotspots while preserving enough freedom for the router to learn useful assignments.

11. Expert Capacity and Token Overflow

Accelerators work best with bounded tensor shapes. Training systems therefore often give each expert a fixed number of assignment slots for one batch. If there are $T$ tokens, $k$ choices per token, $E$ experts, and capacity factor $c$, a common capacity rule is:

12. Auxiliary Load-Balancing Loss

One widely used strategy adds a small auxiliary objective. Let $f_i$ be the fraction of hard token assignments sent to expert $i$, and let $P_i$ be that expert's average router probability. A Switch-style form is:

Concentrating both assignments and probability on a few experts raises this term. The coefficient must be small enough that language modeling remains the main objective. Some newer systems use constrained assignment, routing biases, or auxiliary-loss-free balancing because a global penalty can interfere with the main task.

13. Sparse Gradients and Uneven Learning

An expert receives ordinary backpropagation through the tokens it processed. An unselected expert has no task-dependent computation for that token, so it receives no task gradient from that path. Across a sufficiently large and balanced batch, every expert should still receive useful work.

Sparse activation reduces expert arithmetic, but training must still maintain all expert weights, gradients, and optimizer state. Checkpointing and optimizer memory therefore follow total parameters much more closely than active parameters.

14. Shared Experts and Routed Experts

Some architectures include one or more shared experts that run for every token alongside sparsely selected routed experts. The shared path can learn broadly useful computation while routed experts spend more capacity on conditional behavior.

Shared experts add guaranteed computation, so they reduce sparsity. They can still be valuable when the architecture benefits from separating common transformations from routed capacity. As with routed experts, human-readable labels in a diagram are explanatory examples rather than guaranteed learned roles.

15. Expert Parallelism Across GPUs

Large expert banks do not fit efficiently on one accelerator. Expert parallelism places different experts on different devices. After routing, tokens travel to the devices that own their selected experts, then expert outputs travel back.

Real training jobs combine several forms of parallelism. The main MoE-specific risk is that useful FLOPs wait behind network transfers or one overloaded expert. Fast interconnects, local expert placement, balanced batches, token permutation kernels, and grouped GEMMs are therefore part of the architecture in practice.

16. What Changes During MoE Inference

At inference time, the router still makes a decision for every token at every MoE layer. There is no need for backpropagation or optimizer state, but the full expert bank must remain available and decode batches can be small or irregular.

17. How Real MoE Architectures Differ

MoE names a family of designs rather than one fixed architecture. The examples below show how routing choices changed across influential systems. Parameter figures use each paper's own reporting convention, so they should not be treated as perfectly standardized comparisons.

The progression is not a simple march toward more experts. Top-k, expert size, shared capacity, balancing method, placement, and communication topology are co-designed. A routing method that looks elegant in an equation may perform poorly when it creates tiny matrix multiplications or excessive network traffic.

18. Common MoE Misconceptions

19. A Minimal MoE Implementation

This framework-neutral pseudocode captures the forward pass. Production code replaces the Python loop with token permutation, capacity handling, grouped matrix multiplication, and collective communication kernels.

def sparse_moe(x, router, experts, top_k=2):
    # x: [tokens, d_model]
    router_logits = router(x)                 # [tokens, num_experts]
    router_probs = softmax(router_logits, -1)

    weights, expert_ids = topk(router_probs, top_k)
    weights = weights / weights.sum(-1, keepdim=True)

    output = zeros_like(x)
    for expert_id, expert in enumerate(experts):
        token_ids, slots = where(expert_ids == expert_id)
        if len(token_ids) == 0:
            continue

        expert_input = x[token_ids]
        expert_output = expert(expert_input)
        output[token_ids] += weights[token_ids, slots, None] * expert_output

    return output

A useful test starts with one expert and top-1 routing. That result should match a dense FFN with the same weights. Next, use two identical experts and verify that weighted combination still matches. Only then introduce distinct experts, capacity limits, and distributed dispatch.

Primary Sources and Further Reading

The architecture descriptions and reported model figures in this guide are grounded in the original papers and technical reports below.

• Adaptive Mixtures of Local Experts, Jacobs et al. (1991)
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, Shazeer et al. (2017)
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, Lepikhin et al. (2020)
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, Fedus et al. (2021)
• ST-MoE: Designing Stable and Transferable Sparse Expert Models, Zoph et al. (2022)
• Mixtral of Experts, Jiang et al. (2024)
• DeepSeek-V3 Technical Report, DeepSeek-AI (2024)