详解Mixtral-8x7B背后的MoE！

高端的模型往往只需最朴素的发布方式。

这个来自欧洲的大模型团队在12月8日以一条磁力链接的方式发布了Mixtral-8x7B，这是一种具有开放权重的「高质量稀疏专家混合模型」(SMoE)。

该模型在大多数基准测试中都优于Llama2-70B，相比之下推理速度快了6倍，同时在大多数标准基准测试中匹配或优于GPT-3.5。

之后，Mistral AI将模型权重推送至HuggingFace，并一起推送了Mixtral-8x7B-Instruct。该模型已通过监督微调和直接偏好优化(DPO)进行优化，更加遵循指令。在MT-Bench上，它达到了8.30的分数，使其成为最好的开源模型，性能可与GPT3.5相媲美。

Mixtral-8x7B共有46.7B个参数，但每个token仅使用12.9B个参数。也就是说该模型可以每次只需要120亿参数参与推理就可以达到700亿的LLaMA2、1750亿的GPT-3.5的水平，可以说是成本/性能权衡方面的最佳模型。

Mixtral-8×7B并不是8个7B参数模型的集合，而是Transformer中的前馈块有不同的8份。其背后的一个重要的技术就是「专家混合」(Mixture of Experts，MoE)。

什么是专家混合（MoE）？

模型的规模对于提升其质量至关重要。在有限的计算资源下，相较于用更多步骤训练一个小型模型，训练一个大型模型即便步骤更少效果通常更好。

MoE让模型以「远低于传统密集模型的计算成本」进行预训练，这意味着你可以在相同的计算预算下显著扩大模型或数据集的规模。特别是在预训练阶段，MoE模型能比其同等规模的密集型模型更快地达到相同的性能水平。

那么，MoE究竟是什么呢？在Transformer模型的背景下，MoE主要由两个部分组成：

• 稀疏MoE层代替了传统的密集前馈网络(FFN)层。MoE层包含若干“专家”(例如8个)，每个专家都是一个独立的神经网络。实际上，这些专家通常是FFN，但它们也可以是更复杂的网络，甚至可以是MoE本身，形成一个层级结构的MoE。
• 门控网络或路由器，用于决定哪些Token分配给哪个专家。值得注意的是，一个Token可以被分配给多个专家。如何高效地将Token分配给合适的专家，是使用MoE技术时需要考虑的关键问题之一。路由器由一系列可学习的参数构成，它与模型的其他部分一起进行预训练。

那么，简单回顾一下，MoE的设计思路是这样的：在Transformer模型中，将每一个FFN层替换为MoE层，由一个门控网络和若干“专家”组成。

虽然MoE为我们带来了一些优势，比如更高效的预训练和相较于密集模型更快的推理速度，但同时它也带来了一些挑战：

• 训练：MoEs在预训练阶段的计算效率极高，但在微调时往往难以适应新场景，容易造成过拟合现象。

• 推理：尽管MoE模型可能包含大量参数，但在推理过程中只有部分参数被使用，这使得它的推理速度远快于参数数量相同的密集模型。但这也导致了一个问题：所有参数都需加载到内存中，因此对内存的需求相当大。比如，对于Mixtral-8x7B这样的MoE，我们需要足够的VRAM来支持一个有47B参数的密集型模型。

MoEs简史

MoEs的概念最早出现在1991年的论文Adaptive Mixture of Local Experts中。这一理念与集成方法相似，目的是通过监督程序管理一个由不同网络构成的系统，每个网络处理训练样本的一部分。每个单独网络或“专家”，都在输入空间的不同区域有其特长。由单独的门控网络决定每个专家的权重，在训练过程中，同时对专家和门控网络进行训练。

在2010至2015年间，两个不同的研究领域推动了MoE的进一步发展：

• 将专家作为组件：在传统的MoE结构中，系统由一个门控网络和多个专家组成。MoEs作为整体模型已在SVM、高斯过程等方法中得到应用。Eigen等人的研究将MoEs作为更深层网络的一部分进行探索。这意味着MoE可以作为多层网络中的一层，使模型在大规模和高效率之间达到平衡。
• 条件计算：传统网络会将所有输入数据通过每一层。在此期间，Yoshua Bengio探索了一种基于输入Token动态激活或停用网络组件的方法。

这些研究促进了在自然语言处理领域对混合专家模型的探索。具体来说，Shazeer等人(2017年的研究，团队成员包括Geoffrey Hinton和Jeff Dean)将这一理念应用到了一个137B的LSTM(当时的NLP主要架构)上，通过引入稀疏性概念，即使在大规模应用中也能保持快速的推理速度。这项工作主要关注翻译领域，但也面临着高通信成本和训练不稳定等挑战。

