群魔乱舞：MoE大模型详解

GPT-4是8个2200亿MoE模型

GPT-4远不止1万亿，甚至，还是8个2200亿参数组成的混合专家模型（MoE）。

2023年6月，美国知名骇客George Hotz在接受采访时透露，GPT-4由8个220B模型组成。这么算来，8 x 220B = 1.76万亿。就连PyTorch的创建者Soumith Chintala对此也深信不疑。

下面这张8头怪，看起来就像现在的GPT-4。

MoE 应用于大模型，GPT-4并不是第一个。在2022年的时候，Google 就提出了MoE大模型Switch Transformer，模型大小是1571B，Switch Transformer在预训练任务上显示出比 T5-XXL（11B） 模型更高的样本效率。在相同的训练时间和计算资源下，Switch Transformer 能够达到更好的性能。

除了GPT-4和Switch Transformer，国内的团队DeepSeek 也开源了国内首个 MoE 大模型 DeepSeekMoE。

• DeepSeekMoE 2B可接近2B Dense，仅用了17.5%计算量。
• DeepSeekMoE 16B性能比肩 LLaMA2 7B 的同时，仅用了40%计算量。
• DeepSeekMoE 145B 优于Google 的MoE大模型GShard，而且仅用 28.5%计算量即可匹配 67B Dense 模型的性能。

一时间，国内大模型开始朝着MoE方向大步前进，估计在2024年，会有越来越多大模型选择MoE架构。

那么，究竟什么是MoE大模型？MoE大模型具备哪些优势？本文就带你一探究竟。

什么是MoE大模型？

MoE，全称为Mixed Expert Models，翻译过来就是混合专家模型。MoE并不是什么最新技术，早在1991年的时候，论文Adaptive Mixture of Local Experts就提出了MoE。

我们知道，模型规模是提升模型性能的关键因素之一，这也是为什么今天的大模型能取得成功。在有限的计算资源预算下，用更少的训练步数训练一个更大的模型，往往比用更多的步数训练一个较小的模型效果更佳。

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

MoE基于Transformer架构，主要由两部分组成：

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

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

MoE大模型具备哪些优势？

MoE的最大优势就是与Dense模型相比，在相同计算资源下，训练速度更快，而且可以训练更大的模型。比如Google的Switch Transformer，模型大小是T5-XXL的15倍，在相同计算资源下，Switch Transformer模型在达到固定困惑度 PPL 时，比T5-XXL模型快4倍。

相同计算资源下，Google的MoE大模型能够在相同计算资源下，以更快的速度达到相同的PPL，而且模型是T5的15倍；DeepSeek的16B MoE大模型，仅在40% 的计算量的情况下，性能和LLaMA 2 7B效果比肩。

总结MoE大模型优点，主要有以下3点：

1. 训练速度更快，效果更好。
2. 相同参数，推理成本低。
3. 扩展性好，允许模型在保持计算成本不变的情况下增加参数数量，这使得它能够扩展到非常大的模型规模，如万亿参数模型。
4. 多任务学习能力：MoE在多任务学习中具备很好的新能（比如Switch Transformer在所有101种语言上都显示出了性能提升，证明了其在多任务学习中的有效性）。

而MoE大模型的缺点，主要有以下4点：

1. 训练稳定性：MoE在训练过程中可能会遇到稳定性问题。
2. 通信成本：在分布式训练环境中，MoE的专家路由机制可能会增加通信成本，尤其是在模型规模较大时。
3. 模型复杂性：MoE的设计相对复杂，可能需要更多的工程努力来实现和优化。
4. 下游任务性能：MoE由于其稀疏性，使得在Fine-tuning过程中容易出现过拟合。

接下来，我们就介绍下MoE的主要原理。通过后面的介绍，我们主要需要回答以下3个问题：

1. MoE为什么可以实现更大模型参数、更低训练成本？
2. MoE如何解决训练稳定性问题？
3. MoE如何解决Fine-Tuning过程中的过拟合问题？

一、Adaptive mixtures of local experts

Adaptive mixtures of local experts，这是大多数MoE论文都引用的最早的一篇文章，发表于1991年，作者中有两个大家熟知的大佬：Michael Jordan 和 Geoffrey Hinton。

论文介绍了一种新的监督学习过程，用于由多个独立网络组成的系统，每个网络处理训练集合的子集。这种新方法可以看作是多层监督网络的模块化版本，或者是竞争性学习的关联版本，因此提供了这两种看似不同的方法之间的新联系。

