一文弄懂Mixture of Experts (MoE)的前世今生

“ 连续写了好几天的MOE相关的了，主要是外网的热度真的很高，

        
          
https://huggingface.co/blog/moe  
  
@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 }  
}  

      

模型的参数是提高模型效果的最重要因素之一。在固定的计算资源下，用较少的步骤训练一个更大的模型要比用更多的步骤训练一个较小的模型更好。

专家混合模型使得模型可以用更少的计算资源进行预训练，这意味着你可以在相同的计算预算下大幅扩展模型或数据集的规模。特别是，在预训练期间，MoE模型应该比其密集对应模型更快地达到相同的质量水平。

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

• 使用稀疏的MoE层代替密集的前馈网络（FFN）层。MoE层有一定数量的“专家”（例如8个），其中每个专家都是一个神经网络。在实践中，这些专家是FFN，但它们也可以是更复杂的网络，甚至是一个MoE本身，从而形成分层MoE！
• 一个门控网络或路由网络，确定将哪些token发送给哪个专家。例如，在下图中，“More”token被发送给第二个专家，“Parameters”token被发送给第一个网络。可以将一个token发送给多个专家。如何将token路由到专家是使用MoE时的重要决策之一，路由网络与其余部分同时进行预训练。

在MoEs中，我们将transformer模型的每个FFN层替换为一个MoE层，该层由一个门控网络和一定数量的专家组成。

尽管MoEs提供了诸如高效预训练和与密集模型相比更快的推理等好处，但它们也带来了挑战：

• 训练：MoE 可以显著提高计算效率，但在微调期间过去一直存在泛化困难，导致过拟合。
• 推理：尽管 MoE 可能有许多参数，但在推理期间只使用其中一些参数。这导致推理速度比具有相同参数数量的密集模型要快得多。然而，所有参数都需要加载到 RAM 中，因此内存需求很高。例如，给定一个 Mixtral 8x7B 这样的 MoE，我们需要有足够的 VRAM 来容纳一个密集的 47B 参数模型。为什么是 47B 参数而不是 8 x 7B = 56B？因为在 MoE 模型中，只有 FFN 层被视为单独的专家，其余的模型参数是共享的。同时，假设每个token只使用两个专家，则推理速度（FLOPs）就像使用一个 12B 模型（而不是 14B 模型），因为它计算了 2x7B 的矩阵乘法，但某些层是共享的。

MoEs的根源可以追溯到1991年的论文《 Adaptive Mixture of Local Experts》。这个想法与集成方法类似，目的是为由不同子网络组成的系统提供一种监督过程。每个单独的网络或专家都专门处理输入空间的不同区域。那么如何选择专家呢？一个门控网络确定每个专家的权重。在训练过程中，专家和门控都会被训练。

在2010年至2015年期间，两个不同的研究领域对MoE的进一步发展做出了贡献：

• 专家作为组件：在传统的MoE设置中，整个系统包括一个门控网络和多个专家。MoEs作为整个模型已经在支持向量机、高斯过程和其他方法中进行了探索。《Learning Factored Representations in a Deep Mixture of Experts》工作探索了MoEs作为更深层网络组件的可能性。这使得模型可以同时变得庞大和高效。
• 条件计算：传统的网络通过每一层处理所有输入数据。在这个时期，Yoshua Bengio研究了根据输入token动态激活或停用组件的方法。

这些工作促使人们在自然语言处理的背景下探索专家混合模型。具体而言，在《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》中通过引入稀疏性，将这个想法扩展到了一个137B LSTM，从而实现了高规模下的快速推理。这项工作主要关注翻译，但面临许多挑战，如高通信成本和训练不稳定性。

在密集模型中，所有参数都用于所有输入，而稀疏性允许我们只运行整个系统的某些部分。 条件计算的思想（网络的某些部分根据每个示例的情况作用）使得模型的规模可以扩展而不增加计算量，因此，在每个MoE层中使用了数千个专家。

这种设置引入了一些挑战。例如，尽管较大的批量大小通常对性能更好，但在MOEs中，批量大小会随着数据通过活动的专家而有效减少。例如，如果我们批量输入包含10个token，其中五个token可能会结束在一个专家中，另外五个token可能会结束在五个不同的专家中，导致不均匀的批量大小和低利用率。

