MoE 模型的的可视化指南：揭秘 MoE 在大型语言模型中的角色

原文地址：A Visual Guide to Mixture of Experts (MoE)

📅 作者：Maarten Grootendorst

📆 日期：2024 年 10 月 7 日

探索语言模型：混合专家模型（MoE）可视化指南

图示内容：在这张图中，可以看到一个典型 MoE 结构的两个主要组成部分：Experts（专家）和 Router（路由器或门控网络）。图中显示了一个 Router，以及下方并列的多个 Experts，表明在 LLM 架构中，MoE 会将输入根据需要路由到合适的专家。
图 1 详细说明：

Router：决定将输入（例如 token）发送给哪一个或哪几个专家。
Experts：若干个不同的子模型（通常是 FFNN 结构），每个专家可能在不同方面具有专长。
工作流程：输入先通过 Router，再被分配到不同的专家进行处理，最后汇总结果。

什么是混合专家（MoE）模型？

Mixture of Experts (MoE) 是一种技术，它使用许多不同的子模型（或“experts”）来提升大型语言模型的质量。

在 MoE 中，有两个主要组件：

Experts
- 每个 FFNN 层都不再是一个单独的网络，而是有一组“专家”可供选择。
- 这些“专家”通常也是 FFNN（Feedforward Neural Network）结构。
Router 或 gate network
- 负责决定哪些 tokens 被发送到哪些专家。

在一个带有 MoE 的 LLM 的每一层，我们都能看到（在某种程度上）有所专门化的专家：

图示内容：展示了在 LLM 的每一层都可以拥有多个 Experts。它强调了这些专家在不同的上下文中能够处理不同的输入 token。
图2详细说明：

层结构：图中用不同的层级（Layer 1、Layer 2、Layer 3……）表示多层模型。
Experts：在每一层，都有若干个专家（Expert 1、Expert 2、Expert 3、Expert 4），这些专家并行存在。
目标：强调专家在特定上下文或特定输入时更具备“专业性”，从而被选中来处理该输入。

尽管 MoE 并不会在特定领域（如心理学或生物学）上专门训练专家，但它们仍可能在词法或句法级别上形成一定的偏向：

MoE 专家可能学习到不同的语言特征
- Expert 1 处理标点符号（Punctuation）：如 , . : & - ? 等。
- Expert 2 处理动词（Verbs）：如 said, read, miss 等。
- Expert 3 处理连接词（Conjunctions）：如 the, and, if, not 等。
- Expert 4 处理视觉描述词（Visual Descriptions）：如 dark, outer, yellow 等。

更具体地说，他们的专长是在特定上下文中处理特定的标记（tokens）。

Router (gate network) 选择最适合给定输入的专家或专家组合：

图示内容：展示了 Router 如何在每一层根据输入选择合适的专家。图中高亮了被选中的专家，以及输入 token 的流动过程。
图3详细说明：

输入：图顶部的 Input 代表模型接收到的 token 或向量表示。
Router：位于网络结构中，起到决策作用。
专家选择：被选中的专家会接收输入，其余专家则不被激活。
输出：来自被激活专家的结果被汇总或继续流向下游层。

需要注意的是，每个专家并不是整个 LLM，而是 LLM 架构中的一个子模型部分。

Experts

为了理解专家（Experts）是什么以及它们如何工作，我们先来看看 MoE 希望替代的东西：dense layers。

Dense Layers

所有的 Mixture of Experts (MoE) 都基于 LLM 中一个相对基础的功能：**Feedforward Neural Network (FFNN)**。

回忆一下，一个标准的 decoder-only Transformer 架构中，FFNN 通常是在 layer normalization 之后应用的：

图示内容：展示了一个典型的 decoder 结构，每个 decoder block 包含 Masked Self-Attention 和 FFNN（中间会有 Layer Norm）。

Position Embedding：在输入 token 之前或同时加入位置编码信息。
Decoder Block：包含 Masked Self-Attention、Layer Norm 和 FFNN。
FFNN：在图中用紫色方块表示，是该层对输入进一步变换以捕捉更复杂关系的关键组件。