MoEs的应用使得训练具有数万亿参数的模型成为可能，比如公开的1.6T参数的Switch Transformers等。除此之外，MoEs在计算机视觉领域也有所探索，不过本文将重点讨论NLP领域的应用。

什么是稀疏性？

稀疏性基于条件计算的概念。不同于密集模型中所有参数对所有输入都有效，稀疏性让我们能够只激活系统的部分区域。条件计算(即网络的某些部分仅针对特定样本激活)使得在不增加计算量的情况下扩大模型规模成为可能，从而在每层MoE中使用了数千名专家。

这种方法也带来了挑战。比如，虽然大批量处理通常能提高性能，但在MoE中，当数据通过活跃的专家时，实际的批量大小会减小。例如，如果我们的批量输入包含10个Token，可能有5个Token由一个专家处理，另外5个Token分别由5个不同的专家处理，这导致批量大小不均匀，资源利用率低下。

那我们该如何解决这些问题呢？让我们深入探讨Shazeer在翻译领域对MoE的研究。

通过一个学习型的门控网络(G)，决定将输入的哪些部分分配给哪些专家(E)：

在这种设置中，所有专家都参与处理所有输入——这是一种加权乘法过程。但如果G的值为0呢？这种情况下，就无需计算相应专家的操作，从而节约了计算资源。那么，典型的门控函数是什么样的呢？在传统设置中，我们通常使用一个简单的网络配合softmax函数。这个网络会学习如何选择最合适的专家处理输入。

Shazeer的研究还探索了其他类型的门控机制，如带噪声的Top-K门控。这种方法加入了一些可调节的噪声，然后只保留最高的k个值。具体来说：

• 添加噪音

• 仅保留前k个值

• 应用softmax函数

这种稀疏性带来了一些有趣的特性。如果使用较低的k值(比如一到两个)，我们可以比激活许多专家时更快地进行训练和推理。为什么不只选择最顶尖的专家呢？最初的假设是，为了让门控学习如何路由到不同的专家，需要路由到一个以上的专家，因此至少需要选择两个专家。

我们为什么要加入噪声？这是为了实现负载均衡！

MoEs的负载均衡tokens

正如之前所讨论的，如果所有的tokens都被发送到少数几个受欢迎的专家，这将导致训练效率低下。在标准的多专家系统训练中，门控网络倾向于主要激活相同的几位专家。这会形成自我加强的循环，因为得到优先训练的专家会被更频繁地选择。为了减轻这种情况，引入了一种辅助损失来鼓励平等对待所有专家。这种损失确保所有专家获得大致相同数量的训练样本。下文还将探讨「专家容量」的概念，这涉及到一个专家能处理的tokens数量上限。在transformers中，这种辅助损失可以通过aux_loss参数来调节。

MoEs和Transformers

Transformers模型展示了一个明显的趋势：「增加参数的数量可以显著提高性能」。Google的GShard项目正是在这方面进行了深入探索，试图将Transformers模型扩展到超过6000亿个参数。

在GShard中，编码器和解码器里的部分FFN层被MoE层替代，并采用了一种称为「top-2」的门控机制。下图显示了这种设计在编码器部分的应用。这种设计对大规模计算尤其有利：当模型扩展到多个设备时，MoE层在这些设备间共享，而其他层则在每个设备上独立存在。

为了在大规模应用中保持效率和均衡的负载，GShard团队在设计上做了一些创新，除了引入了类似前一节提到的辅助损失机制外，还包括：

• 随机路由机制：在top-2设计中，我们始终选择表现最优的专家，但第二选择的专家则根据其权重以一定概率被选中。
• 专家处理能力限制：我们可以设定一个专家能处理的Token数量的上限。如果两个专家的处理能力都已达到上限，那么这个Token就会被认为是多余的，并通过残差连接传递到下一层，或在某些情况下被直接丢弃。这一概念在MoEs的应用中非常关键。为什么这样做？因为在模型编译时所有的张量形状都是静态确定的，但我们无法预先知道每个专家将处理多少Token，因此需要设定一个固定的处理能力上限。

在模型推理过程中，只有部分专家会被激活。同时，一些计算过程如自注意力机制会被所有Token共享。因此，尽管一个拥有8个专家的470亿参数模型听起来庞大，但实际上它的计算需求相当于一个120亿参数的密集型模型。如果采用top-2机制，模型会涉及约140亿参数，但由于注意力等操作是共享的，实际上模型真正使用的参数量仍然是120亿。

Switch Transformers

