大模型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给出的解决方案见下图：