Mixture of Experts (MoE)

Overview

Mixture of Experts (MoE) is an architectural pattern where a model contains multiple expert networks and a router network that learns to dispatch different inputs to different experts. Only a subset of experts are active per token, enabling trillion-parameter models to run efficiently.

• Foundational Paper: "Outrageously Large Neural Networks for Efficient Conditional Computation" (Shazeer et al., 2017)
• Modern Scale: "Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity" (Lepikhin et al., 2021)
• Key Insight: Use sparsity to scale parameter count without scaling computation
• Impact: Can build 1.6T parameter models with less compute than 100B dense models
• Adoption: Google Switch Transformer, Mistral MoE, Llama 3.1-405B

---

The Problem: Parameter vs Computation Trade-off

Dense Models

Traditional Scaling:
 - 7B model: 7B multiply-adds per token
 - 70B model: 70B multiply-adds per token
 - 1T model: 1T multiply-adds per token
 - Problem: Computation grows linearly with parameters!

Reality: Training 1T parameter model requires 10x compute
What if we could add parameters without adding compute?

The Insight

Modern observation:
"Not every token needs all parameters"

Question: 
Could we have 1T parameters, but only use 100B per token?

Answer:
Yes! With routing and expert selection.

Architecture:
 - 1000 experts of 1B parameters each (1T total)
 - Clever router learns which 2-4 experts to use per token
 - Only 2-4B parameters active per token
 - Total compute: Similar to 70B dense model
 - Total parameters: 1T (14x more!)

This is Mixture of Experts!

-

How MoE Works

Architecture

Input: (batch, seq_len, hidden_dim)
 ↓
- ┌─────────────────┐
 - Router Network │ (simple: hidden_dim → num_experts)
 - ┘
 ↓
 Routing scores: (batch, seq_len, num_experts)
 ↓
 Top-k selection: Keep only top-2 or top-4 experts
 ↓
- ┌──────────┬──────────┬──────────┬──────────┐
 - Expert 1 │ Expert 2 │ Expert 3 │ Expert 4 │ (only active ones)
 - ┴──────────┴──────────┴──────────┘
 ↓ ↓ ↓ ↓
 Route subset of tokens to each expert
 ↓
- ┌─────────────────────────────────┐
 - Combine outputs (weighted sum) │
 - ┘
 ↓
 Output: (batch, seq_len, hidden_dim)

Simple Example

class SimpleGatedMixtureOfExperts(nn.Module):
 def __init__(self, hidden_dim, num_experts, expert_dim, top_k=2):
 super().__init__()
 self.num_experts = num_experts
 self.top_k = top_k

 # Router: Maps hidden_dim → num_experts (scores)
 self.router = nn.Linear(hidden_dim, num_experts)

 # Experts: Each expert is a small FFN
 self.experts = nn.ModuleList([
 nn.Sequential(
 nn.Linear(hidden_dim, expert_dim),
 nn.ReLU(),
 nn.Linear(expert_dim, hidden_dim)
)
 for _ in range(num_experts)
])

 def forward(self, x):
 # x shape: (batch, seq_len, hidden_dim)
 batch, seq_len, hidden_dim = x.shape

 # Flatten for routing
 x_flat = x.view(-1, hidden_dim) # (batch*seq_len, hidden_dim)

 # Route
 router_scores = self.router(x_flat) # (batch*seq_len, num_experts)
 router_probs = softmax(router_scores, dim=-1)

 # Select top-k experts
 top_k_probs, top_k_indices = torch.topk(router_probs, self.top_k, dim=-1)

 # Normalize top-k probabilities
 top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

 # Apply experts
 output = torch.zeros_like(x_flat)
 for i in range(self.top_k):
 expert_idx = top_k_indices[:, i] # Which expert for each token
 expert_prob = top_k_probs[:, i] # Routing probability

 # Apply each expert to its tokens
 for e in range(self.num_experts):
 mask = (expert_idx == e)
 if mask.sum() > 0:
 expert_out = self.experts[e](x_flat[mask])
 output[mask] += expert_prob[mask].unsqueeze(-1) * expert_out

 # Reshape back
 output = output.view(batch, seq_len, hidden_dim)
 return output

# Usage:
moe = SimpleGatedMixtureOfExperts(
 hidden_dim=2048,
 num_experts=128,
 expert_dim=8192,
 top_k=2
)
output = moe(x) # Only 2 experts active per token!

-

The Challenge: Load Balancing

The Problem

Naive MoE Issue:
 - Router learns to prefer certain experts
 - Token distribution becomes unbalanced
 - Example: 120/128 tokens go to Expert 1
 - 127 experts are mostly idle
 - Wasted computation!

Why it happens:
 - Gradient update makes router prefer rewarded experts
 - Positive feedback loop
 - Model converges to using only 1-2 experts
 - Defeats the purpose of 128 experts!

Result:
 - Dense model efficiency (compute grows with used experts)
 - Sparse model overhead (unused experts in memory)
 - Worst of both worlds!

Solutions: Load Balancing Losses

1. Auxiliary Loss (Switch Transformer approach)