如果一个多层网络用来训练不同的子任务，通常会有强烈的干扰效应，这会导致学习过程变慢和泛化能力差。这种干扰效应的原因在于，当网络试图同时学习多个子任务时，不同任务的学习过程可能会相互干扰。例如，学习一个子任务时对权重的调整可能会影响其他子任务的学习效果，因为这些权重变化会改变其他子任务的loss。这种相互影响使得网络在处理每个子任务时都试图最小化所有其他子任务的loss。

为了解决这个问题，论文提出了使用多个模型（即专家，expert）去学习，使用一个门控网络（gating network）来决定每个数据应该被哪个模型去训练，这样就可以减轻不同类型样本之间的干扰。

对于一个样本 , 第 个 expert 的输出为 , 期望的输出向量为 , 那么损失函数就这么计算:

其中 是门控网络分配给第 个expert的权重。

在论文中，作者提到这个损失函数可能会导致专家网络之间的强烈耦合，因为一个专家网络的权重变化会影响到其他专家网络的loss。这种耦合可能会导致多个专家网络被用于处理每条样本，而不是专注于它们各自擅长的子任务。为了解决这个问题，论文提出了重新定义损失函数的方法，以鼓励专家网络之间的相互竞争。

就是先让不同的expert单独计算loss，然后再加权求和得到总体的loss。这意味着，每个expert在处理特定样本的目标是独立于其他expert的权重。尽管仍然存在一定的间接耦合（因为其他expert权重的变化可能会影响门控网络分配给expert的score）。如果门控网络和expert都使用这个新的loss进行梯度下降训练，系统倾向于将每个样本分配给一个单一expert。当一个expert在给定样本上的的loss小于所有expert的平均loss时，它对该样本的门控score会增加；当它的表现不如平均loss时，它的门控score会减少。这种机制鼓励expert之间的竞争，而不是合作，从而提高了学习效率和泛化能力。下面是一个示意图：

在这种设计下，我们将 Experts 和 Gating Network 一起进行训练，最终的系统就会倾向于让一个expert去处理一个样本。公式（1）鼓励expert之间的相互合作，而公式看（2）鼓励expert之间的竞争。

论文还提到对上面的公式（2）的损失函数进行修改，效果会更好，新的损失函数如下：

对比一下可以看出，在计算每个 expert 的损失之后，先把它给指数化了再进行加权求和，最后取了log。这也是一个我们在论文中经常见到的技巧。这样做有什么好处呢，我们可以对比一下二者在反向传播的时候有什么样的效果，使用 对 第 个 expert 的输出求导，分别得到：

可以看到, 公式 (4) 的导数, 只会跟当前expert有关, 而后者公式 (2) 还考虑其他experts跟当前样本 的匹配程度。换句话说，如果当前样本跟其他的expert也比较匹配，那么 对第 个 expert的输出的导数也会相对更小。

二、Sparsely-Gated MoE

在 2010 至 2015 年间，两个独立的研究领域为混合专家模型 (MoE) 的后续发展做出了显著贡献：

1. 组件专家：在传统的 MoE 设置中，整个系统由一个门控网络和多个专家组成。在支持向量机 (SVMs) 、高斯过程和其他方法的研究中，MoE 通常被视为整个模型的一部分。然而，Eigen、Ranzato 和 Ilya 的研究 探索了将 MoE 作为更深层网络的一个组件。这种方法允许将 MoE 嵌入到多层网络中的某一层，使得模型既大又高效。
2. 条件计算（Conditional Computation）：传统的神经网络通过每一层处理所有输入数据。在这一时期，Yoshua Bengio 等研究人员开始探索基于输入 token 动态激活或停用网络组件的方法。

在 2017 年，Shazeer 等人(团队包括 Geoffrey Hinton 和 Jeff Dean，后者有时被戏称为“谷歌的 Chuck Norris”) 将这一概念应用于 137B 的 LSTM 。通过引入稀疏性，这项工作在保持极高规模的同时实现了快速的推理速度。在牺牲极少的计算效率的情况下，把模型规模提升1000多倍。

这项工作被发表在论文Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer中，和 1991 年Adaptive mixtures of local experts的工作对比，这里的 Sparsely-Gated MoE 主要有两个区别：

