Mixture of Experts (MoE) Explained

Mixture of Experts (MoE) enables transformer models to be pretrained significantly faster than dense models and achieve faster inference speeds relative to their total parameter count. This is achieved by replacing dense feed-forward network (FFN) layers with sparse MoE layers that activate only a few "experts" for any given input token.

Core Architecture of Mixture of Experts

MoE models replace standard FFN layers with a combination of sparse MoE layers and a routing mechanism. This architecture consists of two primary components:

Sparse MoE Layers: Instead of a single dense layer, these layers contain multiple "experts," where each expert is typically a neural network (often an FFN). These can be arranged in hierarchical structures where experts are themselves MoEs.
Gate Network (Router): A learned router determines which tokens are sent to which expert. The router is pretrained alongside the rest of the network and can be configured to send a token to one or more experts.

While MoEs offer compute efficiency, they introduce specific trade-offs. They require high VRAM because all experts must be loaded into memory, even though only a few are active during a single forward pass. For example, Mixtral 8x7B requires enough VRAM to hold a 47B parameter model, though its inference compute (FLOPs) is closer to that of a 12B model when using two experts per token.

Sparsity and Routing Mechanisms

Sparsity allows a model to scale its size without a proportional increase in computation by using conditional computation—activating only parts of the system per example.

Gating Functions

Routing is handled by a gating network $G$ that decides which experts $E$ to activate. In traditional setups, a softmax function is used. More advanced approaches, such as Noisy Top-k Gating, introduce tunable noise and select the top $k$ values to ensure better exploration and load balancing during training.

Load Balancing and Expert Capacity

To prevent a "winner-take-all" scenario where a few experts are over-utilized and others are ignored, MoEs employ:

Auxiliary Loss: An additional loss function that encourages the router to distribute tokens equally across all experts.
Expert Capacity: A threshold that limits the number of tokens a single expert can process. If an expert reaches capacity, tokens may be dropped or sent via residual connections to the next layer.

Evolution of MoE Implementations

GShard and Switch Transformers

Google's GShard scaled transformers beyond 600 billion parameters by using top-2 gating and introducing random routing (where the second expert is picked proportionally to its weight).

Switch Transformers further simplified this by introducing a single-expert strategy. By routing each token to only one expert, Switch Transformers reduced router computation, halved expert batch sizes, and lowered communication costs while preserving quality. This architecture achieved a 4x pre-train speed-up over T5-XXL.

Stability and Specialization

Training stability is managed through techniques like Router z-loss (introduced in ST-MoE), which penalizes large logits entering the gating network to reduce roundoff errors.

Research into expert specialization shows that encoder experts often specialize in shallow concepts (e.g., punctuation or proper nouns), whereas decoder experts show less specialization. In multilingual setups, experts do not typically specialize in a single language due to theL load balancing requirements.

Fine-Tuning and Instruction Tuning

Sparse models are more prone to overfitting than dense models. Effective fine-tuning strategies include:

Higher Regularization: Using higher dropout rates within sparse layers.
Selective Freezing: Freezing MoE layers while updating non-MoE layers can preserve quality while speeding up training.
Hyperparameter Adjustment: Sparse models often benefit from smaller batch sizes and higher learning rates.

Recent research indicates that MoEs benefit significantly more from instruction tuning than dense models. When comparing a MoE to a T5 equivalent, the MoE's performance gain from instruction tuning (Flan-MoE vs MoE) was larger than the dense model's gain (Flan T5 vs T5).

Computational Efficiency and Parallelism

Parallelism Strategies

To handle the massive parameter counts of MoEs, different parallelism techniques are used:

Expert Parallelism: Experts are distributed across different workers. For MoE layers, tokens are routed to the worker where the specific expert resides; for non-MoE layers, it behaves like data parallelism.

Hardware Optimization

Capacity Factor: A trade-off between quality and communication cost. A higher capacity factor increases quality but increases memory for activations and inter-device communication.
Block-Sparse Operations: Projects like MegaBlocks replace traditional batched matrix multiplication with block-sparse operations to handle imbalanced token assignments without dropping tokens, leading to significant speedups.

Comparison: Sparse MoE vs. Dense Models

Feature	Sparse MoE	Dense Model
Pretraining Compute	More efficient / faster	Less efficient
Inference Speed	Faster (relative to total params)	Slower (relative to total params)
VRAM Requirement	High (must load all experts)	Lower (relative to total params)
Best Use Case	High throughput, many machines	Low throughput, limited VRAM

Open Source MoE Ecosystem

Several frameworks and models have brought MoEs to the open community:

Frameworks: MegaBlocks, Fairseq, and OpenMoE.
Models: Google's Switch Transformers (up to 1.6T parameters), Meta's NLLB MoE (54B), and Mistral's Mixtral 8x7B, which outperforms Llama 2 70B with faster inference.