Syed Anzar

Posted on Oct 26

Mixture of Experts: What Actually Happens When a Token Hits the Router

#ai #machinelearning #llm #transformer

Mixture of Experts: What Actually Happens When a Token Hits the Router

You've seen the model names: Mixtral 8x7B, DeepSeek-V3 671B-A37B, Nemotron 3 Ultra 550B-A55B. The first number is total parameters. The second is active parameters. A 671B model that only "uses" 37B per token.

Developers nod along. "Oh, sparse activation. Efficient." But ask what actually happens between the attention output and the next residual add, and the hand-waving starts.

Let's trace a single token through a Mixture of Experts layer. No analogies. The real machinery.

The Setup: What Gets Replaced

A standard Transformer block:

Attention → Add & Norm → FFN → Add & Norm

The FFN is typically the parameter-heavy part. In Llama-3-70B, each FFN expands d_model=8192 → 28672 → 8192 (GELU/SiLU in between). That's ~186M parameters per layer just for the FFN. With 80 layers, FFNs dominate the parameter count.

MoE replaces only the FFN. Attention stays dense. The residual stream, layer norms, the output head — all shared.

# Dense FFN (what you're used to)
class DenseFFN(nn.Module):
    def __init__(self, d_model, d_hidden):
        super().__init__()
        self.gate_proj = nn.Linear(d_model, d_hidden, bias=False)
        self.up_proj   = nn.Linear(d_model, d_hidden, bias=False)
        self.down_proj = nn.Linear(d_hidden, d_model, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

# MoE FFN (what actually runs)
class MoEFFN(nn.Module):
    def __init__(self, d_model, d_hidden, num_experts, top_k):
        super().__init__()
        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.experts = nn.ModuleList([
            DenseFFN(d_model, d_hidden) for _ in range(num_experts)
        ])
        self.top_k = top_k

Every expert is structurally identical to the dense FFN it replaces. Same d_model, same d_hidden, same activation. The difference: there are N of them, and each token sees only k.

Step 1: The Router Scores Every Expert

The token vector x (shape [d_model]) hits the router — a single linear layer, no bias, output dimension num_experts.

router_logits = x @ W_router^T        # [num_experts]

W_router is tiny. For Mixtral 8x7B: d_model=4096, num_experts=8 → 32K parameters. Negligible.

The router is learned jointly with the experts. It's not a classifier trained separately. During pretraining, gradients flow through the selected experts back into the router weights. The router learns: "token x needs expert 3 and 7."

Step 2: Top-K Selection (The Sparsity)

router_logits → softmax → probabilities over all experts.

probs = softmax(router_logits)        # [num_experts], sums to 1

Then: keep top-k, zero the rest, renormalize.

topk_probs, topk_indices = topk(probs, k=2)   # e.g., [0.62, 0.38], [3, 7]

Only experts 3 and 7 compute. The other 6 do zero FLOPs for this token.

This is where the compute savings come from. A dense FFN does 2 * d_model * d_hidden multiply-adds per token. MoE with top-2 does 2 * 2 * d_model * d_hidden_expert. If d_hidden_expert = d_hidden_dense / 4, you get ~4x more total parameters for the same per-token compute.

But wait — the checkpoint still stores all experts. VRAM doesn't shrink. Only FLOPs.

Step 3: Dispatch — The Part Nobody Shows You

This is where implementations diverge. The naive loop:

# Naive: iterate experts, mask tokens
final_output = torch.zeros_like(x_flat)      # [num_tokens, d_model]

for expert_idx in range(num_experts):
    token_mask = (topk_indices == expert_idx).any(dim=1)   # [num_tokens]
    if not token_mask.any():
        continue
    expert_input = x_flat[token_mask]       # [num_assigned, d_model]
    expert_output = self.experts[expert_idx](expert_input)
    # Weight by routing probs...

This works but kills GPU utilization. Real systems (Megablocks, DeepSpeed-MoE, vLLM) permute:

Sort tokens by assigned expert → contiguous blocks per expert
Each expert processes its block as a single dense matmul
Scatter results back to original positions

The permutation is an all-to-all communication if experts live on different GPUs. This is why MoE training needs fast interconnects (NVLink, InfiniBand). At inference with tensor parallelism, it's the same shuffle.

Step 4: Expert Computation (Identical to Dense FFN)

Each selected expert runs its own FFN on its assigned tokens.

Expert_i(x) = down_proj( silu(gate_proj(x)) * up_proj(x) )

Experts are independent. No cross-talk. They can be on different devices, different memory regions, even different architectures (though nobody does that in practice — hurts load balancing).

The outputs are weighted by the router probabilities:

output = Σ (weight_i * expert_output_i)

Then add residual, layer norm, next block.

The Load Balancing Problem

If the router sends 90% of tokens to expert 0, you've recreated a dense model with wasted capacity. Worse: expert 0 becomes a bottleneck, others atrophy.

Two main solutions:

Auxiliary Loss (Switch Transformers, Mixtral)

Add to training loss:

L_aux = α * N * Σ (f_i * p_i)

f_i = fraction of tokens dispatched to expert i
p_i = mean router probability for expert i over the batch
α = hyperparameter (typically 0.01–0.1)

Minimized when f_i = p_i = 1/N — uniform load.

Router Bias / Expert Choice (DeepSeek-V3, newer)

Per-expert bias terms added to router logits before top-k, updated during training to equalize load. No extra loss term, faster convergence.

# Bias added to logits, not part of model parameters
router_logits = x @ W_router^T + expert_bias   # [num_experts]

The bias is adjusted online: if expert i is overloaded, its bias decreases.

Expert Specialization: It's Messier Than You Think

"Expert 3 does code, expert 7 does reasoning" — convenient mental model, often wrong.

In practice:

Specialization emerges but overlaps heavily
Same expert can handle different domains at different layers
Token position matters (early vs late in sequence)
Shared experts (always-active) handle common patterns

Mixtral 8x7B: no clean semantic split visible in routing heatmaps. DeepSeek-V3's 256 experts show some clustering (math, code, language) but it's fuzzy.

Don't design prompts expecting to "trigger the code expert." The router is opaque.

The Two Parameter Numbers Explained

DeepSeek-V3 671B-A37B:

671B total: Every weight in the checkpoint. All 256 experts × FFN params + shared attention + embeddings + router + output head.
37B active: What actually computes per token. 8 experts active (top-8? no — DeepSeek uses top-8 routed + 1 shared = 9 experts per token? Let me check... actually DeepSeek-V3 uses 256 routed experts, top-8, plus 1 shared expert. Active ≈ 8 * (FFN params per expert) + shared attention params).

The "active" number is an approximation. It ignores:

Router computation
Permutation/scatter overhead
Attention (always dense)
KV cache (same as dense)
Embedding/output projection

A 37B-active MoE is not the same latency as a dense 37B model. Usually slower due to memory bandwidth (loading expert weights) and dispatch overhead.

Inference: The VRAM Trap

You need all experts in VRAM simultaneously. Even if top-2 activates 2/8 experts, the other 6 must be resident — the next token might need them.

Mixtral 8x7B: ~47B params in VRAM (not 56B — attention is shared). Needs ~90 GB VRAM (bf16). Two H100s or four A100-80GB.

Quantization helps (4-bit → ~24 GB), but expert count still forces high memory.

What Developers Get Wrong

Misconception	Reality
"MoE = faster inference always"	Only if memory bandwidth isn't the bottleneck. Small batch? Dense often wins.
"Active params = latency equivalent"	Ignores dispatch, expert loading, all-to-all comms.
"Experts specialize cleanly"	They don't. Routing is fuzzy, overlapping, layer-dependent.
"I can run 671B on consumer GPU"	You need the total params in VRAM. Active params only save FLOPs.
"Top-1 vs top-2 doesn't matter much"	Top-1 saves compute but hurts quality. Top-2 is standard; top-4+ rarely worth it.
"Router is just a classifier"	It's trained end-to-end with experts. Gradients flow through selected experts only.

The Mental Model You Should Keep

Token → Attention → Router → [Expert A, Expert B] → Weighted Sum → Residual → Next Layer
             ↑                              ↑
        (shared, dense)              (sparse, conditional)

Router: Learned linear projection, ~0.1% of params, decides where compute goes
Experts: Independent FFNs, ~99% of MoE params, do the actual work
Top-k: Hard sparsity — unselected experts do zero compute
Load balancing: Auxiliary loss or bias terms prevent collapse
VRAM: All experts must fit. Active params ≠ memory footprint.

Why This Matters for You

If you're serving: MoE models need high VRAM, tensor/expert parallelism, and batching to amortize dispatch cost. vLLM, SGLang, and TensorRT-LLm all have MoE-specific kernels now.

If you're fine-tuning: LoRA on MoE is trickier. Do you adapt routers? All experts? Just the shared expert? Current best practice: LoRA on all linear layers including router, or use methods like MoRA (Mixture of LoRAs).

If you're choosing a model: Don't chase the big total-parameter number. A dense 32B often beats a 100B-A12B MoE on consumer hardware. MoE shines when you have the VRAM and want maximum quality per FLOP.

One More Thing: Shared Experts

Modern MoEs (DeepSeek-V3, Nemotron 3 Ultra, OLMoE) add shared experts — always-active FFNs that run for every token alongside the routed ones.

output = shared_expert(x) + Σ routed_expert_i(x)

Why? The router can't perfectly capture "common knowledge every token needs." Shared experts handle universal patterns (basic syntax, common reasoning steps) while routed experts specialize.

Nemotron 3 Ultra: 550B total, 55B active, 1 shared expert always on.

TL;DR

Mixture of Experts replaces one big FFN with many small ones + a learned traffic cop. The cop (router) sends each token to 1–2 experts. The other experts sleep. Total parameters explode; per-token FLOPs stay flat.

The magic isn't the architecture — it's that the router learns useful routing purely from next-token prediction gradients. No supervision. No labels. Just: "this expert helped predict the next token, so route similar tokens there next time."

That's the whole trick. Everything else is engineering to make it fast.

Top comments (1)

Agentic Affairs • Oct 26

The 'checkpoint still stores all experts, VRAM doesn't shrink, only FLOPs do' line is the correction most MoE explainers skip — sparse activation saves compute, not memory. The dispatch section is where it earns its keep: permuting tokens into per-expert blocks is the difference between a textbook MoE and one that actually runs fast. One pushback on load balancing: an auxiliary loss that forces uniform routing can also fight specialization — have you seen cases where the balancing term keeps experts too general to be worth having?