漫画混合专家（MoE）

当查看大型语言模型（LLMs）的最新发布时，你经常会看到标题中带有“MoE”。这个“MoE”代表什么？为什么这么多LLMs都在使用它？

在这份视觉指南中，我们将花时间探索这个重要组件——专家混合（MoE），通过超过 50 个可视化进行深入探讨！

在此视觉指南中，我们将介绍 MoE 的两个主要组件，即专家和路由器，以及它们在典型的基于LLM架构中的应用。

查看与LLMs相关的更多可视化内容，并支持此通讯，请参阅我写的关于大型语言模型的书籍！

什么是混合专家（MoE）？

混合专家（MoE）是一种使用许多不同的子模型（或“专家”）来提高LLMs质量的技术。

两个主要组件定义了 MoE：

• 专家 - 每个 FFNN 层现在都有一个“专家”集合，其中可以选择子集。这些“专家”通常是 FFNN 本身。
• 路由器或网关网络 - 确定哪些令牌发送给哪些专家。

在具有 MoE 的LLM的每一层中，我们发现（某种程度上是专业的）专家：

“专家”并非在特定领域如“心理学”或“生物学”中专业化。最多，它只是词法方面的专家：

更具体地说，他们的专长在于处理特定语境中的特定标记。

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

每位专家不是一个完整的LLM，而是LLM架构的子模型部分。

专家Expert, MoE中的E

为了探索专家所代表的内容以及他们的工作方式，让我们首先考察 MoE 应该取代的内容；密集层。

密集层Dence layer

混合专家（MoE）都始于相对基本的功能LLMs，即前馈神经网络（FFNN）。

请记住，标准解码器 Transformer 架构在层归一化之后应用 FFNN：

一个 FFNN 允许模型使用由注意力机制创建的上下文信息，进一步将其转化为捕捉数据中更复杂关系的能力。

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

稀疏层

在传统的 Transformer 中，FFNN 被称为密集模型，因为所有参数（其权重和偏差）都被激活。没有任何东西被留下，所有内容都用于计算输出。

如果我们仔细观察密集模型，注意输入如何在一定程度上激活所有参数：

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

为了说明，我们可以将我们的密集模型切割成多个部分（所谓专家），重新训练它，并在给定时间只激活一组专家：

底层理念是每个专家在训练期间学习不同的信息。然后，在运行推理时，仅使用与特定任务最相关的特定专家。

当被问及问题时，我们可以选择最适合执行特定任务的专家：

专家学到了什么？

如我们所见，专家比整个领域学习到更细粒度的信息。1 因此，有时将他们称为“专家”被视为误导。

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

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

一个很好的例子可以在 Mixtral 8x7B 论文中找到，其中每个标记都涂上了第一个专家的选择。

此视觉图还表明，专家往往关注语法而非特定领域。

因此，尽管解码器专家似乎没有特定的专长，但他们似乎一直被用于某些类型的标记。

专家架构

尽管将专家视为密集模型中切割成块的隐藏层听起来很酷，但他们本身往往是完整的全连接神经网络（FFNN）：

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

所选专家可能在标记之间有所不同，从而导致采取不同的“路径”：

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

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

路由机制

现在我们有一批专家，模型是如何知道使用哪些专家的？

在专家之前，添加了一个路由器（也称为网关网络），该路由器经过训练，能够选择为给定令牌选择哪个专家。

路由器

路由器（或网关网络）也是一种前馈神经网络（FFNN），用于根据特定输入选择专家。它输出概率，并使用这些概率来选择最佳匹配的专家：

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

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

给定 MoE 层有两种大小，要么是稀疏的专家混合，要么是密集的专家混合。

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

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

在当前状态LLMs下，当你看到“MoE”时，它通常是一个稀疏 MoE，因为它允许你使用专家子集。这在计算上更便宜，这是LLMs的一个重要特性。

专家选择

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

在其最基本的形式中，我们将输入（x）乘以路由权重矩阵（W）：

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

路由器使用这种概率分布来选择给定输入的最佳匹配专家。

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

让我们把所有内容整合起来，探索输入如何通过路由器和专家流动：

路由的复杂性

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

不仅选定的专家分布不均，而且一些专家几乎都没有经过培训。这导致在训练和推理过程中都存在问题。

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

负载均衡

为了平衡专家的重要性，我们需要将路由器视为主要组件，以决定在特定时间选择哪些专家。

保持顶部 K

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

然后，除了您想要激活的前 k 个专家（例如 2 个）之外，其余专家的权重将被设置为-∞：

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

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

Token 选择

The KeepTopK 策略将每个令牌路由到几位选定的专家。这种方法称为 Token Choice 3，允许将特定令牌发送给一位专家（top-1 路由）：

或超过一个专家（top-k 路由）：

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

辅助损失

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

它增加了一个约束，迫使专家们具有同等重要性。

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

