最近DeepSeek火出圈了，我们也凑热闹学习一下DeepSeek背后的技术。MoE 是 DeepSeek 实现高效、高性能模型的关键技术之一，我们首先来了解一下什么是MoE。

在本篇博文中，我们将深入探讨 MoEs 的核心组件、训练方法，以及在推理过程中需要考量的各种因素。

1 概述

模型规模是提升模型性能的关键因素之一。在有限的计算资源预算下，用更少的训练步数训练一个更大的模型，往往比用更多的步数训练一个较小的模型效果更佳。

混合专家模型 (Mixed Expert Models，简称 MoEs) 的一个显著优势是能够在远少于稠密模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，可以显著扩大模型或数据集的规模。特别是在预训练阶段，与稠密模型相比，混合专家模型通常能够更快地达到相同的质量水平。

那么，究竟什么是一个混合专家模型 (MoE) 呢？作为一种基于 Transformer 架构的模型，混合专家模型主要由两个关键部分组成:

• 稀疏 MoE 层: 这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个)，每个专家本身是一个独立的神经网络。在实际应用中，这些专家通常是前馈网络 (FFN)，但它们也可以是更复杂的网络结构，甚至可以是 MoE 层本身，从而形成层级式的 MoE 结构。
• 门控网络或路由: 这个部分用于决定哪些令牌 (token) 被发送到哪个专家。例如，在下图中，“More”这个令牌可能被发送到第二个专家，而“Parameters”这个令牌被发送到第一个专家。有时，一个令牌甚至可以被发送到多个专家。令牌的路由方式是 MoE 使用中的一个关键点，因为路由器由学习的参数组成，并且与网络的其他部分一同进行预训练。

如下图所示：

总结来说，在混合专家模型 (MoE) 中，我们将传统 Transformer 模型中的每个前馈网络 (FFN) 层替换为 MoE 层，其中 MoE 层由两个核心部分组成: 一个门控网络和若干数量的专家。

尽管混合专家模型 (MoE) 提供了若干显著优势，例如更高效的预训练和与稠密模型相比更快的推理速度，但它们也伴随着一些挑战:

• 训练挑战: 虽然 MoE 能够实现更高效的计算预训练，但它们在微调阶段往往面临泛化能力不足的问题，长期以来易于引发过拟合现象。
• 推理挑战: MoE 模型虽然可能拥有大量参数，但在推理过程中只使用其中的一部分，这使得它们的推理速度快于具有相同数量参数的稠密模型。然而，这种模型需要将所有参数加载到内存中，因此对内存的需求非常高。以 Mixtral 8x7B 这样的 MoE 为例，需要足够的 VRAM 来容纳一个 47B 参数的稠密模型。之所以是 47B 而不是 8 x 7B = 56B，是因为在 MoE 模型中，只有 FFN 层被视为独立的专家，而模型的其他参数是共享的。此外，假设每个令牌只使用两个专家，那么推理速度 (以 FLOPs 计算) 类似于使用 12B 模型 (而不是 14B 模型)，因为虽然它进行了 2x7B 的矩阵乘法计算，但某些层是共享的。

2 稀疏性

稀疏性的概念采用了条件计算的思想。在传统的稠密模型中，所有的参数都会对所有输入数据进行处理。相比之下，稀疏性允许我们仅针对整个系统的某些特定部分执行计算。这意味着并非所有参数都会在处理每个输入时被激活或使用，而是根据输入的特定特征或需求，只有部分参数集合被调用和运行。

条件计算的概念 (即仅在每个样本的基础上激活网络的不同部分) 使得在不增加额外计算负担的情况下扩展模型规模成为可能。这一策略在每个 MoE 层中实现了数以千计甚至更多的专家的有效利用。

这种稀疏性设置确实带来了一些挑战。例如，在混合专家模型 (MoE) 中，尽管较大的批量大小通常有利于提高性能，但当数据通过激活的专家时，实际的批量大小可能会减少。比如，假设我们的输入批量包含 10 个令牌， 可能会有五个令牌被路由到同一个专家，而剩下的五个令牌分别被路由到不同的专家。这导致了批量大小的不均匀分配和资源利用效率不高的问题。在接下来的部分中，将会讨论 让 MoE 高效运行 的其他挑战以及相应的解决方案。

那我们应该如何解决这个问题呢？一个可学习的门控网络 (G) 决定将输入的哪一部分发送给哪些专家 (E):

在这种设置下，虽然所有专家都会对所有输入进行运算，但通过门控网络的输出进行加权乘法操作。但是，如果 G (门控网络的输出) 为 0 会发生什么呢？如果是这种情况，就没有必要计算相应的专家操作，因此我们可以节省计算资源。那么一个典型的门控函数是什么呢？一个典型的门控函数通常是一个带有 softmax 函数的简单的网络。这个网络将学习将输入发送给哪个专家。

