Mixture-of-Experts (MoE)：条件计算的诞生与崛起【上篇】

大型语言模型（LLM）的现代进步主要是缩放定律的产物[6]。假设模型是在足够大的数据集上训练出来的，那么随着底层模型规模的增加，我们会看到性能的平滑提升。这种扩展规律最终促使我们创建了 GPT-3 以及随后的其他（更强大的）LLM。在扩展这些高密度 LLM 时，由于内存成本的增加以及模型的计算占用空间与参数总数的关系，我们最终会达到硬件规定的极限。因此，计算和数据已成为训练更好的 LLM 的主要瓶颈–如果能获得更多的计算和数据，创建一个更好的模型就相对简单了。在本综述中，我们将研究如何通过训练稀疏激活的语言模型来避免这一限制。

“由于巨型密集模型的训练受到了当今硬件资源可用性和能力的限制，专家混合物（MoE）模型与质量相当的密集模型相比，训练成本大幅降低，因此已成为最有前途的模型架构之一”。–来自[12]

专家混合层（MoE）非常简单，它允许我们在不相应增加计算量的情况下增加语言模型的规模或容量。我们只需将模型的某些层替换为多个拥有各自参数的 "专家 "层副本。然后，我们可以使用门控机制来（稀疏地）选择用于处理每个输入的专家。这一想法源于 20 世纪 90 年代[15,30]早期的条件计算研究，它允许我们以可控的方式训练庞大的模型，这对语言建模等受益于额外容量模型的领域很有帮助。在这里，我们将研究 MoE 及其起源，以及它在过去二十年中是如何发展的。

什么是专家混合体（MoE）？

在这里插入图片描述

Decoder-only transformer 架构

上图所示为大多数生成式 LLM 使用的标准纯解码器变换器架构；有关该架构的深入概述，请参见此处。在 LLM 中，MoE 对这一架构进行了简单的修改–用 MoE 层取代前馈子层！这个 MoE 层由多个专家组成（即从几个专家 [13] 到数千个专家 [5]），其中每个专家都是自己的前馈子层，具有独立的参数集；具体描述见下文。

在这里插入图片描述
同样，我们也可以将编码器-解码器转换器（如 T5 [8]）转换为 MoE 模型，方法是将编码器和解码器中的前馈子层都替换为 MoE 层。然而，我们通常只将这些子层中的一部分（如每隔一层）转换为 MoE 层。在本概述中，我们将主要概述基于编码器-解码器转换器的 MoE 模型的工作，而基于 MoE 的自回归 LLMS 的工作将在以后的文章中介绍。

“ST-MoE模型有32个专家，专家层频率为1/4（每4个FFN层被一个MoE层取代）”。–来自[24]

稀疏专家。这种方法看似有问题，因为它给模型增加了大量参数。 MoE 模型在变压器的每个前馈子层中都有多个独立的神经网络（即，而不是单个前馈神经网络）。不过，我们在前向传递中只使用了每个 MoE 层的一小部分专家！给定一个令牌列表作为输入，我们使用路由机制稀疏地选择一组专家，将每个令牌发送给他们。因此，与参数数量相同的密集模型相比，MoE 模型前向传递的计算成本要低得多。

MoE 层的组成部分。当应用于变压器模型时，MoE 层有两个主要组成部分：

稀疏 MoE 层：用一个由多个结构类似的 "专家 "组成的稀疏层取代变压器中的密集前馈层。
路由器：决定 MoE 层中哪些标记被发送给哪些专家。

稀疏 MoE 层中的每个专家都是一个前馈神经网络，有自己独立的参数集。每个专家的结构都模仿了标准转换器结构中使用的前馈子层。路由器将每个标记作为输入，并产生一个专家概率分布，决定将每个标记发送给哪个专家；见下文。

在这里插入图片描述
路由器有自己的参数集，并与网络的其他部分共同训练。每个令牌可以发送给许多专家，但我们只将令牌发送给前 K 位专家，从而实现了稀疏性。例如，许多模型设置 k=1 或 k=2，这意味着每个标记分别由一个或两个专家处理。

更大的容量，固定的计算量。如果我们能获得足够大的训练数据集，那么增加语言模型的规模和容量是提高性能的主要途径之一。然而，训练成本会随着基础模型的大小而增加，这使得大型模型在实践中负担沉重。 MoE 模型在推理过程中只使用模型参数的一个子集，从而避免了这种花费。

