Motivation

• The scale of a model is one of the most important axes for better model quality. 

• Given a fixed computing budget, training a larger model for fewer steps is better than training a smaller model for more steps.

Objectives

• MoEs:

  1. Enable models to be pretrained with far less compute.

  2. Have faster inference compared to a model with the same number of parameters.

  3. Require high VRAM as all experts are loaded in memory.

  4. Face many challenges in fine-tuning, but recent work with MoE instruction-tuning is promising.

• In MoEs we replace every FFN layer of the transformer model with an MoE layer, which is composed of a gate network and a certain number of experts.

Related Works

Adaptive Mixture of Local Experts (1991)

• Idea: Akin to ensemble methods:

  1. Have a supervised procedure for a system composed of separate networks, each handling a different subset of the training cases.

  2. Each separate network, or expert, specializes in a different region of the input space. 

• A gating network determines the weights (to choose) for each expert. 

• During training, both the expert and the gating are trained.

Experts as Components

• In the traditional MoE setup, the whole system comprises a gating network and multiple experts. 

• MoEs as the whole model have been explored in SVMs, Gaussian Processes, and other methods. 

• The work by Eigen, Ranzato, and Ilya explored MoEs as components of deeper networks. This allows having MoEs as layers in a multilayer network, making it possible for the model to be both large and efficient simultaneously.

Conditional Computation

• Traditional networks process all input data through every layer. 

• In this period, Yoshua Bengio researched approaches to dynamically activate or deactivate components based on the input token.

Mixture of Experts in NLP

• Shazeer et al. scaled this idea to a 137B LSTM (the de-facto NLP architecture back then, created by Schmidhuber) by introducing sparsity, allowing to keep very fast inference even at high scale. 

• This work focused on translation but faced many challenges, such as high communication costs and training instabilities.

Idea

• Mixture of Experts enable models to be pretrained with far less compute, which means you can dramatically scale up the model or dataset size with the same compute budget as a dense model. 

Architecture

• 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: 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.

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 the model has 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).

Sparsity

• Sparsity uses the idea of conditional computation that allows us to only run some parts of the whole system.

• In Shazeer's exploration of MoEs for translation, the idea of conditional computation (parts of the network are active on a per-example basis) allows one to scale the size of the model without increasing the computation, and hence, this led to thousands of experts being used in each MoE layer.

• Challenges:

  1. For example, although large batch sizes are usually better for performance, batch sizes in MOEs are effectively reduced as data flows through the active experts. 

  2. For example, if our batched input consists of 10 tokens, five tokens might end in one expert, and the other five tokens might end in five different experts, leading to uneven batch sizes and underutilization. 

Gating Network

• A gating mechanism to solve above problem.

• A learned gating network (G) decides which experts (E) to send a part of the input:

• In this setup, all experts are run for all inputs - it’s a weighted multiplication. 

• What happens if G is 0? ==>  there’s no need to compute the respective expert operations and hence we save compute. 

• What’s a typical gating function? ==> a simple network with a softmax function. The network will learn which expert to send the input.

Noisy Top-K Gating

• This gating approach introduces some (tunable) noise and then keeps the top k values. 

1. Add some noise:

2. Only pick the top k:

3. Apply the soft-max:

• By using a low enough k (e.g. one or two), we can train and run inference much faster than if many experts were activated. 

• Why not just select the top expert? The initial conjecture was that routing to more than one expert was needed to have the gate learn how to route to different experts, so at least two experts had to be picked.

• Python code: The forward method calculates the logits, adds Gaussian noise, selects the top k logits and their indices, creates a sparse output, and applies softmax to produce the gating values.

import torch

import torch.nn as nn

import torch.nn.functional as F

