Mixture of Experts (MoE)

Mixture of Experts (MoE) is a transformer design that replaces each FFN with E independent experts and a router that picks k of them per token. Total parameters are E × FFN_size. Active parameters per token are k × FFN_size. Storage scales with E, compute scales with k.

That split matters because a dense transformer's FLOPs at inference equal its parameter count (times 2 for the forward pass). Scale a dense model up and each token pays the full bill. By 2024 the frontier was hitting a compute wall: to be meaningfully smarter, you needed exponentially more FLOPs per token. MoE breaks the link between parameter count and per-token compute. A typical 2026 configuration is E=256, k=8.

The 2026 frontier is almost entirely MoE: DeepSeek-V3 (671B total / 37B active), Mixtral 8×22B, Qwen2.5-MoE, Llama 4, Kimi K2, gpt-oss. On Artificial Analysis's independent leaderboard, the top 10 open-source models are all MoE.

The FFN swap

Dense transformer block:

h = x + attn(norm(x))
h = h + FFN(norm(h))

MoE block:

h = x + attn(norm(x))
scores = router(norm(h))              # (N_tokens, E)
top_k = argmax_k(scores)              # pick k of E per token
h = h + sum_{e in top_k}(
        gate(scores[e]) * Expert_e(norm(h))
    )

Each expert is an independent FFN, typically SwiGLU. The router is a single linear layer. A token picks its own k experts and gets a gated mixture of their outputs.

The load-balancing problem

Suppose the router sends 90% of tokens through expert 3. The other experts starve. Three fixes have been tried:

1. Auxiliary load-balancing loss (Switch Transformer, Mixtral). Add a penalty proportional to the variance in expert usage. Works, but adds a hyperparameter and a second gradient signal.
2. Expert capacity + token dropping (early Switch). Each expert processes at most C × N/E tokens, and overflow tokens skip the layer. Hurts quality.
3. Auxiliary-loss-free balancing (DeepSeek-V3). Add a learned per-expert bias that shifts the router's top-k selection. The bias is updated outside the training loss, so the main objective carries no penalty. That was 2024's big unlock.

DeepSeek-V3's approach: after each training step, check whether each expert's usage is above or below the target, and nudge its bias by ±γ. Selection uses scores + bias. The expert probabilities used for gating are the raw scores, unchanged. Routing and expression stay decoupled.

Shared experts

DeepSeek-V2/V3 also split experts into shared and routed. All tokens pass through the shared experts, while routed experts are picked via top-k. Shared experts capture common knowledge and routed experts specialize. V3 runs 1 shared expert plus top-8 of 256 routed.

Fine-grained experts

Classic MoE (GShard, Switch) makes each expert as wide as a full FFN. E is small (8–64) and k is small (1–2).

Modern fine-grained MoE (DeepSeek-V3, Qwen-MoE) makes each expert narrower (1/8 FFN size). E is large (256+) and k is larger (8+). Total parameters stay the same, but the number of expert combinations grows much faster: C(256, 8) = 400 trillion possible "experts" per token. Quality goes up while latency stays flat.

The cost profile

Per token, per layer:

DeepSeek-V3 beats Llama 3 70B (dense) on almost every benchmark while doing fewer active FLOPs per token. More parameters means more knowledge, and more active FLOPs means more compute per token. MoE decouples the two.

The catch: memory

All experts live on GPU regardless of which ones fire. A 671B model needs ~1.3 TB of VRAM for fp16 weights. Frontier MoE deployment requires expert parallelism: shard experts across GPUs and route tokens across the network. Latency is dominated by the all-to-all communication, not the matmul.

The full version is code/main.py, a compact MoE layer in pure stdlib with:

• n_experts=8 toy experts, each one linear map followed by silu (only the parameter count in Step 3 uses a ×3 SwiGLU-like factor)
• top-k=2 routing
• softmax-normalized gating weights
• auxiliary-loss-free balancing via per-expert bias

The forward pass has no shared expert. The shared expert from the Concept section only shows up in the Step 3 parameter count.

Step 1: the router