我们如何解决这个问题？一个学习到的门控网络（G）决定将输入的一部分发送给哪些专家（E）。

在这个设置中，所有的专家都会对所有的输入进行运算 - 这是一个加权乘法。但是，如果 G 是 0 的话会发生什么？如果是这种情况，就没有必要计算相应的专家操作，从而节省计算资源。那么什么是典型的门控函数呢？在最传统的设置中，我们只是使用一个简单的网络和一个 softmax 函数。这个网络将学习哪个专家来处理输入。

Shazeer的工作还探索了其他门控机制，比如噪声Top-K门控。这种门控方法引入一些（可调节的）噪声，然后保留前k个值。也就是说：

1. 添加一些噪声
2. 选topk
3. softmax

这种稀疏性引入了一些有趣的特性。通过使用足够低的k值（例如1或2），我们可以比激活多个专家时更快地进行训练和推断。为什么不只选择得分最高专家呢？最初的猜想是，需要将路由到多个专家，以使门学会如何路由到不同的专家，因此至少需要选择两个专家。《Switch Transformers》部分将重新验证这个决定。

如前所述，如果我们的所有token都发送给几个概率高的专家，那么将使训练效率低下。在正常的MoE训练中，门控网络会收敛到大多数激活相同的几个专家。这种自我强化会使得受欢迎的专家更快地得到训练，因此被选择的次数也更多。为了缓解这种情况，添加了一个辅助损失以鼓励赋予所有专家相等的重要性。这个损失确保所有专家接收大致相等数量的训练样例。以下部分还将探讨专家容量的概念，它引入了一个专家可以处理多少token的阈值。在transformers中，通过aux_loss参数公开辅助损失。

Transformer是一个非常典型的例子，即增加参数数量可以提高性能，因此谷歌通过GShard探索了将Transformer的参数扩展到超过6000亿的规模。GShard使用top-2 gating在编码器和解码器中将每个FFN层替换为一个MoE层。下图显示了编码器部分的情况。这种设置对大规模计算非常有益：当我们扩展到多个设备时，MoE层在设备之间共享，而所有其他层都会被复制。

为了在规模上保持平衡的负载和效率，GShard的作者除了在前一节讨论的辅助损失之外，还引入了一些改变：

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

GShard论文通过表达适用于MoEs的并行计算模式做出了贡献，但讨论这些内容超出了本博文的范围。

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

尽管MoEs显示出了很多潜力，但它们在训练和微调时仍面临困难。Switch Transformers是一项非常令人兴奋的工作，深入探讨了这些主题。作者甚至在Hugging Face上发布了一个拥有2048个专家和1.6万亿参数的MoE，您可以使用transformers运行。Switch Transformers实现了比T5-XXL快4倍的预训练加速。就像在GShard中一样，作者们用MoE层替换了FFN层。Switch Transformers论文提出了一个Switch Transformer层，它接收两个输入（两个不同的token）并具有四个专家。 与最初使用至少两个专家的想法相反，Switch Transformers使用了简化的单专家策略。这种方法的影响是：

• 路由计算减少了
• 每个专家的批量可以至少减半
• 通信成本降低了
• 质量得到了保证

Switch Transformers还探索了专家能力的概念。

上面建议的容量将批次中的token数量均匀分配给专家的数量。如果我们使用大于1的容量因子，我们为token不完全平衡时提供了缓冲。增加容量将导致更昂贵的设备间通信，因此这是需要权衡的。特别是，Switch Transformer在低容量因子（1-1.25）下表现良好。 Switch Transformer作者还重新审视和简化了在章节中提到的负载平衡损失。对于每个Switch层，在训练期间，辅助损失被添加到总模型损失中。这个损失鼓励均匀路由，并可以使用超参数进行加权。

作者还尝试了选择性精度，例如使用bfloat16训练专家，同时对于其他计算使用完整精度。较低的精度减少了处理器之间的通信成本、计算成本和存储张量的内存。最初的实验中，专家和网关网络都是用bfloat16进行训练，但导致训练不稳定。这主要是由于路由器计算：由于路由器具有指数函数，因此具有更高的精度很重要。为了减轻不稳定性，路由也使用完整精度。

