Mixture of Experts

The most important insight behind MoE fits in one sentence: not every input needs every parameter.

A standard neural network applies all its weights to every input. A 70B dense model runs 70 billion multiply-adds for every single token. A MoE model instead partitions its FFN parameters into E independent "expert" subnetworks and, for each token, activates only a small subset — typically 2 out of 8. Same inference FLOP as a ~17B dense model, but 8× more parameters (and therefore 8× more stored knowledge).

The Human Analogy

Think of a hospital. A patient arrives with a heart condition. The hospital doesn't consult every specialist simultaneously — the cardiologist, the radiologist, and the surgical team are activated. The dermatologist and the pediatric oncologist are not. The total institutional knowledge is vast, but each case activates only the relevant experts.

For language: a token about quantum mechanics routes to "physics" experts; a Python code token routes to "programming" experts. The model learns this routing — not by explicit labeling but entirely through gradient descent on the task objective.

Parameters vs FLOPs: Two Axes of Scale

Dense scaling conflates two things that MoE separates:

MoE is not a 2023 invention. The core idea is over 30 years old. What changed is scale, hardware, and the realization that sparsity is essential.

MoE only changes one thing in the standard transformer: the feed-forward network (FFN) layer. Everything else — attention, layer norms, residual connections, tokenization — stays identical.

Standard Transformer FFN (what you have now)

In a standard transformer, every FFN layer has two weight matrices: W₁ ∈ ℝ^d×d_ff and W₂ ∈ ℝ^d_ff×d where d_ff = 4d typically.

MoE Layer (the replacement)

Replace the single FFN with E identical-architecture FFNs (the "experts"), plus a small router network.

The Key Numbers

The router is the brain of a MoE model. It is a tiny linear layer (d → E parameters) trained end-to-end with the rest of the model. At inference, it produces one score per expert and selects the top K.

Step-by-Step Routing for One Token

Expert Choice Routing (alternative)

Standard top-K routing lets tokens choose experts. Expert Choice routing (Zhou et al., 2022) inverts this: each expert picks its top-C tokens from the batch. This guarantees perfect load balance by design — every expert processes exactly C tokens. The downside: a token might not be selected by any expert, requiring a pass-through with just the residual.

Without explicit regularization, MoE training collapses: a few popular experts attract all the tokens, get the most gradient signal, improve the most, attract even more tokens. This positive feedback loop — the Matthew effect — leads to expert collapse within the first few thousand steps.

The Auxiliary Load Balancing Loss

The Switch Transformer introduced the standard fix. Let f_i = fraction of tokens routed to expert i in the batch, and P_i = mean router probability assigned to expert i. The auxiliary loss penalizes uneven distributions:

Two signals penalized simultaneously: f_i is based on hard routing decisions (non-differentiable), P_i is differentiable through the softmax. The product connects a gradient path to the otherwise hard routing decision.

Expert Capacity

A second mechanism: each expert has a fixed capacity — the maximum tokens it can process per batch.

If more tokens are routed to an expert than its capacity allows, the excess tokens are dropped — they pass through via the residual connection without any expert processing. This adds a small error but prevents any single expert from becoming a bottleneck. Monitoring the "overflow rate" (fraction of dropped tokens) is a key training health metric.

DeepSeek's Device-Level Balancing

Standard auxiliary loss balances experts globally. But if 8 experts run on 8 different GPUs, what matters for efficiency is that each GPU has roughly equal work. DeepSeek-V3 adds a device-level auxiliary loss that balances token counts per device, not just per expert. This prevents communication bottlenecks in distributed training even when global expert balance looks good.

Three Architectural Families

DeepSeek-MoE: Fine-Grained + Shared Experts

DeepSeek's key innovations deserve detail. Instead of 8 large experts, they use 64 small experts with d_ff = d_ff_dense / 8. Plus 2 "shared" experts that always activate regardless of routing. The objective:

Why fine-grained helps: With 8 large experts, each token gets 2 experts = 25% of FFN capacity. With 64 small experts (each 1/8 the size), selecting K=6 gives 6/64 = 9.4% — but more flexibly combined from a richer set of specialists. The router can pick the precise combination of micro-specialties needed.

Does routing create genuine specialization, or is it arbitrary? Research consistently shows: yes, experts develop meaningful semantic specializations — even though no labels or explicit objectives guide this.

What Specialization Has Been Observed

1. Instability and the Router Cold Start

Early in training, the router is undertrained. It routes randomly, which means expert outputs have high variance. Two mitigations used in practice:

2. Expert Collapse

Even with auxiliary loss, expert collapse can happen. Signs: some experts have near-zero utilization after 10k+ steps. Causes: learning rate too high (router commits too early), auxiliary loss coefficient too low, or experts initialized identically (router has no initial preference to exploit).

Fix: Initialize expert weights with small noise (not identically). Use different random seeds per expert for the first few layers. This breaks symmetry and gives the router something to differentiate on.

3. Gradient Imbalance

Experts that receive more tokens get more gradient signal and train faster. This compounds the load imbalance problem — experts that happen to get more tokens early continue to attract them. The auxiliary loss mitigates but doesn't eliminate this. At very large batch sizes, even small imbalances in f_i compound significantly.

4. Communication Overhead in Distributed Training

