Unlocking Mixture of Experts : From 1 Know-it-all to group of Jedi Masters — Pranjal Biyani

Mixture of Experts (MoE) architecture splits complex problems into smaller tasks handled by specialized expert networks, improving efficiency and performance

The feed-forward neural network layer in transformers is replaced with MoE layers, where experts are selected dynamically by a gating router for processing specific tokens

MoE models are more FLOP-efficient than traditional dense models, providing 6x faster inference while maintaining performance (e.g., Mixture-8x7B vs LLaMA 70B)

Key challenges include:

• Unstable training
• Load balancing across experts
• Communication costs between distributed experts
• Expert capacity management
• Non-deterministic behavior

Experts specialize in processing specific tokens rather than topics/domains, contrary to common misconceptions

Training larger MoE models for fewer steps shows better empirical results than training smaller models for more steps

Model requires larger VRAM for loading all expert parameters, even though only a subset of experts is used during inference

Load balancing is achieved through auxiliary loss components that encourage uniform distribution of tokens across experts

Top-K routing strategy (usually K=2) helps balance between specialization and model performance

GPT-4 is rumored to be an 8x220B parameter MoE model, demonstrating the architecture’s effectiveness at scale