Abhik Sarkar's avatar

Mixture of Experts (MoE)

Sparse models that route inputs to specialized expert networks for efficient scaling

Sparse Activation: Only 2 out of 8 experts are active per token, reducing computation by 75% while maintaining model capacity.

What is Mixture of Experts?

Mixture of Experts (MoE) is a neural network architecture that uses conditional computation to scale model capacity without proportionally increasing computational cost. Instead of processing every input through all parameters, MoE models route inputs to a subset of specialized "expert" networks.

Core Concepts

1. Sparse Activation

Unlike dense models where all parameters are active for every input, MoE models activate only a small fraction:

Dense Model: All parameters active (100%)
Sparse MoE: Top-K experts active (e.g., 2/8 = 25%)

Computation savings = 1 - (K/N) where K=active, N=total experts

2. Expert Networks

Each expert is typically a feed-forward network (FFN):

class Expert(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.activation = nn.ReLU()
    
    def forward(self, x):
        return self.w2(self.activation(self.w1(x)))

3. Gating Network (Router)

The gating network decides which experts to activate:

class Router(nn.Module):
    def __init__(self, d_model, num_experts):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts)
    
    def forward(self, x):
        # Compute routing probabilities
        logits = self.gate(x)
        probs = F.softmax(logits, dim=-1)
        
        # Select top-k experts
        top_k_probs, top_k_indices = torch.topk(probs, k=2)
        
        # Normalize selected probabilities
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)
        
        return top_k_indices, top_k_probs

Architecture Variants

Switch Transformer

  • Top-1 routing: Each token to single expert
  • Simplified design: No token dropping
  • Capacity factor: Limits tokens per expert
  • Scale: Up to 1.6T parameters with 2048 experts

Mixtral Architecture

  • Top-2 routing: Better quality vs efficiency trade-off
  • 8 experts per layer: Practical for deployment
  • 46.7B total, 12.9B active: Efficient inference

GShard

  • Top-2 with auxiliary loss: Balanced expert usage
  • Token dropping: Handles capacity overflow
  • 600B parameters: Early large-scale MoE

Load Balancing Strategies

1. Auxiliary Loss

Encourages uniform expert utilization:

def auxiliary_loss(gate_probs, expert_mask):
    # Fraction of tokens per expert
    f = expert_mask.float().mean(dim=0)
    
    # Average probability per expert
    P = gate_probs.mean(dim=0)
    
    # Auxiliary loss encourages balance
    aux_loss = alpha * torch.sum(f * P)
    return aux_loss

2. Capacity Factor

Limits maximum tokens per expert:

Expert Capacity = (tokens_per_batch / num_experts) × capacity_factor

Typical capacity_factor = 1.25 (allows 25% overflow)

3. Random Routing

Add noise to routing decisions for exploration:

# Add noise during training
if training:
    noise = torch.randn_like(logits) * noise_scale
    logits = logits + noise

Training Considerations

Initialization

# Router initialization
router.gate.weight.data.normal_(mean=0, std=0.02)
router.gate.bias.data.zero_()

# Expert initialization (same as dense FFN)
for expert in experts:
    expert.w1.weight.data.normal_(mean=0, std=0.02)
    expert.w2.weight.data.normal_(mean=0, std=0.02/sqrt(2*num_experts))

Gradient Stability

  • Router z-loss: Prevents router logits from growing too large
  • Expert dropout: Randomly drops experts during training
  • Gradient clipping: Essential for stable training

Distributed Training

# Expert parallelism
class DistributedMoE(nn.Module):
    def __init__(self, num_experts, expert_parallel_size):
        self.ep_size = expert_parallel_size
        self.num_local_experts = num_experts // ep_size
        
        # Each device holds subset of experts
        self.local_experts = nn.ModuleList([
            Expert() for _ in range(self.num_local_experts)
        ])
    
    def forward(self, x, indices):
        # All-to-all communication for expert dispatch
        x_dispatched = all_to_all(x, indices)
        
        # Process with local experts
        outputs = [expert(x_e) for expert, x_e in zip(self.local_experts, x_dispatched)]
        
        # All-to-all to combine results
        return all_to_all(outputs, reverse=True)

