混合专家模型（MoE）通过稀疏路由机制，在相同计算资源下实现更大参数规模与更高训练效率。Switch Transformer 作为典型 MoE 架构，采用 Top-1 路由策略简化门控网络，显著降低通信成本。文章详细解析了 MoE 的核心优势，包括训练速度提升、推理成本降低及多任务学习能力，同时探讨了其面临的训练稳定性、通信开销及过拟合挑战。针对这些问题，Switch Transformer 提出了混合精度训练、小参数初始化及专家 Dropout 等优化方案，有效提升了模型收敛速度与微调效果。此外，还介绍了数据并行、模型并行及专家并行等高效训练策略，为构建万亿参数大模型提供了技术参考。
GPT-4 远不止 1 万亿，甚至是由 8 个 2200 亿参数组成的混合专家模型（MoE）。
2023 年 6 月，美国知名骇客 George Hotz 在接受采访时透露，GPT-4 由 8 个 220B 模型组成。这么算来，8 x 220B = 1.76 万亿。就连 PyTorch 的创建者 Soumith Chintala 对此也深信不疑。
MoE 应用于大模型，GPT-4 并不是第一个。在 2022 年的时候，Google 就提出了 MoE 大模型 Switch Transformer，模型大小是 1571B，Switch Transformer 在预训练任务上显示出比 T5-XXL（11B）模型更高的样本效率。在相同的训练时间和计算资源下，Switch Transformer 能够达到更好的性能。
除了 GPT-4 和 Switch Transformer，国内的团队 DeepSeek 也开源了国内首个 MoE 大模型 DeepSeekMoE。
MoE，全称为 Mixed Expert Models，翻译过来就是混合专家模型。MoE 并不是什么最新技术，早在 1991 年的时候，论文 [Adaptive Mixture of Local Experts] 就提出了 MoE。
模型规模是提升模型性能的关键因素之一，这也是为什么今天的大模型能取得成功。在有限的计算资源预算下，用更少的训练步数训练一个更大的模型，往往比用更多的步数训练一个较小的模型效果更佳。
MoE 的一个显著优势是它们能够在远少于 Dense 模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，可以显著扩大模型或数据集的规模。特别是在预训练阶段，与稠密模型相比，混合专家模型通常能够更快地达到相同的质量水平。
MoE 基于 Transformer 架构，主要由两部分组成：
![MoE 结构示意图]
Google Switch Transformer 论文中的 MoE 结构
总结来说，在混合专家模型 (MoE) 中，将传统 Transformer 模型中的每个前馈网络 (FFN) 层替换为 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 大模型优点，主要有以下几点点：
而 MoE 大模型的缺点，主要有以下 4 点：
接下来，就介绍下 MoE 的主要原理。通过后面的介绍，主要需要回答以下 3 个问题：
尽管 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 倍的加速。
![Switch Transformer 对比图]
左图：增加模型稀疏性（更多专家），loss 逐渐降低（图中 256e 表示 256 个 experts）
右图：Switch Transformers 和 T5-Base 的对比（相同计算资源）
上图中，模型参数随着专家数量的增加而增加，但保持了相同的计算成本（FLOPs per token）。这表明模型在保持计算效率的同时，能够利用更多的参数来提高性能。
右图中比较了使用相同计算资源下，Switch Transformer 和 T5-Base 的 PPL。可以看到 Switch Transformer 模型在保持相同计算资源的情况下，相对于 T5-Base 有显著的提升，而且专家数越多（模型参数越多、模型更稀疏），效果越好。
Swith Transformer 在论文中提到其设计的指导原则是——尽可能地把 Transformer 模型的参数量做大！（同时以一种简单高效的实现方式）
和其他 MoE 模型的一个显著不同就是，Switch Transformer 的门控网络每次只路由到 1 个 expert，也就是每次只选取 top1 的专家，而其他的模型都是至少 2 个。这样就是最稀疏的 MoE 了，因此单单从 MoE layer 的计算效率上讲是最高的了。下图是 Switch Transformer 的模型结构。
![Switch Transformer Encoder 模块]
Switch Transformer encoder 模块，使用稀疏的 Switch FFN 替换原来 Dense FFN
与最初使用至少两个专家的想法相反，Switch Transformer 采用了简化的单专家策略，每次只选择一个专家。这种方法的效果包括：
上面提到了 Switch Transformer 可以降低每个专家的专家容量，那么什么是专家容量？
专家容量（Expert Capacity） 是指每个专家在模型中处理的 token 数。专家容量的计算方式如下：
Expert Capacity = (tokens per batch / number of experts) × capacity factor
这里为什么要计算一个专家容量？这个专家容量又有什么作用？
在编译时，所有 tensor 的形状都是静态确定的。这意味着在编译阶段，模型的架构和数据布局已经被定义，包括模型的层数、每层的输入和输出维度等。
尽管 tensor 的形状是静态的，但在训练和推理过程中，模型的计算是动态的。这是因为模型中的路由器（门控网络）会根据输入数据动态地将 token 分配给不同的专家。这种动态性要求模型能够在运行时灵活地处理数据分布。
而这个专家容量的作用就是将 batch 中的总 token 数平均分配给所有专家。然后，为了应对 token 分布不均的情况，会通过一个容量因子（capacity factor）来扩展每个专家的容量。
容量因子是一个大于 1.0 的数，它的作用是为每个专家提供额外的缓冲空间，以容纳可能超出平均分配的 token。这样，即使某些专家接收到的 token 数量超过了平均值，也能够处理这些额外的 token，而不会因为容量不足而导致计算跳过。
下图是不同容量因子下的动态路由。
不同容量因子下的动态路由
如图所示，容量因子 Capacity Factor = 1.0的时候，输入 6 个 token，那么每个专家的专家容量等于 2（Expert Capacity = 6/3 * 1 = 2），Expert 1 被分配到了 3 个 token，超出了专家容量，这些超出的 token 被称为'溢出 token'（图（左）中的虚线部分）。对于这些溢出的 token，模型会跳过计算，直接将 token 的表示通过残差连接传递到下一层。
而如果容量因子 Capacity Factor = 1.5，这时专家容量等于 3，每个专家就能处理 3 个 token（图（右））。
虽然增加容量因子可以减少 token 溢出，但是它也有缺点。如果容量因子设置得过高，会导致计算资源和内存的浪费，因为模型会为可能永远不会用到的 token 分配额外的资源。在论文中，Switch Transformers 在低容量因子 (例如 1 至 1.25) 下表现出色。
下表是不同容量因子的效果对比。
表 1：Switch Transformer 和 MoE 的效果对比
| 模型 | 容量因子 | 训练 100k steps 后的负对数困惑度（越大越好） | 模型到达指定负对数困惑度（-1.5）所需时间（单位小时） | 训练速度（每秒处理的样本数） |
|---|---|---|---|---|
| T5-Base | -1.731 | 没有达到 | 1600 | |
| T5-Large | -1.550 | 131.1 | 470 | |
| MoE-Base | 2.0 | -1.547 | 68.7 | 840 |
| Switch-Base | 2.0 | -1.554 | 72.8 | 860 |
| MoE-Base | 1.25 | -1.559 | 72.8 | 790 |
| Switch-Base | 1.25 | -1.553 | 65.0 | 910 |
| MoE-Base | 1.0 | -1.572 | 80.1 | 860 |
| Switch-Base | 1.0 | -1.561 | 62.8 | 1000 |
| Switch-Base+ | 1.0 | -1.534 | 67.6 | 780 |
Switch Transformer 的作者还重新审视并简化了负载均衡损失。通过合理设置负载均衡损失的系数，可以在训练过程中实现专家之间的良好负载分布。下面介绍下具体实现。
在稀疏模型中，专家的数量通常分布在多个设备上，每个专家负责处理一部分输入数据。理想情况下，每个专家应该处理相同数量的数据，以实现资源的均匀利用。然而，在实际训练过程中，由于数据分布的不均匀性，某些专家可能会处理更多的数据，而其他专家可能会处理较少的数据。这种不均衡可能导致训练效率低下，因为某些专家可能会过载，而其他专家则可能闲置。为了解决这个问题，论文中引入了一种辅助损失函数，以促进专家之间的负载均衡。
给定 N 个专家，索引为 i=1 到 N，以及一个包含 T 个 token 的 batch B，辅助 loss 计算为向量 f 和 P 的缩放点积。表示如下：
auxiliary loss = α · N · Σ(i=1 to N) fi · Pi
其中，fi 是 batch 中分配给专家 i 的 token 占比，计算方式为 batch 中被路由到专家 i 的 token 数除以总 token 数，表示如下：
fi = (1/T) · Σ(x∈B) I{argmax p(x)=i}
Pi 是所有输入 token 被路由到专家 i 的概率，表示如下：
Pi = (1/T) · Σ(x∈B) pi(x)
其中 pi(x) 是给定 token x 被路由到专家 i 的概率。
由于希望 batch 中的所有 token 能够均匀地分配给 N 个专家。这意味着每个专家应该处理相同数量的 token，即每个专家处理的 token 比例应该是 1/N。
通过最小化公式的辅助 loss，可以鼓励这种均匀路由。当 token 均匀分布时，这个损失会被最小化。
最终的 loss 被乘以专家数量 N，这样即使专家数量变化，loss 也能保持恒定。这是因为在均匀路由情况下 Σ(i=1 to N)(fi·Pi) = Σ(i=1 to N)(1/N·1/N) = 1/N。
α 是一个超参数，用于调整辅助 loss 的权重。论文中选择了 α = 10^-2，这个值足够大，可以确保负载均衡，同时又足够小，不会压倒主要的交叉熵目标（即主要的训练损失）。论文实验了从 10^-1 到 10^-5 的 α 值范围，发现 $10^{-2}$ 的值可以快速平衡负载，同时不会干扰训练损失。
前面介绍了 Switch Transformer 的主要优化：稀疏路由和负载均衡损失。下面介绍一下将这两项优化合并在一起的实验效果。
实验设置如下：
实验结论如下（可以参见前面的表 1）：
作者还尝试了混合精度的方法，例如用 bfloat16 精度训练专家，同时对其余计算使用全精度进行。较低的精度可以减少处理器间的通信成本、计算成本以及存储 tensor 的内存。然而，在最初的实验中，当专家和门控网络都使用 bfloat16 精度训练时，出现了不稳定的训练现象。这种不稳定性主要是由路由计算引起的，因为路由涉及指数函数等操作，这些操作对精度要求较高。因此，为了保持计算的稳定性和精确性，保持更高的精度是重要的。为了减轻不稳定性，路由过程也使用了全精度。
下面的表 2 显示了混合精度训练的效果，将路由器输入转换为 float32，同时保持其他部分的精度为 bfloat16。这种策略允许模型在几乎与 bfloat16 精度相同的训练速度下，实现与 float32 训练相当的稳定性。
表 2：不同精度效果对比
| 模型精度选择 | 效果（负对数困惑度） | 训练速度（每秒处理样本数） |
|---|---|---|
| Switch-Base (float32) | -1.718 | 1160 |
| Switch-Base (bfloat16) | -3.780 | 1390 |
| Switch-Base (混合精度) | -1.716 | 1390 |
实验表明，使用混合精度的 Switch-Base 在固定步数的早期训练中，其效果（以负对数困惑度为衡量标准）与使用 float32 训练的模型相似，同时速度接近 bfloat16。
在深度学习中，适当的权重初始化对于模型的成功训练至关重要。作者观察到，在 Switch Transformer 模型中，这一点尤其明显。
为了提高模型的稳定性，作者建议减少默认的 Transformer 初始化规模。在 Transformer 模型中，权重矩阵通常是从一个截断的正态分布，其均值为 0，标准差由一个超参数 s 决定。作者建议将这个初始化超参数 s 从默认值 1.0 减少 10 倍，即 s = 0.1。这种较小的初始化规模有助于提高模型效果和减少训练过程中的不稳定性。
表 3：减小参数初始化规模可以提升训练稳定性
| 权重初始化规模 | 负对数困惑度 | 负对数困惑度标准差 |
|---|---|---|
| 0.1 倍初始化 | -2.72 | 0.01 |
| 1.0 倍初始化 | -3.60 | 0.68 |
表 3 中的数据表明，通过减少初始化规模，模型效果和稳定性得到了提升。这种改进对于大模型，如 Switch Transformer，尤其重要。
为了解决 Fine-Tuning 过程中的过拟合问题，作者提出了增加 dropout 的策略，特别是在专家层（expert layers）中。他们称之为'expert dropout'，即在 Fine-Tuning 时只在专家层增加 dropout 率。
表 4 显示了在 Fine-Tuning Switch Transformer 时，不同 dropout 率的实验结果。这些模型是在 C4 数据集上预训练的，然后进行了 Fine-Tuning。
表 4：Fine-Tuning 过程中正则化效果
| 模型（dropout） | GLUE | CNNDM | SQuAD | SuperGLUE |
|---|---|---|---|---|
| T5-Base (d=0.1) | 82.9 | 19.6 | 83.5 | 72.4 |
| Switch-Base (d=0.1) | 84.7 | 19.1 | 83.7 | 73.0 |
| Switch-Base (d=0.2) | 84.4 | 19.2 | 83.9 | 73.2 |
| Switch-Base (d=0.3) | 83.9 | 19.6 | 83.4 | 70.7 |
| Switch-Base (d=0.1, expert d=0.4) | 85.2 | 19.6 | 83.7 | 73.0 |
通过这种 expert dropout 策略，有效地减少了过拟合的风险，同时保持了模型在下游任务上的性能。这种正则化方法对于处理具有大量参数的稀疏模型特别有用，因为它可以帮助模型更好地泛化到未见过的数据。
任意增加专家数量会导致收益递减（如图 3 所示）。这意味着在某个点之后，继续增加专家数量不会显著提高模型性能。但是可以通过增加模型的维度，如模型的隐藏层大小（dmodel）或前馈网络的维度（dff）来继续提升模型效果。但是这样又会导致显存和内存开销增加，这时候就可以通过并行技术，解决高效训练问题。
这里补充一下关于各种并行的方法的解释。标准的数据并行的定义是一个 batch 的数据在不同的 device 上并行处理，这时每一个 device 上都保存了模型的一份完整拷贝，前向计算完进行梯度汇总和更新。模型并行表示模型不同的参数（层、组件）分配到不同的 device 上，处理一个 batch 的数据。
图（a）模型权重的分配方式
（b）数据的分配方式
图（a）表示模型权重的分配方式，图（b）表示数据的分配方式，一种颜色表示一个矩阵（a unique weight matrix）。其中每一个方格表示一个 core。
第一列表示数据并行，模型权重拷贝 16 份，16 个同一种颜色矩阵分别表示一个完整的模型，图（b）则是一个完整的矩阵，这里可以理解为 16 个模型计算完成后由于存在梯度汇总再更新的步骤，所以整体更新的是一个 batch，因此这里 Data Parallelism 是一个唯一的矩阵。简单来说就是模型复制，数据并行。
数据并行
模型并行部分从模型侧看出来，16 个 cores 维护的是一个整体的模型，但是每一个 core 只分配到其中部分模型参数（图（a）），同一个 batch 数据在所有的 core 上计算（图（b）），由于 1 个 core 中分布了不同的模型权重，每次计算完都需要和其他的 core 进行通信。
模型并行
总共有 N 个 cores，其中 N = n × m，n 代表数据并行维度上的分割因子，m 代表模型并行维度上的分割因子。现在每个 core 处理的是 B/n 个 token 以及 dff/m 个权重。
模型和数据并行
每个专家分配到一个 core 上，同时数据也切分成 16 份，如下图所示：
专家和数据并行
最后将专家、模型、数据并行合并在一起，如下图所示：
模型、专家和数据并行
本文系统性地介绍了混合专家模型（MoE），主要介绍了针对 MoE 的高效训练方法，以及如何提升训练和 Fine-Tuning 的效果。现在我们回答下开篇提出的三个问题。
第一个问题：MoE 为什么能够实现在低成本下训练更大的模型。
这主要是因为稀疏路由的原因，每个 token 只会选择 top-k 个专家进行计算。同时可以使用模型并行、专家并行和数据并行，优化 MoE 的训练效率。而负载均衡损失可提升每个 device 的利用率。
第二个问题：MoE 如何解决训练稳定性问题？
可以通过混合精度训练、更小的参数初始化，以及 Router z-loss 提升训练的稳定性。
第三个问题：MoE 如何解决 Fine-Tuning 过程中的过拟合问题？
可以通过更大的 dropout（主要针对 expert）、更大的学习率、更小的 batch size。目前看到的主要是预训练的优化，针对 Fine-Tuning 的优化主要是一些常规的手段。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online