全面指南！掰细了讲混合专家模型MoE的内部结构

什么是混合专家模型MoE呢？

混合专家模型MoE是一种能让大型语言模型（LLM）质量变好的技术。

它主要由两部分组成：

专家——现在每个前馈神经网络（FFNN）层都有好多个“专家”，每次可以用其中一部分。“专家”一般也是前馈神经网络。

门控网络——它能决定把哪些tokens送到哪些专家那儿。

在带混合专家模型MoE的大型语言模型的每一层，都能看到这些有点专门作用的专家:

要知道，这里的‘专家’可不是像‘心理学’‘生物学’这种特定领域的专家哦。最多就是学学词语层面的语法信息：

说得更细点儿，它们擅长处理特定情境下的特定词。

门控网络会挑出最适合给定输入的专家：

每个专家也不是一个完整的大型语言模型，只是大型语言模型架构里的一部分小模型。

专家

要搞清楚专家代表啥、咋工作的，咱们先瞅瞅混合专家模型MoE是想替代啥；就是密集层。

密集层

混合专家模型MoE都是从大型语言模型（LLM）一个比较基础的功能开始的，就是前馈神经网络（FFNN）。

得记得，标准的只解码器的Transformer架构里，前馈神经网络是在层归一化之后用的。

前馈神经网络（FFNN）能让模型利用注意力机制产生的上下文信息，进一步转换它，从而捕捉数据里更复杂的关系。

不过，前馈神经网络的大小可是会迅速变大的。为了学习这些复杂关系，它通常会扩充它接收到的输入：

稀疏层

传统Transformer里的前馈神经网络（FFNN）被称为密集模型，因为它所有的参数（权重和偏置）都被激活了。啥都不落下，全用来算输出。

要是仔细瞅瞅密集模型，你会发现输入会不同程度地激活所有参数：

相比之下，稀疏模型只激活一部分总参数，和混合专家模型MoE关系很近。

举个例子，我们可以把密集模型切块（这些块就是所谓的专家），重新训练它，然后在特定时间只激活一小部分专家：

背后的想法是，每个专家在训练时学不同的信息。然后在推理时，只用特定的专家，因为它们对特定任务最相关。

当被问问题时，我们可以挑出最适合特定任务的专家：

专家都学了些啥呢？

之前说过，专家学的不是整个领域的知识，而是更细致的信息，所以有时候把它们叫‘专家’还挺容易让人误解的。

在ST-MoE论文里，编码器模型的专家有专门的分工。不过，解码器模型的专家好像就没这种专门分工了。但这并不意味着所有专家都一样。

Mixtral 8x7B论文里有个很好的例子，每个词都用第一个选中的专家来上色。

这个图也说明了专家更关注语法，而不是特定领域。所以，虽然解码器专家好像没有专门的领域，但它们在处理某些类型的词时确实会一直被用到。

专家的架构

虽然把专家想象成密集模型里切块的隐藏层挺形象的，但它们通常本身就是完整的前馈神经网络（FFNN）：

因为大多数大型语言模型（LLM）有好几个解码器模块，一段文本在生成之前要经过多个专家处理：

选中的专家在不同词之间可能不一样，这就导致走不同的“路径”：

要是我们更新一下解码器模块的可视化，那现在里面就会有更多前馈神经网络（每个专家一个）：

解码器模块现在有多个前馈神经网络（每个都是一个“专家”），在推理时可以用它们。

路由机制

既然我们有了这么一帮专家，那模型咋知道该用哪个专家呢？在专家前面加了个门控网络，它会学着挑出给定词用哪个专家。

路由器

路由器（门控网络）也是前馈神经网络，根据特定输入来选专家。它会输出概率，用这些概率来挑出最匹配的专家：

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

路由器加上专家（只选几个）就组成了专家混合（MoE）层：

一个MoE层有两种尺寸，要么是稀疏的，要么是密集的专家混合。

两者都用路由器选专家，但稀疏MoE只挑几个，密集MoE虽然全选，但可能按不同分布来选。

比如，给定一串词，MoE会把词分给所有专家，而稀疏MoE就只挑几个专家。

在目前大型语言模型（LLM）的状况下，看到‘MoE’一般指稀疏MoE，因为它能用专家的一个子集。这样计算成本低，对大型语言模型来说很重要。

专家的选择

门控网络可以说是专家混合（MoE）里最重要的部分，因为它不仅决定了推理时用哪些专家，训练时也一样。

最基础的形式下，我们把输入（x）和路由器的权重矩阵（W）相乘：

然后，对输出应用 SoftMax 函数，为每个专家生成一个概率分布

