Core ML

Mixture of Experts

Replace the single feed-forward layer in a transformer with a bank of N expert networks and a router that picks only k of them per token — total parameter count scales with N, but the compute per token stays constant because most experts are idle.

A mixture of experts replaces a single feed-forward layer with N expert networks; for each token the router activates only the top-k experts, so total parameters can be very large while active parameters per token stay constant

Why Does This Exist?

Scaling language models consistently improves quality, but scaling a dense transformer scales compute proportionally to parameters — double the parameters and you roughly double the FLOPs per token. At some point you hit a wall where the compute budget makes further scaling impractical.

Mixture of Experts (MoE) breaks that coupling. The idea is decades old (Jacobs et al., 1991) but was revived at scale by Google's Switch Transformer (Fedus et al., 2021) and has since become the architecture behind Mixtral, GPT-4 (reportedly), and many production models. Instead of one FFN that all tokens pass through, you have N expert FFNs and a router that selects k of them per token. Total parameters: $N \times \text{expert\_params}$ . Active parameters per token: $k \times \text{expert\_params}$ . Set $k=2$ and $N=8$ and you have 8× the parameters at ~2× the compute — a far better parameter-efficiency ratio than scaling a dense model.

Think of It Like This

A hospital with specialists, not generalists

A hospital with one doctor who treats every condition would work, but a hospital with specialists — cardiologists, neurologists, oncologists — can treat the same patient population with higher quality, because each specialist is deeply trained in one domain. A patient arrives and is routed to the right specialist, not every specialist. The hospital has far more total expertise than any single doctor, but each patient only activates one or two specialists.

Mixture of experts works the same way. Each expert learns to handle specific types of content (code, languages, formal reasoning, factual knowledge — though the split is learned, not designed). Each token is routed to its relevant experts and processed by only those.

How It Actually Works

Where MoE goes in the transformer

A standard transformer block has two sublayers: multi-head attention and a feed-forward network. MoE replaces the FFN with an expert pool and a router. Every token passes through attention normally (the full attention layer is dense, not sparse), then through the router to select k experts from the pool, then through those k expert FFNs. The expert outputs are weighted by router scores and summed.

The computation

For a token $x$ with router producing scores $g_1, \ldots, g_N$ (softmax over $N$ experts), let $\mathcal{T}$ be the top-k expert indices. The MoE output is:

y = \sum_{i \in \mathcal{T}} g_i \cdot \text{FFN}_i(x)

Only the $k$ selected expert FFNs actually run. The rest are skipped entirely — they receive no input and produce no output for this token.

Expert specialisation

Experts do develop soft specialisation during training — you can find that certain experts tend to activate for code tokens, others for mathematical reasoning, others for specific languages — even though nothing in the training objective explicitly encourages this. It emerges because routing tokens to an expert that handles them well is exactly what lowers loss.

Show Me the Code

Implementing a minimal top-2 MoE layer and measuring the compute ratio against a dense equivalent.

1import numpy as np2
3rng = np.random.default_rng(17)4
5n_tokens, d_model, d_ff, n_experts, k = 8, 64, 256, 8, 26
7# Router: linear projection from token representation to expert logits8W_router = rng.standard_normal((d_model, n_experts)) * 0.029# Expert FFN weights (simplified: one weight matrix per expert)10W_experts = rng.standard_normal((n_experts, d_model, d_ff)) * 0.0211W_out = rng.standard_normal((n_experts, d_ff, d_model)) * 0.0212
13X = rng.standard_normal((n_tokens, d_model))   # input tokens14
15logits = X @ W_router                            # (n_tokens, n_experts)16logits -= logits.max(-1, keepdims=True)17probs = np.exp(logits) / np.exp(logits).sum(-1, keepdims=True)18
19# Top-k selection20top_k_idx = np.argsort(probs, axis=-1)[:, -k:]  # (n_tokens, k)21outputs = np.zeros_like(X)22
23for i in range(n_tokens):24    for j, expert_id in enumerate(top_k_idx[i]):25        g = probs[i, expert_id]26        h = np.maximum(0, X[i] @ W_experts[expert_id])   # ReLU FFN27        outputs[i] += g * (h @ W_out[expert_id])28
29# Compute ratio: active vs total FFN parameters30active_params = k * (d_model * d_ff + d_ff * d_model)31total_params = n_experts * (d_model * d_ff + d_ff * d_model)32print(f"Total expert params: {total_params:,}")33print(f"Active params/token: {active_params:,}")34print(f"Ratio: {total_params // active_params}× more params, same compute")35# Total expert params: 524,28836# Active params/token:  131,07237# Ratio: 4× more params, same compute (with top-2 of 8)

At k=2, N=8, you get 4× the parameters for the same per-token compute.

Watch Out For

Assuming load balancing is solved by the router alone

Without an explicit load-balancing mechanism, most tokens route to the same few experts — those experts become over-loaded and the rest are under-utilised. This is called expert collapse. The switch transformer introduced an auxiliary loss that penalises uneven routing; Mixtral uses a similar approach. The loss term adds a small overhead but is essential for MoE to work at scale. Training without it produces a model where most experts are essentially unused.

Confusing active parameter count with memory requirement

All N experts have to be stored in memory, even though only k are active per token. A Mixtral 8×7B model has 46.7B total parameters but only ~12.9B active per token. The inference compute matches a dense 12.9B model, but the memory requirement is 46.7B parameters. On a single GPU, you need to fit all experts at once, which means MoE models need more GPU RAM than their active-parameter count suggests.

The Quick Version

MoE replaces the FFN in each transformer layer with N expert FFNs and a learned router that selects k experts per token.
Total parameters scale with N; compute per token stays at $k \times \text{expert\_params}$ — a direct decoupling of parameters from compute.
Experts develop soft specialisation during training (code, languages, formal reasoning) even without explicit labelling.
Expert collapse (all tokens routing to few experts) is a real failure mode, addressed by an auxiliary load-balancing loss.
All N experts must be stored in memory even though only k are active, so memory scales with total parameters, not active ones.