专家混合（MoE）详解

随着 Mixtral 8x7B（公告，模型卡）的发布，一类 Transformer 模型成为开放 AI 社区最热门的话题：专家混合（Mixture of Experts），简称 MoE。在这篇博客文章中，我们将深入探讨 MoE 的基本组成部分、它们的训练方式以及在推理时需要考虑的权衡。

让我们深入了解！

目录

• 什么是专家混合？
• MoE 的简要历史
• 什么是稀疏性？
• MoE 的令牌负载均衡
• MoE 与 Transformers
• Switch Transformers
• 通过路由器 Z-损失稳定训练
• 专家学到了什么？
• 扩展专家数量如何影响预训练？
• MoE 的微调
• 何时使用稀疏 MoE 与密集模型？
• 让 MoE 加速运行
  • 专家并行
  • 容量因子和通信成本
  • 服务技术
  • 高效训练
• 开源 MoE
• 令人兴奋的工作方向
• 一些资源

摘要

MoE 模型

• 比密集模型预训练速度快得多
• 与相同参数数量的模型相比，推理速度更快
• 需要高 VRAM，因为所有专家都加载到内存中
• 在微调方面面临诸多挑战，但近期关于 MoE 指令微调的研究很有前景

让我们深入了解！

什么是专家混合（MoE）？

模型规模是提高模型质量最重要的因素之一。在固定的计算预算下，训练一个更大的模型，步数更少，比训练一个更小的模型，步数更多，效果更好。

专家混合（MoE）使得模型能够以更少的计算量进行预训练，这意味着在相同的计算预算下，你可以大幅扩展模型或数据集的大小，而使用密集模型则无法做到。具体而言，一个 MoE 模型在预训练期间应该比其密集模型更快地达到相同的质量。

那么，MoE 到底是什么？在 Transformer 模型的背景下，MoE 由两个主要元素组成：

• 稀疏 MoE 层用于代替密集的全连接网络（FFN）层。MoE 层包含一定数量的“专家”（例如 8 个），每个专家都是一个神经网络。实际上，这些专家是 FFN，但它们也可以是更复杂的网络，甚至是 MoE 本身，从而形成分层 MoE！
• 一个门控网络或路由器，它决定哪些令牌被发送到哪个专家。例如，在下图中，“More”令牌被发送到第二个专家，而“Parameters”令牌被发送到第一个网络。正如我们稍后将探讨的，我们可以将一个令牌发送到多个专家。如何将令牌路由到专家是使用 MoE 时的一个重要决策——路由器由可学习参数组成，并与网络的其余部分同时进行预训练。

总而言之，在 MoE 中，我们用 MoE 层取代了 Transformer 模型的每个 FFN 层，该 MoE 层由一个门控网络和一定数量的专家组成。

尽管 MoE 提供了高效预训练和比密集模型更快的推理等优势，但它们也带来了挑战：

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

现在我们对 MoE 有了一个大致的了解，接下来让我们看看导致其发明的研究进展。

MoE 的简要历史

MoE 的根源可以追溯到 1991 年的论文《局部专家的自适应混合》。其思想类似于集成方法，即为由独立网络组成的系统提供一个监督程序，每个网络处理训练案例的不同子集。每个独立的网络或专家专注于输入空间的不同区域。如何选择专家？一个门控网络决定每个专家的权重。在训练期间，专家和门控都被训练。

在 2010-2015 年间，两个不同的研究领域为 MoE 的后续发展做出了贡献：

• 专家作为组件：在传统的 MoE 设置中，整个系统由一个门控网络和多个专家组成。MoE 作为整体模型已被应用于 SVM、高斯过程和其他方法。Eigen、Ranzato 和 Ilya 的工作探讨了将 MoE 作为更深层网络的组件。这使得在多层网络中拥有 MoE 层成为可能，从而使模型能够同时兼具大型和高效的特点。
• 条件计算：传统网络通过每一层处理所有输入数据。在此期间，Yoshua Bengio 研究了根据输入令牌动态激活或停用组件的方法。