Switch Transformers在编码器-解码器设置中使用了T5的MoE对应物。GLaM论文通过训练一个与GPT-3质量相匹配的模型，使用了1/3的能量来探索提高这些模型的规模（是的，多亏了训练MoE所需的更少计算量，他们可以将碳足迹降低到一个数量级）。作者们专注于仅解码器模型以及少样本和一次样本评估，而不是微调。他们使用了Top-2路由和更大的容量因子。此外，他们还探索了容量因子作为一个在训练和评估过程中可以根据需要改变的度量标准，取决于想要使用多少计算资源。

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

在ST-MoE中引入的z-loss通过惩罚进入门控网络的大logits显著提高了训练稳定性而不降低质量。由于这种损失鼓励值的绝对大小较小，可以减少舍入误差，这对于门控等指数函数可能产生很大影响。

ST-MoE的作者观察到编码器专家专门研究一组token或浅层概念。例如，我们可能会得到一个标点符号专家、一个专有名词专家等等。另一方面，解码器专家的专业化程度较低。作者还在多语言设置下进行了训练。虽然可以想象每个专家都专门研究一种语言，但相反的情况发生了：由于token路由和负载平衡，没有单一的专家专门研究任何特定的语言。

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

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

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

Switch Transformers观察到，在固定的预训练困惑度下，稀疏模型在下游任务中表现比密集对应模型差，尤其是在需要推理的任务（如SuperGLUE）中。另一方面，对于知识密集型任务（如TriviaQA），稀疏模型的表现却相对较好。作者还观察到，在微调中使用较少的专家有所帮助。另一个证实了泛化问题的观察是，该模型在较小的任务中表现较差，但在较大的任务中表现良好。

可以尝试冻结所有非专家权重进行实验。由于MoE层对应于网络的大部分，因此预计会导致巨大的性能下降。我们可以尝试相反的方法：只冻结MoE层中的参数，结果证明几乎与更新所有参数一样好。这可以帮助加快微调速度并减少内存使用。

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

此时，您可能会有点失望，因为人们一直在努力微调MoE。令人兴奋的是，最近的一篇论文《MoEs Meets Instruction Tuning》（2023年7月）进行了实验：

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

当作者微调MoE和T5等价物时，T5等价物表现更好。当作者微调Flan T5（T5指令等价物）MoE时，MoE表现显著更好。不仅如此，Flan-MoE相对于MoE的改进比Flan T5相对于T5更大，表明MoEs可能比密集模型更受指令微调的益处。MoEs受益于更多的任务数量。与先前讨论建议关闭辅助损失函数不同，该损失实际上可以防止过拟合。

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

注意：无法直接比较稀疏和密集模型之间的参数数量，因为两者代表着显著不同的内容。

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

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

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

在专家并行性中，专家被放置在不同的工作器上，每个工作器都采用不同的训练样本批次。对于非MoE层，专家并行性的行为与数据并行性相同。对于MoE层，序列中的token被发送到所需专家所在的工作器。

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

MoE的一个重大缺点是参数数量很大。对于本地用例，人们可能想使用较小的模型。让我们快速讨论一些可以帮助提供服务的技术：

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

FasterMoE（2022年3月）分析了MoE在高效分布式系统中的性能，并分析了不同并行策略的理论极限，以及偏斜专家流行度的技术，减少延迟的细粒度通信调度，以及根据最低延迟选择专家的调整拓扑感知门，导致17倍加速。

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

github

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

      

paper

        
          
Adaptive Mixture of Local Experts (1991)  
Learning Factored Representations in a Deep Mixture of Experts (2013)  
Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (2017)  
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (Jun 2020)  
GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (Dec 2021)  
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (Jan 2022)  
ST-MoE: Designing Stable and Transferable Sparse Expert Models (Feb 2022)  
FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models(April 2022)  
MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (Nov 2022)  
Mixture-of-Experts Meets Instruction Tuning:A Winning Combination for Large Language Models (May 2023)  
Mixtral-8x7B-v0.1, Mixtral-8x7B-Instruct-v0.1.