Skip to main content

Chapter 1: Mixture of Experts (MoE)

The MoE Concept​

Instead of one large FFN, use multiple "expert" FFNs with a router:

Input → Router → Select top-k experts → Weighted combination → Output

Mixtral Architecture​

mixtral_config = {
    "num_experts": 8,
    "num_experts_per_token": 2,  # Sparse: only 2 active
    "total_params": "46.7B",
    "active_params": "12.9B",   # Per forward pass
}

Router Mechanism​

def router(hidden_states, num_experts, top_k=2):
    # Compute expert scores
    scores = router_weights @ hidden_states
    
    # Select top-k experts
    top_k_scores, top_k_indices = torch.topk(scores, top_k)
    
    # Normalize weights
    weights = F.softmax(top_k_scores, dim=-1)
    
    return weights, top_k_indices

Benefits​

AspectDenseMoE
ParametersAll activeSubset active
ComputeO(params)O(active_params)
CapacityFixedHigher for same compute

Challenges​

  • Load balancing across experts
  • Communication overhead in distributed setting
  • Expert collapse (all routing to few experts)