A visual guide to MoE models: Uncovering the role of MoE in large language models

Original address：A Visual Guide to Mixture of Experts (MoE)

📅 Author: Maarten Grootendorst

📆 Date: October 7, 2024

---

Exploring Language Models: A Visual Guide to Mixed Expert Models (MoE)

Directory

• A visual guide to MoE models: Uncovering the role of MoE in large language models
• Exploring Language Models: A Visual Guide to Mixed Expert Models (MoE)
  • Directory
  • What is a Mix of Experts (MoE) model?
  • Experts
    • Dense Layers
    • Sparse Layers
    • What does an Expert Learn?
    • Architecture of Experts

When we look at the latest releases of large language models (LLMs, Large Language Models), you often see "MoE". What does this “MoE” stand for? Why are so many LLMs using it?

In this visual guide, we'll step through more than 50 visualizations to explore this key component: **Mixture of Experts (MoE)**.

Illustrated content: In this picture, you can see a typical MoE Two main components of the structure:Experts(expert) and Router(router or gated network). The picture shows a Router, and several juxtaposed below Experts, indicating that in LLM In the architecture, the MoE routes inputs to the appropriate experts as needed.
Figure 1 Detailed description：

1. Router: Decide which expert or experts to send input (such as token) to.
2. Experts: several different submodels (usually FFNN structure), each expert may have expertise in different aspects.
3. Workflow: Input passes first Router, and then assigned to different experts for processing, and finally summarize the results.

What is a Mix of Experts (MoE) model?