def compute_load_balancing_loss(router_probs, top_k_indices, num_experts):
 """
 Encourages uniform distribution of tokens across experts

 Intuition:
 - Soft constraint: "Please use experts evenly"
 - Not hard constraint (allows natural specialization)
 - Auxiliary loss added to main loss
 """
 batch_size = router_probs.shape[0]

 # Importance: How much each token uses each expert
 # (batch*seq_len, num_experts)
 importance = router_probs.sum(dim=0) # Sum across tokens

 # Load: How many tokens are routed to each expert
 load_one_hot = torch.nn.functional.one_hot(
 top_k_indices[:, 0], num_experts
).float() # (batch*seq_len, num_experts)
 load = load_one_hot.sum(dim=0) # (num_experts,)

 # Auxiliary loss encourages load ≈ importance
 # If load is uniform but importance is biased, loss pushes back
 loss = torch.mean(load * importance) * num_experts / batch_size
 return loss

# In training loop:
output = moe(x)
main_loss = compute_loss(output, target)
aux_loss = compute_load_balancing_loss(router_probs, top_k_indices, 128)
total_loss = main_loss + 0.01 * aux_loss # Auxiliary loss term
total_loss.backward()

2. Expert Capacity

def apply_expert_capacity(top_k_indices, expert_capacity):
 """
 Hard constraint: Limit tokens per expert

 Intuition:
 - "Each expert can handle max_tokens tokens per batch"
 - Ensures load balancing through capacity constraint
 - Excess tokens overflow to second-choice expert
 """
 batch_size, seq_len = top_k_indices.shape[:2]
 num_tokens = batch_size * seq_len
 max_tokens_per_expert = expert_capacity * num_tokens

 # Track tokens assigned to each expert
 tokens_per_expert = torch.zeros(num_experts)

 for i in range(seq_len * batch_size):
 for k in range(top_k):
 expert = top_k_indices[i, k]
 if tokens_per_expert[expert] < max_tokens_per_expert:
 tokens_per_expert[expert] += 1
 break # Token assigned to this expert
 # else: Try next expert in top-k

 return tokens_per_expert

# Typical capacity factor:
# capacity_factor = 1.25 means each expert handles 1.25x average
# Prevents concentration, ensures load balancing

---

MoE in Modern Models

Switch Transformers (2021)

Parameters: 1.6 Trillion
Architecture:
 - 2048 experts
 - Top-1 routing (only 1 expert per token!)
 - Reduced computation vs 1.6T dense
 - Quality comparable to 100-200B dense models

Performance:
 - Training: 4x faster than 1.6T dense
 - Inference: Similar cost to 100B dense
 - Quality: Near 1.6T dense model performance
 - Achievement: Massive parameter scaling with modest compute!

Mistral MoE & Llama MoE

Mistral Mixtral 8x7B:
 - 46.7B total parameters
 - 8 experts of 7B each
 - Top-2 routing
 - Sparse computation: ~13B FLOPs per token (vs 47B dense)
 - 5-10x speedup in inference vs dense 46B
 - Competitive with 70B models!

Llama 3.1 405B:
 - MoE variant in development
 - 16-32 experts
 - Mixture of dense + sparse blocks
 - Efficient scaling to trillion+ parameters

-

Pros and Cons

Advantages

Massive Parameter Count
 - 1T+ parameters with reasonable compute
 - Each domain/task can have expert specialization

Compute Efficiency
 - Only subset of experts active per token
 - Can use top-2 instead of computing all experts
 - 3-4x computation reduction vs dense

Specialization
 - Different experts specialize in different domains
 - Emergent task-specific routing
 - Better generalization across diverse data

Scaling Laws
 - Better compute-optimal frontier than dense models
 - Reach quality of larger dense models with less compute

Challenges

Training Instability
 - Load imbalance causes training oscillations
 - Collapse to single expert without careful tuning
 - Requires specialized load balancing techniques

Inference Complexity
 - Router overhead (overhead to decide which expert)
 - Expert memory: All experts loaded in memory
 - Not all hardware optimized for sparse operations

Communication Overhead
 - Distributed training: Routing decisions across devices
 - All-reduce for distributed expert placement
 - Reduced efficiency compared to dense distributed training

Generalization Concerns
 - Some evidence of worse generalization (earlier research)
 - Potential for expert collapse in specific domains
 - Hyperparameter sensitivity (top-k, capacity factor, loss weight)

-

When to Use MoE

Use MoE when:
Parameter count is constraint
 - Want very large models but limited training compute
Inference cost is constraint
 - Can deploy billion-token daily with modest resources
Domain diversity is high
 - Multi-domain tasks benefit from specialization
You have training compute budget
 - MoE training is complex; needs careful engineering

Avoid MoE when:
Model size doesn't matter
Inference latency is critical
 - Routing decision adds latency
Distributed training not feasible
 - MoE requires careful distributed implementation
You want simple, stable models
 - Dense models simpler to train/deploy

---

Key Takeaways

MoE enables 10-100x parameter scaling without proportional compute increase Trade-off: Sparse computation vs. training complexity Load balancing is critical—improper routing kills efficiency Top-k routing (2-4 experts) performs better than top-1 Future of scaling: Dense + MoE hybrid architectures

-

Related Notes

• [Scaling Laws & Optimal Allocation](/01-modeling/02-training/02-scaling/(scaling-laws-optimal-allocation/) - Why MoE scales differently
• Distributed Training - How to train MoE across devices
• [Load Balancing & Request Routing](/01-modeling/04-production/01-deployment/(load-balancing-request-routing/) - Similar concepts in inference
• Llm Inference Optimization - Complete inference stack