Shazeer 等人的工作还探索了其他的门控机制，其中包括带噪声的 TopK 门控 (Noisy Top-K Gating)。这种门控方法引入了一些可调整的噪声，然后保留前 k 个值。具体来说:

这种稀疏性引入了一些有趣的特性。通过使用较低的 k 值 (例如 1 或 2)，我们可以比激活多个专家时更快地进行训练和推理。为什么不仅选择最顶尖的专家呢？最初的假设是，需要将输入路由到不止一个专家，以便门控学会如何进行有效的路由选择，因此至少需要选择两个专家。Switch Transformers 就这点进行了更多的研究。

我们为什么要添加噪声呢？这是为了专家间的负载均衡！

3 微调

3.1 令牌的负载均衡

正如之前讨论的，如果所有的令牌都被发送到只有少数几个受欢迎的专家，那么训练效率将会降低。在通常的混合专家模型 (MoE) 训练中，门控网络往往倾向于主要激活相同的几个专家。这种情况可能会自我加强，因为受欢迎的专家训练得更快，因此它们更容易被选择。为了缓解这个问题，引入了一个 辅助损失，旨在鼓励给予所有专家相同的重要性。这个损失确保所有专家接收到大致相等数量的训练样本，从而平衡了专家之间的选择。接下来的部分还将探讨专家容量的概念，它引入了一个关于专家可以处理多少令牌的阈值。在 transformers 库中，可以通过 aux_loss 参数来控制辅助损失。

3.2 使用HuggingFace微调MoE模型

这个 Jupyter Notebook 展示了如何对 Switch Transformers 进行微调以进行摘要生成的详细指南。

3.3 专家如何学习

ST-MoE 的研究者们发现，编码器中不同的专家倾向于专注于特定类型的令牌或浅层概念。例如，某些专家可能专门处理标点符号，而其他专家则专注于专有名词等。与此相反，解码器中的专家通常具有较低的专业化程度。此外，研究者们还对这一模型进行了多语言训练。尽管人们可能会预期每个专家处理一种特定语言，但实际上并非如此。由于令牌路由和负载均衡的机制，没有任何专家被特定配置以专门处理某一特定语言。

3.4 专家的数量

增加更多专家可以提升处理样本的效率和加速模型的运算速度，但这些优势随着专家数量的增加而递减 (尤其是当专家数量达到 256 或 512 之后更为明显)。同时，这也意味着在推理过程中，需要更多的显存来加载整个模型。值得注意的是，Switch Transformers 的研究表明，其在大规模模型中的特性在小规模模型下也同样适用，即便是每层仅包含 2、4 或 8 个专家。

3.5 微调

稠密模型和稀疏模型在过拟合的动态表现上存在显著差异。稀疏模型更易于出现过拟合现象，因此在处理这些模型时，尝试更强的内部正则化措施是有益的，比如使用更高比例的 dropout。例如，我们可以为稠密层设定一个较低的 dropout 率，而为稀疏层设置一个更高的 dropout 率，以此来优化模型性能。

在微调过程中是否使用辅助损失是一个需要决策的问题。ST-MoE 的作者尝试关闭辅助损失，发现即使高达 11% 的令牌被丢弃，模型的质量也没有显著受到影响。令牌丢弃可能是一种正则化形式，有助于防止过拟合。

Switch Transformers 的作者观察到，在相同的预训练困惑度下，稀疏模型在下游任务中的表现不如对应的稠密模型，特别是在重理解任务 (如 SuperGLUE) 上。另一方面，对于知识密集型任务 (如 TriviaQA)，稀疏模型的表现异常出色。作者还观察到，在微调过程中，较少的专家的数量有助于改善性能。另一个关于泛化问题确认的发现是，模型在小型任务上表现较差，但在大型任务上表现良好。

一种可行的微调策略是尝试冻结所有非专家层的权重。实践中，这会导致性能大幅下降，但这符合我们的预期，因为混合专家模型 (MoE) 层占据了网络的主要部分。我们可以尝试相反的方法: 仅冻结 MoE 层的参数。实验结果显示，这种方法几乎与更新所有参数的效果相当。这种做法可以加速微调过程，并降低显存需求。

在微调稀疏混合专家模型 (MoE) 时需要考虑的最后一个问题是，它们有特别的微调超参数设置——例如，稀疏模型往往更适合使用较小的批量大小和较高的学习率，这样可以获得更好的训练效果。

