A Mixture of Experts: A revolutionary technique to boost generative AI performance?

Mixture of experts (MoE) stands as a machine learning methodology wherein numerous expert networks (learners) collaborate to partition a problem domain into homogeneous regions. This approach diverges from ensemble techniques in that, typically, only one or a select few expert models are employed for each input in MoE, whereas ensemble techniques involve running all models on every input.

The MoE model is based on the concept of leveraging expertise from multiple specialized models, to get a superior performance. In the realm of machine learning, this principle is applied by utilizing multiple expert networks, each focusing on distinct facets of the data, to enhance overall performance. The essence of MoE lies in establishing a dynamic system where the unique strengths of different experts are utilized depending on the input data. This enables the generation of more adaptable and precise predictions compared to what a single model could deliver.

The foundation of Mixture of Experts (MoEs) can be traced back to the 1991 paper "Adaptive Mixture of Local Experts." Similar to ensemble methods, the concept was to employ a supervised approach for a system composed of individual networks, each handling a specific subset of the training cases. These separate networks, or experts, specialize in different regions of the input space. The selection of the expert is determined by a gating network, which assigns weights to each expert. During training, both the expert and the gating network are trained.

Between 2010 and 2015, advancements in two distinct research areas contributed to the later development of MoEs:

·?????? Experts as components: Traditionally, MoEs consisted of a gating network and multiple experts forming the entire system. However, researchers like Eigen, Ranzato, and Ilya explored MoEs as components within deeper networks. This approach allowed MoEs to function as layers within a multilayer network, enabling models to be both large and efficient simultaneously.

·?????? Conditional Computation: Traditional networks process all input data through every layer. In this period, Yoshua Bengio investigated methods to dynamically activate or deactivate components based on the input token.

These advancements laid the groundwork for exploring mixtures of experts within the realm of Natural Language Processing (NLP). Notably, the work by Shazeer et al. (2017), which included prominent figures like Geoffrey Hinton and Jeff Dean, along with Google's Chuck Norris, scaled this idea to a 137B LSTM (Long Short-Term Memory) architecture, which was the prevailing NLP architecture at the time, originally proposed by Schmidhuber. This work introduced sparsity to maintain fast inference even at high scale. While focusing on translation tasks, it encountered challenges such as high communication costs and training instabilities.

How does MoE work with transformers?

In the context of transformer models, a MoE consists of two main elements:

1.???? Sparse MoE layers are used instead of dense feed-forward network (FFN) layers. MoE layers have a certain number of “experts” (e.g. 8), where each expert is a neural network. In practice, the experts are FFNs, but they can also be more complex networks or even a MoE itself, leading to hierarchical MoEs!

2.???? A gate network or router, that determines which tokens are sent to which expert. For example, in the image below, the token “More” is sent to the second expert, and the token "Parameters” is sent to the first network. As we’ll explore later, we can send a token to more than one expert. How to route a token to an expert is one of the big decisions when working with MoEs - the router is composed of learned parameters and is pretrained at the same time as the rest of the network.

Although MoEs provide benefits like efficient pretraining and faster inference compared to dense models, they also come with challenges:

·?????? Training: MoEs enable significantly more compute-efficient pretraining, but they’ve historically struggled to generalize during fine-tuning, leading to overfitting.

·?????? Inference: Although a MoE might have many parameters, only some of them are used during inference. This leads to much faster inference compared to a dense model with the same number of parameters. However, all parameters need to be loaded in RAM, so memory requirements are high. For example, given a MoE like Mixtral 8x7B, we’ll need to have enough VRAM to hold a dense 47B parameter model. Why 47B parameters and not 8 x 7B = 56B? That’s because in MoE models, only the FFN layers are treated as individual experts, and the rest of the model parameters are shared. At the same time, assuming just two experts are being used per token, the inference speed (FLOPs) is like using a 12B model (as opposed to a 14B model), because it computes 2x7B matrix multiplications, but with some layers shared (more on this soon).

What are the business benefits?

Mixture of Experts (MoE) LLM models have the potential to significantly impact businesses by improving both efficiency and performance in various tasks. Here's a breakdown of the key implications:

Reduced Training Costs: MoE models can achieve similar performance to monolithic LLMs with a fraction of the parameters, leading to faster and cheaper training processes. This translates to cost savings for businesses looking to develop or utilize LLMs.

Lower Computational Requirements: By distributing tasks among specialized experts, MoE models require less computational power for inference. This makes them ideal for deployment on resource-constrained environments or for real-time applications.

Task Specialization: Experts within the MoE architecture can be trained for specific tasks, leading to better performance in those areas compared to a general-purpose LLM. Businesses can leverage this to create LLMs tailored to their specific needs.

Flexibility and Adaptability: New experts can be added to an MoE model to address new tasks or improve performance in existing ones. This allows businesses to adapt their LLMs as their needs evolve.