Mixture of Experts (MoE)

Why this matters

The largest production LLMs in 2026 — Mixtral, DeepSeek-V3, GPT-4-class models, Llama-4 — are almost all Mixture-of-Experts. MoE is how you scale a model's parameters (and capacity) without proportionally scaling its active compute. A 671B-parameter DeepSeek-V3 only activates 37B parameters per token, costing roughly the same as a dense 37B model at inference.

If you go to ML interviews and cannot describe sparse routing, load balancing, expert parallelism, and the auxiliary loss, you are missing the architectural story of the last two years.

Learning objectives

1. Distinguish dense vs sparse (MoE) transformers.
2. Implement a top-k router with auxiliary load-balancing loss.
3. Reason about parameters vs active parameters and what each costs.
4. Describe expert parallelism and its bandwidth implications.
5. Recognise modern MoE designs (Mixtral, DeepSeek, Snowflake Arctic, Llama-4).

1. The big idea

In a dense Transformer, every parameter is used for every token.

FFN: x -> W2 (act(W1 x))         # ~2/3 of all params; 100% active per token

In a Sparse Mixture-of-Experts, you replace the single FFN with E parallel FFNs (experts). A small router picks the top-k (typically k=1 or k=2, sometimes k=8 for fine-grained MoE) experts for each token. Only the chosen experts are computed; the rest sit idle.

Total params: ~ E × FFN params       (huge)
Active params per token: k × FFN params      (much smaller)

For Mixtral 8×7B: E=8, k=2 → ~13B active out of ~47B total. Inference cost similar to a dense 13B but quality closer to a 47B.

2. The router

Per token x ∈ R^d:

gate_logits = W_g x                                # (E,)
weights, expert_ids = top_k(softmax(gate_logits), k)
output = sum_i weights[i] * expert_{ids[i]}(x)

A Linear(d, E) layer with no bias. After softmax, each token has a distribution over experts; we keep top-k.

class MoEFFN(nn.Module):
    def __init__(self, d, d_ff, E=8, k=2):
        super().__init__()
        self.gate = nn.Linear(d, E, bias=False)
        self.experts = nn.ModuleList([SwiGLU(d, d_ff) for _ in range(E)])
        self.k = k

    def forward(self, x):
        # x: (B*T, d)
        logits = self.gate(x)                              # (N, E)
        topk_w, topk_idx = logits.topk(self.k, dim=-1)
        topk_w = F.softmax(topk_w, dim=-1)                 # weights re-norm over top-k

        out = torch.zeros_like(x)
        for i, expert in enumerate(self.experts):
            mask = (topk_idx == i)                         # (N, k) booleans
            if not mask.any(): continue
            # rows that route to expert i (in any of their k slots):
            row_idx = mask.any(dim=-1).nonzero().squeeze(-1)
            x_sub = x[row_idx]
            y_sub = expert(x_sub)
            # weight: sum of weights for slot i (could appear in either of k slots)
            w = (topk_w * mask.float()).sum(-1)[row_idx, None]
            out[row_idx] += w * y_sub
        return out

Real implementations are much more efficient — they group tokens per expert, use grouped GEMM kernels, and often dispatch via all-to-all across GPUs.

3. Load balancing — the dirty secret

Without intervention, the router collapses to "always pick expert 3." Then expert 3 trains, the others wither, the model becomes effectively dense-with-extra-junk. To force balanced usage, every MoE adds an auxiliary loss.

Switch / Mixtral-style auxiliary

For each batch:

• f_i = fraction of tokens routed to expert i.
• P_i = average gate probability mass on expert i.

aux_loss = E * sum_i (f_i * P_i)

This loss is small when both f and P are roughly uniform. Add it to the main loss with a small coefficient (e.g., 0.01).

Z-loss

Penalises the magnitude of the router logits to prevent extreme overconfidence:

z_loss = mean( logsumexp(logits)^2 )

Both losses are cheap and standard. PyTorch implementations: see transformers/modeling_mixtral.py.

Capacity factor

Each expert has a fixed capacity C = capacity_factor * tokens_per_expert. If too many tokens want one expert, the surplus get dropped (zero output) or rerouted. Capacity factor of 1.25-2.0 is typical.

4. Expert parallelism (EP)

In a 100B+ MoE, experts don't fit on one GPU. They are placed across GPUs:

GPU 0: experts 0,1
GPU 1: experts 2,3
GPU 2: experts 4,5
GPU 3: experts 6,7

Per layer, tokens all-to-all to the GPU holding their expert, run that expert's FFN, and all-to-all back.

