「混合专家模型」可视化指南：A Visual Guide to MoE

2024-11-15
湖南
本文字数：7564 字
阅读完需：约 25 分钟

编者按： 对大语言模型进行扩展的过程中，如何在保持高性能的同时有效控制计算资源消耗？混合专家模型（MoE）作为一种新兴的架构设计方案正在得到越来越多的关注，但它究竟是如何工作的？为什么越来越多的大语言模型都在采用这种架构？
这篇文章通过 50 多张精心绘制的示意图，系统拆解了 MoE 的核心原理。从最基础的专家网络、路由机制，到负载均衡、容量控制等进阶话题，每个概念都配有直观的可视化讲解。作者还介绍了如何将 MoE 应用于视觉模型，例如 Vision-MoE 和 Soft-MoE。

作者 | Maarten Grootendorst

编译 | 岳扬

在查看最新推出的大语言模型（LLMs）时，你可能会注意到标题中频繁出现的“MoE”字样。这个“MoE”究竟是什么，它为何会成为众多 LLMs 的标配？

在这份可视化指南中，我们将借助超过 50 张图表来详细了解这一关键组件，逐步深入混合专家模型（Mixture of Experts, MoE）的工作原理！

在这份可视化指南里，我们将逐一剖析 MoE 的核心要素：专家网络（Experts）和路由器（Router），并展示它们在常见的大语言模型（LLM）架构中是如何发挥作用的。

01 什么是混合专家模型？

混合专家模型（MoE）技术通过使用众多不同的子模型（即“专家网络”）来提升大语言模型（LLMs）的性能。

MoE 的两大核心组成部分如下：

专家网络 - 当前，每个前馈神经网络层（FFNN）都配备有一组“专家网络”，可以根据需要选择其中的子集。这些“专家网络”通常就是 FFNN 本身。
路由器或门控网络 - 它负责决定哪些 tokens 应该被传递给哪些专家网络。

在集成了 MoE 的 LLM 的每一个模型层中，我们都能找到（具有一定专业性的）专家网络：

需要注意的是，“专家网络”并非专注于某个特定领域（如“心理学”或“生物学”）的专家，它们更多的是在单词层面上学习句法信息：

更准确地说，它们的专长在于在特定的上下文中处理特定的词汇单元。

路由器（门控网络）会为每个输入挑选出最合适的专家网络：

每个专家网络并非一个完整的 LLM，而是 LLM 架构中的一个组成部分。

02 专家网络

为了深入了解专家网络的含义及其运作机制，我们先来考察混合专家模型（Mixture of Experts，MoE）试图优化的结构 —— 密集层（dense layers）。

2.1 密集层（Dense Layers）

混合专家模型（Mixture of Experts，MoE）的概念源自大语言模型（LLMs）的一个基本组件，也就是前馈神经网络（Feedforward Neural Network，FFNN）。

我们需要知道，在标准的仅包含解码器的 Transformer 架构中，前馈神经网络通常应用于层归一化（layer normalization）之后：

前馈神经网络利用注意力机制生成的上下文信息，进一步处理以捕捉数据中的复杂关系。

但是，前馈神经网络的规模增长速度很快。为了掌握这些复杂的关系，它往往需要对输入信息进行扩展处理：

2.2 稀疏层（Sparse Layers）

在传统的 Transformer 中，FFNN 被称作密集模型（dense model），因为它的所有参数（包括权重和偏置项）都会被激活。没有任何参数被闲置，每一个都参与到输出的计算中。

如果我们仔细观察密集模型，会发现输入数据在某种程度上激活了所有参数：

与此相对的是，稀疏模型（sparse models）只激活了其参数总数的一部分，这与混合专家模型（Mixture of Experts）密切相关。

打个比方，我们可以把密集模型分解成若干小块（即所谓的“专家网络”），重新训练它们，并且在任何特定时刻只激活一部分“专家网络”：

这种方法的核心理念是，每个“专家网络”在训练过程中会学习到不同的信息。而在推理阶段，只有那些与当前任务最相关的“专家网络”会被调用。