“假设每个标记只使用两个专家，推理速度就像使用 12B 模型（而不是 14B 模型），因为它计算的是 2x7B 矩阵乘法，只是共享了一些层”。

例如，假设我们有一个 7B 参数的 LLM，并将其每个前馈子层替换为由八个专家组成的 MoE 层，其中每个标记激活两个专家。这正是 Mixtral-8x7B [13] 所使用的架构，它是 Mistral-7B [14] 的 MoE 变种。完整模型有大约 47B 个参数2 ，所有参数都必须加载到内存中。但是，该模型的推理成本与 14B 参数的模型相当。处理每个标记只需要两个专家，这就产生了 ~2 x 7B 矩阵乘法3。这样，我们就能实现 ~50B 参数模型的容量，而不会产生成本！

优点和缺点。 MoE 模型由于能够以固定的计算预算训练更大的模型而被广泛使用，但使用 MoE 式 LLM 有利有弊！与密集模型相比，MoE 模型的预训练速度更快；与参数数量相同的密集模型相比，MoE 模型的推理速度也快得多。 MoE 使我们能够提高模型容量，同时保持较低的计算量。不过，MoE 式 LLM 也会消耗更多的 VRAM，因为所有专家都必须加载到内存中。此外，MoE 模型容易过度拟合，而且众所周知难以微调，这使得它们在实际应用中比密集模型更加复杂。

专家混合制 Mixture-of-Experts 的起源

虽然 MoEs 在最近的人工智能研究中非常流行，但这种想法由来已久。事实上，条件计算的概念–或者说动态开启/关闭神经网络的部分功能–起源于 20 世纪 90 年代早期的工作！在本节中，我们将探讨条件计算的早期工作，并研究它是如何演变成 "MoE "这一概念的。最终，这些想法被应用于训练早期基于 LSTM [4] 的语言模型（取得了一定的成功）。

条件计算的早期工作

“有条件的计算，即网络的某些部分在每个实例的基础上都处于活动状态，在理论上被认为是在不相应增加计算量的情况下显著提高模型容量的一种方法”。–来自 [1]

MoE 的概念源于 Geoffrey Hinton 在 20 世纪 90 年代初[15]的研究成果，该成果提出了一个由多个网络组成的监督学习框架。这些网络中的每一个都要处理训练数据的不同子集，而选择使用哪个网络则由一个门控机制来处理；见下文。换句话说，该系统中的每个网络都是专家，专门处理输入数据的某个领域（或区域）。

在这里插入图片描述

自[15]提出以来，已有多部作品探索并扩展了这一条件计算思想。例如，[30] 中的作者（也写于 20 世纪 90 年代初！）提出了一种分层（树状结构）的 MoE，可以使用期望最大化算法以监督的方式进行训练；见下文。深入研究这一主题的此类著作还有很多。

在这里插入图片描述
通过随机神经元估计或传播梯度，实现条件计算 [16]。这项工作探索了四种可能的技术，用于估计随机神经元和硬阈值激活函数的梯度；例如，通过使用 REINFORCE、直通估计器或加法/乘法噪声。虽然这些术语可能看起来很陌生，但条件计算就是硬阈值激活函数的一个例子–某些神经元的激活在计算中被完全消除（或设置为零）。要训练使用条件计算的网络，就必须使用 [16] 中提出的估计器。文献[16]中探讨的神经网络结构由许多计算单元组成，这些单元通过分布式门网络连接起来，可以用来消除计算块。通过关闭该网络的部分功能，我们可以大大降低大型神经网络的计算成本。

在这里插入图片描述
深度神经网络中条件前馈计算的低秩近似[17]。这项工作探索了一种神经网络设置，在这种设置中，网络节点辅以门控单元，门控单元决定节点是否应被计算；见上文。由于 ReLU 会将神经网络中所有负激活值设为零，因此任何具有负预激活值的节点都可以完全从计算中移除。作者在 [17] 中扩展了这一想法，指出可以通过计算权重矩阵的低秩近似值来估计 ReLU 之前隐藏激活的符号，从而在网络稀疏时显著提高效率。

在这里插入图片描述

在深度混合专家中学习因子表征 [18]。这项研究考虑了由多个专家网络组成的 MoE 层，这些专家网络专门处理输入空间的不同区域。通过学习门控机制将输入映射到这些专家网络，从而可以稀疏地计算更大的网络。更进一步，[18] 中的作者考虑了具有多个连续 MoE 层（即 "堆叠 "MoE）的更深层网络，如上图所示。在 MNIST 数据集上测试这些网络时，我们发现深度 MoE 在每一层都能学习独特的特征；例如，第一层学习与位置相关的特征，而第二层则学习特定类别的特征。

