大模型Transformer的MOE架构介绍及方案整理

前言：DeepSeek模型最近引起了NLP领域的极大关注，也让大家进一步对MOE（混合专家网络）架构提起了信心，借此机会整理下MOE的简单知识和对应的大模型。本文的思路是MOE的起源介绍、原理解释、再到现有MOE大模型的整理。

一、MOE的起源和架构

MoE的概念最早由MIT等人在论文中指出：混合专家网络可以看作是多层监督网络的模块化版本。比如元音识别任务，可以分解为多个子任务，每个子任务可以由一个非常简单的专家网络解决。

图1-1：最早的MOE模型（经典之作，其思想沿用至今）-框架图

从让专家之间学会合作->过渡到让专家之间学会竞争：在合作时，各个专家之间是强耦合的，导致解决方案中使用多个专家；当转换为竞争后，将可以得到少数专家活跃的解决方案。这可以通过修改误差函数实现，见图1-2。

图1-2：最早的MOE模型-损失函数

随着稀疏门控MoE的出现（Sparsely-Gated Mixture-of-Experts），特别是在基于Transformer的LLM中成功地集成（Gshard），为这一30年历史的技术注入了新的活力。

小结——MoE框架基于一个简单而强大的理念：模型的不同部分（称为专家）专注于不同的任务。在这种范式下，只有与给定输入相关的专家会被激活，从而使得模型在具备海量专业知识的同时，保持计算成本的可控性。

二、MOE的分类

根据激活专家情况，可以把MOE模型分为Dense MoE和Sparse MoE，接下来分别展开介绍。

图2-1：MOE模型的分类（根据激活专家情况）

2.1 Dense MoE

Dense MoE在每次迭代中激活所有专家网络，优缺点如下：

• 优点：通常能够提供更高的预测准确性
• 缺点：会显著增加计算开销

Dense MoE层的输出可以表示为：

图2-2：Dense MoE层的输出计算

2.2 Sparse MoE

为了解决Dense MoE的"显著增加计算开销"这一问题，谷歌等人提出了Sparse MoE层，即在每次前向传播过程中仅激活选定的一部分专家，GShard便是其中的经典之作。这一策略通过计算加权和的前 k个专家的输出，而不是聚合所有专家的输出，从而实现了稀疏性。稀疏MoE层的结构如图2-1。稀疏门控机制的公式可以修改为：

图2-3：Sparse MoE层的输出计算

尽管稀疏门控显著扩展了模型的参数空间而不增加计算成本，但它可能导致负载均衡问题：即专家之间工作负载分布不均，某些专家频繁使用，而其他专家很少或从未使用。

为了解决这一问题，每个MoE层都引入了一个辅助负载均衡损失（Auxiliary load balancing loss），以促进每个batch中各专家之间token的均匀分布：

图2-4：Sparse MoE引入的辅助负载均衡损失的公式

通过引入辅助loss，模型保持了所有专家之间的平衡，以促使所有时间内专家的工作负载满足均匀分布。

三、MOE各系列大模型技术点汇总

基于MOE思想构建大模型，自2018的提出->到2022年底chatGPT的出现->再到如今DeepSeek大火，已经经历了七年之久，模型更新脉络如下图3-1所示，本文会将代表性MOE（热度高/效果好）大模型总结在本章节。

图3-1：基于MOE的LLM汇总

3.1 Mistral-MOE

Mixtral 8x7B：一种稀疏混合专家（SMoE）语言模型。它具有与Mistral 7B（其结构可参考笔者另一篇文章）相似的架构，不同之处在于每一层由8个FFN模块（即专家）组成。对于每个token，在每一层，路由网络会选择两个专家（topk=2）来处理当前状态并整合它们的输出。尽管每个token只看到2个专家，但选择的专家在每个时间步可能不同。因此，每个token可以访问47B参数，但在推理过程中只使用13B活跃参数。Mixtral使用32k个token的上下文长度进行训练，并在所有评测基准上优于或等于Llama2-70B和GPT-3.5。

图3-2：Mistral-MOE的架构参数

参考： https://arxiv.org/pdf/2401.04088、 假如给我一只AI：LLM开源大模型汇总-截止0218

3.2 LLaMA-MOE

基于LLaMA2-7B 模型（其结构可参考笔者另一篇文章），作者通过"专家构建"和"持续预训练"这两步就获得了 MoE 模型。

图3-3：LLaMA-MOE模型的两步操作——专家构建和持续预训练

