Standard decoder only architecture

Modern LLMs, such as GPT-style models, use a decoder-only Transformer architecture, unlike the original encoder-decoder Transformer architecture introduced for machine translation. The encoder-decoder design was optimized for tasks where an input sequence needs to be transformed into a different output sequence, with the decoder attending to the encoder’s representation of the input.

For text generation, however, a separate encoder is not required. The model’s objective is to predict the next token based only on the previous tokens. Decoder-only Transformers achieve this by using masked self-attention, allowing each token representation to incorporate information from all preceding tokens while preventing access to future tokens.

But the overall steps stay the same:

Decoder Block without MoE

In a standard decoder block, after the self-attention, a Feed-Forward Neural Network (FFNN) uses the contextual information that was created by the self-attention to capture complex relationships in the data. The FFNN is a dense layer, meaning all neurons are activated.

In MoE, a sparse model is used instead of a dense model for the FFNN layer in the decoder. Instead of a single FFNN, multiple FFNNs are used, and a router in front of the so called “experts” decides which FFNN is activated. The sparse MoE approach used in modern LLMs was introduced in Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (Shazeer et al., 2017).

Dense vs Sparse MoE

Each expert is a FFNN itself, and each decoder block has its own router and its own FFNNs acting as experts. Through the decoder blocks, the experts might differ for a given token, resulting in a different path of experts for a given inference.

Meaning that for a given inference, a set of experts is chosen given the input. The router in each decoder block decides which FFNN expert is chosen. The expert FFNNs and the router together form the MoE layer in the decoder block.

The router itself is a linear layer with a Softmax function at the end, creating probabilities for each expert. The final output is then generated by multiplying the router probability with the expert output, resulting in a weighted activation.

You can think of the MoE architecture as multiple FFNNs acting as experts and a router choosing the right experts to activate given the input tokens.

The term “experts” in Mixture-of-Experts (MoE) models can be somewhat misleading. Unlike human experts, these components do not typically specialize in complete domains, such as computer science, medicine, or law. Instead, individual experts often learn more fine-grained patterns and representations within the data. This specialization can include linguistic structures, semantic concepts, or task-specific features, but it does not mean that an expert has a deep understanding of an entire domain. The Mixtral paper (Jiang et al., 2024) shows this nicely: when the authors analyzed routing decisions, they found no obvious specialization by topic, but experts did show patterns on the syntax level, for example consecutive tokens or indentation in code often going to the same expert.

There are multiple architecture choices for MoE. For example, a dense MoE architecture would activate all experts and a sparse one would only select a few. When multiple experts are enabled, the outputs are weighted by the router probabilities and then aggregated.

Below is the architecture of the decoder block with the MoE layer:

  1. The token embedding for the current position enters the block and starts flowing down the residual stream.
  2. Masked self-attention lets the token pull in information from earlier tokens only (never future ones), and the result is added back onto the stream.
  3. The router reads the token and scores the experts, deciding which one or few should handle it.
  4. The selected expert (one of 4a to 4d) runs its own feed-forward network on the token while the rest stay idle, and its output is added back onto the stream.
  5. The LM Head turns the final vector into a probability distribution over the vocabulary to predict the next token.

MoE Decoder Block

Choosing an expert

A common challenge for training MoE architectures is that some experts might learn faster than others, which results in some experts always being chosen over others.

We want equal importance among experts during training and inference, which can be called also load balancing. In a way, it’s to prevent overfitting on the same experts.

Load Balancing

Instead of just multiplying the input with the router weights and then using the softmax function, noise is introduced so that some experts might get randomly lower scores, which prevents them from being overchosen. Then sparsity is introduced by setting the weights of all but the top k experts to minus infinity. The softmax then sets the probability of those weights to zero. This allows the model to select experts which are undertrained so they can catch up. It also means that a given token can be routed to multiple experts.

The KeepTopK strategy routes each token to a few experts. A major benefit is that it allows the experts’ respective contributions to be weighed and integrated. Both the noisy gating and the KeepTopK selection come from Shazeer et al. (2017).

Auxiliary loss

To get a more even distribution of experts during training, the auxiliary loss (also called load balancing loss) was added to the network’s regular loss. The auxiliary loss is introduced in the loss function, resulting in a more stable training procedure where all experts are given a more equal chance of being chosen.

Expert Capacity

Imbalance is not only in the experts that are chosen for a token, but also in the distribution of the tokens that are sent to experts. Some experts might get only one input token sent to them, while others get all input tokens sent to them. Here, it is not just about which experts are used but how much they are used.

The experts receiving fewer tokens are undertrained. To prevent this problem, we limit how many tokens an expert can process, which is called expert capacity. For example, with an expert capacity of 3, an expert can only process 3 tokens at once. Other tokens will then be routed to the next likely expert. If all experts reached capacity, any new token will not be routed to an expert and is instead sent to the next layer. This is called token overflow. Expert capacity and token overflow were introduced in GShard (Lepikhin et al., 2020), the first work to scale MoE Transformers to over 600 billion parameters.

Computational Requirements

The main advantage of Mixture-of-Experts (MoE) architectures over classical dense Transformer architectures is that they increase the total number of parameters while keeping the computation per token relatively constant. Instead of activating all parameters for every input, an MoE model uses a router to select only a subset of experts for each token. The Switch Transformer paper (Fedus et al., 2021) demonstrated this trade-off at scale: by routing each token to just a single expert, they trained models with up to 1.6 trillion parameters and reached the same quality as a dense T5 baseline up to 7 times faster.

During inference, the full set of model parameters may still need to be available in memory, but only the selected experts are activated for each token prediction. This allows MoE models to provide the capacity of a much larger model while requiring significantly fewer active parameters during computation. But the practical inference speed depends on factors such as hardware, memory bandwidth, and expert routing overhead. Mixtral 8x7B is a good concrete example: it has 47B total parameters but only uses about 13B active parameters per token, since each token is routed to 2 of the 8 experts per layer.

MoE Parameter Counts

MoE for vision models

For text based transformers, text is a sequence split into tokens. For images, similar to text split up into tokens, images are split up into patches (tokens). Patches are then flattened into a sequence of patches. This is passed into a linear projection to create embeddings, one for each patch.

Unlike GPT models using decoders, a Vision Transformer processes all image patches simultaneously, that’s why they have blocks called encoders. There is no causal masking, self-attention is bidirectional, and the full context is available. GPT style architectures do not have an encoder because they are designed for text generation, not for encoding an entire input sequence into a representation, which is important for vision models.

For vision models, the FFN in the encoder block is then replaced with a MoE layer, similar to text based transformers. This was done in V-MoE (Riquelme et al., 2021), which scaled Vision Transformers to 15B parameters and matched dense models at roughly half the inference compute.

References