深度揭秘爆火 MoE：GPT-4 关键架构，成开源模型逆袭杀手锏

上周末，Mistral 甩出的开源 MoE 大模型，震惊了整个开源社区。MoE 究竟是什么？它又是如何提升了大语言模型的性能？

Mistral 上周末丢出的磁力链接震惊了开源圈子，这个 7B×8E 的开源 MoE 大模型性能已经到达了 LLaMA2 70B 的级别！

而根据 Jim Fan 猜测，如果 Mistral 内部训练了 34B×8E 或者甚至 100B+×8E 级别的模型，那他们的能力很有可能已经无限接近 GPT-4 了。

而在之前对于 GPT-4 结构的曝料中，大部分的信息也指向 GPT-4 很可能是由 8 个或者是 16 个 MoE 构成。

为什么 MoE 能成为高性能大模型的必选项？

简单来说，MoE 是一种神经网络架构设计，在 Transformer 模块中集成了专家 / 模型层。

当数据流经 MoE 层时，每个输入 token 都会动态路由到专家子模型进行处理。当每个专家专门从事特定任务时，这种方法可以实现更高效的计算并获得更好的结果。

MoE 最关键的组件：

- 专家（Expert）：MoE 层由许多专家、小型 MLP 或复杂的 LLM（如 Mistral 7B）组成。

- 路由器（Router）：路由器确定将哪些输入 token 分配给哪些专家。

路由策略有两种：token 选择路由器或路由器选择 token。

路由器使用 softmax 门控函数通过专家或 token 对概率分布进行建模，并选择前 k 个。

MoE 能够带来的好处：

- 每个专家都可以专门处理不同的任务或数据的不同部分。

- MoE 构架能向 LLM 添加可学习参数，而不增加推理成本。

- 可以利用稀疏矩阵的高效计算。

- 并行计算所有专家层，以有效利用 GPU 的并行能力。

- 帮助有效地扩展模型并减少训练时间。以更低的计算成本获得更好的结果！

MoE：建立高效大模型的关键技术

大型语言模型（LLM）席卷了机器学习领域，而随着现代数据集的复杂性和规模的增加，每个数据集都包含不同的模式，特征和标注之间的关系截然不同。

—— 这时，就需要 MoE 出手了。

专家混合（MoE）就像神经网络世界中的一种团队合作技术。

想象一下，将一项大任务分解成更小的部分，并让不同的专家处理每个部分。然后，有一个聪明的法官会根据情况，决定遵循哪位专家的建议，所有这些建议都混合在一起。

—— 就像你利用不同的口味，组合成一道好吃的菜。

对于复杂的数据集，可以划分为局部子集（local subsets），同样，将需要预测的问题划分为子任务（采用领域知识或者无监督聚类算法）。

然后，针对每个数据子集训练专家模型（Expert Models），专家模型可以是任何模型，比如支持向量机 （SVM） 或者神经网络，每个专家模型接收相同的输入模式并进行预测。

MoE 还包含门控模型（Gating Model），用于解释每个专家做出的预测，并根据输入选择信任哪个专家。

最后，MoE 需要一种聚合机制（Pooling Method），根据门控模型和专家的输出进行预测。

在现实的应用中，研究人员提出一种称为「稀疏门控专家混合层」的方法，作为原始 MoE 的迭代，这个方法提供了一个通用的神经网络组件，可以适应不同类型的任务。

稀疏门控专家混合层（Sparsely-Gated Mixture-of-Experts Layer）由许多专家网络组成，每个专家网络都是一个简单的前馈神经网络和一个可训练的门控网络。门控网络负责选择这些专家的稀疏组合来处理每个输入。

这里的重点是在门控功能中使用稀疏性，—— 这意味着对于每个输入实例，门控网络只选择少数专家进行处理，其余的专家保持非活动状态。

这种稀疏性和专家选择是针对每个输入动态实现的，整个过程具有高度的灵活性和适应性，而由于不需要处理网络的非活动部分，计算效率大大提高。

—— 简单来说就是，算得快、消耗少、省钱。

MoE 层可以分层堆叠，其中主 MoE 选择稀疏加权的「专家」组合。每个组合都使用一个 MoE 图层。

此外，研究人员还提出了一种名为「Noisy Top-K Gating」的创新技术。

这种机制在门控函数中增加了一个可调的高斯噪声，只保留前 K 个值，并将其余值分配给负无穷大，从而转换为零门控值。

这种方法确保了门控网络的稀疏性，同时保持了对门控函数输出中潜在不连续性的鲁棒性。另外，它还有助于跨专家网络进行负载平衡。

MoE 和 Transformer

下面我们来看一下 MoE 在 Transformer ，也就是当下大火的大语言模型中发挥了怎样的作用。