class NoisyTopKGating(nn.Module):

  def __init__(self, input_dim, num_experts, k):

    super().__init__()

    self.k = k

    self.linear = nn.Linear(input_dim, num_experts)

    self.noise_scale = nn.Parameter(torch.zeros(1))  # Learnable noise scale

  def forward(self, x):

    #logits_i = x ⋅ W_g

    logits = self.linear(x) 

    #StandardNormal() ⋅ Softplus(x ⋅ W_noise)

    noise = torch.randn_like(logits) * self.noise_scale 

    #H(x)_i = logits_i + StandardNormal() ⋅ Softplus(x ⋅ W_noise)

    noisy_logits = logits + noise 

    #KeepTopK(v, k)_i = { v_i if v_i is in the top k elements of v, −∞ otherwise }

    top_logits, top_indices = torch.topk(noisy_logits, self.k, dim=-1)

    # Create a sparse output by setting non-top k values to -inf

    zeros = torch.full_like(noisy_logits, float('-inf'))

    sparse_logits = zeros.scatter(-1, top_indices, top_logits)

    #G(x) = Softmax(KeepTopK(H(x), k))

    gates = F.softmax(sparse_logits, dim=-1)  # Softmax over top k logits

    return gates, top_indices

• Noisy Top-K Gating offers several advantages over traditional softmax gating in MoE:

  1. Improved Load Balancing: By adding noise and selecting the top k experts, it prevents over-reliance on a few experts and distributes the workload more evenly . This ensures that all experts contribute to the model's performance and prevents some from becoming redundant.   

  2. Reduced Computational Cost: Sparsity in the gating output means that only a subset of experts needs to be activated for each input. This reduces the computational cost compared to traditional softmax gating, where all experts are involved in processing every input .   

Load Balancing Tokens for MoEs

• Problem: if all our tokens are sent to just a few popular experts, that will make training inefficient. (In a normal MoE training, the gating network converges to mostly activate the same few experts. This self-reinforces as favored experts are trained quicker and hence selected more).

• To mitigate this, an auxiliary loss is added to encourage giving all experts equal importance. This loss ensures that all experts receive a roughly equal number of training examples. 

• In transformers, the auxiliary loss is exposed via the aux_loss parameter.

MoEs and Transformers

• Google explored this (scaling up the number of parameters improves the performance) with GShard, which explores scaling up transformers beyond 600 billion parameters.

• GShard replaces every other FFN layer with an MoE layer using top-2 gating in both the encoder and the decoder. 

• This setup is quite beneficial for large-scale computing: when we scale to multiple devices, the MoE layer is shared across devices while all the other layers are replicated. 

• To maintain a balanced load and efficiency at scale, the GShard authors introduced a couple of changes in addition to an auxiliary loss similar to the one discussed in the previous section:

  1. Random routing: in a top-2 setup, we always pick the top expert, but the second expert is picked with probability proportional to its weight.

  2. Expert capacity: we can set a threshold of how many tokens can be processed by one expert. If both experts are at capacity, the token is considered overflowed, and it’s sent to the next layer via residual connections (or dropped entirely in other projects). This concept will become one of the most important concepts for MoEs. (Why is expert capacity needed? Since all tensor shapes are statically determined at compilation time, but we cannot know how many tokens will go to each expert ahead of time, we need to fix the capacity factor).

• The GShard paper has contributions by expressing parallel computation patterns that work well for MoEs, but discussing that is outside the scope of this blog post.

Switch Transformers [Code]

• Problem: MoEs still struggles with training and fine-tuning instabilities.

• Switch Transformers solved those problems. 

  1. The authors even released a 1.6 trillion parameters MoE on Hugging Face with 2048 experts, which you can run with transformers. 

  2. Switch Transformers achieved a 4x pre-train speed-up over T5-XXL.

• Switch Transformers:

  1. The authors also replaced the FFN layers with a MoE layer. 

  2. Proposes a Switch Transformer layer that receives two inputs (two different tokens) and has four experts.

• Switch Transformers uses a simplified single-expert strategy. The effects are:

  1. The router computation is reduced.

  2. The batch size of each expert can be at least halved.

  3. Communication costs are reduced.

  4. Quality is preserved.

