50张图解读MoE在LLM中的作用 精华

来源：https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts

如您关注最新发布的大型语言模型（LLM），可能会频繁在标题中看到“MoE”这个词。那么，“MoE”究竟是什么？为何众多LLM都开始采用它呢？

在这份图解中，我们将通过50多张可视化图，详细探索这个重要组件——混合专家模型（Mixture of Experts, MoE）！

在这份图解中，我们将深入探讨MoE的两个主要组成部分，即专家（Experts）和路由器（Router），及其在典型LLM架构中的应用。

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

混合专家模型（MoE）是一种利用多个不同子模型（或“专家”）来提升LLM质量的技术。

MoE由两个主要组件定义：

• 专家（Experts）- 每个前馈神经网络（FFNN）层现在拥有一组“专家”，可以从中选择一个子集。这些“专家”通常本身就是FFNN。
• 路由器（Router）或门控网络（Gate Network）- 决定哪些token被发送到哪些专家。

在采用MoE的LLM的每个层中，我们都会发现（某种程度上专业化的）专家：

请注意，“专家”并非专攻“心理学”或“生物学”等特定领域。它们至多学习词汇层面的句法信息：

更具体地说，它们的专长在于处理特定上下文中的特定token。

路由器（Router）（门控网络）选择最适合给定输入的专家：

每个专家都不是一个完整的LLM，而是LLM架构中的一个子模型。

专家（The Experts）

为了探究专家代表什么以及它们如何工作，我们首先需要了解MoE旨在替代什么：稠密层（Dense Layers）。

稠密层（Dense Layers）

混合专家模型（MoE）都源于LLM的一个相对基本的功能，即前馈神经网络（Feedforward Neural Network, FFNN）。

回想一下，标准的仅解码器Transformer架构在层归一化后应用FFNN：

FFNN支持模型利用注意力机制创建的上下文信息，对其进行进一步转换以捕捉数据中更复杂的关系。

然而，FFNN的规模增长迅速。为了学习这些复杂关系，它通常会扩展其接收到的输入：

稀疏层（Sparse Layers）

传统Transformer中的FFNN被称为稠密模型，因为它的所有参数（权重和偏置）都被激活。没有任何部分被遗漏，所有参数都用于计算输出。

如果我们仔细观察稠密模型，会注意到输入在某种程度上激活了所有参数：

相比之下，稀疏模型只激活其总参数的一部分，并且与混合专家模型密切相关。

举例来说，我们可以将稠密模型切分成多个部分（即所谓的专家），重新训练它，并在给定时间仅激活这些专家中的一个子集：

其基本思想是，每个专家在训练过程中学习不同的信息。然后，在推理时，只使用特定的专家，因为它们与给定任务最相关。

当被问及一个问题时，我们可以选择最适合给定任务的专家：

专家学习什么？

正如我们之前所见，专家学习的信息比整个领域更细粒度¹。因此，称它们为“专家”有时被认为是具有误导性的。

ST-MoE论文中编码器模型的专家专业化示例

然而，解码器模型中的专家似乎没有相同的专业化类型。但这并不意味着所有专家都是平等的。

Mixtral 8x7B论文中有一个很好的例子，其中每个token都用其第一个专家选择进行着色。

此图也表明，专家倾向于关注语法而非特定领域。

因此，尽管解码器专家似乎没有特定专长，但它们确实似乎持续用于某些类型的token。

专家的系统架构

虽然将专家可视化为稠密模型的隐藏层切片很方便，但它们本身通常是完整的FFNN：

由于大多数LLM具有多个解码器块，给定文本在生成之前会通过多个专家：

所选专家可能因token而异，从而导致采取不同的“路径”：

如果我们更新解码器块的可视图，它现在将包含更多的FFNN（每个专家一个），而不是一个：

解码器块现在拥有多个FFNN（每个都是一个“专家”），可以在推理期间使用。

路由机制

现在我们有了一组专家，模型如何知道该使用哪些专家呢？

在专家层之前，会添加一个路由器（Router）（也称为门控网络），它经过训练，可以为给定的token选择专家。

路由器（The Router）

路由器（Router）（或门控网络）也是一个FFNN，用于根据特定输入选择专家。它输出概率，并利用这些概率选择最匹配的专家：

专家层返回所选专家的输出乘以门控值（选择概率）的结果。

路由器与专家（其中只有少数被选中）共同构成了MoE层：

给定的MoE层有两种大小：稀疏（sparse）或稠密（dense）混合专家模型。

两者都使用路由器来选择专家，但稀疏MoE只选择少数，而稠密MoE选择所有专家，但可能采用不同的分布。

例如，给定一组token，MoE会将其token分配给所有专家，而稀疏MoE只会选择少数专家。

对于当前的LLM，当您看到“MoE”时，它通常指的是稀疏MoE，因为它支持您使用专家子集。这在计算上成本较低，而这对于LLM来说是一个重要特性。

专家选择

