Mixture of Experts 详解

随着 Mixtral 8x7B 的发布（公告，模型卡），一种类别的 transformer 已成为开放人工智能社区最热门的话题：Mixture of Experts，简称 MoE。 在这篇博文中，我们将介绍 MoE 的构成要素、如何训练 MoE 以及在使用 MoE 进行推理时需要考虑的权衡因素。

让我们深入了解一下！

目录

• 什么是 Mixture of Experts？
• MoEs 简史
• 什么是稀疏性？
• MoE 的负载平衡 tokens
• MoEs 和 Transformers
• Switch Transformers
• 使用 router Z-loss 让训练稳定
• 专家能学到什么？
• 增加专家数对预训练有何影响？
• 微调 MoEs
• 何时使用稀疏模型，何时使用密集模型？
• 让 "MoEs "沸腾起来
  • 专家并行
  • 容量系数和通信开销
  • 推理服务用到的技术
  • 高效训练
• 开源 MoEs
• 令人兴奋的工作方向
• 一些资源

简要说明

MoEs:

• 与密集模型相比，训练速度要快得多
• 与参数数量相同的模型相比，推理更快
• 要求 高 VRAM，因为所有专家都在内存中加载
• 面临很多微调方面的挑战，但是 MoE 最近的指令微调工作很有希望

让我们深入了解一下！

什么是专家混合体（MoE）？

模型的比例是提高模型质量的最重要轴心之一。 在计算预算固定的情况下，用较少的步骤训练较大的模型，比用较多的步骤训练较小的模型要好。

专家混合模型能以更少的计算量对模型进行预训练，这意味着你可以用与密集模型相同的计算预算来大幅扩大模型或数据集的规模。 特别是，在预训练过程中，MoE 模型应能更快地达到与密集模型相同的质量。

那么，MoE 究竟是什么呢？ 在 transformer 模型中，MoE 由两个主要元素组成：

• 使用稀疏 MoE 层代替密集前馈网络（FFN）层。 MoE 层有一定数量的 "专家"（如 8 个），每个专家是一个神经网络。 在实践中，专家是 FFNs，但它们也可以是更复杂的网络，甚至是 MoE 本身，从而形成分层的 MoE！
• 门网络或路由网络，决定将哪些 tokens 发送给哪些专家。 例如，在下图中，词符 "More" 被发送到第二个专家，词符 "Parameters "被发送到第一个网络。 正如我们稍后将探讨的，我们可以将一个词符发送给多个专家。 如何将词符路由到专家是使用 MoE 时的重大决定之一 — 路由网络由学习到的参数组成，并与网络的其他部分同时进行预训练。

因此，概括地说，在 MoE 中，我们用 MoE 层取代了 transformer 模型的每个 FFN 层，后者由一个门网络和一定数量的专家组成。

虽然 MoE 与密集模型相比具有高效预训练和快速推理等优点，但它们也面临着挑战：

• 训练： MoE 可以大大提高预训练的计算效率，但在微调过程中，MoE 一直难以实现泛化，从而导致过度拟合。
• 推理： 虽然一个 MoE 可能有很多参数，但在推理过程中只使用其中的一部分。 与参数数量相同的密集模型相比，推理速度要快得多。 不过，所有参数都需要加载到 RAM 中，因此对内存的要求很高。 例如，考虑到 Mixtral 8x7B 这样的 MoE，我们需要有足够的 VRAM 来容纳密集的 47B 参数模型。 为什么是 47B 参数，而不是 8 x 7B = 56B？ 这是因为在 MoE 模型中，只有 FFN 层被视为单独的专家，其余的模型参数都是共享的。 同时，假设每个 Token 只使用两个专家，那么推理速度（FLOPs）就相当于使用了 12B 模型（而不是 14B 模型），因为它计算的是 2x7B 矩阵乘法，只是共享了一些层（稍后将详细介绍）。

既然我们已经对 "MoE "有了一个大致的了解，那就让我们来看看导致发明出来 "MoE "的研发过程吧。

MoEs 简史

