Mixtral-8x7B 背后的 MoE 技术
目录
- 什么是专家混合（MoE）？
- MoE 的历史
- 什么是稀疏性？
- MoE 和 Transformers
- Switch Transformers
- 微调 MoE
- 加速 MoE 的运行
- 服务技巧
- 开源 MoE 项目
正文
什么是专家混合（MoE）？
模型的规模对于提升其质量至关重要。在有限的计算资源下，相较于用更多步骤训练一个小型模型，训练一个大型模型即便步骤更少效果通常更好。MoE 让模型以「远低于传统密集模型的计算成本」进行预训练，这意味着你可以在相同的计算预算下显著扩大模型或数据集的规模。特别是在预训练阶段，MoE 模型能比其同等规模的密集型模型更快地达到相同的性能水平。
那么，MoE 究竟是什么呢？在 Transformer 模型的背景下，MoE 主要由两个部分组成：
- 稀疏 MoE 层代替了传统的密集前馈网络 (FFN) 层。MoE 层包含若干“专家”，例如 8 个，每个专家都是一个独立的神经网络。实际上，这些专家通常是 FFN，但它们也可以是更复杂的网络，甚至可以是 MoE 本身，形成一个层级结构的 MoE。
- 门控网络或路由器，用于决定哪些 Token 分配给哪个专家。值得注意的是，一个 Token 可以被分配给多个专家。
如何高效地将 Token 分配给合适的专家，是使用 MoE 技术时需要考虑的关键问题之一。路由器由一系列可学习的参数构成，它与模型的其他部分一起进行预训练。
MoE 的历史
MoE 的概念最早出现在 1991 年的论文《Adaptive Mixture of Local Experts》中。这一理念与集成方法相似，目的是通过监督程序管理一个由不同网络构成的系统，每个网络处理训练样本的一部分。每个单独网络或“专家”，都在输入空间的不同区域有其特长。由单独的门控网络决定每个专家的权重，在训练过程中，同时对专家和门控网络进行训练。
在 2010 至 2015 年间，两个不同的研究领域推动了 MoE 的进一步发展：
- 将专家作为组件：在传统的 MoE 结构中，系统由一个门控网络和多个专家组成。MoE 作为整体模型已在 SVM、高斯过程等方法中得到应用。Eigen 等人的研究将 MoE 作为更深层网络的一部分进行探索。这意味着 MoE 可以作为多层网络中的一层，使模型在大规模和高效率之间达到平衡。
- 条件计算：传统网络会将所有输入数据通过每一层。在此期间，Yoshua Bengio 探索了一种基于输入 Token 动态激活或停用网络组件的方法。这些研究促进了在自然语言处理领域对混合专家模型的探索。具体来说，Shazeer 等人 (2017 年的研究，团队成员包括 Geoffrey Hinton 和 Jeff Dean) 将这一理念应用到了一个 137B 的 LSTM (当时的 NLP 主要架构) 上，通过引入稀疏性概念，即使在大规模应用中也能保持快速的推理速度。这项工作主要关注翻译领域，但也面临着高通信成本和训练不稳定等挑战。
什么是稀疏性？
稀疏性基于条件计算的概念。不同于密集模型中所有参数对所有输入都有效，稀疏性让我们能够只激活系统的部分区域。条件计算 (即网络的某些部分仅针对特定样本激活) 使得在不增加计算量的情况下扩大模型规模成为可能，从而在每层 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 团队在设计上做了一些创新，除了引入了类似前一节提到的辅助损失机制外，还包括：
- 随机路由机制：在 top-2 设计中，我们始终选择表现最优的专家，但第二选择的专家则根据其权重以一定概率被选中。
- 专家处理能力限制：我们可以设定一个专家能处理的 Token 数量的上限。如果两个专家的处理能力都已达到上限，那么这个 Token 就会被认为是多余的，并通过残差连接传递到下一层，或在某些情况下被直接丢弃。这一概念在 MoE 的应用中非常关键。为什么这样做？因为在模型编译时所有的张量形状都是静态确定的，但我们无法预先知道每个专家将处理多少 Token，因此需要设定一个固定的处理能力上限。在模型推理过程中，只有部分专家会被激活。同时，一些计算过程如自注意力机制会被所有 Token 共享。因此，尽管一个拥有 8 个专家的 470 亿参数模型听起来庞大，但实际上它的计算需求相当于一个 120 亿参数的密集型模型。如果采用 top-2 机制，模型会涉及约 140 亿参数，但由于注意力等操作是共享的，实际上模型真正使用的参数量仍然是 120 亿。
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 模型进行微调时，我们需要特别注意以下几点：
- 选择合适的专家数量。 专家数量过多会导致模型过拟合，而专家数量过少会导致模型欠拟合。
- 选择合适的门控函数。 门控函数决定了哪些 Token 分配给哪个专家。不同的门控函数会对模型的性能产生不同的影响。
- 调整专家容量。 专家容量是指每个专家能处理的 Token 数量。专家容量过大或过小都会影响模型的性能。
- 使用辅助损失。 辅助损失可以鼓励模型平等地对待所有专家，从而提高模型的性能。
加速 MoE 的运行
为了加速 MoE 模型的运行，我们可以采用以下方法：
- 使用稀疏性。 稀疏性可以减少模型的计算量，从而提高模型的运行速度。
- 使用随机路由。 随机路由可以减少模型的通信成本，从而提高模型的运行速度。
- 使用专家处理能力限制。 专家处理能力限制可以防止模型过拟合，从而提高模型的运行速度。
- 使用选择性的精确度。 选择性的精确度可以降低模型的计算成本，从而提高模型的运行速度。
服务技巧
在对 MoE 模型进行服务时，我们需要特别注意以下几点：
- 选择合适的硬件。MoE 模型对硬件的要求很高，因此我们需要选择合适的硬件来运行模型。
- 优化模型。 我们可以对模型进行优化，以提高模型的性能。
- 监控模型。 我们需要监控模型的运行情况，以确保模型正常运行。
开源 MoE 项目
目前，有许多开源的 MoE 项目，例如：
这些项目提供了 MoE 模型的实现，我们可以直接使用这些项目来训练和微调 MoE 模型。
总结
MoE 是一种强大的技术，可以显著提高模型的性能。在本文中，我们介绍了 MoE 的基本原理、MoE 的历史、MoE 的应用以及 MoE 的加速方法。我们还提供了几个开源的 MoE 项目，供读者参考。