Tutorial 11: Mixture of Experts (MoE)¶
A standard transformer processes every token through every layer's FFN with the same weights. Mixture of Experts replaces the single FFN with multiple specialized "expert" networks and a learned router that selects which experts to use for each token. This increases model capacity without proportionally increasing compute — only a subset of experts is activated per token.
1. The Core Idea¶
Instead of one FFN with intermediate_size neurons, MoE has num_experts
parallel FFN networks, each with the same architecture. A router (a small
linear layer) examines each token and picks the top-k experts for that token:
Token → Router → "This token should use experts 2 and 5"
Token → Expert_2 processing → partial output
Token → Expert_5 processing → partial output
Final = score_2 × expert_2_output + score_5 × expert_5_output
The Tradeoff¶
- Parameters:
num_experts × (gate + up + down)— much larger total model - Compute: Only
num_experts_per_tok × (gate + up + down)per token — same as a smaller dense model - Result: Large model capacity, moderate compute cost
For example, Qwen3-30B-A3B has 128 experts but only activates 8 per token. The total parameter count is 30B, but the compute per token is equivalent to a 3B model.
2. The Router¶
The router is a simple linear layer that maps each token's hidden state to a score for each expert:
def route_topk(x, w_router, top_k, norm_topk_prob=False):
"""
Route each token to its top-k experts.
Args:
x: Token representations. Shape: (batch, seq_len, hidden_size)
w_router: Router weights. Shape: (num_experts, hidden_size)
top_k: Number of experts to select per token.
norm_topk_prob: If True, normalize selected expert probabilities to sum to 1.
Returns:
router_probs: Softmax over all experts. Shape: (batch, seq_len, num_experts)
expert_ids: Selected expert indices. Shape: (batch, seq_len, top_k)
expert_scores: Scores for selected experts. Shape: (batch, seq_len, top_k)
"""
# Step 1: Compute router logits (which expert does each token "want"?)
router_logits = quantized_linear(x, w_router) # (B, L, num_experts)
# Step 2: Convert to probabilities
router_probs = mx.softmax(router_logits, axis=-1, precise=True)
# Step 3: Select top-k experts by index
# argpartition is O(n) instead of O(n log n) — faster than sorting
expert_ids = mx.argpartition(-router_probs, kth=top_k - 1, axis=-1)[..., :top_k]
# Step 4: Get the scores (probabilities) for the selected experts
expert_scores = mx.take_along_axis(router_probs, expert_ids, axis=-1)
# Step 5: Optionally normalize selected scores to sum to 1
if norm_topk_prob:
expert_scores = expert_scores / mx.sum(expert_scores, axis=-1, keepdims=True)
return router_probs, expert_ids, expert_scores
Why argpartition Instead of argsort?¶
argsort fully sorts the array — O(n log n). argpartition only moves the
top-k elements to the front — O(n). We only need the top-k indices, so
partition is faster.
Why Normalize Selected Probabilities?¶
Without normalization, the raw softmax probabilities for the top-k experts might sum to less than 1 (because we excluded the bottom experts). Normalizing ensures the weighted sum uses the full probability mass:
Before: probs = [0.3, 0.25, 0.15, 0.12, 0.08, 0.05, 0.03, 0.02]
Top-2: experts = [0, 1], scores = [0.3, 0.25], sum = 0.55
After normalization: scores = [0.545, 0.455], sum = 1.0
3. Expert Computation¶
Once we know which experts each token uses, we run the FFN for those experts. But looping over tokens and experts one by one would be slow. Instead, we use a grouped computation:
Step 1: Group tokens by expert¶
# Sort tokens by their assigned expert ID
flat_expert_ids = expert_ids.reshape(-1) # (B × L × top_k,)
sort_idx = mx.argsort(flat_expert_ids) # indices that would sort by expert
# Rearrange tokens so tokens for the same expert are adjacent
grouped_x = flat_x[sort_idx] # tokens sorted by expert assignment
Step 2: Run each expert's FFN on its group¶
# Use a specialized grouped matrix multiply that processes all experts at once
gate = grouped_expert_linear(expanded_x, w_gate, flat_expert_ids)
up = grouped_expert_linear(expanded_x, w_up, flat_expert_ids)
expert_output = grouped_expert_linear(silu(gate) * up, w_down, flat_expert_ids)
Step 3: Unsort and weight by scores¶
# Reverse the sort to restore original token order
output = expert_output[inv_sort_idx]
# Reshape and weight by router scores
output = output.reshape(B, L, top_k, D)
weighted = mx.sum(output * mx.expand_dims(expert_scores, -1), axis=-2)
4. The Complete MoE Layer¶
class Moe:
def __init__(self, w_router, w_gate, w_up, w_down,
num_experts_per_tok, norm_topk_prob=False):
"""
Args:
w_router: Router weights. Shape: (num_experts, hidden_size)
w_gate: Expert gate weights. Shape: (num_experts, intermediate, hidden)
w_up: Expert up weights. Shape: (num_experts, intermediate, hidden)
w_down: Expert down weights. Shape: (num_experts, hidden, intermediate)
num_experts_per_tok: How many experts each token uses (top-k).
norm_topk_prob: Whether to normalize selected expert scores.
"""
self.w_router = w_router
self.w_gate = w_gate
self.w_up = w_up
self.w_down = w_down
self.num_experts_per_tok = num_experts_per_tok
self.norm_topk_prob = norm_topk_prob
def __call__(self, x):
B, L, D = x.shape
# Step 1: Route each token to its top-k experts
_, expert_ids, expert_scores = route_topk(
x, self.w_router,
top_k=self.num_experts_per_tok,
norm_topk_prob=self.norm_topk_prob,
)
# Step 2: Expand x for top-k expert selection
# Each token now appears top_k times (once per selected expert)
expanded_x = mx.broadcast_to(
mx.expand_dims(x, -2), # (B, L, 1, D)
(B, L, self.num_experts_per_tok, D), # (B, L, top_k, D)
).reshape(-1, D) # (B × L × top_k, D)
flat_expert_ids = expert_ids.reshape(-1) # (B × L × top_k,)
# Step 3: Run expert FFNs (grouped by expert assignment)
gate = grouped_expert_linear(expanded_x, self.w_gate, flat_expert_ids)
up = grouped_expert_linear(expanded_x, self.w_up, flat_expert_ids)
# Step 4: SiLU-gated FFN (same as standard MLP, but per-expert)
expert_output = grouped_expert_linear(
silu(gate) * up,
self.w_down,
flat_expert_ids,
).reshape(B, L, self.num_experts_per_tok, D)
# Step 5: Weight by router scores and sum across experts
return mx.sum(
expert_output * mx.expand_dims(expert_scores, -1),
axis=-2,
)
5. When MoE Layers Are Used¶
Not every transformer block is an MoE layer. The model config specifies which layers are MoE:
def is_moe_layer(self, layer_idx):
return (
self.num_experts > 0 # MoE is enabled
and layer_idx not in self.mlp_only_layers # not in the exclusion list
and (layer_idx + 1) % self.decoder_sparse_step == 0 # hits the frequency
)
For example, if decoder_sparse_step=4, every 4th layer is an MoE layer.
The other layers use a standard dense FFN. This is a common pattern — mixing
dense and sparse layers keeps the model efficient while adding capacity where
it matters most.
6. Memory Considerations¶
MoE models are parameter-heavy. Qwen3-30B-A3B has 30B parameters but only activates 3B per token. At 4-bit quantization:
Dense 8B model: ~2.5 GB
MoE 30B-A3B: ~10 GB (all expert weights in memory)
Per-token compute: ~1 GB (only 8 experts active)
The bottleneck is memory, not compute. All expert weights must be in memory even though only a few are used per token. This is why MoE models need more RAM but not proportionally more GPU time.
7. Summary¶
| Concept | Description |
|---|---|
| Router | Linear layer that assigns tokens to experts |
| Expert | A standard gated FFN (gate_proj + up_proj + down_proj) |
| Top-k selection | Each token uses k experts (typically 4-8) |
| Grouped computation | Tokens are sorted by expert for efficient batched matmul |
| Weighted sum | Expert outputs are combined using router scores |
| Sparse activation | Only k out of N experts compute per token |
MoE is how models get "smart" without getting "slow" — more parameters for knowledge, fewer for compute per token.