门控网络可以说是任何MoE最重要的组成部分，因为它不仅决定了在推理期间选择哪些专家，还决定了在训练期间选择哪些专家。

在其最基本的形式中，我们将输入 (x) 与路由器权重矩阵 (W) 相乘：

然后，我们对输出应用SoftMax，为每个专家创建概率分布G(x)：

路由器使用此概率分布为给定输入选择最匹配的专家。

最后，我们将每个路由器的输出与每个选定的专家相乘，并对结果求和。

让我们把所有内容放在一起，探究输入如何流经路由器和专家：

路由的复杂性

然而，这个简单的函数常常导致路由器选择相同的专家，因为某些专家可能比其他专家学习得更快：

这不仅会导致专家选择分布不均，还会导致一些专家几乎得不到训练。这在训练和推理过程中都会导致问题。

相反，我们希望在训练和推理过程中专家之间具有同等的重要性，这被称为负载均衡（load balancing）。从某种意义上说，这是为了防止对相同的专家过度拟合。

负载均衡

为了平衡专家的重要性，我们需要关注路由器，因为它是决定在给定时间选择哪些专家的主要组件。

KeepTopK

一种对路由器进行负载均衡的方法是通过一个直接的扩展，称为KeepTopK²。通过引入可训练的（高斯）噪声，我们可以防止总是选择相同的专家：

然后，除了你想要激活的K个顶部专家（例如2个）之外，所有其他专家的权重都将被设置为-∞：

通过将这些权重设置为-∞，SoftMax对这些权重输出的概率将为0：

KeepTopK策略是许多LLM仍在使用的策略，尽管存在许多有前景的替代方案。请注意，KeepTopK也可以在不添加额外噪声的情况下使用。

Token选择（Token Choice）

KeepTopK策略将每个token路由到几个选定的专家。这种方法被称为Token Choice³，它支持将给定token发送到一个专家（top-1路由）：

或发送到多个专家（top-k路由）：

一个主要的好处是它允许对专家各自的贡献进行加权和整合。

辅助损失（Auxiliary Loss）

为了在训练期间获得更均匀的专家分布，网络常规损失中增加了辅助损失（auxiliary loss）（也称为负载均衡损失）。

它增加了一个约束，强制专家具有同等重要性。

这个辅助损失的第一个组成部分是将整个批次中每个专家的路由器值求和：

这为我们提供了每个专家的重要性分数（importance scores），它代表了无论输入如何，给定专家被选中的可能性。

我们可以用它来计算变异系数（coefficient variation, CV），它告诉我们专家之间重要性分数的差异程度。

例如，如果重要性分数差异很大，CV将很高：

相反，如果所有专家具有相似的重要性分数，CV将很低（这是我们的目标）：

利用这个CV分数，我们可以在训练期间更新辅助损失，使其旨在尽可能降低CV分数（从而给予每个专家同等的重要性）：

最后，辅助损失作为单独的损失在训练期间进行优化。

专家容量（Expert Capacity）

不平衡不仅存在于被选择的专家中，还存在于发送给专家的token分布中。

例如，如果输入token不成比例地发送给一个专家而不是另一个专家，那么这也可能导致训练不足：

在这里，重要的不仅是使用了哪些专家，还在于使用了多少。

解决这个问题的一个方案是限制给定专家可以处理的token数量，即专家容量（Expert Capacity）⁴。当一个专家达到容量时，剩余的token将被发送到下一个专家：

如果两个专家都已达到容量，则token将不会被任何专家处理，而是发送到下一个层。这被称为token溢出（token overflow）。

使用Switch Transformer简化MoE

首批处理MoE训练不稳定性问题（如负载均衡）的基于Transformer的MoE模型之一是Switch Transformer⁵。它极大地简化了架构和训练过程，同时提高了训练稳定性。

切换层（The Switching Layer）

Switch Transformer是一个T5模型（编码器-解码器），它将传统的FFNN层替换为切换层（Switching Layer）。切换层是一个稀疏MoE层，为每个token选择一个专家（Top-1路由）。

路由器在计算选择哪个专家时没有特殊技巧，它将输入与专家权重相乘后进行softmax（与我们之前做的相同）。

这种架构（top-1路由）假设路由器只需要1个专家就可以学习如何路由输入。这与我们之前看到的假设相反，我们之前假设token应该被路由到多个专家（top-k路由）以学习路由行为。

容量因子（Capacity Factor）

容量因子是一个重要值，因为它决定了专家可以处理多少token。Switch Transformer通过引入直接影响专家容量的容量因子来扩展这一点。

专家容量的组成部分很简单：

如果我们增加容量因子，每个专家将能够处理更多的token。

然而，如果容量因子过大，我们会浪费计算资源。相反，如果容量因子过小，模型性能将因token溢出而下降。

辅助损失（Auxiliary Loss）

为了进一步防止token丢失，引入了简化版的辅助损失。

这种简化的损失不再计算变异系数，而是权衡了调度token的比例与每个专家路由器概率的比例：

由于目标是实现token在$N$个专家间的均匀路由，我们希望向量$P$和$f$的值均为1/N。