“有人提出，条件计算是一种在不增加所需计算量的情况下提高深度神经网络容量的方法，其方法是按需激活某些参数，以每个实例为基础”。 - 摘自[19]

指数级提高深度学习中条件计算的容量计算比[19]。众所周知，更大的模型性能更好4。如果计算量不是问题，我们可以使用更大的数据集和模型来提高泛化效果，但深度网络的计算量会随着参数数量的增加而增加。为了避免这一问题并更好地利用条件计算，作者在 [19] 中提出了一种新颖的权重矩阵参数化方法，当观察到特定的隐藏激活模式时，就会关闭参数组5。这种方法可以成倍增加参数与计算的比率。

神经网络中的条件计算以提高模型速度[20]。这项研究探索了基于强化学习的训练策略（即 REINFORCE），用于使用条件计算的神经网络。我们发现，这种使用策略梯度算法的方法可用于训练网络，既能保持准确性，又能最大限度地提高计算速度。此外，我们还可以利用正则化使网络中的条件计算模式多样化。实验表明，这种学习策略可以生成在计算成本和性能之间达到良好平衡的网络。

在这里插入图片描述
动态容量网络 (DCN) [21] 通过定义低容量和高容量子网络，自适应地为输入数据的不同部分分配容量。对于大部分数据来说，应用的是低容量子网络。不过，可以使用基于梯度的关注机制，根据网络对数据中这一区域的敏感度来选择输入的一部分，并将其应用到高容量子网络中；见上文。我们在实验中看到，DCN 可以减少计算量，同时获得与传统卷积神经网络（CNN）相当（或更好）的性能。

深度序列神经网络[22]摒弃了神经网络架构中使用的 "层 "这一传统概念，而是选择构建一整套候选变换，并有选择地应用于每一层。在处理输入时，每一层都会从候选集中选择一个元素，从而形成一个基于变换序列的类似 DAG 的架构。根据输入数据的属性，所选的变换序列可能会有所不同，这与传统的多层网络相比，提高了模型的表达能力。与之前的工作类似，这种序列架构也是采用策略梯度法进行训练的。

超大型神经网络：稀疏门控的专家混合层 [1]

“我们的模型容量提高了 1000 倍以上，而计算效率只略有下降，大大推进了公共语言建模和翻译数据集的最新成果”。 - 摘自 [1]

条件计算的理念前景广阔。特别是在语言建模领域，训练数据集往往是海量的6，在不相应增加计算量的情况下提高模型的底层能力至少可以说是非常诱人的。尽管 MoEs 已被研究了二十多年，但由于各种技术难题，MoEs 的前景并不乐观。在[1]中，作者提出了 “稀疏门控专家混杂层”（Spararsely-Gated Mixture-Of-Experts layer），试图克服其中的一些挑战，表明 MoEs 可应用于语言建模和翻译领域。

关于 "模态建模 "的先前问题。尽管[1]中的工作是最早将 MoE 应用于语言建模的工作之一，但 MoE 的概念可以追溯到 20 世纪 90 年代初。有鉴于此，我们不禁要问：是什么阻碍了建模引擎被更广泛地采用？是什么阻碍了移动电子设备被更广泛地采用？有几个技术障碍导致了缺乏采用：

GPU 擅长高效地执行算术运算，但在分支（即条件计算的主要组成部分）方面却不尽如人意。
要高效地训练神经网络，需要较大的批次规模，但 MoE 可减少批次规模（即每个专家只接收输入批次的一部分）。
在研究具有较大训练数据集的领域时，提高模型容量的影响最大，而 MoE 主要是在计算机视觉领域研究训练数据集不够大的问题。

在这里插入图片描述
什么是稀疏 MoE？ 大多数现代 LLM 所使用的 MoE 层都类似于 [1] 中提出的稀疏门控专家混合层（如上图所示），它是一个通用的神经网络模块，由两个部分组成：

专家：每一层都有几个 “专家”，它们是独立的神经网络模块或层，具有独立的参数集。
路由器：一种参数化（可学习）的门控机制，用于选择一组（稀疏的）专家来处理每个输入。

在 [1] 中，MoE 中的每个专家都是一个具有相同架构的前馈神经网络。不过，我们也可以使用更复杂的架构；例如，我们甚至可以将每个专家作为另一个 MoE 来实现，从而创建 "分层 "的 MoE 模块！通过梯度下降法，专家和门控机制与其他网络参数一起接受联合训练。

