Mixture of Experts (MoE) in AI Models Explained

The Mixture of Experts (MoE) is offering a unique approach to efficiently scaling models while maintaining, or even improving, their performance.

Traditionally, the trade-off in model training has been between size and computational resources. Larger models typically promise better performance but at the cost of greater computational demand. MoE challenges this norm by enabling the pretraining of models with substantially less compute. This approach allows for a dramatic increase in model or dataset size within the same compute budget as a traditional dense model. Consequently, a MoE model is expected to reach comparable quality to its dense counterpart much faster during the pretraining phase.

1. Understanding Mixture of Experts

Definition and Components

At its core, a MoE, particularly in the context of transformer models, consists of two primary elements: Sparse MoE layers and a gate network (or router).

Role of Sparse MoE Layers and Experts

Unlike the dense feed-forward network (FFN) layers typically used in transformer models, MoE employs sparse MoE layers. Each layer houses several “experts,” with each expert being a neural network, often in the form of FFNs. These experts can vary in complexity and, intriguingly, can even encompass MoEs themselves, leading to the creation of hierarchical MoE structures.

Gate Network Functionality

The gate network plays a crucial role in determining the routing of tokens to appropriate experts. For instance, in a given scenario, the token “More” might be directed to one expert, while “Parameters” to another. This routing is not just pivotal in the functioning of MoEs but also brings in the complexity of decision-making about token routing, where the router itself is a learned entity that evolves during the pretraining of the network.

2. Challenges and Solutions

Training and Inference Challenges

While MoEs offer efficiency in pretraining and faster inference, they are not without their challenges.

Training Challenges

A significant obstacle has been in generalizing the MoE during fine-tuning, where it can tend toward overfitting.

Inference Challenges

Despite their large parameter size, only a subset of these parameters are active during inference, leading to faster processing. However, this also implies substantial memory requirements, as all parameters must be loaded into RAM regardless of their active status during inference.

Solutions and Strategies

To address these challenges, various strategies are employed. These include load balancing to prevent the overuse of certain experts and the incorporation of an auxiliary loss to ensure equitable training across all experts.

3. Historical Context and Evolution

The concept of MoEs dates back to 1991, with the paper “Adaptive Mixture of Local Experts .” This early work laid the foundation for MoEs by proposing a system where separate networks (experts) handle different subsets of training cases, guided by a gating network.

Advancements in NLP and Beyond

The period between 2010 and 2015 saw significant advancements that contributed to the development of MoEs. These include the exploration of MoEs as components within deeper networks and the introduction of conditional computation by Yoshua Bengio, which dynamically activates network components based on the input data.

4. The Principle of Sparsity

Concept of Sparsity

Sparsity, as introduced in Shazeer’s exploration of MoEs for translation, is based on the principle of conditional computation. This allows scaling the model size without proportionally increasing the computation, leading to the use of thousands of experts in each MoE layer.

Gating Mechanisms

Various gating mechanisms, such as the Noisy Top-K Gating, have been explored. This approach adds noise to the routing process and then selects the top ‘k’ values, creating a balance between efficiency and diversity in expert utilization.

5. MoEs in Transformers

GShard’s implementation of MoEs in transformers is a notable example of large-scale application. It introduces novel concepts like random routing and expert capacity, ensuring balanced load and efficiency at scale.

6. Breakthrough with Switch Transformers

Switch Transformers represent a significant advancement in the MoE domain. They simplify the routing process and reduce the communication costs, all while preserving the quality of the model. The concept of expert capacity is further refined here, striking a balance between token distribution and computational efficiency.

7. Fine-Tuning MoEs

Fine-tuning MoEs poses unique challenges, particularly regarding overfitting. Strategies like higher regularization within experts and adjustments to the auxiliary loss have been employed to mitigate these issues. Additionally, selective freezing of MoE layer parameters during fine-tuning has shown promise in maintaining performance while streamlining the process.

8. Practical Applications and Future Directions

MoEs have found applications in various fields, notably in language translation and large-scale models. The potential for MoEs in AI is vast, with ongoing research exploring new domains and applications.

• Further experiments on distilling a sparse MoE back to a dense model with less parameters but similar number of parameters. (Hugging Face, December 11, 2023)
• Another area will be quantization of MoEs. QMoE (Oct. 2023) is a good step in this direction by quantizing the MoEs to less than 1 bit per parameter, hence compressing the 1.6T Switch Transformer which uses 3.2TB accelerator to just 160GB. (Hugging Face, December 11, 2023)

9. Open Source and Accessibility

The accessibility of MoEs has been enhanced by several open-source projects. These resources facilitate the training and implementation of MoEs, contributing to a collaborative and progressive AI research community.

There are nowadays several open source projects to train MoEs:

• Megablocks: https://github.com/stanford-futuredata/megablocks
• Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm
• OpenMoE: https://github.com/XueFuzhao/OpenMoE

In the realm of released open access MoEs, you can check:

• Switch Transformers (Google) : Collection of T5-based MoEs going from 8 to 2048 experts. The largest model has 1.6 trillion parameters.
• NLLB MoE (Meta) : A MoE variant of the NLLB translation model.
• OpenMoE : A community effort that has released Llama-based MoEs.
• Mixtral 8x7B (Mistral) : A high-quality MoE that outperforms Llama 2 70B and has much faster inference. A instruct-tuned model is also released. Read more about it in the announcement blog post .

10. Conclusion

The Mixture of Experts model represents a significant leap in the field of artificial intelligence, offering a scalable, efficient approach to building large and powerful AI models. As research and development in this area continue to evolve, the potential applications and advancements of MoEs in various domains of AI are boundless.

Follow me on social media

• https://twitter.com/nifty0x
• https://vidrihmarko.medium.com/
• https://www.dhirubhai.net/in/marko-vidrih/

Project I’m currently working on