• Sparsely-Gated：不是所有expert都会起作用，而是极少数的expert会被使用来进行推理。这种稀疏性，也使得我们可以使用海量的experts来把模型容量做的超级大。
• token-level：前面那个文章，是 sample-level 的，即不同的样本，使用不同的专家，但是这篇则是 token-level 的，一个句子中不同的 token 使用不同的专家。

下图是Sparsely-Gated MoE的模型结构，被运用于RNN结构中：

如图 2-1 所示，每个 token，都会有一个 MoE Layer，每个 MoE layer 中包含了一堆的 experts，每个 expert都是一个小型的 FFN，还有一个 Gating Network 会根据当前 token，选择少数几个 expert 来进行计算。

2.1 门控网络（Gating Network）

门控网络（Gating Network）的设计和实现，这是Sparsely-Gated MoE 层的核心组成部分。门控网络负责为每个输入 token 选择一个稀疏的专家组合。

设 和 分别是门控网络和第 个 expert 的输出, 那么对于在当前的输入x, 输出就是所有 experts 的加权和:

在公式（6）的设置下，所有 expert 都会对所有输入进行运算，再通过门控网络的输出进行加权和，如果 experts 的数量太大，就会导致计算量非常大。但是，如果 G (门控网络的输出) 为 0 会发生什么呢？如果是这种情况，就没有必要计算相应的 E(x)，因此我们可以节省计算资源。

那么一个典型的门控网络是什么呢？一个典型的门控网络通常是一个带有 softmax 函数的简单的网络。这个网络将学习将输入发送给哪个 expert。如下所示：

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

1. 添加一些噪声

1. 选择保留前 k 个值

1. 应用 Softmax 函数

总而言之，Sparsely-Gated MoE 是通过 TopK 采样的方式实现的。对于非 TopK 的部分，由于值是负无穷，这样在经过 Softmax 之后就会变成 0，不会被选中。noise 项则可以使得不同 expert 的负载更加均衡。在具体实验中，作者使用的 K=2~4。

2.2 平衡专家利用率（Balancing Expert Utilization）

论文指出，门控网络倾向于收敛到一种状态，总是为相同的几个专家产生大的权重。这种不平衡是自我强化的，因为受到青睐的专家训练得更快，因此被门控网络更多地选择。这种不平衡可能导致训练效率低下，因为某些专家可能从未被使用过。

为了解决这个问题, 论文提出了一种软约束方法。作者定义了专家相对于一批训练样本的重要性 Importance , 就是该专家在这批样本中门控值的总和。然后, 他们定义了一个额外的损失函数 , 这个损失函数被添加到模型的整体损失函数中。这个损失函数等于重要性值集合的CV (coefficient of variation) 平方，乘以一个手动调整的缩放因子 。这个额外的损失鼓励所有专家具有相等的重要性，具体计算公式如下所示：

三、GShard

上面的两篇，都是 MoE 系列的基础工作，而且都没有在大模型上得到广泛应用。接下来介绍的工作，都是近几年在大模型上应用上比较出色的工作。

GShard 是谷歌 2021 年在论文 GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding 中提出的使用 GShard 实现 MoE 跨设备分片的方法。

按照文章的说法，GShard 是第一个将 MoE 的思想拓展到 Transformer 上的工作。具体的做法是，把Transformer 的 encoder 和 decoder 中，每隔一个的 FFN 层，替换成 MoE 层，使用的都是 Top-2 门控网络。

GShard 将在编码器和解码器中的每个 FFN 层替换为使用 Top-2 门控的 MoE 层。下图展示了编码器部分的结构。这种架构对于大规模计算非常有效：当扩展到多个设备时，MoE 层在不同设备间共享，而其他所有层则在每个设备上复制。

图中3-1中，分为标准Transformer（a）、MoE Transformer（b）以及GShard Transformer：

1. 标准 Transformer（a）：是标准的Transformer编码器，其中每个 token 通过一个标准的 FFN。
2. MoE Transformer（b）：将每隔一个的 FFN 层替换为 MoE 层。这意味着在编码器中，不再是每个 token 都通过相同的 FFN，而是通过一个由多个专家组成的 MoE 层。
3. MoE跨设备分片（c）：它展示了 MoE 层是如何在多个设备上进行分片的。GShard MoE 层中的专家网络（experts）被分布在不同的设备上。每个专家网络负责处理一部分输入数据，并且每个 token 根据门控机制的输出被分配到一个或两个专家网络中。这样，整个 MoE 层的计算被分散到了多个设备上，每个设备负责处理一部分计算任务。