当研究者们对 MoE 和对应性能相当的 T5 模型进行微调时，他们发现 T5 的对应模型表现更为出色。然而，当研究者们对 Flan T5 (一种 T5 的指令优化版本) 的 MoE 版本进行微调时，MoE 的性能显著提升。更值得注意的是，Flan-MoE 相比原始 MoE 的性能提升幅度超过了 Flan T5 相对于原始 T5 的提升，这意味着 MoE 模型可能从指令式微调中获益更多，甚至超过了稠密模型。此外，MoE 在多任务学习中表现更佳。与之前关闭 辅助损失 函数的做法相反，实际上这种损失函数可以帮助防止过拟合。

4 提速

最初的混合专家模型 (MoE) 设计采用了分支结构，这导致了计算效率低下。这种低效主要是因为 GPU 并不是为处理这种结构而设计的，而且由于设备间需要传递数据，网络带宽常常成为性能瓶颈。

4.1 并行计算

并行计算的几种形式:

• 数据并行: 相同的权重在所有节点上复制，数据在节点之间分割。
• 模型并行: 模型在节点之间分割，相同的数据在所有节点上复制。
• 模型和数据并行: 我们可以在节点之间同时分割模型和数据。注意，不同的节点处理不同批次的数据。
• 专家并行: 专家被放置在不同的节点上。如果与数据并行结合，每个节点拥有不同的专家，数据在所有节点之间分割。

在专家并行中，专家被放置在不同的节点上，每个节点处理不同批次的训练样本。对于非 MoE 层，专家并行的行为与数据并行相同。对于 MoE 层，序列中的令牌被发送到拥有所需专家的节点。

4.2 容量因子和通信开销

提高容量因子 (Capacity Factor, CF) 可以增强模型的性能，但这也意味着更高的通信成本和对保存激活值的显存的需求。在设备通信带宽有限的情况下，选择较小的容量因子可能是更佳的策略。一个合理的初始设置是采用 Top-2 路由、1.25 的容量因子，同时每个节点配置一个专家。在评估性能时，应根据需要调整容量因子，以在设备间的通信成本和计算成本之间找到一个平衡点。

Switch Transformers 也对 专家容量 这个概念进行了研究。

上述建议的容量是将批次中的令牌数量均匀分配到各个专家。如果我们使用大于 1 的容量因子，我们为令牌分配不完全平衡时提供了一个缓冲。增加容量因子会导致更高的设备间通信成本，因此这是一个需要考虑的权衡。特别值得注意的是，Switch Transformers 在低容量因子 (例如 1 至 1.25) 下表现出色。

4.3 部署技术

部署混合专家模型 (MoE) 的一个关键挑战是其庞大的参数规模。对于本地使用情况，我们可能希望使用更小的模型。为了使模型更适合部署，下面是几种有用的技术:

• 预先蒸馏实验: Switch Transformers 的研究者们进行了预先蒸馏的实验。他们通过将 MoE 模型蒸馏回其对应的稠密模型，成功保留了 30-40%的由稀疏性带来的性能提升。预先蒸馏不仅加快了预训练速度，还使得在推理中使用更小型的模型成为可能。
• 任务级别路由: 最新的方法中，路由器被修改为将整个句子或任务直接路由到一个专家。这样做可以提取出一个用于服务的子网络，有助于简化模型的结构。
• 专家网络聚合: 这项技术通过合并各个专家的权重，在推理时减少了所需的参数数量。这样可以在不显著牺牲性能的情况下降低模型的复杂度。

4.4 高效训练

FasterMoE (2022 年 3 月) 深入分析了 MoE 在不同并行策略下的理论性能极限，并且探索了一系列创新技术，包括用于专家权重调整的方法、减少延迟的细粒度通信调度技术，以及一个基于最低延迟进行专家选择的拓扑感知门控机制。这些技术的结合使得 MoE 运行速度提升高达 17 倍。

Megablocks (2022 年 11 月) 则专注于通过开发新的 GPU kernel 来处理 MoE 模型中的动态性，以实现更高效的稀疏预训练。其核心优势在于，它不会丢弃任何令牌，并能高效地适应现代硬件架构 (支持块稀疏矩阵乘)，从而达到显著的加速效果。Megablocks 的创新之处在于，它不像传统 MoE 那样使用批量矩阵乘法 (这通常假设所有专家形状相同且处理相同数量的令牌)，而是将 MoE 层表示为块稀疏操作，可以灵活适应不均衡的令牌分配。

5 MoE和稠密模型对比

混合专家模型（MoE）和稠密模型（Dense Model）的对比表格：