在这里插入图片描述
要计算 MoE 模块的输出，我们采用专家输出的加权组合，其中权重由路由器提供；见上文。路由器输出一个 N 维的权重向量（即 N 为专家人数）。虽然这种方法最初看起来似乎没什么用，但当路由器的输出稀疏时，神奇的事情就发生了–在计算 MoE 的输出时，权重为零的专家将不再被考虑。因此，MoE 可以让我们创建、训练和使用超大型神经网络，而无需大量计算，因为在任何时候都只需使用部分模型参数。

在这里插入图片描述
门控机制。对于 MoE 内部的路由选择，人们提出了许多不同的策略。最简单的方法是将我们的输入乘以权重矩阵，然后应用 softmax（见上文）。然而，这种方法并不能保证专家的选择是稀疏的。为了解决这个问题，[1] 中的作者提出了一种改进的门控机制，在这种简单的 softmax 门控机制的基础上增加了稀疏性和噪声；见下文。

上述门控机制与软最大门控机制执行路由的方式类似，但我们增加了两个步骤：

在应用 softmax 之前，会在路由器的输出中加入一定量的高斯噪声。
除了前 K 位专家的输出外，其他专家的输出都会被屏蔽（即设置为 -∞），以确保专家的选择是稀疏的。

平衡专家。 MoEs 的一个问题是，网络在训练过程中倾向于反复使用相同的几个专家。门控机制不会学会统一使用所有专家，而是会趋同于每次输入时都选择同一组专家的状态。这是一个自我实现的循环–如果一个专家被选择的频率最高，那么它的训练速度就会更快，因此会继续被选择，而不是其他专家！之前的研究提出了几种解决这个问题的方法 [2，3]，但我们在 [1] 中看到，可以通过在训练损失中添加一个简单的 "软 "约束来平衡专家；见下文。

在这里插入图片描述
我们首先定义了每位专家在一批输入数据中的 "重要性 "得分，计算方法是将每位专家在整批数据中的门限值相加。简单地说，在批次中被频繁选中的专家将具有较高的重要性得分。然后，我们可以利用专家重要性得分的平方变异系数（CV）来计算辅助损失函数；见上文。如果专家的重要性得分都非常相似，那么 CV 就会很小，反之亦然。这一损失项可以添加到模型的训练目标中，以鼓励专家在每一批中获得相同的重要性。

经验评估。为了测试稀疏 MoE 层的性能，作者在 [1] 中训练了一个语言模型，在 LSTM8 的堆叠层之间插入了一个 MoE 模块。 MoE 模块以卷积方式应用，即一次（批量）处理多个 LSTM 时间步。此外，我们还提出了一种新的分布式训练策略，将模型和数据并行训练结合起来，以消除 MoE 中批处理规模缩小的问题。这种策略加上其他一些系统优化措施，可以在合理的计算要求下训练多达 137B 个参数的语言模型！

在这里插入图片描述
这些模型在不同规模的语言建模数据集（即 10、100 和 1000 亿个单词）上进行了评估。我们使用不同数量的专家对扁平式和分层式 MoE 进行了测试，并对所有模型进行了训练，使其总计算成本大致相同（即活跃专家的数量 k 相同）9。即使在研究最小的数据集时，我们也发现最大的 MoE 模型也能明显改善复杂性（越低越好）；见上文。

在这里插入图片描述
如上图所示，如果训练数据集的规模足够小，那么通过增加专家来增加容量的收益就会递减。然而，我们发现在更大的数据集上，性能会继续提高，直至达到 68B 个参数的规模（即 65,536 个专家）！这一发现强调了 MoE 层与语言建模领域之间的协同作用–在足够大的训练语料库中，增加的模型容量是有帮助的。文献[1]中的研究首次表明，在不明显增加计算成本的情况下，这些优势在实践中是可以实现的，这为 MoE 成为 LLM 的流行工具奠定了基础。

