Mixture of Experts (MoE) is an innovative architectural technique that is transforming the development of large language models (LLMs). MoE-based LLMs, such as GPT-4 and Mixtral 8x7B, achieve new levels of performance and efficiency by partitioning the computational workload across specialized “expert” networks. Each expert network specializes in different types of inputs, allowing the model to leverage their strengths more effectively than a single generalist network. This article delves into the mechanics of MoE, its benefits, and provides insights from experiments conducted with the Mixtral 8x7B model, illustrating how MoE contributes to the advancement of LLM technology.
What is Mixture of Experts (MoE)?
MoE is an architectural pattern for neural networks that splits the computation of a layer or operation into multiple “expert” subnetworks. These subnetworks independently perform their own computations, and their results are combined to create the final output of the MoE layer. MoE architectures can be either dense or sparse:
- Dense MoE: All experts are used for every input, with a gating mechanism determining the weighting of each expert’s contribution.
- Sparse MoE: Only a subset of experts is used for each input, reducing computational cost.
Importance of Mixture of Experts in LLMs
Model Capacity
Model capacity refers to the complexity a model can understand or express. Historically, models with more parameters have proven to have larger capacity. MoE models effectively increase capacity by replacing layers of the model with MoE layers, where expert subnetworks are the same size as the original layer. This allows MoE models to achieve greater complexity without a proportional increase in resource demands.
Training Efficiency
Sparse MoE models are more flop-efficient per parameter used, allowing them to process more tokens and achieve better training results under fixed compute cost constraints. For instance, Mixtral 8x7B uses eight experts, with only two experts activated for each token. This means fewer parameters are used per token, reducing computational costs and making it more efficient than similarly-sized dense models.
Reduced Latency
MoE architectures can deliver decreased first-token serving latency, especially in use cases like retrieval-augmented generation (RAG) and autonomous agents that require many calls to the model. This decreased latency is crucial as it compounds over multiple calls, significantly improving performance in practical applications.
How MoE Architectures Work
Expert Subnetworks
MoE architectures consist of “expert” subnetworks that compose the mixture. These subnetworks are used in both dense and sparse MoE models. The experts independently compute their outputs, which are then combined using averaging or summation.
Routing Algorithms
Routing algorithms are critical in sparse MoE models. They determine which experts process which tokens, ranging from simple uniform selection to complex mechanisms designed to maximize accuracy while ensuring load balance. A well-designed routing algorithm strikes a balance between accuracy and flop efficiency.
Application in Transformers
MoE techniques are often applied to Multi-Layer Perceptrons (MLPs) within transformer blocks. The MLP in the transformer block is replaced with a set of expert MLP subnetworks. Recent research suggests that MoE can also be applied to other parts of the transformer architecture, such as the projection layers for Q, K, and V matrices and the attention heads themselves.
Experimenting with the Mixtral Model
A possible interpretation of the Mixtral 8x7B model; Source NVIDIA Blog
Experiment Setup
To understand how experts specialize, NVIDIA Researchers designed an experiment using the Mixtral 8x7B model. This model has 32 sequential transformer blocks, with each MLP layer replaced by a sparse MoE block containing eight experts, two of which are activated for each token. They ran all samples of the Massive Multitask Language Understanding (MMLU) benchmark through the model, recording the token-expert assignment for each of the eight experts on layers 1, 16, and 32.
Observations
Simplified Mixtral 8x7B model architecture; Source NVIDIA Blog
Load Balancing
Despite the load-balancing algorithm, the researchers observed that experts received equalized loads, but the busiest expert still handled 40-60% more tokens than the least busy one. This imbalance can affect inference efficiency, as some experts finish their work early while others are overloaded.
Domain-Expert Assignment
Certain domains activated specific experts more frequently. For example, in layer 32, abstract algebra activated experts three and eight much more than others. This suggests that while the load is balanced overall, domain-specific specializations emerge within the experts.
Benefits of Mixture of Experts in LLMs
Improved Performance
By activating only the most relevant experts for each input, MoE models achieve higher accuracy and better generalization across diverse datasets and tasks.
Enhanced Efficiency
The sparse activation of experts in MoE architectures significantly reduces computational cost and memory footprint, making them more scalable and deployable.
Increased Capacity
MoE allows LLMs to scale their capacity by adding more expert networks without proportionally increasing resource demands.
Final Words
Mixture of Experts in LLM models provide significant benefits in terms of training efficiency, reduced latency, and increased capacity. By leveraging specialized expert subnetworks and efficient routing algorithms, MoE models like Mixtral 8x7B can achieve competitive performance under fixed compute budgets. Our experiments show how tokens are assigned to experts and highlight the balance and specialization within the model. As MoE continues to evolve, it holds promise for developing more sophisticated and capable AI systems.