这为我们提供了每位专家的重要性得分，它代表了无论输入如何，某位专家被选中的可能性。

我们可以用这个来计算变异系数（CV），它告诉我们专家之间的重要性分数差异有多大。

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

相比之下，如果所有专家的重要性评分都相似，则 CV 值将较低（这正是我们追求的目标）：

使用此 CV 得分，我们可以在训练过程中更新辅助损失，使其尽可能降低 CV 得分（从而给予每个专家同等的重要性）：

最后，将辅助损失作为一个单独的损失项用于训练优化。

专家能力

不平衡不仅存在于被选中的专家中，也存在于发送给专家的代币分布中。

例如，如果输入标记不成比例地发送给某个专家而不是另一个，这也可能导致欠训练：

这里，不仅仅是关于使用哪些专家，而是他们被使用了多少。

该问题的解决方案是限制特定专家可以处理的令牌数量，即专家容量 4。当专家达到容量时，产生的令牌将被发送到下一个专家：

如果两位专家都已达到其容量，则令牌不会被任何专家处理，而是发送到下一层。这被称为令牌溢出。

简化 MoE 与 Switch Transformer

之一是处理 MoE（如负载均衡）训练不稳定问题的第一个基于 transformer 的 MoE 模型——Switch Transformer。5 它简化了大部分架构和训练过程，同时提高了训练稳定性。

交换层

开关变换器是一种 T5 模型（编码器-解码器），它用切换层替换了传统的 FFNN 层。切换层是一个稀疏 MoE 层，为每个标记选择一个专家（Top-1 路由）。

路由器在计算选择哪个专家时没有特殊技巧，只是将输入乘以专家的权重后取 softmax（与我们之前做的一样）。

这种架构（top-1 路由）假设路由器只需要 1 位专家就能学习如何路由输入。这与我们之前所见的假设不同，我们之前假设应该将令牌路由到多个专家（top-k 路由）以学习路由行为。

容量系数

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

专家能力组成部分简单明了：

如果我们提高容量系数，每位专家将能够处理更多的标记。

然而，如果容量系数过大，我们会浪费计算资源。相比之下，如果容量系数过小，由于标记溢出，模型性能将下降。

辅助损失

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

而不是计算变异系数，这种简化的损失衡量了分配的标记分数与每个专家的路由概率分数之比：

由于目标是让 N 个专家之间具有统一的令牌路由，我们希望向量 P 和 f 的值为 1/N。

α 是一个超参数，我们可以用它来微调训练过程中这个损失函数的重要性。值过高会超过主损失函数，值过低对负载均衡的帮助很小。

专家混合模型

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

作为快速回顾，ViT（视觉 Transformer）是一种将图像分割成类似标记处理的补丁的架构。6

这些补丁（或令牌）随后被投影到嵌入（带有额外的位置嵌入）中，然后再输入到常规编码器中：

这些补丁进入编码器的那一刻，它们就像标记一样被处理，这使得这种架构非常适合 MoE。

Vision-MoE

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

这允许 ViT 模型，通常比语言模型尺寸更小，通过添加专家进行大规模扩展。

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

为了保持容量低，网络将重要性分数分配给补丁，并首先处理这些补丁，以便溢出的补丁通常不太重要。这被称为批量优先路由。

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

优先路由通过关注最重要的部分，允许处理更少的补丁。

在 V-MoE 中，优先级评分器帮助区分重要程度不同的补丁。然而，补丁被分配给每个专家，未处理补丁中的信息丢失。

软 MoE 旨在通过混合补丁从离散补丁（标记）分配过渡到软补丁（标记）分配。8

在第一步中，我们将输入 x（补丁嵌入）与可学习的矩阵Φ相乘。这为我们提供了路由信息，告诉我们某个标记与特定专家的相关程度。

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

更新的补丁嵌入实际上是所有补丁嵌入的加权平均值。

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

路由矩阵影响令牌级别的输入和专家级别的输出。

因此，我们得到“软”补丁/令牌，而不是处理离散输入。

活跃参数与稀疏参数在 Mixtral 8x7B 中

莫埃（MoE）引人入胜的一大原因是其计算需求。由于在特定时间只使用专家子集，我们能够访问比实际使用的更多参数。

尽管给定的 MoE 有更多参数需要加载（稀疏参数），但由于我们只在推理时使用一些专家（激活参数），因此激活的参数较少。

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

让我们通过一个例子来探讨稀疏参数与活跃参数的数量，Mixtral 8x7B。9

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

我们不得不加载 8x5.6B（46.7B）个参数（包括所有共享参数），但推理时我们只需要使用 2x5.6B（12.8B）个参数。

结论

这标志着我们与专家混合模型的旅程结束！希望这篇帖子能让你更好地理解这种有趣技术的潜力。现在，几乎所有模型集都至少包含一个 MoE 变体，感觉它将在这里停留。

本文转载自​​AI大模型世界​​，作者： Maarten ​​​​