• 最终效果：LLaMA-MoE 模型能够保持语言能力，并将输入的 token 路由到特定的专家，且部分参数被激活。
• 实验表明：通过训练 200B token，LLaMA-MoE-3.5B 模型在性能上显著优于包含类似激活参数的Dense模型。

1）专家构建：将原始FFN层的参数分割成多个专家

图3-4：LLaMA-MOE的专家构建流程梳理

2）持续预训练：进一步训练转换后的 MoE 模型和额外的门网络

在经历"专家构建"后，原始LLaMA模型结构会被重新组织为MoE架构，为了恢复其语言建模能力，作者采用"持续预训练"策略进一步训练LLaMA-MoE模型（该策略使用的训练目标与 LLaMA2 相同）。为了提高训练效率，作者探索了不同的'数据采样策略"和"数据过滤策略"。

如果要采用"持续预训练策略"，可能遇到问题见表3-1：

表3-1：持续预训练可能遇到的问题

文章具体采用的方法：1）采用"数据过滤"，得到去噪且流畅性高的数据；2）对比四种"数据采样策略"，实验对比哪种好选择哪种即可。具体总结如下表3-2：

表3-2：4种采样策略和2种数据过滤策略

参考： https://arxiv.org/pdf/2406.16554、 Swish激活函数、 LLaMA2论文)

3.3 Deepseek-MOE

1）DeepSeek-MoE（V1版模型）

解决当前MOE模型存在的两方面问题：

• 专家数量小但token信息丰富：将多样的知识分配给有限的专家，有概率导致专家"试着在有限的参数中学习大量不同类型的知识"，然而这些知识又难以同时利用，最终会降低专家的专业性。
• 多个专家之间存在知识冗余：在传统路由策略中，分配token给不同专家时可能需要一些"共享知识"。因此，多个专家可能在各自参数中"收敛于共享知识"，这就会导致专家参数冗余。

DeepSeek-MoE给出的解决方案见下图：

图3-5：DeepSeek-MoE的细粒度专家和共享专家

在此基础上，DeepSeek-MOE也具有考虑了负载平衡：即自动学习的路由策略可能会遇到负载不平衡的问题，这会导致两个显著的缺陷：[A] 存在路由崩溃的风险，即模型始终选择少数几个专家，其他专家缺乏充分训练；[B] 如果专家分布在多个设备上，负载不平衡会加剧计算瓶颈。

解决2个问题，分别提出了专家级负载loss和设备级负载loss，问题->解决->公式的解释如下图：

图3-6：DeepSeek-MoE的专家级负载和设备级负载，公式推导见https://zhuanlan.zhihu.com/p/18565423596

2）DeepSeek-V2模型

在DeepSeek-MoE的基础上，新增了一个路由机制和两个负载均衡方法，即设备受限的专家路由机制、通信负载均衡loss和设备级Token丢弃策略，它们的问题->解决->公式的解释如下两图：

图3-7a：DeepSeek-V2的设备受限的专家路由机制

图3-7b：DeepSeek-V2的通信负载均衡和设备级Token丢弃策略

3）DeepSeek-V3模型

相比DeepSeek-V2，DeepSeek-V3在MOE架构上的改进有三点：

• 使用 sigmoid 函数计算亲和度，并对所有选定的亲和度进行归一化以产生门值（图3-8a）。
• 提出了无辅助Loss的负载均衡技术和sequence粒度的负载均衡Loss（图3-8b）。
• 接入了节点限制的路由和无token丢弃策略（图3-8c）。

图3-8a：DeepSeek-V3的亲和度计算公式

图3-8b：DeepSeek-V3的无辅助Loss的负载均衡技术和sequence粒度的负载均衡Loss

图3-8c：DeepSeek-V3的节点限制的路由和无token丢弃策略

代码学习：DeepSeek-MoE源码、DeepSeek-V3源码

参考： DeepSeek-MOE论文、 DeepSeek-V2论文、 DeepSeek-V3论文

3.4 Qwen-MOE

【持续更新】

https://qwenlm.github.io/blog/qwen-moe/

3.5 Nvidia-MOE

【持续更新】

3.6 Grok-MOE

【持续更新】

3.7 Skywork-MOE

【持续更新】

https://arxiv.org/pdf/2406.06563

四、参考文献

• MOE综述：https://arxiv.org/pdf/2407.06204
• 姜富春：deepseek技术解读(3)-MoE的演进之路
• Gshard：https://arxiv.org/pdf/2006.16668
• https://arxiv.org/pdf/1701.06538
• Mistral-moe：https://arxiv.org/pdf/2401.04088