Performance Analysis

Computational Efficiency

FLOPs comparison (for 8 experts, top-2):
- Dense model: 1.0×
- Sparse MoE: 0.25× (routing) + 0.01× (gating) = 0.26×

Memory usage:
- Parameters: 8× (all experts stored)
- Activations: ~0.25× (only active experts)

Scaling Laws

MoE scaling: L = A × (C_active)^(-α) × (N_total)^(-β)

Where:
- C_active: Active compute per token
- N_total: Total model parameters
- α ≈ 0.07, β ≈ 0.05 (better than dense β ≈ 0.08)

Implementation Example

Simple MoE Layer

class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, num_experts, top_k):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff) for _ in range(num_experts)
        ])
        self.router = Router(d_model, num_experts)
    
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)
        
        # Get routing decision
        indices, weights = self.router(x_flat)
        
        # Dispatch to experts
        output = torch.zeros_like(x_flat)
        for i in range(self.top_k):
            expert_idx = indices[:, i]
            expert_weight = weights[:, i:i+1]
            
            for e in range(self.num_experts):
                mask = (expert_idx == e)
                if mask.any():
                    expert_input = x_flat[mask]
                    expert_output = self.experts[e](expert_input)
                    output[mask] += expert_weight[mask] * expert_output
        
        return output.view(batch_size, seq_len, d_model)

Optimization Techniques

1. Expert Caching

# Cache frequently used expert combinations
@lru_cache(maxsize=128)
def get_expert_combination(expert_indices):
    return [experts[i] for i in expert_indices]

2. Batched Expert Execution

# Group tokens by selected experts
def batch_by_expert(tokens, indices):
    expert_batches = defaultdict(list)
    for token, expert_idx in zip(tokens, indices):
        expert_batches[expert_idx].append(token)
    return expert_batches

3. Dynamic Capacity

# Adjust capacity based on load
def dynamic_capacity(base_capacity, expert_loads):
    imbalance = expert_loads.std() / expert_loads.mean()
    adjusted_capacity = base_capacity * (1 + 0.5 * imbalance)
    return adjusted_capacity

Common Pitfalls & Solutions

Expert Collapse

Problem: All tokens routed to same expert
Solution: Auxiliary loss, noise injection, expert dropout

Load Imbalance

Problem: Some experts overloaded, others idle
Solution: Capacity constraints, load-aware routing

Training Instability

Problem: Router gradients explode or vanish
Solution: Gradient clipping, careful initialization, router z-loss

Memory Overhead

Problem: All experts stored in memory
Solution: Expert pruning, parameter sharing, hierarchical experts

Advanced Techniques

Hierarchical MoE

Level 1: Coarse routing (4 super-experts)
Level 2: Fine routing (4 experts per super-expert)
Total: 16 experts with 2-level hierarchy

Soft MoE

Instead of hard top-K selection, use weighted combination of all experts with sparsity regularization.

Expert Pruning

Remove underutilized experts post-training for efficient deployment.

Best Practices

  1. Start Simple: Begin with fewer experts (4-8) and top-2 routing
  2. Monitor Metrics: Track expert utilization, routing entropy, auxiliary loss
  3. Gradual Scaling: Increase experts gradually, ensure stable training
  4. Profile Performance: Measure actual speedup vs theoretical
  5. Consider Deployment: Plan for memory constraints and latency requirements

Conclusion

Mixture of Experts represents a powerful paradigm for scaling neural networks efficiently. By activating only a subset of parameters per input, MoE models achieve better scaling laws than dense models while maintaining computational efficiency. As models continue to grow, sparse architectures like MoE will become increasingly important for practical deployment of large-scale AI systems.

Further Reading

Mastodon