实现 MoE 跨设备分片的关键技术是模型并行化（model parallelism）和数据并行化（data parallelism）的结合。在模型并行化中，模型的不同部分（在这里是 MoE 层的专家网络）被分配到不同的设备上。在数据并行化中，输入数据（token）被分割成多个部分，每个部分被分配给不同的设备进行处理。

为了实现这种分片，论文中提到的 GShard 模块提供了一套 API 和编译器扩展，允许用户在模型代码中简单地注释关键张量，指定它们应该如何在设备集群上进行分片。这样，编译器就可以自动地将计算图（computation graph）转换为可以在多个设备上并行执行的程序，而不需要用户手动处理复杂的数据分片和通信逻辑。

由于专家被分配到不同设备，可以并行计算，因此大大提升了模型的计算效率，这也解释了为什么 MoE 可以实现更大模型参数、更低训练成本。

为了保持负载平衡和训练效率，GShard 的作者除了引入上节 Sparsely-Gated MoE 中的辅助 loss 外，还引入了一些关键变化：

• 随机路由： 在 Top-2 设置中，GShard 始终选择排名最高的专家，但第二个专家是根据其权重比例随机选择的。
• 专家容量： 我们可以设定一个阈值，定义一个专家能处理多少 token。如果两个专家的容量都达到上限，token 就会溢出，并通过残差连接传递到下一层，或在某些情况下被完全丢弃。专家容量是 MoE 中最重要的概念之一。为什么需要专家容量呢？因为所有张量的形状在编译时是静态确定的，我们无法提前知道多少 token 会分配给每个专家，因此需要一个固定的容量因子。

注意： 在推理过程中，只有部分专家被激活。同时，有些计算过程是共享的，例如自注意力 (self-attention) 机制，它适用于所有 token。这就解释了为什么我们可以使用相当于 12B Dense 模型的计算资源来运行一个包含 8 个专家的 47B 模型。如果我们采用 Top-2 门控，模型会使用高达 14B 的参数。但是，由于自注意力操作 (专家间共享) 的存在，实际上模型运行时使用的参数数量是 12B。

四、Switch Transformers

尽管 MoE 显示出了很大的潜力，但是由于复杂性、通信成本以及训练和微调过程的不稳定性，模型广泛采用仍需要优化。

而在2022年，Google 提出的 Switch Transformers 一定程度缓解了这些问题。Switch Transformers 是一项非常激动人心的工作，它深入研究了这些话题。作者在 Hugging Face 上发布了一个 1.6 万亿参数的 MoE，拥有 2048 个专家，你可以使用 transformers 库来运行它。Switch Transformers 实现了与 T5-XXL 相比 4 倍的预训练速度提升。

Switch Transformers 简化了 MoE 路由算法，设计了直观的改进模型，降低了通信和计算成本。Switch Transformers 的训练方法减轻了不稳定性，并且首次展示了用较低精度（bfloat16）格式训练大型稀疏模型的可能性。

和 T5 Base、T5 Large 相比，Switch Transformers 在相同计算资源情况下获得了高达 7 倍的预训练速度。在多语言实验中，Switch Transformers 在所有 101 种语言测试中都取得了提升。Switch Transformers 通过在爬虫语料库上预训练了一个高大万亿参数规模的模型，实现了与 T5-XXL 相比4倍的加速。

上图中，模型参数随着专家数量的增加而增加，但保持了相同的计算成本（FLOPs per token）。这表明模型在保持计算效率的同时，能够利用更多的参数来提高性能。

上图 4-3 比较了使用相同计算资源下，Switch Transformer 和 T5-Base 的 PPL。可以看到 Switch Transformer 模型在保持相同计算资源的情况下，相对于 T5-Base 有显著的提升，而且专家数越多（模型参数越多、模型更稀疏），效果越好。

4.1 Switch Transformer 主要优化

Swith Transformer 在论文中提到其设计的指导原则是——尽可能地把 Transformer 模型的参数量做大！（同时以一种简单高效的实现方式）

和其他 MoE 模型的一个显著不同就是，Switch Transformer 的门控网络每次只路由到 1 个 expert，也就是每次只选取 top1 的专家，而其他的模型都是至少 2 个。这样就是最稀疏的 MoE 了，因此单单从 MoE layer 的计算效率上讲是最高的了。下图 4-4 是 Switch Transformer 的模型结构。