特性	混合专家模型（MoE）	稠密模型（Dense Model）
模型结构	由多个专家网络和一个门控网络组成，门控网络决定输入数据由哪些专家处理，每个专家网络专注于特定任务或数据子集。	所有神经元在每一层都参与计算，数据通过全连接层逐层传递。
计算效率	较高，门控网络只激活部分专家网络，适合大规模数据和任务。	较低，所有神经元都参与计算，计算量随模型规模增加而显著增长。
参数共享	专家网络之间参数不共享，门控网络参数共享。	所有参数在每一层共享，所有神经元都参与计算。
应用场景	适合大规模数据和复杂任务，如大规模语言模型、推荐系统。	适合中小规模数据和任务，如图像分类、文本分类。
训练难度	较复杂，需同时训练门控网络和多个专家网络，容易出现专家网络之间的不平衡问题。	相对简单，只需优化单一网络参数，训练过程更稳定。
扩展性	扩展性强，可通过增加专家网络处理更大规模的任务。	扩展性有限，增加模型规模会显著增加计算量。
资源需求	需要更多内存和计算资源来存储和训练多个专家网络。	资源需求相对较低，但随模型规模增加而增加。
灵活性	更灵活，专家网络可以针对不同任务进行专门设计。	灵活性较低，所有任务共享同一网络结构。
训练稳定性	训练过程可能不稳定，容易出现专家网络之间的不平衡问题。	训练过程更稳定，参数共享和统一优化目标有助于收敛。
适用任务类型	适合多任务学习、大规模数据分布不均匀的任务。	适合单一任务或数据分布相对均匀的任务。

总结：

• MoE：适合大规模、复杂任务，计算效率高，但训练复杂且资源需求大。
• Dense Model：适合中小规模任务，训练简单且稳定，但计算效率较低。

6 为什么是替换FFN层

在Transformer架构中，用MoE（混合专家模型）层替换FFN（前馈神经网络）层是一种常见的设计选择，这主要是基于以下几个方面的考虑：

6.1 FFN层的角色与特性

• FFN的作用：在Transformer中，FFN层是每个Transformer块的重要组成部分，负责对输入特征进行非线性变换和特征提取。它的计算量通常占整个Transformer块的很大一部分（例如，70%以上）。
• FFN的结构：FFN通常由两个全连接层和一个激活函数组成，结构相对简单，且与MoE的专家网络结构相似（都是多层感知机，MLP）。
• 计算密集性：FFN层的计算是密集的，所有神经元都参与计算，导致计算成本较高。

由于FFN层的计算密集性和其在Transformer中的重要性，用MoE层替换FFN层可以显著提高计算效率，同时保留其功能。

6.2 MoE的优势与FFN的适配性

• MoE的计算效率：MoE通过门控网络动态选择部分专家网络进行计算，而不是像FFN那样对所有输入进行密集计算。这种稀疏性可以大幅减少计算量。
• 专家网络与FFN的相似性：MoE中的每个专家网络通常是一个MLP（多层感知机），与FFN的结构非常相似。因此，用MoE替换FFN可以无缝集成到Transformer架构中，而不需要对其他部分（如自注意力层）进行修改。
• 任务适配性：FFN层主要负责特征变换，而MoE的专家网络可以针对不同输入数据动态选择不同的特征变换路径，从而更好地适配多样化的任务。

6.3 自注意力层的特殊性

• 自注意力层的角色：自注意力层是Transformer的核心，负责捕捉输入序列中不同位置之间的依赖关系。它的计算复杂度较高（(O(n^2))，其中 (n) 是序列长度），但它的作用是全局的，难以通过MoE的稀疏机制来优化。
• 自注意力层的不可替代性：自注意力层的全局交互特性是Transformer的关键优势，如果用MoE替换自注意力层，可能会破坏这种全局建模能力。
• 计算瓶颈：虽然自注意力层的计算复杂度高，但在实际应用中，FFN层的计算量通常更大（尤其是在大模型中），因此优化FFN层对提升整体效率更有意义。

6.4 工程实现的便利性

• 模块化设计：Transformer的FFN层是一个独立的模块，用MoE替换FFN层可以在不改变其他部分的情况下实现，工程实现较为简单。
• 训练稳定性：自注意力层的训练已经比较复杂，如果再用MoE替换自注意力层，可能会进一步增加训练难度。而FFN层的训练相对稳定，用MoE替换后可以通过辅助损失（如负载均衡损失）来稳定训练过程。

6.5 实验验证

• 研究支持：许多研究表明，用MoE替换FFN层可以在保持模型性能的同时显著减少计算量。例如，Google的Switch Transformer和GShard等模型都采用了这种设计，并取得了显著的效果。
• 性能提升：实验证明，MoE层在保持模型性能的同时，可以大幅减少计算量，尤其是在大规模模型和大规模数据集上。

因此，用MoE替换FFN层是一种基于计算效率、结构适配性和工程实现的综合考虑。

---

参考资料：

1. https://huggingface.co/blog/zh/moe
2. https://colab.research.google.com/drive/1aGGVHZmtKmcNBbAwa9hbu58DDpIuB5O4?usp=sharing
3. Switch Transformers