50张图，直观理解混合专家(MoE)大模型

OSCHINA

的高效训练与推理效果曾引发AI社区对混合专家（MoE）模型的广泛关注，后来居上的国产开源大模型以及（基于Transformer的最大MoE模型）也选择了MoE框架路线。为何大语言模型总是离不开MoE的身影？

借助50多个图例，数据科学家Maarten Grootendorst由浅 入深多维度剖析了MoE模型，从基础概念出发，逐步介绍MoE核心组件专家和路由机制，以及它们在典型LLM架构中的应用。

（本文经作者授权后由OneFlow编译发布。原文：https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts）

作者 | Maarten Grootendorst

OneFlow编译

翻译｜张雪聃、林心宇
题图由平台生成

1

什么是混合专家（MoE）？

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

MoE的两个主要组成部分是：

• 专家：每个前馈神经网络（FFNN）层现在都有一组“专家”，可以选择其中的一部分。这些“专家”通常也是FFNN。

• 路由或门控网络：决定哪些词元发送到哪些专家。

在每个具有MoE的模型层中，我们会找到（相对专业化的）专家：

需要注意的是，“专家”并不专注于特定领域，如“心理学”或“生物学”。专家在学习过程中最多只能掌握关于单词层面的句法信息：

更具体地说，专家的专长是在特定上下文中处理特定词元。

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

单个专家并不是整个LLM，而是LLM架构中的一个子模型部分。

2

专家

为了探讨专家的含义及其工作方式，我们首先需要了解MoE所替代的内容：密集层。

密集层

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

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

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

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

稀疏层

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

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

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

为了说明这一点，我们可以将密集模型分割成片段（即专家），重新训练，并在给定时间仅激活一部分专家：

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

当收到问题时，我们可以选择最适合特定任务的专家：

专家学习的内容

如前所述，专家学习到的信息比整个领域的信息更加精细。因此，称它们为“专家”有时被视为误导。

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

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