In expert parallelism, different experts live on different GPUs. Routing a token to an expert on a different GPU requires all-to-all communication — every GPU sends some tokens to every other GPU. At scale, this communication becomes a significant overhead.

DeepSeek addresses this with expert grouping: ensure experts are grouped so that popular token types route to experts on the same GPU (group intra-GPU experts). Also: limit K to prevent communication from dominating.

Full Training Recipe

from transformers import MixtralConfig, MixtralForCausalLM
import torch

# ── MoE Configuration ──────────────────────────────────────────────────────
config = MixtralConfig(
    num_local_experts=8,              # E: total experts per layer
    num_experts_per_tok=2,            # K: experts activated per token
    router_aux_loss_coef=0.02,        # alpha: auxiliary loss weight
    hidden_size=4096,                 # d: model dimension
    intermediate_size=14336,          # d_ff: expert FFN hidden dim
    num_hidden_layers=32,
    num_attention_heads=32,
)

# ── Loss Computation ────────────────────────────────────────────────────────
def compute_moe_loss(model_output):
    """Total loss = task loss + auxiliary balancing loss."""
    task_loss = model_output.loss

    # Collect auxiliary losses from all MoE layers
    aux_losses = []
    for layer in model.model.layers:
        if hasattr(layer, 'block_sparse_moe'):
            router_logits = layer.block_sparse_moe.router_logits  # (seq_len, E)

            # f_i: fraction of tokens routed to expert i
            routing_weights = torch.softmax(router_logits, dim=-1)
            _, selected = torch.topk(router_logits, k=config.num_experts_per_tok)
            tokens_per_expert = torch.zeros(config.num_local_experts)
            for expert_idx in range(config.num_local_experts):
                tokens_per_expert[expert_idx] = (selected == expert_idx).sum().float()
            f_i = tokens_per_expert / tokens_per_expert.sum()         # normalize

            # P_i: mean router probability for expert i
            P_i = routing_weights.mean(dim=0)                         # (E,)

            # Auxiliary loss: E × Σ f_i × P_i (minimized at uniform distribution)
            aux_loss = config.num_local_experts * (f_i * P_i).sum()
            aux_losses.append(aux_loss)

    total_aux = config.router_aux_loss_coef * sum(aux_losses)
    return task_loss + total_aux

# ── Monitoring ──────────────────────────────────────────────────────────────
def log_routing_stats(router_logits, step):
    """Track expert utilization — critical for diagnosing MoE health."""
    _, selected = torch.topk(router_logits, k=2)
    counts = torch.zeros(8)
    for e in range(8):
        counts[e] = (selected == e).float().mean()

    print(f"Step {step} | Expert utilization: {counts.tolist()}")
    print(f"  Max/min ratio: {counts.max()/counts.min():.2f}x "
          f"(healthy: <3x, concerning: >10x)")

    # Overflow rate: tokens dropped due to capacity limits
    overflow = getattr(model, '_last_overflow_rate', 0)
    print(f"  Overflow rate: {overflow:.1%} (healthy: <2%, problem: >5%)")

MoE models have a fundamental inference challenge: all expert weights must live in memory even though only K/E are used per token. You pay full memory cost for the parameter advantage, but only partial compute cost.

Memory vs Compute

Batch Size and Throughput

MoE models benefit disproportionately from larger batch sizes. When batch = 1 (chat), each expert processes at most a handful of tokens — very inefficient. When batch = 512 (serving), each expert processes many tokens in parallel. The K/E throughput advantage is fully realized only at large batch sizes. For interactive applications, MoE is less efficient than its FLOP numbers suggest.

This is where the RL guides and MoE connect. The same PPO/GRPO algorithms apply to MoE models, but there are specific complications.

The Routing Instability Problem in GRPO

Recall from the Beyond PPO guide: GRPO's probability ratio is computed at the token level:

For a dense model, this ratio is stable across PPO epochs — the same weights produce the same logits for the same input. For a MoE model, this ratio has an additional source of variance: routing decisions can change between the numerator and denominator evaluations.

When π_θ and π_old route a token to different experts, the numerator and denominator compute through different parameter subsets. The probability ratio can be large even when the intended "change in probability" is small. This is why GSPO was invented.

Expert Load Balance During RL

A second MoE-specific challenge during RL training: the policy update can shift the input distribution enough that the routing balance, carefully tuned during pre-training, breaks. Experts that were balanced during pre-training may become unbalanced after 1000 RL steps because the policy now generates different token distributions.

Fix: Keep the auxiliary load-balancing loss active during RL training. Use a lower coefficient than pre-training (α = 0.001–0.005 vs 0.01–0.02 during pre-training). Occasionally monitor expert utilization during RL and reset the router if collapse occurs.

Knowledge Retention and Expert Forgetting

MoE experts are less likely to suffer from catastrophic forgetting during RL compared to dense models. Because each expert specializes in a subset of the input space, a gradient update on "LLM-specific" tokens mainly updates the relevant domain experts — not the language, math, or code experts. This is a practical benefit: MoE models tend to retain broader capabilities better during task-specific RL fine-tuning.

A MoE architecture is a natural fit for a multimodal LLM. response sequences have distinctly different token types that benefit from different processing strategies.

Why MoE Suits text Specifically

Suggested Expert Configuration for a text MoE

MoE + GRPO/DAPO for text

Combining MoE with the RL training pipeline from the previous guides: