Mixture of Experts (MoE) — conditional computation at scale

A standard transformer block runs every token through the same feed-forward network (FFN). In a Mixture of Experts block, that single FFN is replaced by N parallel FFN "experts" and a lightweight router that selects which experts process each token. Only k experts (typically k=2) activate per token, so the model stores far more parameters than it uses on any given forward pass. The result: model capacity scales with total parameters, but inference cost scales with active parameters.

The router

The router is a single linear layer that takes a token's hidden state h (dimension d_model) and produces a score for each expert. For N experts, the router weight matrix is W_r of shape (d_model, N). The routing scores for a token are:

\text{scores} = \text{softmax}(h @ W_{r})

The router then selects the top-k experts by score. Each selected expert processes the token independently, and the expert outputs are combined as a weighted sum using the softmax scores as weights:

python

import torch
import torch.nn as nn
import torch.nn.functional as F

class Router(nn.Module):
    def __init__(self, d_model, num_experts, top_k=2):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        self.top_k = top_k

    def forward(self, hidden_states):
        # hidden_states: (batch, seq_len, d_model)
        logits = self.gate(hidden_states)  # (batch, seq_len, num_experts)
        scores = F.softmax(logits, dim=-1)

        top_k_scores, top_k_indices = torch.topk(scores, self.top_k, dim=-1)
        # Renormalize the top-k scores so they sum to 1
        top_k_scores = top_k_scores / top_k_scores.sum(dim=-1, keepdim=True)

        return top_k_scores, top_k_indices

The router itself is tiny — for Mixtral 8x7B with $d_{\text{model}} = 4096$ and 8 experts, W_r has $4096 \times 8 = 32, 768$ parameters. The routing decision adds negligible compute overhead.

Mixtral 8x7B — the canonical open MoE

Mistral AI released Mixtral 8x7B in December 2023. Each of its 32 transformer layers contains a standard grouped-query attention block (shared across all tokens) and 8 expert FFN blocks, of which 2 are selected per token.

The parameter breakdown:

Each expert FFN — a SwiGLU network with $d_{\text{model}} = 4096$ and intermediate size 14,336. Three weight matrices per expert: W_gate (4096 × 14336), W_up (4096 × 14336), W_down (14336 × 4096). That's 3 × 4096 × 14336 ≈ 176M parameters per expert.
8 experts per layer — 8 × 176M ≈ 1.41B FFN parameters per layer.
Shared attention per layer — 32 query heads, 8 KV heads (GQA with 4 groups), $d_{\text{head}} = 128$ . Roughly 4 × 4096 × 4096 ≈ 67M parameters (Q, K, V, O projections, with KV being smaller due to GQA).
32 layers total — 32 × (1.41B + 67M) ≈ 47B total parameters.
Active per token — 2 experts selected means 2 × 176M ≈ 352M FFN params active, plus the full attention (~67M), per layer. Across 32 layers: roughly 13B active parameters per token.

Mixtral matches or exceeds Llama 2 70B on most benchmarks while using ~13B active parameters per forward pass — inference cost comparable to a dense 13B model.

Expert specialization

Analysis of Mixtral's routing patterns (Jiang et al. 2024) reveals that experts specialize by function rather than by topic. Specific experts consistently handle:

Punctuation and formatting tokens
Code syntax and operators
Mathematical expressions
Named entities and proper nouns
Function words and connectives

This specialization emerges entirely from training — the router learns to dispatch tokens to whichever expert produces the lowest loss for that token type. No supervision or labeling guides which expert handles what. The specialization is also partial: most experts handle a mix of token types, with statistical tendencies rather than hard boundaries.

Load balancing and expert collapse

Without constraints, the router converges to routing most tokens to a small number of "popular" experts, leaving others undertrained. This is expert collapse — a positive feedback loop where better-trained experts attract more tokens, which makes them even better-trained, which attracts even more tokens.

The standard mitigation is an auxiliary load-balancing loss (Fedus et al. 2022, Switch Transformer). For N experts, define f_i as the fraction of tokens routed to expert i and P_i as the mean router probability assigned to expert i across all tokens:

L_{\text{balance}} = \alpha \cdot N \cdot \sum f_{i} \cdot P_{i} \text{for} i \text{in} \text{range}(N)