Expert capacity

• The capacity suggested above evenly divides the number of tokens in the batch across the number of experts. 

• If we use a capacity factor greater than 1, we provide a buffer for when tokens are not perfectly balanced. Increasing the capacity will lead to more expensive inter-device communication, so it’s a trade-off to keep in mind. 

• In particular, Switch Transformers perform well at low capacity factors (1-1.25)

Load Balancing Loss

• For each Switch layer, the auxiliary loss is added to the total model loss during training. This loss encourages uniform routing and can be weighted using a hyperparameter.

Stabilizing training with router Z-loss

• Problem: The balancing loss previously discussed can lead to instability issues. 

• We can use many methods to stabilize sparse models at the expense of quality. 

  1. For example, introducing dropout improves stability but leads to loss of model quality. 

  2. On the other hand, adding more multiplicative components improves quality but decreases stability.

• Router z-loss, introduced in ST-MoE, significantly improves training stability without quality degradation by penalizing large logits entering the gating network. Since this loss encourages absolute magnitude of values to be smaller, roundoff errors are reduced, which can be quite impactful for exponential functions such as the gating. 

What Does An Expert Learn?

• The ST-MoE authors observed that:

  1. Encoder experts specialize in a group of tokens or shallow concepts. For example, we might end with a punctuation expert, a proper noun expert, etc. 

  2. Decoder experts have less specialization. 

• The authors also trained in a multilingual setup. Although one could imagine each expert specializing in a language, the opposite happens: due to token routing and load balancing, there is no single expert specialized in any given language.

How does scaling the number of experts impact pretraining?

• More experts lead to improved sample efficiency and faster speedup, but these are diminishing gains (especially after 256 or 512), and more VRAM will be needed for inference. 

• The properties studied in Switch Transformers at large scale were consistent at small scale, even with 2, 4, or 8 experts per layer.

Fine-tuning MoEs

• The overfitting dynamics are very different between dense and sparse models. Sparse models are more prone to overfitting, so we can explore higher regularization (e.g. dropout) within the experts themselves (e.g. we can have one dropout rate for the dense layers and another, higher, dropout for the sparse layers).

• One decision question is whether to use the auxiliary loss for fine-tuning. The ST-MoE authors experimented with turning off the auxiliary loss, and the quality was not significantly impacted, even when up to 11% of the tokens were dropped. Token dropping might be a form of regularization that helps prevent overfitting.

• Switch Transformers observed that:

  1. At a fixed pretrain perplexity, the sparse model does worse than the dense counterpart in downstream tasks, especially on reasoning-heavy tasks such as SuperGLUE. 

  2. On the other hand, for knowledge-heavy tasks such as TriviaQA, the sparse model performs disproportionately well. 

  3. The authors also observed that a fewer number of experts helped at fine-tuning. 

  4. Another observation that confirmed the generalization issue is that the model did worse in smaller tasks but did well in larger tasks.

• Experiments:

  1. Freezing all non-expert weights. That led to a huge performance drop, which was expected as the MoE layers correspond to most of the network. 

  2. Freezing only the parameters in MoE layers, which turned out to work almost as well as updating all parameters. This can help speed up and reduce memory for fine-tuning.

• Explain: MoE layers often contain a large portion (about 80%) of total model parameters, so intuitively, updating these parameters seems critical.

• One last part to consider when fine-tuning sparse MoEs is that they have different fine-tuning hyperparameter setups - e.g., sparse models tend to benefit more from smaller batch sizes and higher learning rates.

• MoEs Meets Instruction Tuning (July 2023), performs experiments doing:

  1. Single task fine-tuning

  2. Multi-task instruction-tuning

  3. Multi-task instruction-tuning followed by single-task fine-tuning