路由器用这个概率分布来挑出给定输入最适合的专家。

最后，我们把每个路由器的输出和每个选中的专家的输出相乘，再把结果加起来。

咱们把所有步骤串起来，看看输入是怎么通过路由器和专家的：

路由的复杂性

不过，这个简单的函数常常会导致路由器总是选同一个专家，因为有些专家可能学得比其他专家快：

这不仅会导致专家选择的分布不均，还会使一些专家几乎不被训练，从而在训练和推理过程中引发问题。

为了解决这个问题，我们希望在训练和推理过程中，所有专家都能得到平等的重视，这被称为负载均衡。这样可以防止模型过度依赖少数几个专家，从而避免过拟合。

负载均衡

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

KeepTopK

一种负载均衡路由器的方法是通过一个简单的扩展，称为 KeepTopK。通过引入可训练的高斯噪声，有助于打破对特定专家的偏好，使选择更加随机化，这样我们可以防止总是选择相同的专家：

然后，除了你想要激活的顶部 k 个专家（例如 2 个）之外，其他专家的权重会被设置为 −∞：

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

SoftMax 函数的性质： SoftMax 函数将一个向量转换为概率分布，每个元素的值在 0 和 1 之间，且所有元素的和为 1。 SoftMax 函数的公式为：

其中，是输入向量，是向量的长度，是向量的第 个元素。 设置权重为 ：

当某个元素的值为 时，的值为 0。

因此，对于任何包含的向量，SoftMax 函数将这些元素的输出概率设置为 0。 具体步骤

假设我们有中间向量 ，其中某些元素被设置为：

应用 SoftMax 函数：

具体计算：由于e^{-\infty} = 0，上述表达式简化为：\text{SoftMax}(\math

计算具体值：\text{SoftMax}(\mathbf{z}_{\text{top-k}}) = \left, \frac{1.6487}{1.6487 + $2.22$55}, 0, \frac{2.2255}{1.6487 + 2.2255}, 0 \right] = \left[ 0, 0.4256, 0,0.5744, 0 \right]

结果

通过将权重设置为 ，SoftMax 函数确保了这些专家的激活概率为 0，从而只激活顶部 个专家。

优点

负载均衡：确保每个专家都有机会被激活，避免某些专家过度使用而其他专家闲置。

防止过拟合：通过引入噪声，防止模型过度依赖少数几个专家，从而提高模型的泛化能力。

灵活性：可以根据具体需求调整的值，灵活控制专家的选择数量。

KeepTopK 策略是许多大型语言模型（LLM）仍在使用的一种方法，尽管有许多有前景的替代方案。需要注意的是，KeepTopK 也可以在不添加额外噪声的情况下使用。

Token Choice

KeepTopK 策略将每个词路由到几个选定的专家。这种方法称为 Token Choice，允许给定的词被发送到一个专家（top-1 路由）或多个专家（top-k 路由）。

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

辅助损失

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

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

它增加了一个约束，迫使专家具有相等的重要性。 这个辅助损失的第一个组成部分是在整个批次中对每个专家的路由器值进行求和：

这为我们提供了每个专家的重要性分数，这些分数表示无论输入是什么，给定专家被选中的可能性。

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

例如，如果重要性分数差异很大，变异系数（CV）将会很高：

相反，如果所有专家的重要性分数相似，变异系数（CV）将会很低（这是我们所期望的）：

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

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

专家容量

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

例如，如果输入令牌不成比例地发送到一个专家而不是另一个专家，这可能会导致训练不足。

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

一个解决方案是限制给定专家可以处理的令牌数量，即专家容量（Expert Capacity）。当一个专家达到其容量时，后续的令牌将被发送给下一位专家。

如果两个专家都达到容量了，这个词就不会被任何专家处理，而是直接送到下一层。这就叫做词溢出。

简单来说，就好比有两个专家在处理任务，但他们都忙不过来了，手上的活儿已经堆得满满的。这时候再来的新任务，他们就没办法处理了，只能把这个任务往后传，让下一层的专家或者别的机制来处理。这就叫词溢出。

用 Switch Transformer 简化 MoE

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

Switch Transformer

Switch Transformer 是 Google 在 2022 年发表的一篇论文中提出的一种模型，它简化了 MoE 的路由算法，减少了计算量和通信量，同时支持 bfloat16 精度进行训练。基于 T5-Base 和 T5-Large 设计的模型在相同的算力下训练速度提升了 7 倍；同时发布了 1.6 万亿参数的 MoE 模型，相比 T5-XXL 模型训练速度提升了 4 倍。

模型介绍

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

简化稀疏路由