This loss is minimized when routing is perfectly uniform ( $f_{i} = \frac{1}{N}$ for all i). The coefficient alpha controls the tradeoff between routing efficiency and balance — typical values are alpha=0.01 to alpha=0.1. Too high and the router prioritizes balance over quality; too low and experts collapse.

python

def load_balancing_loss(router_logits, top_k_indices, num_experts, alpha=0.01):
    """
    router_logits: (batch * seq_len, num_experts) — raw router outputs
    top_k_indices: (batch * seq_len, top_k) — selected expert indices
    """
    scores = F.softmax(router_logits, dim=-1)
    num_tokens = router_logits.shape[0]

    # f_i: fraction of tokens routed to each expert
    one_hot = F.one_hot(top_k_indices, num_experts).float()  # (tokens, top_k, experts)
    tokens_per_expert = one_hot.sum(dim=1).sum(dim=0)  # (num_experts,)
    f = tokens_per_expert / num_tokens

    # P_i: mean router probability for each expert
    P = scores.mean(dim=0)  # (num_experts,)

    return alpha * num_experts * (f * P).sum()

The Switch Transformer (Fedus et al. 2022) introduced this formulation with k=1 (a single expert per token) and 128 experts, demonstrating 7x speedup over the dense T5-XXL at equivalent quality.

DeepSeek-V3 — scaling MoE to 671B parameters

DeepSeek-V3 (DeepSeek-AI 2024) pushes the MoE architecture to its current extreme. Each layer has 256 fine-grained experts plus 1 shared expert. The router selects 8 of the 256 routed experts per token; the shared expert processes every token unconditionally.

The numbers:

Total parameters: ~671B
Active parameters per token: ~37B (8 routed experts + 1 shared expert + attention)
Training cost: 2,788K H800 GPU-hours — roughly $5.6M at market rates, remarkably efficient for a 671B model
Context length: 128K tokens
Vocabulary: 128K tokens

DeepSeek-V3 introduces two routing innovations:

Auxiliary-loss-free load balancing

Instead of the auxiliary loss (which degrades model quality by competing with the language modeling objective), DeepSeek-V3 adds a per-expert bias term to the routing scores. The bias for over-loaded experts decreases and for under-loaded experts increases, adjusted dynamically during training. This achieves balanced routing without any auxiliary loss term.

Expert-level routing with token dropping

Each expert has a capacity limit — the maximum number of tokens it can process in a single batch. If an expert reaches capacity, additional tokens routed to it are "dropped" (processed by the shared expert only or redistributed). This prevents memory and compute imbalances across GPUs during training, where each expert typically resides on a different accelerator.

The shared expert pattern

DeepSeek-V3's shared expert processes every token, handling "common knowledge" — syntactic patterns, frequent word predictions, and other capabilities that all tokens need. The routed experts specialize in rarer or more complex patterns. This decomposition is more efficient than routing everything: the shared expert's parameters are always active (no routing overhead), and the routed experts can specialize more aggressively because they don't need to redundantly learn basic language competence.

The shared expert also provides a safety net — if the router makes a poor decision, the shared expert still produces a reasonable baseline output, preventing catastrophic failures on individual tokens.

GPT-4 and the rumored MoE architecture

GPT-4 (OpenAI 2023) is widely reported to use a MoE architecture, though OpenAI has never confirmed this officially. Leaked reports (since partially corroborated by OpenAI employees in interviews) describe 16 experts with 2 active per token, approximately 1.8T total parameters, and roughly 220B active parameters per forward pass. Whether these specific numbers are accurate, the inference cost characteristics of GPT-4 (competitive with much smaller dense models in latency) are consistent with MoE.

Expert parallelism — distributing MoE across GPUs

Training and serving MoE models introduces a unique distributed-computing challenge: tokens must be physically routed to the GPU holding the selected expert. In a dense model, every GPU processes the same layers on different data (data parallelism) or different parts of the same layer on the same data (tensor parallelism). In MoE, expert parallelism adds a third dimension:

Each expert resides on a specific GPU (or set of GPUs). With 8 experts and 8 GPUs, each GPU holds one expert.
After the router selects experts, tokens are dispatched via all-to-all communication: each GPU sends tokens destined for other GPUs' experts and receives tokens destined for its own expert.
Each GPU processes tokens through its local expert, then another all-to-all sends results back.