尽管MoEs充满潜力，但它们在训练和微调时面临稳定性挑战。Switch Transformers这项研究深入剖析了这些问题，并发布了一个具有2048个专家和1.6万亿参数的MoE模型。相较于T5-XXL，Switch Transformers的预训练速度提高了四倍。

Switch Transformers提出了一种处理两种不同token的新型Transformer层，包含四个专家。

不同于最初至少使用两个专家的设想，Switch Transformers采用了更简洁的单专家策略。这种策略的影响包括：

• 简化了路由计算
• 每个专家处理的批量至少减少了一半
• 减少了通信成本
• 保持了模型质量

此外，Switch Transformers还探讨了专家容量的概念。专家容量的计算公式是：

每批token数量除以专家数量，再乘以容量因子。按此计算方式，可以均匀分配批次中的Token给每个专家。如果容量因子大于1，可以为Token分配不均的情况提供缓冲。但容量增加会带来更高的设备间通信成本，这是一个需要权衡的问题。Switch Transformers在较低的容量因子下表现优异。

Switch Transformer的研究者还对上文提到的负载均衡损失进行了简化。在训练过程中，每个Switch层的辅助损失会加入到总模型损失中，这种做法促进了均匀的路由分配，并可以通过超参数进行调整。

研究者们还尝试了一种选择性的精确度方法，例如在训练专家系统时使用bfloat16格式，而在其他计算过程中则采用全精度。降低精度能够显著减少处理器间的通信成本、计算成本以及存储数据的内存需求。但初期实验中，无论是专家系统还是门控网络都采用bfloat16进行训练，结果训练过程变得更加不稳定。特别是路由器计算部分，由于其涉及到指数函数，因此更高的精度显得尤为重要。为了缓解这种不稳定性，路由过程最终也采用了「全精度处理」。

使用路由器Z-loss稳定模型训练

上文讨论过的平衡损失可能会引起训练稳定性的问题。为了稳定稀疏模型，我们可以采用多种方法，但这可能会牺牲模型的质量。例如，引入dropout虽然能增强稳定性，却会削弱模型的效果。而增加乘法运算组件虽然能提高模型质量，但又会降低其稳定性。

在ST-MoE研究中提出的路由器z-loss通过对门控网络输入的大数值logits施加惩罚，显著提高了训练的稳定性，同时又不会影响模型的质量。这种方法通过降低数值的绝对大小来减少舍入误差，这对于像门控这样的指数函数来说非常重要。

“专家”在模型学习中的影响

ST-MoE的研究者发现，编码器的专家倾向于专注于特定的Token组或基础概念。例如，可能形成专门处理标点符号或专有名词的专家。而解码器的专家则在专业化方面表现得较为平均。此外，作者还在多语言环境中进行了训练。虽然人们可能会认为每个专家会专注于一种特定语言，但实际情况却恰恰相反：由于Token的路由和负载均衡，没有任何一个专家专门对某一特定语言进行专研。

增加更多的专家可以提高样本效率和加速训练过程，但增益逐渐减少(特别是在达到256或512个专家后)，并且在推理过程中需要更多的VRAM。所以“专家”并不是越多越好。

微调MoE技术

密集型模型和稀疏型模型在过拟合上表现出明显不同的特点。稀疏型模型更易于过拟合，因此我们可以尝试在专家系统内部应用更强的正则化手段，例如「不同层次的dropout率」——对密集层和稀疏层分别设置不同的dropout率。

在微调过程中，一个关键的决策是是否采用辅助损失。ST-MoE的研究人员尝试关闭辅助损失，并发现即使高达11%的Token被丢弃，模型的质量也几乎不受影响。这表明「Token丢弃」可能是一种有效的防止过拟合的正则化策略。

Switch Transformers的研究发现，在预训练阶段达到固定的困惑度时，稀疏模型在下游任务中的表现通常不及密集型模型，特别是在逻辑推理较多的任务，如SuperGLUE上。然而，在知识密集型的任务，比如TriviaQA上，稀疏模型的表现却出奇地好。研究还发现，在微调阶段使用较少数量的专家有助于模型表现。此外，模型在小型任务中表现不佳，但在大型任务中则表现良好，这也证明了其泛化能力的问题。

另一个尝试是冻结所有非MoE层的权重，结果如预期那样导致了性能大幅下降，因为MoE层占据了网络的大部分。相反，仅冻结MoE层的参数几乎能达到更新所有参数的效果。这种方法可以加速微调过程，同时减少内存使用。

在调整稀疏型MoEs时，我们需要特别关注它们独特的微调超参数配置。比如，这类稀疏模型通常更适合较小的批量大小和较高的学习率。

2023年7月的一篇新论文MoEs Meets Instruction Tuning展示了一些有趣的实验：单任务微调、多任务指令微调以及在多任务指令微调后进行单任务微调。

