Mixture of Experts: Core Concepts and Hands-on Implementation
Chapter 1: Foundations of Mixture of Experts Models
Overview of Sparsely-Gated MoE Architecture
The Gating Network: Formulation and Function
Expert Networks: Specialization and Capacity
Mathematical Formulation of the MoE Layer
Load Balancing and Auxiliary Losses
Challenges in MoE Training: Expert Collapse
Comparison with Dense Model Scaling
Hands-on: Implementing a Basic MoE Layer
Chapter 2: Advanced Routing Mechanisms
Analysis of Top-k Gating and its Variants
Noisy Top-k Gating for Load Balancing
Hash-based Routing for Deterministic Selection
Switch Transformers: Simplified Routing
Soft MoE: Differentiable Routing
Analyzing Routing Decisions and Specialization
Hands-on: Implementing Different Routing Strategies
Chapter 3: Training and Optimization of Large-Scale MoEs
Expert Parallelism for Distributed Training
Combining Model, Data, and Expert Parallelism
Capacity Factor and its Impact on Performance
Techniques for Mitigating Router Z-Loss Instability
Precision and its Effects: BFloat16 Training
Fine-tuning Strategies for Pre-trained MoE Models
Practice: Configuring a Distributed Training Job
Chapter 4: Efficient Inference with MoE Models
Inference Challenges: Memory and Latency
Expert Offloading to CPU or NVMe
Batching Strategies for Sparse Activation
Model Distillation for MoE Compression
Quantization Techniques for MoE Layers
Speculative Decoding with MoE Models
Hands-on: Building an Optimized Inference Pipeline
Chapter 5: Integrating MoE into Modern Architectures
Replacing FFNs with MoE Layers in Transformers
Placement of MoE Layers: Frequency and Location
MoE in Vision Transformers (ViT)
MoE in Multi-modal Models
Architectural Variants and their Properties
Analyzing Parameter vs. FLOPs Trade-offs
Practice: Modifying a Transformer to use MoE
Hands-on: Implementing Different Routing Strategies
Was this section helpful?
- Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean, 2017 International Conference on Learning Representations (ICLR) DOI: 10.48550/arXiv.1701.06538 - Introduces the sparsely-gated Mixture-of-Experts layer and the noisy gating mechanism for load balancing, foundational for the Noisy Top-k Router.
- Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, William Fedus, Barret Zoph, Noam Shazeer, 2022 Journal of Machine Learning Research, Vol. 23 (Microtome Publishing) DOI: 10.5555/3545945.3545958 - Presents the Switch Transformer architecture, which uses a simplified k=1 routing strategy, directly relevant to the Switch Router implementation.
- torch.nn.Module, PyTorch Contributors, 2024 (PyTorch Foundation) - Official documentation for PyTorch's base class for all neural network modules, essential for implementing custom layers like the routing strategies shown.
- Designing and Training Sparse Mixture-of-Experts Models for Language, Elena Albarran, William Fedus, Andrew M. Dai, Sharan Narang, Noam Shazeer, 2023 (Google AI Blog) - Provides an accessible overview of Mixture-of-Experts models, including discussions on routing mechanisms and the importance of load balancing.