Mixture of Experts (MoE) is a technique that uses many different submodels (or "experts”) to improve the quality of large language models.

In MoE, there are two main components:

1. Experts
   • each FFNN Each layer is no longer a separate network, but has a group of "experts" to choose from.
   • These "experts" are often also FFNN(Feedforward Neural Network) structure.
2. Router or gate network
   • Responsible for deciding which tokens Which experts are being sent to.

In a LLM At each level, we see experts who are specialized (to some extent):

Illustrated content: Shown in LLM Each layer of can have multiple Experts. It highlights that these experts are able to handle different input tokens in different contexts.
Figure 2 details：

1. layer structure: The figure uses different levels (Layer 1, Layer 2, Layer 3...) to represent the multi-layer model.
2. Experts: At each layer, there are several experts (Expert 1, Expert 2, Expert 3, Expert 4), and these experts exist in parallel.
3. target: Emphasizes that experts are more "professional" in a specific context or specific input, and are thus selected to handle that input.

Although MoEs do not specifically train experts in specific fields (such as psychology or biology), they may still develop certain biases at the lexical or syntactic level:

• MoE experts may learn different language features
  • Expert 1 ProcessPunctuation(Punctuation): such as , . : & - ? Wait.
  • Expert 2 Processverb(Verbs): Such as said, read, miss Wait.
  • Expert 3 ProcessConnective word(Conjunctions): such as the, and, if, not Wait.
  • Expert 4 Processvisual descriptors(Visual Descriptions): Such as dark, outer, yellow Wait.

More specifically, their expertise is in handling specific tokens in specific contexts.

---

Router (gate network) Select the expert or combination of experts that best suits a given input:

Illustrated content: shows Router How to select appropriate experts at each layer based on input. The figure highlights the selected experts and the flow of input tokens.
Figure 3 details：

1. input: The Input at the top of the diagram represents the token or vector representation received by the model.
2. Router: Located in the network structure, it plays a decision-making role.
3. Expert choice: The selected expert will receive input, and the remaining experts will not be activated.
4. output: Results from activated experts are aggregated or continue flowing to downstream layers.

It is important to note that each expert is not the entire LLM; LLM A submodel part of the architecture.

---

Experts

To understand the experts (Experts) and how they work, let’s first take a look at what the MoE hopes to replace:dense layers。

Dense Layers

all Mixture of Experts (MoE) They are all based on a relatively basic function in LLM: **Feedforward Neural Network (FFNN)**.

Recall that a standard decoder-only Transformer In the architecture,FFNN usually in layer normalization Applied afterwards:

Illustrated content: shows a typical decoder structure, each decoder block contains Masked Self-Attention and FFNN(There will be Layer Norm）。

1. Position Embedding: Add position coding information before or at the same time as input token.
2. Decoder Block: Contains Masked Self-Attention、Layer Norm and FFNN。
3. FFNN: Represented by purple squares in the figure, it is a key component of this layer that further transforms the input to capture more complex relationships.

FFNN The contextual information produced by the attention mechanism can be exploited and further transformed to capture more complex relationships in the data.

However, in order to learn these complex relationships,FFNN The size of will grow, usually expanding on the input (for example, the middle layer dimension will become larger):

Illustrated content: shows a FFNN structure, the input is first mapped to a higher dimension and then mapped back to the output dimension.

1. input dimensions: The figure shows 512 input units.
2. Hidden layer: Usually there is a 4x or more expansion (the example in the figure is a 4x expansion to 2048 dimensions).
3. Output dimensions: Then map back to 512-dimensional output.

Sparse Layers

In a traditional Transformer,FFNN called dense model, since all its parameters (weights and biases) will be activated. In other words, all parameters of the model are involved in calculating the output.

If we look carefully dense model, you can see that input will activate all parameters:

Illustrated content: shows a "dense" model, where each neuron in the input layer is connected to all neurons in the hidden layer, and all neurons in the hidden layer are connected to neurons in the output layer.
Figure 6 details：

1. Fully connected: All nodes in the graph are connected to all nodes in the next layer, indicating no sparsity.
2. All parameters are activated: There are no "idle" or "inactive" parameters.

In contrast,sparse models(sparse model) only activates a part of the total parameters, which is consistent with Mixture of Experts closely related.

To illustrate this, we can put dense model split into parts (i.e. experts,experts), retrain it, and activate only part of it during inference:

Illustrated content: Split the original dense model into multiple experts (Expert 1, Expert 2, Expert 3, Expert 4). During the inference phase, only a subset of experts is selected for activation.

1. Model segmentation: The original large network is split into multiple smaller "experts".
2. sparse activation: Not all experts are activated, only some are activated under certain inputs.
3. Benefits: With sparse activation, you can have more potential parameter capacity without significantly increasing computational cost.

The core idea is that during training, each expert learns different information; during inference, only those experts most relevant to the current task are used.

When we ask a question, we select the experts best suited for the task:

Illustrated content: shows an example: when the input is a number-related question like “What is 1 + 1?”, the router activates only number-related experts.

1. input: A sentence or token representing an arithmetic problem.
2. Expert choice: Only experts in the "Numbers" field are activated.
3. output: The expert gives the result "2".

What does an Expert Learn?

As mentioned earlier, experts (Experts) often learn more detailed knowledge than the entire field. One might feel that calling them “experts” might be misleading, but this is because each expert tends to specialize only in certain types of input features or contexts.

Illustrated content: Shows a table or comparison to illustrate that in some cases, different experts may learn different features (such as punctuation marks, verbs, numbers, etc.).

1. Example Expert: Punctuation, Conjunctions, Verbs, Numbers, etc.
2. hierarchical location: Different experts may appear in different layers of the model.
3. allocate: Certain tokens are routed to certain experts for more efficient processing.

in decoder In the model, there may not be such an obvious division of labor between experts. However, this does not mean that all experts are exactly the same.
in Mixtral 8x7B In this paper, there is a good example: each token will be marked with its preferred experts, which do not necessarily correspond to intuitive semantic fields, but show certain tendencies statistically.

This visual example also shows that experts tend to focus on syntax rather than a specific domain. So while decoder experts don't seem to have a clear "specialism", they do seem to be consistently used on certain types of tokens.

In [Figure 1], a sample code or visualization result about MoELayer is shown. The color blocks distinguish different parts and emphasize theExperts and routersrelationship between. It can be seen from the color blocks:

• The experts list (represented by nn.ModuleList in the code) contains multiple sub-networks (i.e. multiple FFNN, Feed-Forward Neural Network, feed-forward neural network).
• The gate (gated network, also called router) is responsible for selecting which experts will be activated.
• Overall, it can be seen that these experts usually focus on the syntactic level of the input sentences rather than on a specific topic or domain.

Architecture of Experts