Mixture of Experts (MoE): Scaling to Trillions of Parameters

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 - Why MoE scales differently
• Distributed Training - How to train MoE across devices
• Load Balancing & Request Routing - Similar concepts in inference
• Llm Inference Optimization - Complete inference stack