这样在面对问题时，我们可以挑选出最适合处理该任务的“专家网络”：

2.3 专家网络学习了什么？

正如我们之前所见，专家网络所学习的信息比泛泛地学习整个领域的所有信息更为精细 1。因此，将它们称为“专家”有时会被认为是一种误导。

在 ST-MoE 论文中提到的编码器模型的专业化

相比之下，解码器模型中的专家网络并没有展现出同样的专业化。但这并不意味着所有专家网络都是等同的。

在 Mixtral 8x7B 论文[1]中提供了一个很好的案例，其中每个 token 都根据首次选择的专家网络进行了着色。

这张图片也表明，专家网络通常更关注句法结构，而不是特定领域的内容。

因此，尽管解码器中的专家网络没有显示出特定的专业领域，但它们似乎确实会一致地用于处理某些特定类型的 tokens 。

2.4 专家网络架构

虽然将专家网络可视化成被分解的密集模型中的隐藏层确实很有帮助，但实际上，他们本身往往就是一个完整的前馈神经网络（FFNN）：

大语言模型（LLM）通常包含多个解码器块，这意味着一段文本在最终生成之前，会经过一系列专家网络的处理：

每个 tokens 可能由不同的专家网络处理，这就形成了多样化的处理“路径”：

如果我们重新构想解码器块的构造，会发现它现在集成了更多的前馈神经网络（每个专家网络对应一个）：

因此，解码器块在推理阶段能够调用多个前馈神经网络（每个都被视为一个“专家网络”）来完成任务。

03 路由机制

现在已经有了一组专家网络，模型如何知道使用哪些专家网络呢？

就在专家网络之前，会添加一个路由器（也称为门控网络），它会被训练来为给定的 token 选择专家网络。

3.1 路由器

路由器（或称门控网络）本身也是一个前馈神经网络（FFNN），其作用是基于特定的输入来挑选合适的专家网络。它会输出一系列概率值，用以确定最匹配的专家网络：

专家层（expert layer）将所选专家网络的输出结果与门控值（即选择概率）相乘后返回。

路由器与被选中的少数专家网络共同组成了 MoE 层：

MoE 层有两种类型，一种是稀疏混合专家（sparse mixture of experts），另一种是密集混合专家（dense mixture of experts）。

这两种模式都通过路由器来挑选专家网络，不同的是，稀疏 MoE 只挑选部分专家网络，而密集 MoE 则会选择所有专家网络，但每个专家网络的参与程度可能不同。

打个比方，面对一组 tokens ，MoE 会将其分配给所有专家网络处理，而稀疏 MoE 则只挑选少数几个专家网络来处理。

在现有的 LLMs 中，提到的 “MoE” 通常指的是稀疏 MoE，因为它只使用部分专家网络，大大降低了计算成本，对于 LLMs 来说，这一点至关重要。

3.2 专家网络的选择

门控网络可以说是任何 MoE 最重要的组成部分，因为它不仅决定推理过程中选择哪些专家网络，还决定训练过程中的选择。

最基本的形式是：我们将输入（x）乘以路由器权重矩阵（W）：

然后，我们对输出应用 SoftMax 操作，为每个专家网络创建一个概率分布 G(x)：

路由器利用该概率分布来为给定输入选择最匹配的专家网络。

最后，我们将每个路由器的输出与每个选定专家网络的输出相乘，并将结果相加。

现在将整个流程组合起来，探索输入是如何流经路由器和专家网络的：

3.3 路由过程的复杂性

然而，这种简单的功能往往使得路由器倾向于选择相同的专家网络，因为某些专家网络的学习速度可能比其他专家网络要快：

这样不仅会造成专家网络选择的不均衡，还可能导致部分专家网络几乎得不到训练。这种现象在训练和推理阶段都会引发问题。

为了解决这个问题，我们希望在训练和推理过程中，所有专家网络都能得到平等的重视，我们称之为“专家网络负载均衡”。这样做一定程度上是为了避免模型对某些专家网络产生过度的依赖，从而避免过拟合。

