Implementing a Sparse Mixture of Experts Language Model from Scratch with PyTorch

In this article, we dive into the implementation of a sparse mixture of experts (MoE) language model using PyTorch. This project, inspired by Andrej Karpathy's 'makemore' and built on its reusable components, aims to provide an intuitive understanding of how MoE architectures work. The code is available in the GitHub repo: https://github.com/AviSoori1x/makeMoE/tree/main.

What Changed and Why It Matters

Sparse MoE models are a hot topic, especially with recent developments like Mixtral and the speculation around GPT-4. These models offer significant computational efficiency by activating only a subset of their experts (sub-models) for each input token. However, training stability remains a challenge, making small-scale, hackable implementations crucial for rapid experimentation.

Key Components

Sparse Mixture of Experts

• Sparse MoE: Unlike the traditional feed-forward neural network in transformers, this model uses a sparse mixture of experts. Only a few experts are activated per token, reducing computational load.
• Top-k Gating and Noisy Top-k Gating: These techniques determine which experts to activate for each input. Top-k gating selects the top k experts based on their scores, while noisy top-k adds noise to these scores for better exploration during training.

Initialization

• Kaiming He Initialization: This is used by default but can be swapped out for Xavier/Glorot initialization or other methods. The flexibility of this implementation allows you to experiment with different initializations.

Unchanged from makemore

• Dataset and Preprocessing: The Shakespeare dataset and tokenization process remain the same.
• Casual Self Attention: The self-attention mechanism is implemented as in 'makemore'.
• Training Loop and Inference Logic: These components are also unchanged, ensuring consistency with the original project.

Implementation Details

Self-Attention

Self-attention is a crucial component of transformers, allowing each token to attend to all other tokens in the sequence. This mechanism helps capture contextual information effectively. Here’s a brief overview:

• Query, Key, and Value Vectors: Each token is transformed into query (Q), key (K), and value (V) vectors.
• Attention Scores: The dot product of Q and K is scaled by the square root of the key dimension to produce attention scores.
• Softmax: These scores are passed through a softmax function to get attention weights.
• Weighted Sum: The final output is a weighted sum of the value vectors using these attention weights.

Mixture of Experts Block

The MoE block introduces the sparse activation mechanism:

• Router: This component decides which experts to activate for each token. It uses top-k gating or noisy top-k gating.
  • Top-k Gating: Selects the top k experts based on their scores.
  • Noisy Top-k Gating: Adds noise to the scores to encourage exploration.
• Experts: Each expert is a small neural network that processes the token. Only the selected experts are activated, reducing computational overhead.
• Output Aggregation: The outputs from the activated experts are combined and fed into the next layer.

Training Stability

Training MoE models can be unstable due to issues like load imbalance (where some experts are overused while others are underutilized) and vanishing gradients. Techniques such as noisy top-k gating and careful initialization help mitigate these problems.

Benchmarks and Performance

While this implementation is primarily for educational purposes, it provides a solid foundation for experimenting with MoE models. The flexibility to swap in different components (like initialization methods) allows researchers to explore various strategies for improving training stability and performance.

Conclusion

By following this implementation, you can gain a deeper understanding of how sparse mixture of experts language models work. The code is designed to be hackable, making it an excellent starting point for your own research and experiments.