“两条线之间的差距越来越大，这表明（毫不奇怪），模型能力的提高对更大的训练集更有帮助”。 - 摘自[1］

[1] Shazeer, Noam, et al. “Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.” arXiv preprint arXiv:1701.06538 (2017).

[2] Eigen, David, Marc’Aurelio Ranzato, and Ilya Sutskever. “Learning factored representations in a deep mixture of experts.” arXiv preprint arXiv:1312.4314 (2013).

[3] Bengio, Emmanuel, et al. “Conditional computation in neural networks for faster models.” arXiv preprint arXiv:1511.06297 (2015).

[4] Hochreiter, Sepp, and Jürgen Schmidhuber. “Long short-term memory.” Neural computation 9.8 (1997): 1735-1780.

[5] Fedus, William, Barret Zoph, and Noam Shazeer. “Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.” Journal of Machine Learning Research 23.120 (2022): 1-39.

[6] Kaplan, Jared, et al. “Scaling laws for neural language models.” arXiv preprint arXiv:2001.08361 (2020).

[7] Hoffmann, Jordan, et al. “Training compute-optimal large language models.” arXiv preprint arXiv:2203.15556 (2022).

[8] Raffel, Colin, et al. “Exploring the limits of transfer learning with a unified text-to-text transformer.” The Journal of Machine Learning Research 21.1 (2020): 5485-5551.

[9] Lepikhin, Dmitry, et al. “Gshard: Scaling giant models with conditional computation and automatic sharding.” arXiv preprint arXiv:2006.16668 (2020).

[10] Lin, Junyang, et al. “M6: A chinese multimodal pretrainer.” arXiv preprint arXiv:2103.00823 (2021).

[11] Riquelme, Carlos, et al. “Scaling vision with sparse mixture of experts.” Advances in Neural Information Processing Systems 34 (2021): 8583-8595.

[12] Rajbhandari, Samyam, et al. “Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale.” International conference on machine learning. PMLR, 2022.

[13] Jiang, Albert Q., et al. “Mixtral of experts.” arXiv preprint arXiv:2401.04088 (2024).

[14] Jiang, Albert Q., et al. “Mistral 7B.” arXiv preprint arXiv:2310.06825 (2023).

[15] Jacobs, Robert A., et al. “Adaptive mixtures of local experts.” Neural computation 3.1 (1991): 79-87.

[16] Bengio, Yoshua, Nicholas Léonard, and Aaron Courville. “Estimating or propagating gradients through stochastic neurons for conditional computation.” arXiv preprint arXiv:1308.3432 (2013).

[17] Davis, Andrew, and Itamar Arel. “Low-rank approximations for conditional feedforward computation in deep neural networks.” arXiv preprint arXiv:1312.4461 (2013).

[18] Eigen, David, Marc’Aurelio Ranzato, and Ilya Sutskever. “Learning factored representations in a deep mixture of experts.” arXiv preprint arXiv:1312.4314 (2013).

[19] Cho, Kyunghyun, and Yoshua Bengio. “Exponentially increasing the capacity-to-computation ratio for conditional computation in deep learning.” arXiv preprint arXiv:1406.7362 (2014).

[20] Bengio, Emmanuel, et al. “Conditional computation in neural networks for faster models.” arXiv preprint arXiv:1511.06297 (2015).

[21] Almahairi, Amjad, et al. “Dynamic capacity networks.” International Conference on Machine Learning. PMLR, 2016.

[22] Denoyer, Ludovic, and Patrick Gallinari. “Deep sequential neural network.” arXiv preprint arXiv:1410.0510 (2014).

[23] Fan, Angela, et al. “Beyond english-centric multilingual machine translation.” Journal of Machine Learning Research 22.107 (2021): 1-48.

[24] Zoph, Barret, et al. “St-moe: Designing stable and transferable sparse expert models.” arXiv preprint arXiv:2202.08906 (2022).

[25] Du, Nan, et al. “Glam: Efficient scaling of language models with mixture-of-experts.” International Conference on Machine Learning. PMLR, 2022.

[26] Vaswani, Ashish, et al. “Attention is all you need.” Advances in neural information processing systems 30 (2017).

[27] Shazeer, Noam. “Glu variants improve transformer.” arXiv preprint arXiv:2002.05202 (2020).

[28] Zhang, Biao, and Rico Sennrich. “Root mean square layer normalization.” Advances in Neural Information Processing Systems 32 (2019).

[29] Clark, Aidan, et al. “Unified scaling laws for routed language models.” International conference on machine learning. PMLR, 2022.

[30] Jordan, Michael I., and Robert A. Jacobs. “Hierarchical mixtures of experts and the EM algorithm.” Neural computation 6.2 (1994): 181-214.

[31] Shazeer, Noam, et al. “Mesh-tensorflow: Deep learning for supercomputers.” Advances in neural information processing systems 31 (2018).