04 负载均衡

为了平衡专家网络的重要性，我们需要重点关注路由器，因为它是决定在特定时间选择哪些专家网络的主要组成部分。

4.1 KeepTopK

对路由器进行负载平衡的一种方法是一个名为 KeepTopK[2]2 的简易扩展策略。通过引入可训练的高斯噪声，避免了总是重复选择同一批专家网络的情况：

接下来，除了我们希望激活的前 k 个专家网络（比如 2 个）之外，其他专家网络的权重将被调整至负无穷大：

这样做之后，这些权重在 SoftMax 函数作用下的输出概率将会是 0：

尽管还存在许多有潜力的替代策略，但许多大语言模型（LLMs）仍然采用 KeepTopK 策略。需要注意的是，即使不添加额外的噪声，KeepTopK 策略同样可以使用。

4.1.1 Token 选择策略

KeepTopK 策略能够将每个 token 分配给若干选定的专家网络。这种技术被称为 Token Choice3，它允许每个 token 可以被路由至一个专家网络（单专家网络路由，即 top-1）：

或者同时发送至多个专家网络（多专家网络路由，即 top-k）：

这种方法的一个主要优点在于，它能够对各个专家的贡献进行评估权衡，并将其整合起来。

4.1.2 辅助损失（Auxiliary Loss）

为了在训练过程中实现专家网络的均衡分配，我们在网络的常规损失基础上增加了辅助损失（也称作负载平衡损失）。

这一措施增加了一个约束条件，旨在强制确保每个专家网络都拥有同等的重要性。

辅助损失的第一部分是对整个批次中每个专家网络的路由器值（router values）进行求和：

这样我们就得到了每个专家网络的重要性得分（importance scores），它代表了在不考虑输入的情况下，给定专家网络被选中的可能性。

通过这些重要性得分（importance scores），我们可以计算出变异系数（CV），它表示不同专家网络之间重要性得分的差异程度。

比如，如果专家网络之间的重要性得分差异显著，CV 值就会偏高：

反之，如果所有专家网络的重要性得分相近，CV 值就会偏低（这正是我们希望达到的效果）：

利用 CV 值，我们可以在训练过程中调整辅助损失，力求最大限度地降低 CV 值（以此确保每个专家网络都具有相同的重要性）：

最终，这个辅助损失将作为一个独立的 loss，参与训练优化。

4.2 专家网络容量

不平衡不仅体现在被选中的专家网络上，还体现在分配给专家网络的 token 分布上。

例如，如果输入的 token 过多地被分配给某一个专家网络，而不是均匀分配，可能会导致某些专家网络得不到充分的训练：

这里的关键不仅仅在于哪些专家网络被使用，更在于它们被使用的程度。

为了解决这个问题，我们可以采用一种方法，即限制每个专家网络能够处理的 token 数量，我们称之为专家网络容量（Expert Capacity4）。当一个专家网络的处理能力达到上限时，新来的 token 将会被转交给下一个专家网络：

如果所有相关专家网络都已达到其处理容量，那么这些 tokens 将不会由任何专家网络处理，而是直接传递到网络的下一层。这种现象我们称之为 token 溢出。

4.3 使用 Switch Transformer 简化 MoE

在处理 MoE 模型训练不稳定问题（例如负载平衡）的早期基于 Transformer 的模型中，Switch Transformer 是其中的一个。5 它通过简化架构和训练流程，有效提升了训练的稳定性。

4.3.1 Switching Layer

Switch Transformer 是一个 T5 模型（编码器-解码器结构），它用 Switching Layer 替换了传统的 FFNN 层。Switching Layer 是一个稀疏的 MoE 层，它为每个 token 选择一个专家网络（Top-1 路由）。

路由器在计算选择哪个专家网络时没有使用特殊技巧，只是将输入乘以专家网络的权重后取 softmax（与我们之前做的一样）。