The all-to-all is bandwidth-heavy: limits scale and latency. NVLink (intra-node) and InfiniBand (inter-node) bandwidth determine practical EP size.

5. Modern MoE designs

Trends:

• More, smaller experts (DeepSeek's 256) → finer specialization.
• Shared experts: a few experts are always on for every token (DeepSeek). They learn common patterns; routing chooses specialised ones.
• Auxiliary-free balancing: DeepSeek-V3 replaces the load-balancing loss with a clever bias-update scheme; cleaner gradients, no aux-loss tuning.
• Per-token vs per-sequence routing: per-token is standard; per-sequence is simpler but less expressive.

6. Why MoE works (intuitively)

• Capacity per token cost: an expert is only computed when "needed." Total memorised facts can scale with E without scaling FLOPs.
• Specialisation: experts spontaneously specialise (some on math, some on code, some on language X). Visible in routing analyses.
• Better scaling laws: at fixed compute, MoE achieves lower loss than equivalent-FLOPs dense, by a factor that depends on E and k.

But:

• Memory at inference is all parameters (you can't free unused experts; they may be picked next token).
• Communication overhead (all-to-all) on multi-GPU is real.
• Routing introduces noise and capacity drops.
• Fine-tuning MoE is harder (you may damage routing).

7. Inference considerations

Memory

You hold all experts in VRAM. A 47B Mixtral 8×7B needs the full 47B in memory at bf16.

Throughput

Each GPU computes a fraction of FFN work (only its experts), so GPU utilisation is bursty. Engines (vLLM, TensorRT-LLM, SGLang) batch tokens by expert via grouped GEMM (grouped_gemm kernels).

Quantization

AWQ, GPTQ, FP8 work for MoE weights. compressed-tensors in vLLM handles MoE-specific quant. Gains scale to make 671B fit on a small cluster.

Routing entropy

Watch the entropy of the gate distribution: too low → collapse; too high → no specialization. Healthy is around log(k) + small.

8. Fine-tuning an MoE — special care

If you fine-tune (SFT/DPO) an MoE:

• Freeze the router initially (requires_grad=False on gate.weight). Otherwise the router shifts away from balanced and you damage rarely-routed experts.
• Use a small subset of experts if compute is tight; LoRA on each expert, or LoRA only on the shared expert (DeepSeek style).
• Re-introduce auxiliary loss with small weight (~1e-3) to maintain balance.

Hands-on lab (4 hours)

moe_lab.ipynb:

1. Implement MoEFFN with E=4, k=2 using SwiGLU experts. Replace the FFN in your nano-GPT (Lesson 2.4).
2. Train on tiny_shakespeare for 3000 steps. Track per-expert routing fractions and aux loss.
3. Set aux_loss coefficient to 0. Show that one expert collapses to 80%+ of routings.
4. Reset and add z-loss; observe more stable routing.
5. Load Mixtral-8x7B-Instruct (or mistralai/Mixtral-8x7B-Instruct-v0.1 if you have the GPU). Inspect its config; compute total vs active params.
6. Plot the gate distribution for several prompts (math, code, English). Discuss specialization.
7. Bonus: implement a shared-expert version where one expert always fires plus 2 specialised; compare.

Common pitfalls

1. Skipping aux loss → expert collapse. Always include.
2. Top-k=E (no sparsity) → you're back to dense, but with extra mediation overhead. Use k ≪ E.
3. Tiny per-expert capacity → frequent token drop → mysterious quality drops on long sequences.
4. Wrong batched implementation — naive Python for expert in experts: is slow. Use grouped GEMM or efficient libraries (fused_moe in vLLM).
5. Fine-tuning router weights with the same lr as everything else → router drifts.

Self-check

1. What is "active params" vs "total params"?
2. Why do MoE models need an auxiliary load-balancing loss?
3. What does expert parallelism communicate over the network?
4. What is a shared expert (DeepSeek-style)?
5. Why is MoE memory at inference still high despite sparsity?

References

• Shazeer et al. (2017), "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer."
• Fedus et al. (2021), "Switch Transformer: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity."
• Du et al. (2022), "GLaM: Efficient Scaling of Language Models with Mixture-of-Experts."
• Jiang et al. (2024), "Mixtral of Experts."
• DeepSeek-AI (2024), "DeepSeek-V2 Technical Report."
• DeepSeek-AI (2025), "DeepSeek-V3 Technical Report."
• Snowflake (2024), "Arctic — A Truly Open, Enterprise-Grade LLM."
• HuggingFace Mixtral source.