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

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模型的训练和部署提供了帮助。