研究者对比了微调后的MoE和T5等效模型，发现后者性能更优。但当微调Flan-T5(T5指令等效模型)MoE时，MoE的表现显著提高。不仅如此，Flan-MoE相比MoE的提升幅度，甚至超过了Flan-T5相比T5的提升，这表明 MoEs 可能从指令微调中获益更大，尤其是在任务数量更多的情况下。这与先前建议关闭辅助损失功能的讨论相反，实际上，这种损失可以帮助防止过拟合。

在多机器、高吞吐量的场景中，MoEs是非常有效的。如果预训练的计算预算有限，那么稀疏模型将是更佳的选择。对于VRAM较少、吞吐量低的情况，稠密模型则更为合适。

注意：我们不能直接比较稀疏和稠密模型之间的参数数量，因为这两种模型代表的是完全不同的概念。

加速MoEs的运行

在最初的MoE研究中，MoE层被设计成分支结构，这导致计算速度较慢，因为 GPU本身并不适合这种设计。同时，由于设备间需要传输信息，网络带宽成为了性能瓶颈。下面我们将探讨一些方法，以提高这些模型在预训练和推理阶段的实用性，使MoEs运行更加高效。

并行处理技术

简要介绍一下并行处理技术：

• 数据并行：相同的权重在所有核心上复制，数据则在核心之间分配。
• 模型并行：模型在各核心之间分配，数据在所有核心上复制。
• 模型和数据并行：可以在核心间分配模型和数据。需要注意的是，不同核心处理的是不同批次的数据。
• 专家并行：将不同的专家部署在不同的处理单元上。如果与数据并行结合，每个核心将配备一个不同的专家，数据则在所有核心间分配。

在专家并行模式下，不同的处理单元部署了不同的专家，每个处理单元处理不同批次的训练样本。对于非MoE层，专家并行的行为类似于数据并行。对于MoE层，序列中的Token被发送到拥有相应专家的处理单元。

容量因子和通信成本

提高容量因子(CF)可以增加模型质量，但同时也会增加通信成本和激活内存的需求。如果全面通信速度较慢，那么使用较小的容量因子将是更好的选择。一个较好的初始设置是使用top-2路由，1.25的容量因子，并且每个核心配置一个专家。在评估阶段，可以调整容量因子以减少计算量。

服务技巧

MoE的一个主要问题是它的参数特别多。如果是在本地环境中使用，可能会更倾向于使用一个体积更小的模型。下面，我们来看看几种有助于优化服务的技巧：

• Switch Transformers的研究者们早期就做了一些模型蒸馏的实验。通过将MoE模型蒸馏成更密集的形式，他们能够保留大约30-40%的稀疏性优势。因此，蒸馏不仅加快了模型的预训练速度，还能在实际应用中使用更小的模型。
• 最新的一些方法对路由机制进行了改进，能将整个句子或特定任务直接指派给某个专家，从而提取出适合服务的子网络。
• 专家聚合技术：这种方法通过合并不同专家的权重，在推理阶段有效减少了模型的参数数量。

关于高效训练的更多讨论

FasterMoE(2022年3月)深入分析了MoE在高效分布式系统中的表现。研究不仅探讨了不同并行处理策略的理论极限，还包括了如何倾斜专家的受欢迎程度、减少延迟的精细通信调度，以及一种新型的拓扑感知门控机制。这种机制通过选择延迟最低的专家来进行决策，从而实现了高达17倍的速度提升。

Megablocks(2022年11月)致力于探索高效的稀疏预训练技术。他们提出了一种新的GPU核心，能够处理MoE中的动态性。这一创新方法不会丢失任何 Token，并且能够高效地适应现代硬件，带来了显著的速度提升。那么，它的独特之处在哪里呢？与传统的MoE使用批量矩阵乘法不同(这种方法假设所有专家的形状和Token数量都一样)，Megablocks使用块稀疏运算来表达MoE层，这使得它能够适应不均匀的任务分配。

开源MoEs项目

1. Megablocks: https://github.com/stanford-futuredata/megablocks
2. Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm
3. OpenMoE: https://github.com/XueFuzhao/OpenMoE
4. Switch Transformers(Google): https://huggingface.co/collections/google/switch-transformers-release-6548c35c6507968374b56d1f
5. NLLB MoE(Meta): https://huggingface.co/facebook/nllb-moe-54b
6. Mixtral-8x7B(Mistral): https://huggingface.co/mistralai

以上就是关于Mistral AI新发布模型背后的MoE技术的详细解释。如果觉得还不错的话，欢迎关注我😘。

个人博客：https://jenqyang.github.io/

公众号：「ChaosstuffAI」