α 是一个超参数，我们可以用它来微调训练期间此损失的重要性。值过高会压倒主损失函数，值过低则对负载均衡几乎没有作用。

视觉模型中的混合专家模型

MoE并非仅适用于语言模型的技术。视觉模型（如ViT）利用基于Transformer的架构，因此有潜力使用MoE。

快速回顾一下，ViT（Vision-Transformer）是一种将图像分割成小块（patches）的架构，这些小块的处理方式类似于token⁶。

这些小块（或token）随后被投影成嵌入（附加位置嵌入），然后被送入常规编码器：

当这些小块进入编码器时，它们的处理方式与token类似，这使得这种架构非常适合MoE。

Vision-MoE

Vision-MoE（V-MoE）是MoE在图像模型中的首批实现之一⁷。它采用我们之前看到的ViT，并将编码器中的稠密FFNN替换为稀疏MoE。

这使得ViT模型（通常比语言模型小）能够通过添加专家实现大规模扩展。

每个专家都使用了预定义的小专家容量，以减少硬件限制，因为图像通常有许多小块。然而，低容量往往会导致小块被丢弃（类似于token溢出）。

为了保持低容量，网络会给小块分配重要性分数，并优先处理那些分数较高的小块，从而使溢出的小块通常不那么重要。这被称为批量优先级路由（Batch Priority Routing）。

因此，即使token的百分比降低，我们仍然应该看到重要的补丁被路由。

优先级路由通过关注最重要的补丁来处理更少的补丁。

从稀疏到软MoE（From Sparse to Soft MoE）

在V-MoE中，优先级评分器帮助区分重要性更高和更低的补丁。然而，补丁被分配给每个专家，未处理补丁中的信息会丢失。

软MoE（Soft-MoE）旨在通过混合补丁，从离散到软补丁（token）分配⁸。

第一步，我们将输入$x$（补丁嵌入）与一个可学习矩阵$\Phi$相乘。这为我们提供了路由器信息，它告诉我们某个token与给定专家的相关程度。

然后，通过对路由器信息矩阵（按列）进行softmax运算，我们更新每个补丁的嵌入。

更新后的补丁嵌入本质上是所有补丁嵌入的加权平均值。

从视觉上看，这就像所有的补丁都被混合在一起了。这些组合后的补丁然后被发送到每个专家。生成输出后，它们再次与路由器矩阵相乘。

路由器矩阵在token层面影响输入，在专家层面影响输出。

因此，我们得到的是“软”补丁/token，而不是离散输入。

Mixtral 8x7B中的激活参数与稀疏参数（Active vs. Sparse Parameters with Mixtral 8x7B）

MoE之所以引人注目，很大程度上在于其计算需求。由于在给定时间只使用专家子集，我们能够访问比实际使用更多的参数。

尽管一个MoE需要加载更多的参数（稀疏参数），但激活的参数较少，因为我们在推理时只使用部分专家（激活参数）。

换句话说，我们仍然需要将整个模型（包括所有专家）加载到设备上（稀疏参数），但当我们运行推理时，我们只需要使用一个子集（激活参数）。MoE模型需要更多的VRAM来加载所有专家，但在推理期间运行速度更快。

让我们以Mixtral 8x7B⁹为例，探究稀疏参数与激活参数的数量。

这里，我们可以看到每个专家的大小是5.6B，而不是7B（尽管有8个专家）。

我们将需要加载8x5.6B (46.7B) 参数（以及所有共享参数），但我们只需要使用2x5.6B (12.8B) 参数进行推理。

结论

至此，我们的混合专家模型之旅告一段落！希望这篇文章能让您更好地理解这项有趣技术的潜力。现在几乎所有模型系列都至少包含一个MoE变体，感觉它将持续存在。

参考文献

1. Zoph, Barret, et al. "St-moe: Designing stable and transferable sparse expert models. arXiv 2022." arXiv preprint arXiv:2202.08906.
2. Shazeer, Noam, et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." arXiv preprint arXiv:1701.06538 (2017).
3. Shazeer, Noam, et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." arXiv preprint arXiv:1701.06538 (2017).
4. Lepikhin, Dmitry, et al. "Gshard: Scaling giant models with conditional computation and automatic sharding." arXiv preprint arXiv:2006.16668 (2020).
5. Fedus, William, Barret Zoph, and Noam Shazeer. "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity." Journal of Machine Learning Research 23.120 (2022): 1-39.
6. Dosovitskiy, Alexey. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).
7. Riquelme, Carlos, et al. "Scaling vision with sparse mixture of experts." Advances in Neural Information Processing Systems 34 (2021): 8583-8595.
8. Puigcerver, Joan, et al. "From sparse to soft mixtures of experts." arXiv preprint arXiv:2308.00951 (2023).
9. Jiang, Albert Q., et al. "Mixtral of experts." arXiv preprint arXiv:2401.04088 (2024).

本文转载自​​​​​Andy730​​​​​，作者：常华​