这些工作促使人们在 NLP 领域探索专家混合。具体来说，Shazeer 等人（2017 年，其中的“等人”包括 Geoffrey Hinton 和 Jeff Dean，Google 的 Chuck Norris）通过引入稀疏性，将这一思想扩展到 137B 的 LSTM（当时事实上的 NLP 架构，由 Schmidhuber 创建），从而即使在大型模型中也能保持极快的推理速度。这项工作主要集中在翻译领域，但也面临许多挑战，例如高通信成本和训练不稳定性。

MoE 使得训练万亿参数模型成为可能，例如开源的 1.6T 参数 Switch Transformers 等。MoE 也已在计算机视觉领域进行了探索，但本博客文章将重点关注 NLP 领域。

什么是稀疏性？

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

让我们深入探讨 Shazeer 对翻译 MoE 的探索。条件计算的思想（网络的一部分根据每个示例动态激活或停用）允许在不增加计算量的情况下扩展模型大小，因此，这使得每个 MoE 层中使用数千个专家成为可能。

这种设置带来了一些挑战。例如，尽管通常大批量大小对性能更好，但 MoE 中的批量大小会随着数据流经活跃专家而有效减少。例如，如果我们的批处理输入包含 10 个令牌，其中五个令牌可能进入一个专家，而另外五个令牌可能进入五个不同的专家，导致批量大小不均匀和利用率不足。下面的“让 MoE 加速运行”部分将讨论其他挑战和解决方案。

我们如何解决这个问题？一个学习到的门控网络（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 的工作还探讨了其他门控机制，例如 Noisy Top-k Gating。这种门控方法引入了一些（可调的）噪声，然后保留了 Top-k 值。也就是说：

1. 我们添加了一些噪声：

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

2. 我们只选择 Top k：

$$\text{KeepTopK}(v, k)_i = \begin{cases} v_i & \text{if } v_i \text{ is in the top } k \text{ elements of } v, \\ -\infty & \text{otherwise.} \end{cases}$$

3. 我们应用 Softmax。

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

这种稀疏性引入了一些有趣的特性。通过使用足够低的 k 值（例如一个或两个），我们可以比激活许多专家时更快地进行训练和推理。为什么不只选择顶级专家？最初的推测是，路由到多个专家是必要的，以使门控网络学习如何路由到不同的专家，因此至少必须选择两个专家。Switch Transformers 部分重新审视了这一决策。

为什么要添加噪声？这是为了负载均衡！

MoE 的令牌负载均衡

如前所述，如果所有令牌都只发送给少数几个热门专家，那会使训练效率低下。在正常的 MoE 训练中，门控网络会收敛到主要激活相同的几个专家。这种现象会自我强化，因为受青睐的专家训练得更快，因此被选择得更多。为了缓解这个问题，添加了一个辅助损失来鼓励所有专家获得同等的重要性。这个损失确保所有专家接收到大致相同数量的训练样本。以下章节还将探讨专家容量的概念，该概念引入了专家可以处理多少令牌的阈值。在 transformers 中，辅助损失通过 aux_loss 参数暴露。

MoE 与 Transformers

Transformer 模型很明显地证明了参数数量的扩展能够提高性能，因此 Google 在 GShard 中探索了这一点，它将 Transformer 模型扩展到超过 6000 亿个参数，这并不令人惊讶。

GShard 在编码器和解码器中都用 MoE 层（使用 top-2 门控）替换了隔层 FFN 层。下图展示了编码器部分的这种设置。这种设置对于大规模计算非常有利：当我们扩展到多个设备时，MoE 层在设备之间共享，而所有其他层都进行复制。这将在“让 MoE 加速运行”部分进一步讨论。

为了在大规模下保持负载均衡和效率，GShard 作者除了引入类似于上一节中讨论的辅助损失之外，还引入了一些更改：

• 随机路由：在 top-2 设置中，我们总是选择排名第一的专家，而第二个专家则以与其权重成比例的概率进行选择。
• 专家容量：我们可以设置一个阈值，限制一个专家可以处理多少个令牌。如果两个专家都达到了容量，则该令牌被视为溢出，并通过残差连接发送到下一层（或者在其他项目中完全丢弃）。这个概念将成为 MoE 最重要的概念之一。为什么需要专家容量？因为所有张量形状在编译时都是静态确定的，但我们无法提前知道有多少令牌会发送到每个专家，所以我们需要固定容量因子。

