什么是MoE？

一、概念

        业内认为，早在1991年，Geoffrey Hinton老爷子就联手Michael I. Jordan发表了《Adaptive Mixtures of Local Experts》奠定了Mixture of Experts（MoE）的理论基础。自本世纪以来，陆续有一些研究在推动着MoE技术的发展与成熟。而真正发挥出MoE架构强大能力的正是近两年的大语言模型时代，结合了MoE的大语言模型训练工作不论是在参数量还是在训练速度上，都远远优于常规的解决方案。

        从概念上看，MoE可以看作是一种集成学习策略，它通过结合多个专家模型来处理不同的输入数据或任务（可以是不同类型的神经网络，例如前馈神经网络或卷积神经网络），每个专家专注于解决特定的子问题。MoE的核心在于其门控机制，该机制负责根据输入数据的特征动态地决定哪个专家模型应该被激活（路由机制）以生成最佳预测。因此，不同于传统集成学习方法，并不是简单地把多个模型生成的结果融合到一起进行输出就能够叫MoE的，而是要在训练阶段实实在在进行模型结构的改进。

二、技术特点

1、稀疏激活

        在MoE模型中，只有一部分专家模型会被激活来处理输入数据，这使得模型能够更高效地利用计算资源。这其中的每个输入数据都会被门控网络评估，以确定哪些专家模型最适合处理该数据，通常可以使用softmax输出一个概率分布，指示每个专家模型被激活的可能性。例如，有的MoE模型设计时激活Top2概率的专家，有的是Top1，这就是所谓的稀疏性。

2、动态路由

        在MoE模型中，动态路由的实现是通过门控网络来完成的。门控网络的作用是根据输入数据的特征动态地决定哪些专家模型应该被激活以生成预测。这种方法允许模型在处理不同输入时能够灵活调整，从而提高整体性能。

3、易于扩展

        MoE模型可以通过增加专家模型的数量来扩展模型的容量，而不需要对现有的模型结构进行大规模的修改。换句话说，想要增加参数量就拼命增加专家模型的数量，调整激活的Top-K专家数值。

4、训练效率

        MoE模型能够在较少的计算资源下进行有效的预训练（因为只激活了一部分专家网络中的参数，而且是一小部分），这意味着在相同的计算预算条件下，可以显著扩大模型或数据集的规模。

5、推理速度

        由于只有部分专家模型被激活，MoE模型在推理时可以更快地处理输入数据，提高了模型的推理速度。

6、多任务学习

        MoE模型在多任务学习中表现出色，能够同时处理多个不同的任务，而不会相互干扰。

三、MoE在大语言模型中的应用

       这里，我们介绍两个比较经典的MoE大模型。

1、GShard

        GShard是一种由Google提出，用于扩展大型模型的混合专家模型（MoE）架构。它的核心思想是在Transformer模型的前馈网络（FFN）层中引入MoE结构，通过动态路由机制，使得每个输入标记（token）可以被分配给不同的专家网络进行处理。

• MoE层设计：在GShard中，每个MoE层包含多个专家网络和一个门控网络。门控网络负责为每个输入token选择两个专家网络进行计算，并通过加权求和的方式得到最终的输出。
• 自动分片：GShard使用自动分片技术将模型参数和数据分布到多个设备上，以支持大规模模型的训练。
• 条件计算：GShard通过引入条件计算机制，使得模型能够根据输入数据的不同，动态地选择不同的专家（即子模型）进行处理。
• 实验结果：GShard在多个自然语言处理任务上取得了显著的性能提升，并展示了MoE架构在扩展模型容量方面的巨大潜力。例如，在多语言机器翻译任务中，GShard成功实现了超过6000亿参数的模型，显著提高了翻译质量 。
• 负载均衡：为了解决MoE训练中的负载不均衡问题，GShard引入了动态负载均衡机制，确保所有专家都能得到充分利用。这种机制通过动态调整专家容量，避免了某些专家过载而其他专家闲置的情况 。
• 计算效率：GShard能够在2048个TPU v3加速器上高效地训练6000亿参数的模型，仅需4天时间！

2、Switch Transformer

        Switch Transformer是一种新型的混合专家模型（MoE），它在多个方面对GShard模型进行了改进：

• 简化稀疏路由：Switch Transformer采用了硬路由策略，只将每个输入标记（token）发送给一个专家处理，而不是像GShard中的top-2专家。这种简化路由减少了路由器的计算量和通信量，同时每个专家对应的batch size减半。
• 高效稀疏路由： 为了解决动态路由导致的输入张量形状不固定的问题，Switch Transformer引入了专家容量（Expert Capacity）的概念。专家容量是根据每个批次的token数量、专家数量和容量因子计算得出的，它决定了每个专家需要处理的token数量。
• 改进的训练技术：Switch Transformer采用了选择性精度训练，即在路由器内部使用全精度（FP32），而其他部分使用低精度（bfloat16），以提高训练稳定性并减少通信成本。
• 容量因子： 容量因子是一个自由超参数，用于权衡计算效率和建模精度，较低的容量因子可以减少因token丢失而造成的模型性能损失。
• 预训练可扩展性： Switch Transformer展示了在预训练中的可扩展性，通过增加专家数量，可以在近似恒定的计算复杂度下扩展模型规模。
• 模型蒸馏：Switch Transformer的大型稀疏模型可以被蒸馏成小而稠密的版本，同时性能保持在合理水平。

四、总结

        在如今各家大模型开发商军备竞赛的时期，MoE架构逐渐成为了主流的预训练架构之一，毕竟对于大语言模型而言，谁使用的预训练数据集更大、谁的参数量更大，那么谁的表现就趋于更优。数据量和参数量上去了，如何高效利用计算资源进行高质量的预训练和微调就成为了关键，而MoE在这种场景就成为了优质答案之一了。