• Results:

  1. When the authors fine-tuned the MoE and the T5 equivalent, the T5 equivalent was better. 

  2. When the authors fine-tuned the Flan T5 (T5 instruct equivalent) MoE, the MoE performed significantly better. 

  3. Not only this, the improvement of the Flan-MoE over the MoE was larger than Flan T5 over T5, indicating that MoEs might benefit much more from instruction tuning than dense models. 

  4. MoEs benefit more from a higher number of tasks. Unlike the previous discussion suggesting to turn off the auxiliary loss function, the loss actually prevents overfitting.

When to use sparse MoEs vs dense models?

• Experts are useful for high throughput scenarios with many machines. Given a fixed compute budget for pretraining ==> a sparse model will be more optimal. 

• For low throughput scenarios with little VRAM ==> a dense model will be better.

• Note: one cannot directly compare the number of parameters between sparse and dense models, as both represent significantly different things.

Making MoEs go brrr

• The initial MoE work presented MoE layers as a branching setup, leading to slow computation as GPUs are not designed for it and leading to network bandwidth becoming a bottleneck as the devices need to send info to others. 

• This section will discuss some existing work to make pretraining and inference with these models more practical. MoEs go brrrrr.

Parallelism

1. Data parallelism: the same weights (of model) are replicated across all cores, and the data is partitioned across cores.

2. Model parallelism: the model is partitioned across cores, and the data is replicated across cores.

3. Model and data parallelism: we can partition the model and the data across cores. Note that different cores process different batches of data.

4. Expert parallelism: experts are placed on different workers. If combined with data parallelism, each core has a different expert and the data is partitioned across all cores

   • With expert parallelism, experts are placed on different workers, and each worker takes a different batch of training samples. 

   • For non-MoE layers, expert parallelism behaves the same as data parallelism. 

   • For MoE layers, tokens in the sequence are sent to workers where the desired experts reside.

Capacity Factor and communication costs

• Increasing the capacity factor (CF) increases the quality but increases communication costs and memory of activations. 

• If all-to-all communications are slow, using a smaller capacity factor is better. A good starting point is using top-2 routing with 1.25 capacity factor and having one expert per core. 

• During evaluation, the capacity factor can be changed to reduce compute.

Serving Techniques

• A big downside of MoEs is the large number of parameters. For local use cases, one might want to use a smaller model. Let's quickly discuss a few techniques that can help with serving:

• The Switch Transformers authors did early distillation experiments. By distilling a MoE back to its dense counterpart, they could keep 30-40% of the sparsity gains. Distillation, hence, provides the benefits of faster pretaining and using a smaller model in production.

• Recent approaches modify the routing to route full sentences or tasks to an expert, permitting extracting sub-networks for serving.

• Aggregation of Experts (MoE): this technique merges the weights of the experts, hence reducing the number of parameters at inference time.

More on efficient training

• FasterMoE (March 2022) analyzes the performance of MoEs in highly efficient distributed systems and analyzes the theoretical limit of different parallelism strategies, as well as techniques to skew expert popularity, fine-grained schedules of communication that reduce latency, and an adjusted topology-aware gate that picks experts based on the lowest latency, leading to a 17x speedup.

• Megablocks (Nov 2022) explores efficient sparse pretraining by providing new GPU kernels that can handle the dynamism present in MoEs. Their proposal never drops tokens and maps efficiently to modern hardware, leading to significant speedups. What’s the trick? Traditional MoEs use batched matrix multiplication, which assumes all experts have the same shape and the same number of tokens. In contrast, Megablocks expresses MoE layers as block-sparse operations that can accommodate imbalanced assignment.

Dataset

Metrics

Experimental Results

Ablations

• Further experiments on distilling a sparse MoE back to a dense model with less parameters but similar number of parameters.

• 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.

• Some interesting areas to explore:

  1. Distilling Mixtral into a dense model

  2. Explore model merging techniques of the experts and their impact in inference time

  3. Perform extreme quantization techniques of Mixtral

• [From DeepSeekV3] MoE has Routing Collapse problem, where the router learns to use only one or very few experts every single time and ignore the rest.