MoE 的起源来自 1991 年的论文 局部专家的适应性混合。 这种想法与集成方法类似，都是为不同网络组成的系统制定一个监督程序，每个网络处理不同的训练样本子集。 每个独立的网络或专家专门负责输入空间的不同区域。 如何选择专家？ 门控网络决定每个专家的权重。 在训练过程中，专家和门控都要接受训练。

2010-2015 年间，两个不同的研究领域促进了后来 MoE 的进步：

• 作为组件的专家：在传统的 MoE 设置中，整个系统由门控网络和多个专家组成。 在 SVM、高斯过程和其他方法中，对作为整体模型的 MoE 进行了探索。 Eigen、Ranzato 和 Ilya 的研究将 MoE 作为更深层网络的组成部分进行了探讨。 这样就可以将 MoE 作为多层网络中的层，使模型既庞大又高效。
• 条件计算：传统网络通过每一层处理所有输入数据。 在此期间，Yoshua Bengio 研究了根据输入 token 动态激活或停用组件的方法。

这些工作促使我们在 NLP 的背景下探索专家的混合。 具体来说 沙泽尔等人 (2017年，"等人 "包括 Geoffrey Hinton 和 Jeff Dean，Google’s Chuck Norris）通过引入稀疏性，将这一想法扩展到137B LSTM（当时事实上的NLP架构，由Schmidhuber创建），从而即使在大规模下也能保持极快的推理速度。 这项工作的重点是翻译，但面临许多挑战，如通信成本高和训练不稳定。

MoE 可以训练数万亿参数模型，如开源的 1.6T 参数 Switch Transformers 等。 计算机视觉领域也对 MoE 进行了探索，但本博文将侧重于 NLP 领域。

什么是稀疏性？

稀疏性使用了条件计算的理念。 在密集模型中，所有参数都用于所有输入，而稀疏性允许我们只运行整个系统的某些部分。

让我们深入了解一下 Shazeer 对翻译的 MoEs 的探索。 有条件计算的理念（网络的某些部分基于每个样本激活）使我们可以在不增加计算量的情况下扩大模型的规模，因此，这导致每个 MoE 层使用了成千上万的专家。

这种设置带来了一些挑战。 例如，虽然大的 batch size 通常更有利于提高性能，但 MOE 中的 batch size 会在数据流过活动专家时而显著减少。 例如，如果我们的批次输入包含 10 个 token，其中 5 个 token 可能在一个专家那里结束，而另外 5 个 token 可能在 5 个不同的专家那里结束，从而导致批次大小不均和利用率不足。 下面的 Making MoEs go brrr 部分将讨论其他挑战和解决方案。

如何解决这个问题？ 可学习门控网络（G）决定向哪些专家（E）发送部分输入：

$$y=\sum _{i=1}^{n}G(x{)}_i{E}_i(x)$$

在这种情况下，所有专家都要对所有输入进行运算，这是一种加权乘法。 但是，如果 G 为 0 会怎样呢？ 如果是这样，就不需要计算相应的专家运算，因此我们节省了计算量。 典型的门控函数是什么样子的？ 在最传统的设置中，我们只使用一个带有 softmax 函数的简单网络。 网络将学习向哪位专家发送输入。

$$G_{\sigma }(x)=\text{Softmax}(x\cdot W_g)$$

Shazeer 的研究还探索了其他门控机制，如噪声 Top-k 门控。 这种门控方法会引入一些（可调整的）噪声，然后保留最高的 k 值。 那就是

1. 我们添加一些噪音

$$H(x{)}_i=(x\cdot W_{\text{g}}{)}_i+\text{StandardNormal()}\cdot \text{Softplus}((x\cdot W_{\text{noise}}{)}_i)$$

2. 我们只选取前 k

$$$$KeepTopK(v,k)i=。{viif vi is in the top k elements of v,vi位于 k 元素的顶部。−∞否则

3. 我们应用 softmax。

$$G(x)=\text{Softmax}(\text{KeepTopK}(H(x),k))$$

这种稀疏性带来了一些有趣的特性。 通过使用足够小的 k（例如 1 或 2），我们可以比激活许多专家更快地完成训练和运行推理。 为什么不直接选择最顶尖的专家？ 最初的猜想是，要让 "闸门 "学会如何将路由选择给不同的专家，就需要将路由选择给不止一位专家，因此至少要选择两位专家。 Switch Transformers 一节重新讨论了这一决定。

为什么要增加噪音？ 这是为了平衡负载！

MoE 的 tokens 负载均衡

如前所述，如果我们的所有 tokens 都只发送给少数几个受欢迎的专家，这将导致训练效率低下。 在正常的 MoE 训练中，门控网络会趋于主要激活相同的几个专家。 这种情况会自我强化，因为受青睐的专家会更快地得到训练，从而被更多地选中。 为了减轻这种情况，我们添加了辅助损失，以鼓励给予所有专家同等的重视。 这种损失可确保所有专家获得大致相同数量的训练示例。 接下来的章节还将探讨专家能力的概念，该概念引入了专家可以处理多少 tokens 的阈值。 在 transformers 中，辅助损失通过 aux_loss 参数开放出来。

MoEs 和 Transformers

Transformers 是一个非常明显的例子，增加参数数量可以提高性能，因此 Google 通过 GShard对此进行探索也就不足为奇了，他们将 transformers 扩展到了 6000 亿个参数。

GShard 用 MoE 层取代其他每个 FFN 层，并在编码器和解码器中都使用 top-2 门控。 下一张图片显示了编码器部分的外观。 这种设置对大规模计算非常有利：当我们扩展到多个设备时，MoE 层会在设备间共享，而所有其他层都会被复制。 "Making MoEs go brrr" 部分将对此作进一步讨论。

为了保持均衡负载和规模效率，GShard 的作者除了采用与上一节讨论的类似的辅助损失外，还引入了一些变化：

• 随机路由：在 top-2 设置中，我们总是挑选第一名专家，但挑选第二名专家的概率与其权重成正比。
• 专家容量：我们可以设定一个专家可以处理多少 tokens 的阈值。 如果两个专家容量都已满，tokens 就会被视为溢出，并通过残差连接发送到下一层（在其他项目中是完全丢弃）。 这一概念将成为 MoEs 最重要的概念之一。 为什么需要专家容量？ 由于所有张量形状都是在编译时静态确定的，但我们无法提前知道每个专家将获得多少 tokens，因此我们需要固定容量因子。

GShard 论文通过表达并行计算模式为 MoE 做出了贡献，但这不在本博文讨论范围之内。

请注意： 当我们运行推理时，只有部分专家会被触发。 同时，还有一些共享计算，如 self-attention，它适用于所有 tokens。 这就是为什么当我们谈论由 8 位专家组成的 47B 模型时，我们可以使用 12B 密集模型的计算量。 如果使用 top-2，将使用 14B 个参数。 但考虑到注意力操作是共享的（除其他外），实际使用的参数数量为 12B。

Switch Transformers

尽管 MoEs 前景广阔，但它们在训练和微调的不稳定性方面却举步维艰。 Switch Transformers是一项深入探讨这些主题的令人兴奋的工作。 作者甚至还在 Hugging Face 上发布了一个 1.6 万亿参数的 MoE，其中有 2048 个专家，你可以使用 transformers 运行。 Switch Transformers 的预训练速度比 T5-XXL 提高了 4 倍。

就像在 GShard 中一样，作者用 MoE 层取代了 FFN 层。 Switch Transformers 论文提出一个 Switch Transformer 层，该层接收两个输入（两个不同的 token），有四个专家。

与最初的至少使用两名专家想法相反，Switch Transformers 采用简化的单一专家策略。 这种方法的效果是

• 路由网络计算量减少
• 每个专家的 batch size 至少可以减半
• 降低通信成本
• 保持质量

Switch Transformers 还探讨了专家容量的概念。

$$\text{Expert Capacity}=\left(\frac{\text{tokens per batch}}{\text{number of experts}}\right)\times \text{capacity factor}$$

上面建议的容量将批次中的 tokens 数量平均分配给专家。 如果我们使用大于 1 的容量系数，就能在 tokens 不完全平衡时提供缓冲。 增大容量会导致设备间通信费用增加，因此要注意权衡利弊。 特别是，Switch Transformers 在低容量系数（1-1.25）条件下表现出色。

Switch Transformer 的作者还重新审视并简化了章节中提到的负载平衡损失。 在训练过程中，每个 Switch 层的辅助损失都会被添加到模型总损失中。 这种损耗鼓励均衡的路由，并可使用超参数进行加权。

作者还尝试了选择性精度，例如使用 bfloat16 训练专家，而在其余计算中使用全精度。 较低的精度可降低处理器之间的通信成本、计算成本和用于存储张量的内存。 在最初的实验中，专家和门控网络都是在 bfloat16 中进行训练的，但训练结果更不稳定。 这主要是由于路由网络的计算：由于路由网络具有指数函数，因此更高的精度非常重要。 为了减少不稳定性，路由也采用了全精度。

这份笔记展示了微调 Switch Transformers 的总结，我们建议首先查看 微调部分。

Switch Transformers 使用一种编码器-解码器设置，他们在其中制作了与 T5 相对应的 MoE。 GLaM论文探讨了通过使用1/3的能源来训练一个与GPT-3质量相匹配的模型，从而扩大这些模型的规模（是的，由于训练MoE所需的计算量较低，它们可以将碳足迹减少多达一个数量级）。 作者们关注的重点是纯解码器模型以及少量和一次评估，而不是微调。 他们使用 Top-2 路由和更大的容量系数。 此外，他们还探索了容量因子，将其作为一个可以在训练和评估过程中根据想要使用的计算量而改变的指标。

使用路由网络 Z 损失进行稳定训练

前面讨论过的平衡损失会导致不稳定问题。 我们可以用很多方法来稳定稀疏模型，但这是以牺牲质量为代价的。 例如，引入 dropout 可以提高稳定性，但会导致模型质量下降。 另一方面，增加更多的乘法成分可以提高质量，但会降低稳定性。

ST-MoE 中引入的路由网络 z 损失通过惩罚进入门控网络的大对数，在不降低质量的情况下显著提高了训练稳定性。 由于这种损失会促使数值的绝对值变小，从而减少舍入误差，这对门控等指数函数的影响很大。 我们建议阅读该文件了解详情。

专家能学到什么？

ST-MoE 的作者观察到，编码专家专门研究一组 tokens 或浅层概念。 例如，我们可能会遇到标点符号专家、专有名词专家等。另一方面，解码专家的专业化程度较低。 作者还在多语言环境中进行了训练。 虽然我们可以想象每个专家都擅长于一种语言，但事实恰恰相反：由于 token 路由和负载平衡，没有一个专家擅长于任何一种语言。

增加专家数对预训练有何影响？

专家数量越多，采样效率就越高，速度也就越快，但收益也就越小（尤其是在 256 或 512 专家之后），而且推理还需要更多的 VRAM。 Switch Transformers 研究表明，这个特性在大规模网络和小规模网络上一致的，即使每层有 2 个、4 个或 8 个专家。

微调 MoEs

Mixtral 受 transformers 4.36.0 版支持。 您可以使用 pip install "transformers==4.36.0 --upgrade 安装它。

密集模型和稀疏模型的过拟合动态截然不同。 稀疏模型更容易出现过度拟合，因此我们可以在专家本身内部探索更高的正则化（如滤除）（例如，我们可以为密集层设定一个滤除率，为稀疏层设定另一个更高的滤除率）。

一个问题是，是否利用辅助损耗进行微调。 ST-MoE 作者曾尝试关闭辅助损耗，结果发现即使丢弃多达 11% 的标记，质量也不会受到明显影响。 丢弃标记可能是正则化的一种形式，有助于防止过度拟合。

Switch Transformers 观察到，在固定的预训练困惑度下，稀疏模型在下游任务中的表现不如密集模型，尤其是在推理繁重的任务中，如 SuperGLUE。 另一方面，对于 TriviaQA 等知识含量高的任务，稀疏模型的表现却好得不成比例。 作者还注意到，较少的专家人数有助于微调。 另一个证实了泛化问题的观察结果是，模型在较小的任务中表现较差，但在较大的任务中表现良好。

我们可以尝试冻结所有非专家权重。 也就是说，我们只更新 MoE 图层。 这会导致性能大幅下降。 我们可以尝试相反的方法：只冻结 MoE 图层中的参数，这与更新所有参数的效果几乎一样好。 这有助于加快微调速度并减少内存。 这可能有点违背直觉，因为 80% 的参数都在 MoE 层（在 ST-MoE 项目中）。 他们对这种结构的假设是，由于专家层每 1/4 层才出现一次，而每个标记每层最多看到两个专家，因此更新 MoE 参数对各层的影响要比更新其他参数小得多。

在微调稀疏 MoE 时需要考虑的最后一个问题是，它们有不同的微调超参数设置--例如，稀疏模型往往更受益于较小的批次规模和较高的学习率。

此时此刻，你可能会为人们一直在努力微调 MoE 而感到有些难过。 令人兴奋的是，最近的一篇论文MoEs Meets Instruction Tuning（2023 年 7 月）进行了实验：

• 单一任务微调
• 多任务指令调整
• 多任务指令调整后进行单任务微调

当作者对 MoE 和 T5 等效值进行微调时，T5 等效值的效果更好。 当作者对 Flan T5（相当于 T5 指示灯）模拟电子显微镜进行微调时，模拟电子显微镜的性能明显提高。 不仅如此，Flan-MoE 相对于 MoE 的改进幅度大于 Flan T5 相对于 T5 的改进幅度，这表明 MoE 可能比密集模型更受益于指令调整。 任务越多，MoEs 受益越大。 与之前建议关闭辅助损失函数的讨论不同，损失实际上是为了防止过度拟合。

何时使用稀疏模型，何时使用密集模型？

专家适用于有许多机器的高吞吐量方案。 在预训练的计算预算固定的情况下，稀疏模型将更为理想。 对于 VRAM 较少的低吞吐量场景，使用密集模型会更好。

请注意： 我们不能直接比较稀疏模型和稠密模型的参数数量，因为二者所代表的内容明显不同。

让 "MoEs "沸腾起来

最初的 MoE 工作将 MoE 层作为一个分支设置，由于 GPU 并非为此而设计，导致计算速度缓慢，并且由于设备需要向其他设备发送信息，导致网络带宽成为瓶颈。 本节将讨论现有的一些工作，以使这些模型的预训练和推理更加实用。 MoEs go brrrrr.

并行性

让我们简要回顾一下并行性：

• 数据并行 在所有内核上复制相同的权重，并在各内核上分割数据。
• 模型并行： 模型在不同内核间分区，数据在不同内核间复制。
• 模型和数据并行 我们可以将模型和数据划分到不同的内核上。 请注意，不同的内核处理不同批次的数据。
• 专家并行：专家被安排在不同的工作者上。 如果与数据并行相结合，每个内核都有不同的专家，数据在所有内核之间进行分割

在专家并行模式下，专家被安排在不同的工作者身上，每个工作者负责不同批次的训练样本。 对于非 MoE 层，专家并行性与数据并行性表现相同。 对于 MoE 层，序列中的令牌会被发送到所需专家所在的工作站。

容量系数和通信费用

提高容量因子 (CF) 可以提高质量，但会增加通信成本和激活记忆。 如果全对全通信速度较慢，使用较小的容量因子会更好。 一个好的起点是使用容量系数为 1.25 的 top-2 路由，每个核心配备一名专家。 在评估过程中，可以更改容量系数以减少计算量。

推理技术

您可以将 mistralai/Mixtral-8x7B-Instruct-v0.1 部署到推理端点。

模拟电子设备的一大缺点是参数过多。 在本地使用情况下，可能需要使用较小的型号。 让我们快速讨论一些有助于推理服务的技巧：

• Switch Transformers 的作者进行了早期的蒸馏实验。 通过将一个 MoE 提炼回其密集的对应物，他们可以保留 30-40% 的稀疏性收益。 因此，蒸馏法的好处是可以更快地进行预处理，并在生产中使用更小的型号。
• 最近的方法修改了路由，将完整的句子或任务路由给专家，允许提取子网络为专家服务。
• 专家汇总 (MoE)：这一技术可合并专家权重，从而减少推理时的参数数量。

关于高效训练的更多信息

FasterMoE（2022 年 3 月）分析了高效分布式系统中 MoE 的性能，并分析了不同并行策略的理论极限，还分析了倾斜专家受欢迎程度的技术、减少延迟的细粒度通信调度，以及根据最低延迟挑选专家的调整拓扑感知门，从而使速度提高了 17 倍。

Megablocks（2022 年 11 月）通过提供能处理 MoEs 中存在的动态性的新 GPU 内核，探索高效的稀疏预训练。 他们的建议从不丢弃令牌，并能有效地映射到现代硬件，从而大大加快了速度。 有什么诀窍？ 传统的 MoE 使用分批矩阵乘法，假定所有专家都具有相同的形状和相同的代币数量。 相比之下，Megablocks 将 MoE 层表示为块稀疏操作，可以适应不平衡分配。

开源 MoE

现在有几个开源项目可以训练 MoEs：

• Megablocks：https://github.com/stanford-futuredata/megablocks
• Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm
• OpenMoE: https://github.com/XueFuzhao/OpenMoE

在已发布的公开获取 MoE 方面，你可以检查：

• Switch Transformers (Google)：基于 T5 的 MoE 系列，专家数从 8 到 2048 不等。 最大的模型有 1.6 万亿个参数。
• NLLB MoE (Meta)：NLLB 翻译模型的 MoE 变体。
• OpenMoE：社区努力发布了基于 Llama 的 MoE。
• Mixtral 8x7B (Mistral)：高质量的 MoE，性能优于 Llama 2 70B，推理速度更快。 此外，还发布了一个指导性调谐模型。 了解更多详情，请参阅公告博文。

令人兴奋的工作方向

进一步实验将稀疏的 MoE 分解为参数较少但参数数量相似的稠密模型。

另一个领域是 MoE 的量化。 QMoE（2023 年 10 月）通过将 MoE 量化到每个参数小于 1 位，从而将使用 3.2TB 加速器的 1.6T 交换变压器压缩到仅 160GB，这是朝这个方向迈出的良好一步。

因此，简而言之，一些有趣的领域值得探索：

• 将 Mixtral 提炼为密集模型
• 探索专家的模型合并技术及其对推理时间的影响
• 执行 Mixtral 的极端量化技术

一些资源

• 适应性地方专家混合物（1991 年）
• 在专家深度混合中学习因子表征（2013 年）
• 超大型神经网络：稀疏门控的专家混合层（2017）
• GShard：利用条件计算和自动分片技术扩展巨型模型（2020 年 6 月）
• GLaM：利用专家混合物高效扩展语言模型（2021 年 12 月）
• 开关变压器：以简单高效的稀疏性扩展到万亿参数模型（2022 年 1 月）
• ST-MoE：设计稳定且可转移的稀疏专家模型（2022 年 2 月）
• FasterMoE：大规模动态预训练模型的建模和优化训练（2022 年 4 月）
• MegaBlocks：利用专家混合物进行高效稀疏训练（2022 年 11 月）
• 专家混合与指令调整：大型语言模型的成功组合（2023 年 5 月）
• Mixtral-8x7B-v0.1、Mixtral-8x7B-Instruct-v0.1。

引用

@misc {sanseviero2023moe,
    author       = { Omar Sanseviero and
                     Lewis Tunstall and
                     Philipp Schmid and
                     Sourab Mangrulkar and
                     Younes Belkada and
                     Pedro Cuenca
                   },
    title        = { Mixture of Experts Explained },
    year         = 2023,
    url          = { https://huggingface.co/blog/moe },
    publisher    = { Hugging Face Blog }
}

Sanseviero, et al., "Mixture of Experts Explained", Hugging Face Blog, 2023.