GShard 论文通过表达适用于 MoE 的并行计算模式做出了贡献，但讨论这一点超出了本博客文章的范围。

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

Switch Transformers

尽管 MoE 展现出巨大的潜力，但它们在训练和微调稳定性方面面临挑战。Switch Transformers 是一项非常激动人心的工作，深入探讨了这些主题。作者甚至在 Hugging Face 上发布了一个拥有 2048 个专家、1.6 万亿参数的 MoE 模型，你可以使用 transformers 库运行它。Switch Transformers 在预训练速度上比 T5-XXL 提升了 4 倍。

正如 GShard 中一样，作者用 MoE 层替换了 FFN 层。Switch Transformers 论文提出了一种 Switch Transformer 层，它接收两个输入（两个不同的令牌）并拥有四个专家。

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

• 路由器计算量减少
• 每个专家的批量大小至少可以减半
• 通信成本降低
• 质量得以保持

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

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

上述建议的容量将批处理中的令牌数量平均分配给专家数量。如果使用大于 1 的容量因子，则在令牌未完美平衡时提供了一个缓冲区。增加容量将导致更昂贵的设备间通信，因此需要权衡。特别是，Switch Transformers 在低容量因子（1-1.25）下表现良好。

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

作者还尝试了选择性精度，例如用 bfloat16 训练专家，而其余计算则使用全精度。较低精度可以降低处理器之间的通信成本、计算成本以及存储张量的内存。最初的实验中，专家和门控网络都用 bfloat16 训练，但导致训练更不稳定。这尤其归因于路由器计算：由于路由器具有指数函数，因此更高的精度很重要。为了缓解不稳定性，路由器也使用了全精度。

这个笔记本展示了如何微调 Switch Transformers 进行摘要，但我们建议首先回顾微调部分。

Switch Transformers 使用编码器-解码器设置，其中他们制作了 T5 的 MoE 版本。GLaM 论文探索了通过训练一个匹配 GPT-3 质量的模型来提高这些模型的规模，该模型仅使用 1/3 的能量（是的，由于训练 MoE 所需的计算量更少，它们可以将碳足迹减少一个数量级）。作者专注于仅解码器模型以及少量样本和单样本评估，而不是微调。他们使用了 Top-2 路由和更大的容量因子。此外，他们探索了容量因子作为一个可以在训练和评估期间根据所需的计算量进行更改的指标。

通过路由器 Z-损失稳定训练

前面讨论的平衡损失可能导致不稳定性问题。我们可以使用许多方法来稳定稀疏模型，但代价是质量下降。例如，引入 dropout 可以提高稳定性，但会导致模型质量损失。另一方面，添加更多的乘性组件可以提高质量，但会降低稳定性。

路由器 Z-损失（由 ST-MoE 引入）通过惩罚进入门控网络的大 logits，显著提高了训练稳定性，而不会降低质量。由于这种损失鼓励值保持较小的绝对幅度，因此舍入误差会减少，这对于指数函数（如门控）可能产生相当大的影响。我们建议查阅论文以获取详细信息。

专家学到了什么？

ST-MoE 的作者观察到，编码器专家专注于一组令牌或浅层概念。例如，我们可能最终得到一个标点专家、一个专有名词专家等。另一方面，解码器专家的专业化程度较低。作者还在多语言设置中进行了训练。尽管人们可能会想象每个专家都专注于一种语言，但情况恰恰相反：由于令牌路由和负载均衡，没有任何一个专家专门针对任何给定语言。

扩展专家数量如何影响预训练？

更多的专家可以提高样本效率和加速，但这些增益是递减的（尤其是在 256 或 512 之后），并且推理需要更多的 VRAM。在 Switch Transformers 中大规模研究的特性在小规模下也是一致的，即使每层只有 2、4 或 8 个专家。

MoE 的微调

Mixtral 支持 transformers 4.36.0 版本。您可以通过 pip install transformers==4.36.0 --upgrade 进行安装。

密集模型和稀疏模型之间的过拟合动态非常不同。稀疏模型更容易过拟合，因此我们可以在专家内部探索更高的正则化（例如 dropout）（例如，我们可以为密集层设置一个 dropout 率，为稀疏层设置另一个更高的 dropout 率）。

