Mixture of Experts (MoE)

The defining characteristic of MoE: only a fraction of the total parameters are activated for any given input token. A gating network (router) selects the top-k experts—typically 1 or 2—from potentially hundreds of available experts. This means a model with 8 total experts where k=2 activates only 25% of its parameters per forward pass.

• Compute efficiency: FLOPs scale with active parameters, not total parameters
• Router z-loss: An auxiliary loss term that penalizes large router logits to stabilize training
• Load balancing: Ensures no single expert becomes a bottleneck or receives zero gradients

Experts naturally specialize in different aspects of the input distribution without explicit programming. In language models, individual experts may focus on syntactic patterns, domain-specific knowledge, or linguistic phenomena like dates and numbers.

• Emergent specialization: Experts self-organize during training based on token-level routing decisions
• Domain affinity: Certain experts become preferentially activated for STEM content, code, or multilingual text
• Redundancy: Multiple experts often develop overlapping competencies, providing resilience against individual expert failure

The gating function is a learned linear layer that computes logits for each expert given the input token representation. A softmax over top-k operation selects the winning experts and assigns normalized weights to their outputs.

• Auxiliary load balancing loss: Penalizes imbalanced expert utilization during training
• Expert capacity factor: Limits tokens per expert to prevent memory overflow in distributed settings
• Random routing: Some implementations add noise or stochasticity to encourage exploration during training
• Switch Transformer variant: Routes to exactly k=1 expert, achieving maximum sparsity with simplified routing

In large-scale deployments, experts are sharded across multiple accelerators. This introduces all-to-all communication overhead, where tokens must be dispatched to the devices hosting their assigned experts and results gathered back.

• Expert parallelism: A distinct distributed strategy orthogonal to data and tensor parallelism
• Capacity factor tuning: Balances memory usage against token overflow risk per expert
• Cross-node communication: Expert routing across nodes introduces latency that must be amortized over large batch sizes
• DeepSpeed-MoE: Microsoft's optimized implementation that co-designs sharding strategies with the ZeRO optimizer

MoE models introduce unique training challenges. The discrete routing decision is non-differentiable, requiring careful auxiliary loss design. Without proper stabilization, routers can collapse to always selecting the same experts.

• Router z-loss: Prevents the router from producing extremely large logits that destabilize softmax
• Expert dropout: Randomly dropping experts during training forces the router to diversify
• Gradient estimation: Straight-through estimators approximate gradients through the discrete top-k selection
• Load balancing loss: Typically a coefficient of 0.01 multiplied by the coefficient of variation of expert assignment counts

While MoE reduces FLOPs per token, it introduces memory overhead because all expert weights must reside in memory simultaneously. This creates a tension between compute efficiency and memory footprint.

• Memory-bound: Expert parameters consume VRAM even when inactive during inference
• Batch size sensitivity: Small batch sizes underutilize experts and amplify routing overhead
• Expert caching: Frequently accessed experts can be kept in fast memory while cold experts reside in slower tiers
• Quantization synergy: Combining MoE with 4-bit quantization (QLoRA-style) enables massive models on single GPUs