MoE 的概念最早出现在 1991 年的论文《Adaptive Mixture of Local Experts》中。这一理念与集成方法相似，目的是通过监督程序管理一个由不同网络构成的系统，每个网络处理训练样本的一部分。每个单独网络或“专家”，都在输入空间的不同区域有其特长。由单独的门控网络决定每个专家的权重，在训练过程中，同时对专家和门控网络进行训练。
$$仅保留前 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}$$
Mixtral-8x7B 背后的 MoE 技术
目录
正文
什么是专家混合（MoE）？
模型的规模对于提升其质量至关重要。在有限的计算资源下，相较于用更多步骤训练一个小型模型，训练一个大型模型即便步骤更少效果通常更好。MoE 让模型以「远低于传统密集模型的计算成本」进行预训练，这意味着你可以在相同的计算预算下显著扩大模型或数据集的规模。特别是在预训练阶段，MoE 模型能比其同等规模的密集型模型更快地达到相同的性能水平。
那么，MoE 究竟是什么呢？在 Transformer 模型的背景下，MoE 主要由两个部分组成：
如何高效地将 Token 分配给合适的专家，是使用 MoE 技术时需要考虑的关键问题之一。路由器由一系列可学习的参数构成，它与模型的其他部分一起进行预训练。
MoE 的历史
MoE 的概念最早出现在 1991 年的论文《Adaptive Mixture of Local Experts》中。这一理念与集成方法相似，目的是通过监督程序管理一个由不同网络构成的系统，每个网络处理训练样本的一部分。每个单独网络或“专家”，都在输入空间的不同区域有其特长。由单独的门控网络决定每个专家的权重，在训练过程中，同时对专家和门控网络进行训练。
在 2010 至 2015 年间，两个不同的研究领域推动了 MoE 的进一步发展：
什么是稀疏性？
稀疏性基于条件计算的概念。不同于密集模型中所有参数对所有输入都有效，稀疏性让我们能够只激活系统的部分区域。条件计算 (即网络的某些部分仅针对特定样本激活) 使得在不增加计算量的情况下扩大模型规模成为可能，从而在每层 MoE 中使用了数千名专家。
这种方法也带来了挑战。比如，虽然大批量处理通常能提高性能，但在 MoE 中，当数据通过活跃的专家时，实际的批量大小会减小。例如，如果我们的批量输入包含 10 个 Token，可能有 5 个 Token 由一个专家处理，另外 5 个 Token 分别由 5 个不同的专家处理，这导致批量大小不均匀，资源利用率低下。
那我们该如何解决这些问题呢？让我们深入探讨 Shazeer 在翻译领域对 MoE 的研究。通过一个学习型的门控网络 (G. ，决定将输入的哪些部分分配给哪些专家 (E)：✅
$$y = \sum_{i=1}^{\text{n}} G(x)_i E_i(x)$$
在这种设置中，所有专家都参与处理所有输入——这是一种加权乘法过程。但如果 G 的值为 0 呢？这种情况下，就无需计算相应专家的操作，从而
节约了计算资源。
那么，典型的门控函数是什么样的呢？在传统设置中，我们通常使用一个简单的网络配合 softmax 函数。这个网络会学习如何选择最合适的专家处理输入。
$$G_\sigma(x) = \text{Softmax}(x \cdot W_g)$$
Shazeer 的研究还探索了其他类型的门控机制，如带噪声的 Top-K 门控。这种方法加入了一些可调节的噪声，然后只保留最高的 k 个值。具体来说：
$$添加噪音H(x)i = (x \cdot W_g)_i + \text{StandardNormal()} \cdot \text{Softplus}((x \cdot W{\text{noise}})_i)$$
$$仅保留前 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}$$
$$应用softmax函数G(x) = \text{Softmax}(\text{KeepTopK}(H(x),k))$$
这种稀疏性带来了一些有趣的特性。如果使用较低的 k 值 (比如一到两个)，我们可以比激活许多专家时更快地进行训练和推理。为什么不只选择最顶尖的专家呢？最初的假设是，为了让门控学习如何路由到不同的专家，需要路由到一个以上的专家，因此至少需要选择两个专家。
我们为什么要加入噪声？这是为了实现负载均衡！
MoE 的负载均衡
正如之前所讨论的，如果所有的 token 都被发送到少数几个受欢迎的专家，这将导致训练效率低下。在标准的多专家系统训练中，门控网络倾向于主要激活相同的几位专家。这会形成自我加强的循环，因为得到优先训练的专家会被更频繁地选择。为了减轻这种情况，引入了一种辅助损失来鼓励平等对待所有专家。这种损失确保所有专家获得大致相同数量的训练样本。
下文还将探讨「专家容量」的概念，这涉及到一个专家能处理的 token 数量上限。在 transformers 中，这种辅助损失可以通过 aux_loss 参数来调节。
MoE 和 Transformers
Transformers 模型展示了一个明显的趋势：「增加参数的数量可以显著提高性能」。Google 的 GShard 项目正是在这方面进行了深入探索，试图将 Transformers 模型扩展到超过 6000 亿个参数。在 GShard 中，编码器和解码器里的部分 FFN 层被 MoE 层替代，并采用了一种称为「top-2」的门控机制。下图显示了这种设计在编码器部分的应用。
这种设计对大规模计算尤其有利：当模型扩展到多个设备时，MoE 层在这些设备间共享，而其他层则在每个设备上独立存在。
为了在大规模应用中保持效率和均衡的负载，GShard 团队在设计上做了一些创新，除了引入了类似前一节提到的辅助损失机制外，还包括：
Switch Transformers
尽管 MoE 充满潜力，但它们在训练和微调时面临稳定性挑战。Switch Transformers 这项研究深入剖析了这些问题，并发布了一个具有 2048 个专家和 1.6 万亿参数的 MoE 模型。相较于 T5-XXL，Switch Transformers 的预训练速度提高了四倍。
Switch Transformers 提出了一种处理两种
不同 token 的新型 Transformer 层，包含四个专家。不同于最初至少使用两个专家的设想，Switch Transformers 采用了更简洁的单专家策略。这种策略的影响包括：
此外，Switch Transformers 还探讨了专家容量的概念。专家容量的计算公式是：
$$\text{Expert Capacity} = \left( \frac{\text{tokens per batch}}{\text{number of experts}} \right) \times \text{capacity factor}$$
每批 token 数量除以专家数量，再乘以容量因子。按此计算方式，可以均匀分配批次中的 Token 给每个专家。如果容量因子大于 1，可以为 Token 分配不均的情况提供缓冲。但容量增加会带来更高的设备间通信成本，这是一个需要权衡的问题。Switch Transformers 在较低的容量因子下表现优异。
Switch Transformer 的研究者还对上文提到的负载均衡损失进行了简化。在训练过程中，每个 Switch 层的辅助损失会加入到总模型损失中，这种做法促进了均匀的路由分配，并可以通过超参数进行调整。研究者们还尝试了一种选择性的精确度方法，例如在训练专家系统时使用 bfloat16 格式，而在其他计算过程中则采用全精度。降低精度能够显著减少处理器间的通信成本、计算成本以及存储数据的内存需求。但初期实验中，无论是专家系统还是门控网络都采用 bfloat16 进行训练，结果训练过程变得更加不稳定。特别是路由器计算部分，由于其涉及到指数函数，因此更高的精度能够显著改善准确性。
微调 MoE
在对 MoE 模型进行微调时，我们需要特别注意以下几点：
加速 MoE 的运行
为了加速 MoE 模型的运行，我们可以采用以下方法：
服务技巧
在对 MoE 模型进行服务时，我们需要特别注意以下几点：
开源 MoE 项目
目前，有许多开源的 MoE 项目，例如：
这些项目提供了 MoE 模型的实现，我们可以直接使用这些项目来训练和微调 MoE 模型。
总结
MoE 是一种强大的技术，可以显著提高模型的性能。在本文中，我们介绍了 MoE 的基本原理、MoE 的历史、MoE 的应用以及 MoE 的加速方法。我们还提供了几个开源的 MoE 项目，供读者参考。