与最初使用至少两个专家的想法相反，Switch Transformer 采用了简化的单专家策略，每次只选择一个专家。这种方法的效果包括：

• 减少了路由计算，一个 token 每次只路由到一个专家
• 每个专家的 batch size（专家容量、Expert Capacity） 至少可以减半
• 简化路由的实现，降低了 MoE 中的通信成本

4.2 Switch Routing

上面提到了 Switch Transformer 可以降低每个专家的专家容量，那么什么是专家容量？

专家容量（Expert Capacity） 是指每个专家在模型中处理的 token 数。专家容量的计算方式如下：

这里为什么要计算一个专家容量？这个专家容量又有什么作用？

在编译时，所有 tensor 的形状都是静态确定的。这意味着在编译阶段，模型的架构和数据布局已经被定义，包括模型的层数、每层的输入和输出维度等。

尽管 tensor 的形状是静态的，但在训练和推理过程中，模型的计算是动态的。这是因为模型中的路由器（门控网络）会根据输入数据动态地将 token 分配给不同的专家。这种动态性要求模型能够在运行时灵活地处理数据分布。

而这个专家容量的作用就是将 batch 中的总 token 数平均分配给所有专家。然后，为了应对 token 分布不均的情况，会通过一个容量因子（capacity factor）来扩展每个专家的容量。

容量因子是一个大于 1.0 的数，它的作用是为每个专家提供额外的缓冲空间，以容纳可能超出平均分配的 token。这样，即使某些专家接收到的 token 数量超过了平均值，也能够处理这些额外的 token，而不会因为容量不足而导致计算跳过。

下图是不同容量因子下的动态路由。

如图4-5所示，容量因子 Capacity Factor = 1.0的时候，输入6个 token，那么每个专家的专家容量等于 2（Expert Capacity = 6/3 * 1 = 2），Expert 1 被分配到了 3 个 token，超出了专家容量，这些超出的 token 被称为“溢出 token”（图4-5（左）中的虚线部分）。对于这些溢出的 token，模型会跳过计算，直接将 token 的表示通过残差连接传递到下一层。

而如果容量因子 Capacity Factor = 1.5，这时专家容量等于 3，每个专家就能处理 3 个 token（图4-5（右））。

虽然增加容量因子可以减少 token 溢出，但是它也有缺点。如果容量因子设置得过高，会导致计算资源和内存的浪费，因为模型会为可能永远不会用到的 token 分配额外的资源。在论文中，Switch Transformers 在低容量因子 (例如 1 至 1.25) 下表现出色。

下表是不同容量因子的效果对比。

表4-1：Switch Transformer 和 MoE 的效果对比

• 以上模型都是在相同的计算资源（32核）和硬件（TPUv3）上进行训练的。
• 所有 MoE 和 Switch Transformer 模型都使用 128 个专家。
• 为了达到负对数困惑度为-1.50，所有模型都需要进行超过 100k steps 的预训练。
• Switch-Base+：对于这个模型，作者增加了模型的大小，直到其训练速度与 MoE 模型相匹配。这通过增加模型的隐藏层大小（从768增加到896）和 head 的数量（从14增加到16）来实现。
• T5-Base 在训练的 100k 步内没有达到这个负对数困惑度：这表示在给定的训练步数内，T5-Base 模型没有达到设定的效果，这可能是由于其性能不如 Switch Transformer 或 MoE Transformer 模型。

Switch Transformer 的作者还重新审视并简化了前面章节中提到的负载均衡损失（公式（12））。通过合理设置负载均衡损失的系数，可以在训练过程中实现专家之间的良好负载分布。下面介绍下具体实现。

4.3 不同的负载均衡损失

在稀疏模型中，专家的数量通常分布在多个设备上，每个专家负责处理一部分输入数据。理想情况下，每个专家应该处理相同数量的数据，以实现资源的均匀利用。然而，在实际训练过程中，由于数据分布的不均匀性，某些专家可能会处理更多的数据，而其他专家可能会处理较少的数据。这种不均衡可能导致训练效率低下，因为某些专家可能会过载，而其他专家则可能闲置。为了解决这个问题，论文中引入了一种辅助损失函数，以促进专家之间的负载均衡。

给定 个专家, 索引为 到 , 以及一个包含 个 token 的 batch , 辅助 loss 计算为向量 和 的缩放点积。表示如下:

其中, 是 batch 中分配给专家 的 token 占比, 计算方式为 batch 中被路由到专家 的 token 数除以总token 数, 表示如下:

是所有输入token 被路由到专家 的概率，表示如下:

其中 是给定 token 被路由到专家 的概率。

由于我们希望 batch 中的所有 token 能够均匀地分配给 个专家。这意味着每个专家应该处理相同数量的token, 即每个专家处理的 token 比例应该是 。

通过最小化公式 (14) 的辅助 loss, 可以鼓励这种均匀路由。当 token 均匀分布时, 这个损失会被最小化。

最终的 loss 被乘以专家数量 , 这样即使专家数量变化, loss 也能保持恒定。这是因为在均匀路由情况下 。 是一个超参数, 用于调整辅助 loss 的权重。论文中选择了 , 这个值足够大, 可以确保负载均衡, 同时又足够小, 不会压倒主要的交叉嫡目标（即主要的训练损失）。论文实验了从 到 的 值范围, 发现 的值可以快速平衡负载, 同时不会干扰训练损失。

4.4 稀疏路由和负载均衡loss的合并效果

前面介绍了 Switch Transformer 的主要优化：稀疏路由和负载均衡损失。下面介绍一下将这两项优化合并在一起的实验效果。

实验设置如下：

• 首先从 C4 数据集上进行预训练，使用 MLM（Masked Language Modeling） 作为预训练目标。在这个任务中，模型被训练来预测被 mask 的 token。
• 对比 Switch Transformer 与 MoE Transformer 以及 T5。Switch Transformer 在计算量（FLOPs）上与 T5-Base 匹配，即每个 token 应用的计算量相同。MoE Transformer 使用 top-2 路由。
• 所有模型都在相同的硬件（TPUv3）上进行了相同数量的步骤训练。

实验结论如下（可以参见前面的表4-1）：

• Switch Transformer 在速度和效果上都优 MoE Transformer。对于固定的计算量和时间，Switch Transformer 实现了最佳结果。
• Switch Transformer 的计算量小于同等参数的 MoE 模型。如果将 Switch Transformer 的规模增加到匹配MoE Transformer 的训练速度，那么它在每个步骤上都优于所有 MoE 模型。
• Switch Transformer 在较低的容量因子（1.0, 1.25）下表现更好。较低的专家容量表明在大模型中，模型内存非常稀缺，容量因子应尽可能小。

4.5 改进训练和Fine-Tuning技术

1. 精度选择

作者还尝试了混合精度的方法，例如用 bfloat16 精度训练专家，同时对其余计算使用全精度进行。较低的精度可以减少处理器间的通信成本、计算成本以及存储 tensor 的内存。然而，在最初的实验中，当专家和门控网络都使用 bfloat16 精度训练时，出现了不稳定的训练现象。这种不稳定性主要是由路由计算引起的，因为路由涉及指数函数等操作，这些操作对精度要求较高。因此，为了保持计算的稳定性和精确性，保持更高的精度是重要的。为了减轻不稳定性，路由过程也使用了全精度。

下面的表 4-2 显示了混合精度训练的效果，将路由器输入转换为 float32，同时保持其他部分的精度为 bfloat16。这种策略允许模型在几乎与 bfloat16 精度相同的训练速度下，实现与 float32 训练相当的稳定性。

表4-2：不同精度效果对比

实验表明，使用混合精度的 Switch-Base 在固定步数的早期训练中，其效果（以负对数困惑度为衡量标准）与使用 float32 训练的模型相似，同时速度接近 bfloat16。

1. 更小的参数初始化

在深度学习中，适当的权重初始化对于模型的成功训练至关重要。作者观察到，在 Switch Transformer 模型中，这一点尤其明显。

为了提高模型的稳定性，作者建议减少默认的 Transformer 初始化规模。在 Transformer 模型中，权重矩阵通常是从一个截断的正态分布，其均值为0，标准差由一个超参数 s 决定。作者建议将这个初始化超参数 s 从默认值1.0 减少 10 倍，即 s = 0.1。这种较小的初始化规模有助于提高模型效果和减少训练过程中的不稳定性。

表4-3：减小参数初始化规模可以提升训练稳定性

表 4-3 中的数据表明，通过减少初始化规模，模型效果和稳定性得到了提升。这种改进对于大模型，如 Switch Transformer，尤其重要。

1. Fine-Tuning 过程正则化

为了解决 Fine-Tuning 过程中的过拟合问题，作者提出了增加 dropout的策略，特别是在专家层（expert layers）中。他们称之为“expert dropout”，即在 Fine-Tuning 时只在专家层增加 dropout 率。