def route(hidden, W_router, top_k, bias):
    scores = [sum(h * w for h, w in zip(hidden, W_router[e])) for e in range(len(W_router))]
    biased = [s + b for s, b in zip(scores, bias)]
    top_idx = sorted(range(len(biased)), key=lambda i: -biased[i])[:top_k]
    # softmax over ORIGINAL scores of the chosen experts
    chosen = [scores[i] for i in top_idx]
    m = max(chosen)
    exps = [math.exp(c - m) for c in chosen]
    s = sum(exps)
    gates = [e / s for e in exps]
    return top_idx, gates

Bias affects selection, not gate weight. That is the DeepSeek-V3 trick: the bias corrects load imbalance without steering the model's predictions.

Step 2: run 1000 tokens through the router

Track how often each expert fires. The target is 250 per expert (1000 tokens × 2 / 8). With random router weights the toy starts close to uniform already: iteration 0 has entropy 2.072 against the ln(8) = 2.079 ceiling, so there is little imbalance to fix. Each iteration applies -γ to over-used experts and +γ to under-used ones. With seed 42 and γ=0.15 the usage does not settle. From iteration 1 it alternates between two fixed patterns (entropy 2.076, then 2.073), because each step overshoots. That is why γ needs tuning.

Step 3: param count comparison

Print the parameter counts for the toy MoE and for a DeepSeek-V3-shaped config: 256 routed + 1 shared experts, d_model=7168, 61 layers, each expert 1/8 of a full FFN. The printout counts the shared expert plus 8 routed experts, so 9 experts, as active. It reports about 806B FFN parameters in total (13.2B per layer × 61) and about 28B active. That is the toy's FFN-only arithmetic, not the published 671B total / 37B active, which comes from the real architecture and includes attention and a different expert layout. Use the published figures for the real model.

The next line prints a "llama-3-70b FFN total" of about 13B. It reuses the d=7168 FFN width for 32 layers and counts only FFN parameters, so it is not Llama 3 70B's real size. The closing "same active FLOPs" line is likewise just the toy's shorthand. Compare against published numbers instead: 37B active for DeepSeek-V3 versus 70B for dense Llama 3 70B, which is fewer active FLOPs per token.

HuggingFace loading:

from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x22B-v0.1")

For 2026 production inference, vLLM supports MoE routing natively, and SGLang has the fastest expert-parallel path. Both handle top-k selection and expert parallelism for you.

When to pick MoE:

• You want frontier quality at lower inference cost per token.
• You have the VRAM / expert-parallel infrastructure.
• Your workload is token-heavy (chat, code) not context-heavy (long docs).

When NOT to pick MoE:

• Edge deployment. You pay full storage for any active FLOP.
• Latency-critical single-user serving. Expert routing adds overhead.
• Small models (<7B). MoE's quality advantage appears above a compute threshold (~6B active params).

The skill in outputs/skill-moe-configurator.md picks E, k, and the shared-expert layout for a new MoE given a parameter budget, training tokens, and a deployment target.

1. Easy. Run code/main.py. Watch the auxiliary-loss-free bias update over its 10 iterations. Usage starts near uniform and then alternates between two patterns instead of settling. Try a smaller gamma and see whether the swings shrink.
2. Medium. Replace the learned router with a hash-based router (deterministic, no learning). Compare quality and balance. Why is the learned router better?
3. Hard. Implement GRPO-style "rollout-matched routing" (DeepSeek-V3.2 trick): log which experts fire during inference, force the same routing during gradient computation. Measure the effect on a toy policy-gradient setup.

• Shazeer et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer: the idea.
• Fedus, Zoph, Shazeer (2022). Switch Transformer: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity: Switch, the classic MoE.
• Jiang et al. (2024). Mixtral of Experts: Mixtral 8×7B.
• DeepSeek-AI (2024). DeepSeek-V3 Technical Report: MLA + auxiliary-loss-free MoE + MTP.
• Wang et al. (2024). Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts: the bias-based balancing paper.
• Dai et al. (2024). DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models: the fine-grained and shared-expert split.
• Kim et al. (2022). DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training: original shared-expert paper.