在传统的 MoE 层中，网络接收一个输入词元 ，输出会路由到 top-K 个专家进行计算。Switch Transformer 采用了一种简化的路由策略，即每次只发给一个专家，这样可以显著减少路由的计算量，同时保证模型的性能。

高效稀疏路由

Switch Transformer 通过以下方式实现高效稀疏路由：

路由计算量减少：只有一个专家被激活，减少了计算量。

专家容量减半：专家中的 batch_size 至少减半，进一步减少了计算量。

简化路由实现：减少了传统 MoE 方法中通信的代价。

分布式 Switch 实现

Switch Transformer 的分布式实现中，所有核心被分配给数据分区维度 ，这也对应于模型中的专家数量。对于每个核心的每个词元，路由器会本地计算分配给专家的分配。输出是一个大小为 的二进制矩阵，该矩阵在第一维度上进行分区，并确定专家的分配。

负载均衡损失

为了促使每个专家都可以拿到近似均匀分布的样本，Switch Transformer 引入了负载均衡损失。 当 时，损失是最小的。

切换层

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

简单来说，就是路由器在挑专家的时候，就是按照老规矩来，把输入数据和专家的权重相乘，然后用 SoftMax 函数处理一下，看看哪个专家最合适。没有啥复杂的操作，就是常规操作。

Switch Transformer 用了一种很简单的办法 来挑专家。以前的方法是让每个词元找好几个专家来处理，这种方法叫 top-k 路由。但 Switch Transformer 说，其实每个词元找一个专家就足够了，这种方法叫 top-1 路由。这样做的好处是，计算量变小了，通信成本也降低了，而且训练起来更稳定。

容量因子

容量因子就像是给每个专家分配任务的“额度”。这个值很重要，因为它决定了每个专家能处理多少个词元。Switch Transformer 在这方面做了改进，通过直接调整容量因子来控制每个专家的处理能力。这样做的好处是，可以更好地平衡每个专家的负载，避免有些专家忙不过来，而有些专家却很闲。

tokens per batch：每个批次中的令牌总数。

number of experts：专家的数量。

capacity factor：容量因子，通常大于 1，以提供额外的缓冲空间。

专家容量的组成部分其实挺简单的。专家容量主要由两个因素决定：每个批次中的令牌数量和专家的数量。

如果我们把容量因子调大，每个专家就能处理更多的词元。这听起来好像挺好的，因为专家们能干更多的活儿。

但是，如果把这个因子调得太大，就会浪费很多计算资源，因为有些专家可能会闲着没事干。反过来，如果容量因子太小，专家们就会忙不过来，处理不了所有的词元，这就叫令牌溢出。结果就是，模型的性能会变差，因为有些词元没被处理好。所以，容量因子得找个合适的值，这样才能让专家们既不过载，也不闲置，模型性能也能保持在最佳状态。

辅助损失

辅助损失这个东西，就是为了防止模型在处理数据时丢掉一些重要的信息。想象一下，每个专家都有自己的工作量，我们希望这些工作量分配得均匀一些。为了达到这个目的，我们引入了一个简化的辅助损失函数。

这个辅助损失函数不是去计算那些复杂的变异系数，而是直接看每个专家被分配到的令牌比例和路由器预测的概率。具体来说，我们希望每个专家处理的令牌数量和它们被选中的概率都差不多，也就是每个专家都能均匀地分到活儿干。

α 这个超参数就像是一个调节旋钮，用来控制这个辅助损失在训练过程中的重要性。如果 α 值设得太高，辅助损失就会太强势，把主要的损失函数都盖过去了；如果 α 值太低，那这个辅助损失就起不到什么作用，专家们的工作量还是不平衡。所以，α 值需要好好调整，才能让模型既不丢信息，又能让每个专家都忙得过来。

混合专家模型MoE在视觉模型中的应用

混合专家模型（MoE）这种技术不仅在语言模型里有用，在视觉模型里也能大显身手。比如 ViT（Vision-Transformer）这种模型，它用的是 Transformer 的架构，所以完全可以把 MoE 用起来。

简单来说，ViT 这个模型是把图像切成很多小块，这些小块就像是语言模型里的词元一样被处理。具体来说，比如一张 224x224 像素的图像，可以切成 16x16 像素的小块，这样就能得到 196 个小块。每个小块再被转换成一维向量，这样就能用 Transformer 来处理了。

这种方法的好处是，可以利用 Transformer 在处理序列数据上的强大能力，来处理图像这种二维数据。这样，视觉模型也能像语言模型一样，通过 MoE 来提高效率和可扩展性。