表 4-4 显示了在 Fine-Tuning Switch Transformer 时，不同 dropout 率的实验结果。这些模型是在 C4 数据集上预训练的，然后进行了 Fine-Tuning。

表4-4：Fine-Tuning 过程中正则化效果

通过这种 expert dropout 策略，有效地减少了过拟合的风险，同时保持了模型在下游任务上的性能。这种正则化方法对于处理具有大量参数的稀疏模型特别有用，因为它可以帮助模型更好地泛化到未见过的数据。

4.6 高效训练：数据、模型、专家并行

任意增加专家数量会导致收益递减（如图4-3所示）。这意味着在某个点之后，继续增加专家数量不会显著提高模型性能。但是可以通过增加模型的维度，如模型的隐藏层大小（dmodel）或前馈网络的维度（dff）来继续提升模型效果。但是这样又会导致显存和内存开销增加，这时候就可以通过并行技术，解决高效训练问题。

这里补充一下关于各种并行的方法的解释。标准的数据并行的定义是一个 batch 的数据在不同的 device 上并行处理，这时每一个 device 上都保存了模型的一份完整拷贝，前向计算完进行梯度汇总和更新。模型并行表示模型不同的参数（层、组件）分配到不同的 device 上，处理一个 batch 的数据。

图4-6（a）表示模型权重的分配方式，图4-6（b）表示数据的分配方式，一种颜色表示一个矩阵（a unique weight matrix）。其中每一个方格表示一个 core。

1. 数据并行（Data Parallelism）

第一列表示数据并行，模型权重拷贝 16 份，16 个同一种颜色矩阵分别表示一个完整的模型，图4-6（b）则是一个完整的矩阵，这里可以理解为 16 个模型计算完成后由于存在梯度汇总再更新的步骤，所以整体更新的是一个batch，因此这里 Data Parallelism 是一个唯一的矩阵。简单来说就是模型复制，数据并行。

2. 模型并行（Model Parallelism）

模型并行部分从模型侧看出来，16个 cores 维护的是一个整体的模型，但是每一个 core 只分配到其中部分模型参数（图4-6（a）），同一个 batch 数据在所有的 core 上计算（图4-6（b）），由于 1 个 core 中分布了不同的模型权重，每次计算完都需要和其他的 core 进行通信。

3. 模型和数据并行

总共有 NN 个 cores，其中 N=n\times mN=n\times m ， nn 代表数据并行维度上的分割因子， mm 代表模型并行维度上的分割因子。现在每个 core 处理的是 B/nB/n 个 token 以及 d_{ff}/md_{ff}/m 个权重。

4. 专家和数据并行

每个专家分配到一个 core 上，同时数据也切分成16份，如下图所示：

5. 专家、模型和数据并行

最后将专家、模型、数据并行合并在一起，如下图所示：

五、GLaM

除了 Switch Transformer，Google还推出另外一个 MoE 模型：GLaM (Generalist Language Model)。

GLaM 比 GPT-3 大三倍，但是由于使用了 Sparse MoE 的设计，训练成本却只有 GPT-3 的 1/3，而且在 29 个NLP 任务上超越了 GPT-3。

下面是 Google Blog 中 GLaM 的模型结构，非常形象。

从图5-1来看，和前面的 GShard 非常像，反正都是出自 Google，也不知道谁借鉴的谁。

表5-1：GLaM 实验

上表展示了 GLaM 跟其他大模型的对比。可以看到，虽然 GLaM 的总参数量有 1.2T，但是在计算中实际激活的参数量只有 96B，所以在 inference 的时候，比 GPT-3 等 dense model 要快得多。

六、ST-MOE

之前讨论的负载均衡损失可能会导致稳定性问题。我们可以使用许多方法来稳定稀疏模型的训练，但这可能会牺牲模型质量。例如，引入 dropout 可以提高稳定性，但会导致模型质量下降。

6.1 用 Router z-loss 稳定模型训练

在论文 ST-MOE: Designing Stable and Transferable Sparse Expert Models 中，作者提出了一种新的辅助损失函数，称为 Router z-loss，用于提高稀疏模型的训练稳定性，同时保持或稍微提高模型质量。这个损失函数是针对稀疏专家模型中的路由器（router）部分设计的，路由器负责将输入的 token 路由到最合适的专家（expert）层。