Mixtral 8x7B论文(https://arxiv.org/pdf/2401.04088) 中有一个很好的例子，其中每个词元都标记了第一个专家的选择。

这一图片也表明，专家往往专注于句法而非特定领域。

因此，尽管解码器专家似乎没有专业化，但它们似乎在特定类型的词元上使用得相对一致。

专家的架构

尽管将专家可视化为切成块的密集模型的隐藏层很不错，但它们通常是完整的FFNN：

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

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

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

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

3

路由机制

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

我们可以在专家层之前添加一个路由（也称为门控网络），它是专门训练用来选择针对特定词元的专家。

路由

路由（或门控网络）也是一个前馈神经网络（FFNN），用于根据特定输入选择专家。它可以输出概率，用于选择最匹配的专家：

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

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

给定的MoE层有两种类型：稀疏混合专家或密集混合专家。

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

例如，给定一组词元，MoE会将词元分配到所有专家，而稀疏MoE仅选择少数专家。

在当前的LLM状态下，当看到“MoE”时，通常指的是稀疏MoE，因为它允许使用一部分专家。这在计算上更为经济（消耗的资源更少），这是LLM的重要特性。

选择专家

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

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

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

路由使用这个概率分布来选择最匹配的专家。

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

我们可以将所有内容结合在一起，探索输入如何通过路由和专家流动：

路由的复杂性

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

这不仅会导致选择的专家分布不均，而且一些专家几乎无法受到训练。这在训练和推理期间都会产生问题。

相反，我们希望在训练和推理期间让专家之间保持均等的重要性，这称为负载均衡。这样可以防止对同一专家的过度拟合。

负载均衡

为平衡专家的重要性，我们需要把关注点放在路由上，因为它是在特定时间决定选用哪些专家的关键组件。

KeepTopK

对路由进行负载均衡的一种方式是借助"KeepTopK"（https://arxiv.org/pdf/1701.06538）直接扩展。通过引入可训练的（高斯）噪声，可以避免重复选择相同的专家。

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

通过将这些权重设置为-∞，这些权重上的SoftMax输出所产生的概率将会是0：

尽管许多替代方案都很有前景，但许多语言模型仍然使用KeepTopK策略。请注意，KeepTopK也可以在不添加额外噪声的情况下使用。

• 词元选择

KeepTopK策略将每个词元路由到少数选定的专家。这种方法称为词元选择，它允许将给定的词元发送给一个专家（top-1路由）：

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

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

• 辅助损失

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

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

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

这为我们提供了每个专家的重要性得分，它代表了在任何输入下，给定专家被选中的可能性。

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

例如，如果重要性得分有很大差异，变异系数就会很高：

相反，如果所有专家的重要性得分相似，变异系数就会很低（这是我们的目标）：

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

最后，辅助损失被单独添加进来，作为一个独立的损失项在训练期间进行优化。

专家容量

不平衡现象不仅存在于被选中的专家中，还存在于发送给专家的词元分布中。

例如，如果输入的词元在分配给不同专家时比例失调，过多地发送给一个专家而较少地发送给另一个专家，那么可能会出现训练不足的问题。

这里，问题不仅仅在于使用了哪些专家，还在于对它们的使用程度。

这个问题的一个解决方案是限制给定专家可以处理的词元数量，即专家容量。当一位专家达到其容量时，后续的词元将被发送给下一位专家：

如果两位专家都达到了他们的容量，那么该词元将不会被任何专家处理，而是被发送到下一层。这被称为词元溢出（token overflow）。

借助Switch Transformer简化MoE

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

• 切换层

Switch Transformer是一个T5模型（编码器-解码器），它用切换层取代了传统的前馈神经网络层。切换层是一个稀疏的MoE层，它为每个词元选择一个专家（Top-1路由）。

路由在计算选择哪个专家时没有特殊技巧，它只是对输入乘以专家权重后的结果取Softmax（与我们之前所做的相同）。

这种架构（Top-1 路由）假定路由只需要一个专家就能学会如何对输入进行路由。这与我们之前看到的情况形成对比，之前我们假设词元应该被路由到多个专家（Top-k 路由）以学习路由行为。

• 容量因子

容量因子是一个重要的值，因为它决定了一个专家能够处理多少个词元。Switch Transformer在此基础上进行了扩展，直接引入了一个容量因子，它对专家容量产生直接影响。

专家容量的组成部分很直接：

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

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

• 辅助损失

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

这个简化后的损失并非去计算变异系数，而是依据每个专家的路由概率所占的比例，来对分配给各个专家的词元的比例进行权衡。

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

α是一个超参数，我们可以在训练期间使用它来微调这个损失的重要性。过高的值将主导主要的损失函数，而过低的值对负载均衡的作用很小。

4

视觉模型中的混合专家

混合专家（MoE）技术并非仅适用于语言模型。视觉模型（例如视觉transformerViT）利用基于transformer的架构，因此也有使用混合专家的潜力。

ViT（视觉transformer）是一种将图像分割成小块（patch）的架构，这些小块的处理方式与词元类似。

这些小块（或词元）随后被映射到嵌入中（带有额外的位置嵌入），然后再输入到常规编码器中：

这些小块一旦进入编码器，就会像词元一样被处理，这使得这种架构非常适合用于混合专家。

视觉混合专家（Vision-MoE）

视觉混合专家（V-MoE）是在图像模型中最早实现混合专家的方法之一。它采用我们之前看到的视觉变换器（ViT），并将编码器中的密集前馈神经网络替换为稀疏混合专家（Sparse MoE）。

这使得通常比语言模型规模更小的ViT模型能够通过添加专家而大规模扩展。

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

为了保持低容量，网络为小块分配重要性得分，并首先处理那些重要的小块，这样溢出的小块通常就不太重要了。这被称为批量优先级路由。

因此，如果词元的百分比降低，我们仍然应该看到重要的小块被路由。

优先级路由通过专注于最重要的小块，使得较少的小块被处理。

从稀疏混合专家到软混合专家

在视觉混合专家（V-MoE）中，优先级评分器有助于区分更重要和不太重要的小块。然而，小块被分配给每个专家，未处理小块中的信息会丢失。

软混合专家（Soft-MoE）旨在通过混合小块将离散的小块（词元）分配转变为软小块（词元）。

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

然后，通过对路由信息矩阵（在列上）取Softmax，我们更新每个小块的嵌入。

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

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

路由矩阵在词元级别影响输入，在专家级别影响输出。

结果，我们得到了被处理的“软”小块/词元，而不是离散的输入。

Mixtral 8x7B的活跃参数与稀疏参数

混合专家之所以有趣的很大一部分原因在于其计算需求。由于在给定时间只使用一部分专家，所以我们可以访问比实际使用更多的参数。

虽然给定的混合专家模型有更多的参数要加载（稀疏参数），但由于在推理期间我们只使用一些专家，所以激活的参数较少（活跃参数）。

换句话说，我们仍然需要将整个模型（包括所有专家）加载到你的设备上（稀疏参数），但当我们进行推理时，我们只需要使用一部分（活跃参数）。混合专家模型需要更多的显存（VRAM）来加载所有专家，但在推理期间运行得更快。

让我们以Mixtral 8x7B来探讨稀疏参数与活跃参数的数量。

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

我们将不得不加载8×5.6B（46.7B）的参数（以及所有共享参数），但在推理时我们只需要使用2×5.6B（12.8B）的参数。

5

总结

至此，我们对混合专家的探索之旅就结束了。希望这篇文章能让你更好地理解这一有趣技术的潜力。如今，几乎所有的模型系列中都至少包含一种混合专家的变体，它将会持续存在下去。