FFNN 可以利用注意力机制产生的上下文信息，对其进行进一步的转换，以捕捉数据中更复杂的关系。

不过，为了学习这些复杂关系，FFNN 的规模会随之增长，通常会在输入上进行扩张（例如，中间层维度会变大）：

图示内容：展示了一个 FFNN 的结构，输入先被映射到更高维度，然后再被映射回输出维度。

输入维度：图中显示有 512 个输入单元。
隐藏层：通常会有 4 倍或更多的扩张（图中示例为 4 倍扩张到 2048 维）。
输出维度：再映射回 512 维的输出。

Sparse Layers

在传统的 Transformer 中，FFNN 称为 dense model，因为它的所有参数（权重和偏置）都会被激活。也就是说，模型的全部参数都参与计算输出。

如果我们仔细观察 dense model，可以看到输入会激活所有的参数：

图示内容：展示了一个“密集”模型，输入层的每个神经元都与隐藏层所有神经元相连，隐藏层所有神经元又与输出层神经元相连。
图6详细说明：

全连接：图中所有节点都连接到下一层的所有节点，表示无稀疏性。
所有参数被激活：没有任何“闲置”或“未激活”的参数。

与之对比，sparse models（稀疏模型）只激活一部分总参数，这与 Mixture of Experts 密切相关。

为了说明这一点，我们可以把 dense model 切分成多个部分（即专家，experts），重新训练它，并且在推理（inference）时只激活其中一部分：

图示内容：将原本的密集模型分割成多个专家（Expert 1、Expert 2、Expert 3、Expert 4）。在推理阶段，只选择一部分专家进行激活。

模型切分：原有的大网络被拆分成多个较小的“专家”。
稀疏激活：并不是所有专家都被激活，只有部分专家在某些输入下被激活。
好处：通过稀疏激活，可以在不显著增加计算成本的情况下，拥有更多的潜在参数容量。

其核心思想是：在训练期间，每个专家学习不同的信息；在推理时，只用到与当前任务最相关的那些专家。

当我们提出一个问题时，就会选择最适合该任务的专家：

图示内容：展示了一个示例：当输入是 “What is 1 + 1?” 这样的数字相关问题时，路由器只激活与数字相关的专家。

输入：一个表示算术问题的句子或 token。
专家选择：只激活 “Numbers” 领域的专家。
输出：专家给出结果 “2”。

What does an Expert Learn?

正如前面所提到的，专家（Experts）往往学习到比整个领域更细致的知识。有人会觉得称它们为“专家”可能会带来误解，但这是因为每个专家往往只专注于某些特定类型的输入特征或上下文。

图示内容：展示了一个表格或对照，说明在某些情况下，不同的专家可能学习到不同的特征（比如标点符号、动词、数字等）。

示例化专家：Punctuation、Conjunctions、Verbs、Numbers 等。
分层位置：不同专家可能出现在模型的不同层。
分配：某些 token 会路由到某些专家，以获得更有效的处理。

在 decoder 模型中，专家之间可能没有那么明显的领域分工。然而，这并不意味着所有专家都完全相同。
在 Mixtral 8x7B 这篇论文中，有一个很好的示例：每个 token 会被标记为其首选专家，这些专家并不一定对应直观的语义领域，但在统计上表现出某些倾向。

这张可视化示例还展示了，experts（专家）更倾向于关注句法（syntax），而不是特定的领域（domain）。因此，虽然 decoder experts（解码器专家）似乎并没有明确的“专业领域（specialism）”，但它们似乎会在某些特定类型的 tokens（标记）上被持续地使用。

在[图1]中，展示了一段关于 MoELayer 的示例代码或可视化结果，色块区分了不同部分，强调了专家（experts）与路由器（router）之间的关系。通过色块可以看出：

experts 列表（在代码中用 nn.ModuleList 表示）包含了多个子网络（即多个 FFNN，Feed-Forward Neural Network，前馈神经网络）。
gate（门控网络，也称 router）负责选择哪些专家会被激活。
整体上可以看到，这些专家通常关注到输入句子的句法层面，而非特定主题或领域。

MoE模型的的可视化指南：揭秘 MoE 在大型语言模型中的角色