在 Vision Transformer（ViT）模型中，我们先把图像切成很多小块，这些小块就像是语言模型里的词元一样。然后，我们把这些小块转换成嵌入向量，这个过程有点像给每个小块加上一个标签，帮助模型更好地理解它们。

这个转换过程包括两个步骤：

线性变换：每个小块被展平成一维向量，然后通过一个线性变换（全连接层）映射到一个固定的维度 D，形成每个小块的嵌入向量。这个嵌入向量的维度就是 Transformer 的输入维度。

位置嵌入：为了帮助模型理解每个小块在图像中的位置，我们还会加上位置嵌入。位置嵌入是一个额外的向量，它包含了每个小块的位置信息。这样，模型不仅知道每个小块的内容，还知道它们在图像中的位置。

最后，这些带有位置嵌入的嵌入向量被输入到常规的 Transformer 编码器中，编码器会通过自注意力机制来处理这些向量，提取图像的全局特征。这个过程有点像语言模型处理词元序列一样，只不过这里是处理图像的小块序列。

当这些小块（也就是图像被切成的小方块）进入编码器的时候，它们的处理方式和语言模型里的词元一样。这就意味着，每个小块都会被当作一个独立的单元来处理，就像处理句子中的每个词一样。这种处理方式让 Vision Transformer（ViT）这种架构非常适合用混合专家（MoE）模型来优化。因为 MoE 模型本来就是设计来处理这种“分而治之”的任务的，每个专家可以专注于处理一部分数据，这样既能提高效率，又能保持高性能。

Vision-MoE

Vision-MoE（V-MoE）是图像模型中第一个实现 MoE 的例子之一。它把 ViT 中的密集前馈神经网络（FFNN）层换成了稀疏 MoE。这样，ViT 模型（通常比语言模型小）就能通过增加专家数量来大规模扩展。

因为图像通常包含很多小块，为了减少硬件限制，每个专家使用了一个较小的预定义专家容量。但是，较低的容量会导致一些小块被丢弃，就像令牌溢出一样。

为了应对这个问题，网络会为小块分配重要性分数，并优先处理这些小块，这样溢出的小块通常就不那么重要了。这个方法叫做批量优先级路由（Batch Priority Routing）。

所以，即使处理的小块数量减少了，重要的小块仍然会被路由。

优先级路由允许模型通过专注于最重要的小块来减少处理的小块数量。这样，模型在保持高性能的同时，还能减少计算资源的消耗。

活跃与稀疏参数：以 Mixtral 8x7B 为例

MoE（混合专家）模型的一个关键优势在于它的计算效率。

以 Mixtral 8x7B 为例，这个模型虽然总参数量很大，但在实际运行时，每次只激活一部分专家。这意味着，虽然模型有 47B 参数，但实际计算量相当于 12B 参数的密集模型。这种设计让模型在保持高性能的同时，大大减少了计算资源的消耗。

具体来说，MoE 模型在训练和推理时有以下优势：

训练速度更快：与相同参数规模的密集模型相比，MoE 可以用更少的计算资源和时间完成预训练。例如，Google 的 Switch Transformer 模型大小是 T5-XXL 的 15 倍，但在相同计算资源下，Switch Transformer 模型在达到固定困惑度 PPL 时，比 T5-XXL 模型快 4 倍。

推理速度快：虽然模型总参数量大，但推理时只激活部分专家网络。对于同等参数规模的密集模型，MoE 推理速度更快。例如，Mixtral 8x7B 虽然有 47B 参数，但推理时计算量相当于 12B 参数的密集模型。

扩展性强：MoE 允许模型在保持计算成本不变的情况下增加参数数量，这使得它能够扩展到非常大的模型规模，如万亿参数模型。

多任务学习能力：MoE 在多任务学习中表现出色，能够同时处理多种不同的任务，每个专家可以专注于特定的任务或数据类型。

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

Mixtral 8x7B 就是这样一个模型，它有很多参数（稀疏参数），但实际推理时只用到其中的一部分（活跃参数）。

具体来说，Mixtral 8x7B 模型有 8 个专家，每个专家有 5.6B 参数。虽然模型总参数量是 46.7B（8 个专家的参数加上共享参数），但在推理时，每个输入只用到 2 个专家，所以实际激活的参数量是 12.8B。这意味着，虽然加载整个模型需要很多显存，但实际运行时效率很高。

这种设计让 Mixtral 8x7B 在处理大规模数据时非常高效，既能保持高性能，又减少了计算资源的消耗。这就是为什么 MoE 模型在大型语言模型和视觉模型中越来越受欢迎的原因。

文章参考：https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts

本文转载自 ​​智驻未来​​，作者： 小智