The all-to-all communication is the bottleneck. For Mixtral 8x7B with 8 experts across 8 GPUs, each routing decision requires two all-to-all operations per layer, 32 layers deep. With NVLink interconnects (~900 GB/s per GPU on H100 SXM), the communication overhead is manageable; on PCIe or cross-node InfiniBand, it can dominate training time.

DeepSeek-V3's 256 experts across thousands of GPUs required careful expert placement and routing-aware batch scheduling to keep communication costs under control. Their training report describes achieving 61% model FLOPs utilization (MFU) on 2048 H800 GPUs — respectable for an MoE model, where 55–65% MFU is typical versus 50–60% for dense models at similar scale.

Implementing a complete MoE layer

python

import torch
import torch.nn as nn
import torch.nn.functional as F

class ExpertFFN(nn.Module):
    """A single SwiGLU expert FFN, same as Llama's FFN."""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, num_experts=8, top_k=2, alpha=0.01):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.alpha = alpha

        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.experts = nn.ModuleList([
            ExpertFFN(d_model, d_ff) for _ in range(num_experts)
        ])

    def forward(self, hidden_states):
        batch, seq_len, d_model = hidden_states.shape
        flat = hidden_states.view(-1, d_model)  # (B*T, d_model)

        # Route
        logits = self.router(flat)               # (B*T, num_experts)
        scores = F.softmax(logits, dim=-1)
        top_k_scores, top_k_idx = torch.topk(scores, self.top_k, dim=-1)
        top_k_scores = top_k_scores / top_k_scores.sum(dim=-1, keepdim=True)

        # Dispatch to experts and combine
        output = torch.zeros_like(flat)
        for k in range(self.top_k):
            expert_indices = top_k_idx[:, k]      # (B*T,)
            expert_weights = top_k_scores[:, k]   # (B*T,)

            for e in range(self.num_experts):
                mask = (expert_indices == e)
                if mask.any():
                    expert_input = flat[mask]
                    expert_output = self.experts[e](expert_input)
                    output[mask] += expert_weights[mask].unsqueeze(-1) * expert_output

        # Compute load balancing loss
        num_tokens = flat.shape[0]
        one_hot = F.one_hot(top_k_idx, self.num_experts).float()
        tokens_per_expert = one_hot.sum(dim=1).sum(dim=0)
        f = tokens_per_expert / num_tokens
        P = scores.mean(dim=0)
        self.aux_loss = self.alpha * self.num_experts * (f * P).sum()

        return output.view(batch, seq_len, d_model)


# Verify shapes and routing
moe = MoELayer(d_model=512, d_ff=2048, num_experts=8, top_k=2)
x = torch.randn(2, 16, 512)
out = moe(x)
print(f"Input:  {x.shape}")    # (2, 16, 512)
print(f"Output: {out.shape}")  # (2, 16, 512)
print(f"Aux loss: {moe.aux_loss.item():.6f}")

total_params = sum(p.numel() for p in moe.parameters())
expert_params = sum(p.numel() for e in moe.experts for p in e.parameters())
print(f"Total params: {total_params:,}")    # ~25M (8 experts × ~3.1M each + router)
print(f"Expert params: {expert_params:,}")  # ~25M
print(f"Active per token: ~{expert_params // 8 * 2 + 512 * 8:,}")  # 2 experts + router

This implementation loops over experts sequentially, which is fine for understanding but inefficient on GPU. Production implementations (Megablocks, Fairseq MoE) use batched matrix multiplications — grouping all tokens destined for the same expert into a single batched GEMM call, eliminating the Python loop.

The economics of MoE

The appeal of MoE is economic. Consider two models with identical benchmark performance:

Dense 70B — 70B active parameters per token. Inference on 2× H100 GPUs (tensor parallelism). At ~1000 tokens/sec throughput, cost per million tokens is ~$1.50.
MoE 47B total / 13B active — equivalent quality, 13B active parameters. Fits on a single H100 with room to spare. At ~3000 tokens/sec, cost per million tokens is ~$0.50.

The tradeoff: the MoE model needs more GPU memory (47B params must be in memory even though only 13B activate), but the compute cost per token is dramatically lower. For inference-heavy workloads — which is most production LLM deployment — MoE dominates dense architectures at equivalent quality.

← Previous