一个问题是是否在微调时使用辅助损失。ST-MoE 的作者尝试关闭辅助损失，即使有高达 11% 的令牌被丢弃，质量也没有受到显著影响。令牌丢弃可能是一种有助于防止过拟合的正则化形式。

Switch Transformers 观察到，在固定的预训练困惑度下，稀疏模型在下游任务中表现不如密集模型，尤其是在 SuperGLUE 等推理密集型任务上。另一方面，对于 TriviaQA 等知识密集型任务，稀疏模型表现得异常出色。作者还观察到，少量专家有助于微调。另一个证实泛化问题的观察是，模型在较小任务中表现较差，但在较大任务中表现良好。

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

微调稀疏 MoE 时要考虑的最后一点是，它们具有不同的微调超参数设置——例如，稀疏模型往往更能受益于较小的批量大小和较高的学习率。

到目前为止，你可能对人们在微调 MoE 方面遇到的困难感到有些沮丧。令人兴奋的是，一篇最近的论文《MoEs 遇见指令微调》（2023 年 7 月）进行了以下实验：

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

当作者微调 MoE 和 T5 等效模型时，T5 等效模型表现更好。当作者微调 Flan T5 (T5 instruct 等效) MoE 时，MoE 模型的表现显著更好。不仅如此，Flan-MoE 相对于 MoE 的提升比 Flan T5 相对于 T5 的提升更大，这表明 MoE 可能比密集模型从指令微调中受益更多。MoE 从更多任务中受益更多。与之前建议关闭辅助损失函数的讨论不同，该损失实际上可以防止过拟合。

何时使用稀疏 MoE 与密集模型？

专家在具有多台机器的高吞吐量场景中非常有用。给定预训练的固定计算预算，稀疏模型将更优化。对于 VRAM 很少的低吞吐量场景，密集模型将更好。

注意：不能直接比较稀疏模型和密集模型的参数数量，因为两者代表着截然不同的事物。

让 MoE 加速运行

最初的 MoE 工作将 MoE 层呈现为分支结构，导致计算速度缓慢，因为 GPU 不是为此设计的，并且导致网络带宽成为瓶颈，因为设备需要向其他设备发送信息。本节将讨论一些现有工作，以使这些模型的预训练和推理更加实用。MoE 加速运行。

并行性

我们来简要回顾一下并行性

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

在专家并行中，专家被放置在不同的工作节点上，每个工作节点处理不同批次的训练样本。对于非 MoE 层，专家并行与数据并行行为相同。对于 MoE 层，序列中的令牌被发送到所需专家所在的工作节点。

容量因子和通信成本

增加容量因子（CF）可以提高质量，但会增加通信成本和激活内存。如果全对全通信速度较慢，使用较小的容量因子会更好。一个好的起点是使用 top-2 路由，容量因子为 1.25，并且每个核心有一个专家。在评估期间，可以更改容量因子以减少计算量。

服务技术

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

MoE 的一个主要缺点是参数数量庞大。对于本地用例，可能需要使用较小的模型。让我们快速讨论一些有助于服务部署的技术：

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

更多高效训练

FasterMoE (2022 年 3 月) 分析了 MoE 在高效分布式系统中的性能，并分析了不同并行策略的理论极限，以及倾斜专家受欢迎度的技术、减少延迟的精细通信调度，以及根据最低延迟选择专家的调整后的拓扑感知门控，从而实现了 17 倍的加速。

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

开源 MoE

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

• 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 Switch Transformer 压缩到仅 160GB。

所以，简而言之，一些值得探索的有趣领域包括：

• 将 Mixtral 蒸馏成密集模型
• 探索专家模型合并技术及其对推理时间的影响
• 对 Mixtral 进行极端量化技术

一些资源

• 局部专家的自适应混合（1991）
• 在深度专家混合中学习因子化表示（2013）
• 惊人庞大的神经网络：稀疏门控专家混合层（2017）
• GShard：使用条件计算和自动分片扩展巨型模型（2020 年 6 月）
• GLaM：使用专家混合高效扩展语言模型（2021 年 12 月）
• Switch Transformers：通过简单高效的稀疏性扩展到万亿参数模型（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.