这种架构（Top-1 路由）假定路由器只需要一个专家网络就能学会如何对输入进行路由。这与我们之前看到的情况截然不同，之前我们认为 token 应该被路由到多个专家网络（Top-k 路由）以学习路由行为。

4.3.2 容量因子（Capacity Factor）

容量因子是一个重要的参数，它决定了专家网络可以处理多少个 token。Switch Transformer 在此基础上进行了扩展，引入了一个直接影响专家网络容量的容量因子。

专家网络容量的组成部分非常简单：

如果我们增加容量因子，则每个专家网络将能够处理更多的 token。

然而，如果容量因子过大，就会浪费计算资源。相反，如果容量因子太小，由于 token 溢出，模型性能将会下降。

4.3.3 辅助损失（Auxiliary Loss）

为了进一步防止 token 被丢弃，Switch Transformer 引入了一个简化版的辅助损失（Auxiliary Loss）。

这个简化版的辅助损失不再计算系数变化，而是根据每个专家网络的路由器概率与分发的 token 数量进行加权比较：

由于我们的目标是让 token 在 N 个专家网络之间均匀路由，因此我们希望向量 P 和 f 的值都是 1/N。

α 是一个超参数，我们可以在训练过程中用它来微调这个辅助损失的重要性。如果 α 值过高，会影响主要损失函数；如果 α 值过低，则无法有效进行负载平衡。

05 视觉模型中的 MoE

MoE 不仅仅是一种仅适用于语言模型的技术。视觉模型（如 ViT）使用基于 transformer 的架构，因此也有潜力使用 MoE。

简单回顾一下，ViT（Vision-Transformer）是一种将图像分割成 patches 的架构，这些 patch 的处理方式与 token 类似。6

这些 patchs（或 tokens）随后会被投影到嵌入向量中（并加上位置嵌入向量），然后被送入常规的编码器中：

一旦这些 patches 进入编码器，它们就像 tokens 一样被处理，这使得这种架构非常适合用于 MoE。

5.1 Vision-MoE

Vision-MoE（V-MoE）是第一次在图像模型中实现 MoE 的尝试。7 它采用了我们之前看到的 ViT，并将编码器中的密集 FFNN 层替换为稀疏 MoE。

这使得 ViT 模型（通常模型规模小于语言模型），可以通过添加专家网络来进行大规模扩展。

为了减少硬件限制，每个专家网络都设置了较小的预定义容量，因为图像通常包含许多 patch 。然而，低容量往往导致 patch 被丢弃（类似于 token 溢出）。

为了保持较低的专家网络容量，网络会为每个 patch 分配重要性分数，并优先处理这些得分较高的 patches，这样使得溢出的 patches 一般就不那么重要了。这种方法被称为批量优先级路由（Batch Priority Routing）。

因此，即使 token 的数量减少，我们仍然应该看到重要的 patch 被成功路由。

优先级路由（priority routing）可以集中处理最重要的 patches，从而减少处理的 patch 数量。

5.2 从稀疏 MoE 到软 MoE

在 V-MoE 中，优先级评分器（priority scorer）能够帮助区分重要和不重要的 patch。然而，patch 被分配给每个专家网络后，未被处理的 patch 中的信息就会丢失。

Soft-MoE 的目标是通过混合 patch 来从离散的 patch（token）分配转变为软分配。8

在第一步中，我们将输入 x（patch 嵌入）与一个可学习的矩阵 Φ 相乘。这将生成路由器信息（router information），告诉我们某个 token 与给定专家网络的相关程度。

然后，通过对路由器信息矩阵（按列）进行 softmax 操作，从而更新每个 patch 的嵌入向量。

更新后的 patch 嵌入本质上是所有 patch 嵌入的加权平均。

从视觉上看，就好像所有的 patch 都被混合在一起。然后，这些合并后的 patch 会被发送给每个专家网络。生成输出后，它们再次与路由器矩阵相乘。

路由器矩阵在 token 层面上影响输入，在专家网络层面上影响输出。