MoE 作为一种神经网络架构设计，可以集成进 Transformer 的结构中。

当数据流经 MoE 层时，每个输入（tokens）都会动态路由到某个专家模型进行计算，这样每个专家都可以专注于特定的任务，更好、更高效地给出结果。

上图展示了具有 MoE 层的 Transformer 编码器的演化（对于解码器的修改也类似），MoE 层取代了 Transformer 的前馈层。

上图左边是标准 Transformer 模型的编码器，包括自注意力层和前馈层，交错着残差连接和归一化层。

中间部分通过用 MoE 层替换其他前馈层，得到了 MoE Transformer Encoder 的模型结构。

上图右侧是当模型扩展到多个设备时的情况，MoE 层将跨设备分片，而所有其他层都会被复制。

我们可以看到 MoE 的关键组件就是各种专家模型和路由模块。

专家模型也可以是小型 MLP，或者复杂的 LLM，如 Mistral 7B。

而路由模块用来确定将哪些输入 tokens 分配给哪些专家。

一般有两种路由策略：token 选择路由器，或路由器选择 token。这里使用 softmax 门控函数，通过专家模型或者 tokens，对概率分布进行建模，并选择 top k。

由此，我们可知 MoE 层在 Transformer 中发挥了重要的作用。

每个专家都可以专门处理不同的任务或数据的不同部分；使用 MoE 可以在不增加推理成本的情况下向 LLM 添加可学习参数；

此外，MoE 还有助于对稀疏矩阵进行高效计算；而 MoE 中的专家层可以并行计算，这样就有效利用了 GPU 的并行能力。

最后，MoE 在帮助减少训练时间的同时，还可以做到有效地扩展模型，以更低的计算成本获得更好的结果。

MoE 开源再受关注

在 Mistral 放出这个开源的 7B×8E 的 MoE 之前，英伟达和谷歌也放出过其他完全开源的 MoE。

曾在英伟达实习的新加坡国立大学博士生 Fuzhao Xue 表示，他们的团队在 4 个月前也开源了一个 80 亿参数的 MoE 模型。

数据来源

- 一半来自 The RedPajama，另一半来自 The Stack Dedup。

- 为提升模型的推理能力，采用了大量的编程相关数据。

模型架构

- OpenMoE 模型基于「ST-MoE」，但采用了 decoder-only 架构。

其它设计

- 采用 umT5 tokenizer。

- 使用 RoPE 技术。

- 采用 SwiGLU 激活函数。

- 设定 2000 token 的上下文长度。

BigBench 评估

团队在 BigBench-Lite 上进行了少样本测试，其中包括与 BIG-G、BIG-G-Sparse 以及 GPT-3 的对比。

通过计算每个词元激活的参数数量和训练词元的数量来大致估计相对成本。图中每个点的大小代表了相应词元激活的参数数量。特别需要指出的是，浅灰色的点表示 MoE 模型的总参数量。

对此，Jim Fan 也表示，MoE 并不新鲜，它只是没有得到那么多关注而已......

比如，谷歌很早之前就开源了基于 T5 的 MoE 模型 ——Switch Transformer。

面临的挑战和机遇

• MoE 基础设施建设

由于 MoE 拥有大量可训练参数，理想的软件环境应该支持灵活组合的专家级、张量级、流水线级和数据并行，无论是节点内还是节点间。

此外，如果能支持简单快速的激活卸载和权重量化，从而减轻 MoE 权重的内存占用，就更好了。

• MoE 指令微调

FLAN-MoE 研究提出：尽管将 MoE 的性能通过特定任务的微调转移到下游任务上存在挑战，但指令微调却能有效地与 MoE 模型协调一致。这展示了基于 MoE 的语言模型巨大的潜力。

• MoE 评估

MoE 模型的归纳偏置（Inductive bias）可能在困惑度（perplexity）之外还有其他效果，就像其他自适应模型（如 Universal Transformer 和 AdaTape）那样。

• 硬件挑战

值得一提的是，GPU 在跨节点通信方面面临挑战，因为每个节点通常只能配备有限数量的 GPU。这使得专家并行中，通信成为瓶颈。

幸运的是，NVIDIA 最近推出了 DGX GH200，将 256 个 NVIDIA Grace Hopper Superchips 集成到一个单一 GPU 中，很大程度上解决了通信带宽问题，为开源领域的 MoE 模型的训练和部署提供了帮助。

参考资料：

• https://twitter.com/sophiamyang/status/1733505991600148892

广告声明：文内含有的对外跳转链接（包括不限于超链接、二维码、口令等形式），用于传递更多信息，节省甄选时间，结果仅供参考，IT之家所有文章均包含本声明。