在 MoE 模型中，每个输入 token 可能被路由到多个专家，但通常只有一个专家层会被激活。为了确保路由器能够稳定地工作并产生高质量的输出，作者引入了 Router z-loss。这个损失函数的目标是鼓励路由器产生较小的logits 值，因为较大的 logits 值在 softmax 激活函数中会导致较大的梯度，这可能会引起训练不稳定。

Router z-loss 的定义如下：

其中, 是 batch 中的 token 数量, 是专家的数量, 是路由器的 logits。这个损失函数通过惩罚较大的 logits 值来工作，因为这些值在 softmax 函数中会导致较大的梯度。通过这种方式，Router z-loss 有助于减少训练过程中的不稳定性，并可能提高模型的泛化能力。

6.2 专家如何学习？

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

6.3 专家的数量对预训练有何影响？

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

6.4 Fine-Tuning MoE 模型

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

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

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

一种可行的 Fine-Tuning 策略是尝试冻结所有非专家层的权重。实践中，这会导致性能大幅下降，我们可以尝试相反的方法：仅冻结 MoE 层的参数。实验结果显示，这种方法几乎与更新所有参数的效果相当。这种做法可以加速 Fine-Tuning 过程，并降低显存需求。

在 Fine-Tuning MoE 时还需要考虑的一个问题是，它们有需要特殊设置的超参数，例如，稀疏模型往往更适合使用较小的 batch size 和较高的学习率，这样可以获得更好的训练效果。

七、开源 MoE 模型

目前已经有一些开源的 MoE 大模型。国内的 MoE 大模型则是最近 DeepSeek 团队开源的 DeepSeekMoE，模型、代码、论文均已同步发布。

• 模型下载：https://huggingface.co/deepseek-ai
• 微调代码：https://github.com/deepseek-ai/DeepSeek-MoE
• 技术报告：https://github.com/deepseek-ai/DeepSeek-MoE/blob/main/DeepSeekMoE.pdf

此外还有一些国外的开源 MoE 模型，开源了训练代码。

• Megablocks：https://github.com/stanford-futuredata/megablocks
• Fairseq：https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm
• OpenMoE：https://github.com/XueFuzhao/OpenMoE

下面是开源了模型，但是没有开源代码：

• Switch Transformers (Google)：基于 T5 的 MoE，专家数量从 8 到 2048。最大的模型有 1.6 万亿个参数。
• NLLB MoE (Meta)：NLLB 翻译模型的一个 MoE 变体。
• OpenMoE：社区对基于 Llama 的模型的 MoE 尝试。
• Mixtral 8x7B (Mistral)：一个性能超越了 Llama 2 70B 的高质量 MoE，并且具有更快的推理速度。此外，还发布了一个经过指令微调的模型。

总结

本文系统性地介绍了混合专家模型（MoE），主要介绍了针对 MoE 的高效训练方法，以及如何提升训练和 Fine-Tuning 的效果。现在我们回答下开篇提出的三个问题。

第一个问题：MoE 为什么能够实现在低成本下训练更大的模型。

这主要是因为稀疏路由的原因，每个 token 只会选择 top-k 个专家进行计算。同时可以使用模型并行、专家并行和数据并行，优化 MoE 的训练效率。而负载均衡损失可提升每个 device 的利用率。

第二个问题：MoE 如何解决训练稳定性问题？

可以通过混合精度训练、更小的参数初始化，以及 Router z-loss 提升训练的稳定性。

第三个问题：MoE 如何解决 Fine-Tuning 过程中的过拟合问题？

可以通过更大的 dropout （主要针对 expert）、更大的学习率、更小的 batch size。目前看到的主要是预训练的优化，针对 Fine-Tuning 的优化主要是一些常规的手段。

参考

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

社区发布 | 深度求索开源国内首个 MoE 大模型，技术报告、模型权重同时发布

Adaptive mixtures of local experts

混合专家模型（MoE）详解

蝈蝈：Mixture-of-Experts (MoE) 经典论文一览

Miller：详解专家混合：MoE模型

朱翔宇：NLP炼丹笔记：Switch Transformers 朴实无华 大招秒杀

AI Box专栏：深入解读首个万亿级语言模型 Switch Transformer

Kimi Chat - 帮你看更大的世界

GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

GLaM: Efficient Scaling of Language Models with Mixture-of-Experts

ST-MOE: Designing Stable and Transferable Sparse Expert Models

Mixtral of experts