因此，我们得到的是经过处理的“软” patch / token，而不是离散的输入。

06 Mixtral 8x7B 的激活与稀疏参数对比

MoE 的一大亮点在于其计算需求。由于在给定时间内只会使用部分专家网络，我们可以访问的参数比我们正在使用的参数要多。

尽管给定的 MoE 有更多的参数需要加载（稀疏参数），但由于在推理过程中只使用部分专家网络（活跃参数），因此实际激活的参数数量较少。

换句话说，我们仍然需要将整个模型（包括所有专家网络）加载到设备上（稀疏参数），但是当我们进行推理时，只需要使用部分参数（活跃参数）。MoE 模型在加载所有专家网络时需要更多的 VRAM，但在推理过程中运行得更快。

让我们以 Mixtral 8x7B 为例，来探讨稀疏参数与活跃参数的数量差异。9

在这里，我们可以看到每个专家网络的参数量是 5.6B，而不是 7B（有 8 个专家网络）。

我们将不得不加载 8x5.6B（46.7B）的参数（以及所有共享参数），但在推理过程中我们只需要使用 2x5.6B（12.8B）的参数。

07 Conclusion

我们的混合专家模型（MoE）学习之旅到此结束！希望这篇文章能让你更好地理解这种有趣技术的潜力。现在，几乎所有的模型系列都至少包含一个 MoE 变体，感觉它已经成为了主流。

想要查看更多与大语言模型（LLMs）相关的可视化内容，请查看作者编写的大语言模型书籍！

书籍相关网站[3]。你可以在亚马逊[4]上订购这本书。所有代码都已上传至 GitHub[5]。

Resources

如果你想要更深入地了解 MoE，我建议你学习以下内容：

这两篇论文[6][7]是关于 MoE 最新创新技术的概述。
这篇关于 expert choice routing[8] 的论文已经引起了业界的一些关注。
一篇很棒的博客文章[9]，梳理了一些关于 MoE 的主要论文（及其发现）。
这篇博客文章[10] 讲述了 MoE 的发展时间线。

1.Zoph, Barret, et al. "St-moe: Designing stable and transferable sparse expert models. arXiv 2022." arXiv preprint arXiv:2202.08906.

2.Shazeer, Noam, et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." arXiv preprint arXiv:1701.06538 (2017).

3.Shazeer, Noam, et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." arXiv preprint arXiv:1701.06538 (2017).

4.Lepikhin, Dmitry, et al. "Gshard: Scaling giant models with conditional computation and automatic sharding." arXiv preprint arXiv:2006.16668 (2020).

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.Dosovitskiy, Alexey. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).

7.Riquelme, Carlos, et al. "Scaling vision with sparse mixture of experts." Advances in Neural Information Processing Systems 34 (2021): 8583-8595.

8.Puigcerver, Joan, et al. "From sparse to soft mixtures of experts." arXiv preprint arXiv:2308.00951 (2023).

9.Jiang, Albert Q., et al. "Mixtral of experts." arXiv preprint arXiv:2401.04088 (2024).

Thanks for reading!

Hope you have enjoyed and learned new things from this blog!

About the authors

Maarten Grootendorst

Data Scientist | Psychologist | Writer | Open Source Developer (BERTopic, PolyFuzz, KeyBERT) | At the intersection of Artificial Intelligence and Psychology

END

本期互动内容 🍻

❓如果让你用一个生活中的比喻来解释 MoE 的工作原理，你会怎么比喻？欢迎在评论区分享你的想法！

🔗文中链接🔗

[1]https://arxiv.org/pdf/2401.04088

[2]https://arxiv.org/pdf/1701.06538

[3]https://www.llm-book.com/

[4]https://www.amazon.com/Hands-Large-Language-Models-Understanding/dp/1098150961

[5]https://github.com/handsOnLLM/Hands-On-Large-Language-Models

[6]https://arxiv.org/pdf/2209.01667

[7]https://arxiv.org/pdf/2407.06204

[8]https://arxiv.org/pdf/2202.09368