MoE
随着 Mixtral 8x7B (announcement (opens new window), model card (opens new window)) 的推出,一种称为混合专家模型 (Mixed Expert Models,简称 MoEs) 的 Transformer 模型在开源人工智能社区引起了广泛关注。
简短总结
混合专家模型 (MoEs):
- 与稠密模型相比, 预训练速度更快
- 与具有相同参数数量的模型相比,具有更快的 推理速度
- 需要 大量显存,因为所有专家系统都需要加载到内存中
- 在 微调方面存在诸多挑战,但 近期的研究 (opens new window) 表明,对混合专家模型进行 指令调优具有很大的潜力。
什么是混合专家模型?
模型规模是提升模型性能的关键因素之一。在有限的计算资源预算下,用更少的训练步数训练一个更大的模型,往往比用更多的步数训练一个较小的模型效果更佳。
混合专家模型 (MoE) 的一个显著优势是它们能够在远少于稠密模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下,您可以显著扩大模型或数据集的规模。特别是在预训练阶段,与稠密模型相比,混合专家模型通常能够更快地达到相同的质量水平。
那么,究竟什么是一个混合专家模型 (MoE) 呢?作为一种基于 Transformer 架构的模型,混合专家模型主要由两个关键部分组成:
- 稀疏 MoE 层: 这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个),每个专家本身是一个独立的神经网络。在实际应用中,这些专家通常是前馈网络 (FFN),但它们也可以是更复杂的网络结构,甚至可以是 MoE 层本身,从而形成层级式的 MoE 结构。
- 门控网络或路由: 这个部分用于决定哪些令牌 (token) 被发送到哪个专家。例如,在下图中,“More”这个令牌可能被发送到第二个专家,而“Parameters”这个令牌被发送到第一个专家。有时,一个令牌甚至可以被发送到多个专家。令牌的路由方式是 MoE 使用中的一个关键点,因为路由器由学习的参数组成,并且与网络的其他部分一同进行预训练。

Switch Transformers paper (opens new window) 论文中的 MoE layer
总结来说,在混合专家模型 (MoE) 中,我们将传统 Transformer 模型中的每个前馈网络 (FFN) 层替换为 MoE 层,其中 MoE 层由两个核心部分组成: 一个门控网络和若干数量的专家。
尽管混合专家模型 (MoE) 提供了若干显著优势,例如更高效的预训练和与稠密模型相比更快的推理速度,但它们也伴随着一些挑战:
- 训练挑战: 虽然 MoE 能够实现更高效的计算预训练,但它们在微调阶段往往面临泛化能力不足的问题,长期以来易于引发过拟合现象。
- 推理挑战: MoE 模型虽然可能拥有大量参数,但在推理过程中只使用其中的一部分,这使得它们的推理速度快于具有相同数量参数的稠密模型。然而,这种模型需要将所有参数加载到内存中,因此对内存的需求非常高。以 Mixtral 8x7B 这样的 MoE 为例,需要足够的 VRAM 来容纳一个 47B 参数的稠密模型。之所以是 47B 而不是 8 x 7B = 56B,是因为在 MoE 模型中,只有 FFN 层被视为独立的专家,而模型的其他参数是共享的。此外,假设每个令牌只使用两个专家,那么推理速度 (以 FLOPs 计算) 类似于使用 12B 模型 (而不是 14B 模型),因为虽然它进行了 2x7B 的矩阵乘法计算,但某些层是共享的。
专家的架构
虽然将专家可视化为密集模型的隐藏层被切分成若干部分很有帮助,但实际上它们通常本身就是完整的FFNN。

由于大多数 LLM 具有多个解码器块(Decoder Blocks),因此给定的文本会在生成之前通过多个专家:

不同 tokens 被选中的专家可能各不相同,这导致了不同的“路径”被选择:

如果我们更新解码器块的可视化,现在它将包含多个 FFNN(每个 FFNN 对应一个“专家”):

解码器块现在包含多个可以在推理时使用的 FFNN(即“专家”)。
混合专家模型简史
混合专家模型 (MoE) 的理念起源于 1991 年的论文 Adaptive Mixture of Local Experts (opens new window)。这个概念与集成学习方法相似,旨在为由多个单独网络组成的系统建立一个监管机制。在这种系统中,每个网络 (被称为“专家”) 处理训练样本的不同子集,专注于输入空间的特定区域。那么,如何选择哪个专家来处理特定的输入呢?这就是门控网络发挥作用的地方,它决定了分配给每个专家的权重。在训练过程中,这些专家和门控网络都同时接受训练,以优化它们的性能和决策能力。
在 2010 至 2015 年间,两个独立的研究领域为混合专家模型 (MoE) 的后续发展做出了显著贡献:
- 组件专家: 在传统的 MoE 设置中,整个系统由一个门控网络和多个专家组成。在支持向量机 (SVMs) 、高斯过程和其他方法的研究中,MoE 通常被视为整个模型的一部分。然而,Eigen、Ranzato 和 Ilya 的研究 (opens new window) 探索了将 MoE 作为更深层网络的一个组件。这种方法允许将 MoE 嵌入到多层网络中的某一层,使得模型既大又高效。
- 条件计算: 传统的神经网络通过每一层处理所有输入数据。在这一时期,Yoshua Bengio 等研究人员开始探索基于输入令牌动态激活或停用网络组件的方法。
这些研究的融合促进了在自然语言处理 (NLP) 领域对混合专家模型的探索。特别是在 2017 年,Shazeer 等人 (opens new window) (团队包括 Geoffrey Hinton 和 Jeff Dean,后者有时被戏称为 “谷歌的 Chuck Norris” (opens new window)) 将这一概念应用于 137B 的 LSTM (当时被广泛应用于 NLP 的架构,由 Schmidhuber 提出)。通过引入稀疏性,这项工作在保持极高规模的同时实现了快速的推理速度。这项工作主要集中在翻译领域,但面临着如高通信成本和训练不稳定性等多种挑战。

Outrageously Large Neural Network 论文中的 MoE layer
混合专家模型 (MoE) 的引入使得训练具有数千亿甚至万亿参数的模型成为可能,如开源的 1.6 万亿参数的 Switch Transformers 等。这种技术不仅在自然语言处理 (NLP) 领域得到了广泛应用,也开始在计算机视觉领域进行探索。
什么是稀疏性?
稀疏性的概念采用了条件计算的思想。在传统的稠密模型中,所有的参数都会对所有输入数据进行处理。

相比之下,稀疏性允许我们仅针对整个系统的某些特定部分执行计算。这意味着并非所有参数都会在处理每个输入时被激活或使用,而是根据输入的特定特征或需求,只有部分参数集合被调用和运行。

让我们深入分析 Shazeer 对混合专家模型 (MoE) 在翻译应用中的贡献。条件计算的概念 (即仅在每个样本的基础上激活网络的不同部分) 使得在不增加额外计算负担的情况下扩展模型规模成为可能。这一策略在每个 MoE 层中实现了数以千计甚至更多的专家的有效利用。
这种稀疏性设置确实带来了一些挑战。例如,在混合专家模型 (MoE) 中,尽管较大的批量大小通常有利于提高性能,但当数据通过激活的专家时,实际的批量大小可能会减少。比如,假设我们的输入批量包含 10 个令牌, 可能会有五个令牌被路由到同一个专家,而剩下的五个令牌分别被路由到不同的专家。这导致了批量大小的不均匀分配和资源利用效率不高的问题。在接下来的部分中,将会讨论 让 MoE 高效运行 的其他挑战以及相应的解决方案。
路由机制
现在我们有了一组专家,模型如何知道该使用哪些专家呢?在专家之前会加入一个可学习的门控网络 (

在这种设置下,虽然所有专家都会对所有输入进行运算,但通过门控网络的输出进行加权乘法操作。但是,如果
Shazeer 等人的工作还探索了其他的门控机制,其中包括带噪声的 TopK 门控 (Noisy Top-K Gating)。这种门控方法引入了一些可调整的噪声,然后保留前
- 添加一些噪声
- 选择保留前
个值
- 应用 Softmax 函数
这种稀疏性引入了一些有趣的特性。通过使用较低的
我们为什么要添加噪声呢?这是为了专家间的负载均衡!
混合专家模型中令牌的负载均衡
正如之前讨论的,如果所有的令牌都被发送到只有少数几个受欢迎的专家,那么训练效率将会降低。在通常的混合专家模型 (MoE) 训练中,门控网络往往倾向于主要激活相同的几个专家。这种情况可能会自我加强,因为受欢迎的专家训练得更快,因此它们更容易被选择。为了缓解这个问题,引入了一个 辅助损失,旨在鼓励给予所有专家相同的重要性。这个损失确保所有专家接收到大致相等数量的训练样本,从而平衡了专家之间的选择。接下来的部分还将探讨专家容量的概念,它引入了一个关于专家可以处理多少令牌的阈值。在 transformers
库中,可以通过 aux_loss
参数来控制辅助损失。
KeepTopK
一种对路由器进行负载平衡的方法是使用一个简单的扩展策略,称为 KeepTopK。通过引入可训练的(高斯)噪声,我们可以防止总是选择相同的专家:

https://arxiv.org/pdf/1701.06538
然后,除希望激活的前

将这些权重设为

尽管存在许多有前景的替代方案,许多 LLM 仍然使用 KeepTopK 策略。需要注意的是,KeepTopK 也可以不使用额外的噪声。
Token 选择策略
KeepTopK 策略会将每个 token 路由到若干选定的专家。这种方法被称为 Token 选择策略(Token Choice),它允许一个给定的token 被路由到一个专家(top-1 路由):

或者被路由到多个专家(top-k 路由):

这种策略的主要优点在于,它可以对各个专家的贡献进行加权,并将其整合起来。
辅助损失(Auxiliary Loss)
为了在训练期间实现专家的均匀分布,网络的常规损失中加入了辅助损失(也称为负载平衡损失)。
辅助损失增加了一个约束,强制专家在训练过程中具有相同的重要性。
辅助损失的第一个组成部分是对整个批次中每个专家的路由值进行求和:

这为我们提供了每个专家的重要性得分,即在不考虑输入的情况下,给定专家被选中的概率。
我们可以使用这些重要性得分计算变异系数(Coefficient of Variation, CV),它表示各个专家的重要性得分之间的差异程度。

例如,如果重要性得分之间的差异较大,那么 CV 值就会较高:

相反,如果所有专家的得分都相似,则 CV 值较低(这是我们期望的情况):

通过使用这个 CV 得分,我们可以在训练过程中更新辅助损失,使其尽可能降低 CV 得分(从而使每个专家具有相同的重要性):

最后,辅助损失将作为一个独立的损失项,参与训练优化。
专家容量(Expert Capacity)
专家的不平衡不仅体现在被选中的专家上,还体现在分配给这些专家的 token 分布上。例如,如果输入 token 被不成比例地分配到某些专家上,而不是平均分配,这可能导致某些专家的训练不足:

这里不仅要考虑使用了哪些专家,还需要关注这些专家被使用的频率。
解决这个问题的方法是限制每个专家能够处理的 token 数量,即专家容量(Expert Capacity)。当一个专家达到其容量时,多余的 token 将被分配到下一个专家:

如果两个专家都达到了其容量,token 将不会被任何专家处理,而是直接传递到下一层。这种情况被称为 token 溢出(token overflow)。

MoEs and Transformers
Transformer 类模型明确表明,增加参数数量可以提高性能,因此谷歌使用 GShard (opens new window) 尝试将 Transformer 模型的参数量扩展到超过 6000 亿并不令人惊讶。
GShard 将在编码器和解码器中的每个前馈网络 (FFN) 层替换为使用 Top-2 门控的混合专家模型 (MoE) 层。下图展示了编码器部分的结构。这种架构对于大规模计算非常有效: 当扩展到多个设备时,MoE 层在不同设备间共享,而其他所有层则在每个设备上复制。“让 MoE 起飞” 部分对这一点进行更详细的讨论。

为了保持负载平衡和训练效率,GShard 的作者除了引入了上一节中讨论的类似辅助损失外,还引入了一些关键变化:
- 随机路由: 在 Top-2 设置中,我们始终选择排名最高的专家,但第二个专家是根据其权重比例随机选择的。
- 专家容量: 我们可以设定一个阈值,定义一个专家能处理多少令牌。如果两个专家的容量都达到上限,令牌就会溢出,并通过残差连接传递到下一层,或在某些情况下被完全丢弃。专家容量是 MoE 中最重要的概念之一。为什么需要专家容量呢?因为所有张量的形状在编译时是静态确定的,我们无法提前知道多少令牌会分配给每个专家,因此需要一个固定的容量因子。
注意: 在推理过程中,只有部分专家被激活。同时,有些计算过程是共享的,例如自注意力 (self-attention) 机制,它适用于所有令牌。这就解释了为什么我们可以使用相当于 12B 稠密模型的计算资源来运行一个包含 8 个专家的 47B 模型。如果我们采用 Top-2 门控,模型会使用高达 14B 的参数。但是,由于自注意力操作 (专家间共享) 的存在,实际上模型运行时使用的参数数量是 12B。
Switch Transformers
尽管混合专家模型 (MoE) 显示出了很大的潜力,但它们在训练和微调过程中存在稳定性问题。Switch Transformers (opens new window) 是一项非常激动人心的工作,它深入研究了这些话题。作者甚至在 Hugging Face 上发布了一个 1.6 万亿参数的 MoE (opens new window),拥有 2048 个专家,你可以使用 transformers
库来运行它。Switch Transformers 实现了与 T5-XXL 相比 4 倍的预训练速度提升。
切换层(Switching Layer)
就像在 GShard 中一样,作者用混合专家模型 (MoE) 层替换了前馈网络 (FFN) 层。Switch Transformers 提出了一个 Switch Transformer 层,它接收两个输入 (两个不同的令牌) 并拥有四个专家。

与最初使用至少两个专家的想法相反,Switch Transformers 采用了简化的单专家策略。这种方法的效果包括:
- 减少门控网络 (路由) 计算负担
- 每个专家的批量大小至少可以减半
- 降低通信成本
- 保持模型质量

路由器在选择专家时并没有使用特殊的方法,只是对输入与专家权重相乘的结果取 softmax(与之前的方法相同)。

这种架构(top-1 路由)假设每个 token 只需要一个专家来学习如何进行路由。这与我们之前讨论的 top-k 路由(将 token 分配给多个专家)有所不同。
容量因子
容量因子是一个重要参数,它决定了每个专家可以处理的 token 数量。Switch Transformer 通过引入直接影响专家容量的容量因子扩展了这一概念。


如果我们增加容量因子,则每个专家能够处理更多的 token。

然而,如果容量因子过大,就会浪费计算资源。相反,如果容量因子过小,模型性能会因为 token 溢出而下降。
上述建议的容量是将批次中的令牌数量均匀分配到各个专家。如果我们使用大于 1 的容量因子,我们为令牌分配不完全平衡时提供了一个缓冲。增加容量因子会导致更高的设备间通信成本,因此这是一个需要考虑的权衡。特别值得注意的是,Switch Transformers 在低容量因子 (例如 1 至 1.25) 下表现出色。
Switch Transformer 的作者还重新审视并简化了前面章节中提到的负载均衡损失。在训练期间,对于每个 Switch 层的辅助损失被添加到总模型损失中。这种损失鼓励均匀路由,并可以使用超参数进行加权。
作者还尝试了混合精度的方法,例如用 bfloat16
精度训练专家,同时对其余计算使用全精度进行。较低的精度可以减少处理器间的通信成本、计算成本以及存储张量的内存。然而,在最初的实验中,当专家和门控网络都使用 bfloat16
精度训练时,出现了不稳定的训练现象。这种不稳定性特别是由路由计算引起的,因为路由涉及指数函数等操作,这些操作对精度要求较高。因此,为了保持计算的稳定性和精确性,保持更高的精度是重要的。为了减轻不稳定性,路由过程也使用了全精度。

使用混合精度不会降低模型质量并可实现更快的训练
这个 Jupyter Notebook (opens new window) 展示了如何对 Switch Transformers 进行微调以进行摘要生成的详细指南。然而,在开始微调 Switch Transformers 之前,强烈建议您先阅读关于 微调混合专家模型 部分的内容。
Switch Transformers 采用了编码器 - 解码器的架构,实现了与 T5 类似的混合专家模型 (MoE) 版本。GLaM (opens new window) 这篇工作探索了如何使用仅为原来 1/3 的计算资源 (因为 MoE 模型在训练时需要的计算量较少,从而能够显著降低碳足迹) 来训练与 GPT-3 质量相匹配的模型来提高这些模型的规模。作者专注于仅解码器 (decoder-only) 的模型以及少样本和单样本评估,而不是微调。他们使用了 Top-2 路由和更大的容量因子。此外,他们探讨了将容量因子作为一个动态度量,根据训练和评估期间所使用的计算量进行调整。
辅助损失
为了进一步防止 token 被丢弃,Switch Transformer 引入了简化版的辅助损失(Auxiliary Loss)。
在简化版的辅助损失中,不再计算变异系数,而是将分配的 token 数量与每个专家的路由概率进行加权比较:

由于目标是希望将 token 在
用 Router z-loss 稳定模型训练
之前讨论的平衡损失可能会导致稳定性问题。我们可以使用许多方法来稳定稀疏模型的训练,但这可能会牺牲模型质量。例如,引入 dropout 可以提高稳定性,但会导致模型质量下降。另一方面,增加更多的乘法分量可以提高质量,但会降低模型稳定性。
ST-MoE (opens new window) 引入的 Router z-loss
在保持了模型性能的同时显著提升了训练的稳定性。这种损失机制通过惩罚门控网络输入的较大 logits
来起作用,目的是促使数值的绝对大小保持较小,这样可以有效减少计算中的舍入误差。这一点对于那些依赖指数函数进行计算的门控网络尤其重要。
专家如何学习?
ST-MoE 的研究者们发现,编码器中不同的专家倾向于专注于特定类型的令牌或浅层概念。例如,某些专家可能专门处理标点符号,而其他专家则专注于专有名词等。与此相反,解码器中的专家通常具有较低的专业化程度。此外,研究者们还对这一模型进行了多语言训练。尽管人们可能会预期每个专家处理一种特定语言,但实际上并非如此。由于令牌路由和负载均衡的机制,没有任何专家被特定配置以专门处理某一特定语言。

ST-MoE 论文中显示了哪些令牌组被发送给了哪个专家的表格
在 Mixtral 8x7B 论文中,每个 token 都被其选择的第一个专家进行了着色,这是一个很好的例子。

https://arxiv.org/pdf/2401.04088
这一可视化结果也表明,专家往往更关注句法,而不是某个特定领域的内容。因此,尽管解码器中的专家似乎没有特定的专业化,但它们确实在某些类型的 tokens 上表现得更加一致。
专家的数量对预训练有何影响?
增加更多专家可以提升处理样本的效率和加速模型的运算速度,但这些优势随着专家数量的增加而递减 (尤其是当专家数量达到 256 或 512 之后更为明显)。同时,这也意味着在推理过程中,需要更多的显存来加载整个模型。值得注意的是,Switch Transformers 的研究表明,其在大规模模型中的特性在小规模模型下也同样适用,即便是每层仅包含 2、4 或 8 个专家。
微调混合专家模型
4.36.0
版本的transformers
库支持 Mixtral 模型。你可以用以下命令进行安装:pip install "transformers==4.36.0 --upgrade
稠密模型和稀疏模型在过拟合的动态表现上存在显著差异。稀疏模型更易于出现过拟合现象,因此在处理这些模型时,尝试更强的内部正则化措施是有益的,比如使用更高比例的 dropout。例如,我们可以为稠密层设定一个较低的 dropout 率,而为稀疏层设置一个更高的 dropout 率,以此来优化模型性能。
在微调过程中是否使用辅助损失是一个需要决策的问题。ST-MoE 的作者尝试关闭辅助损失,发现即使高达 11% 的令牌被丢弃,模型的质量也没有显著受到影响。令牌丢弃可能是一种正则化形式,有助于防止过拟合。
Switch Transformers 的作者观察到,在相同的预训练困惑度下,稀疏模型在下游任务中的表现不如对应的稠密模型,特别是在重理解任务 (如 SuperGLUE) 上。另一方面,对于知识密集型任务 (如 TriviaQA),稀疏模型的表现异常出色。作者还观察到,在微调过程中,较少的专家的数量有助于改善性能。另一个关于泛化问题确认的发现是,模型在小型任务上表现较差,但在大型任务上表现良好。

一种可行的微调策略是尝试冻结所有非专家层的权重。实践中,这会导致性能大幅下降,但这符合我们的预期,因为混合专家模型 (MoE) 层占据了网络的主要部分。我们可以尝试相反的方法: 仅冻结 MoE 层的参数。实验结果显示,这种方法几乎与更新所有参数的效果相当。这种做法可以加速微调过程,并降低显存需求。

在微调稀疏混合专家模型 (MoE) 时需要考虑的最后一个问题是,它们有特别的微调超参数设置——例如,稀疏模型往往更适合使用较小的批量大小和较高的学习率,这样可以获得更好的训练效果。

此时,您可能会对人们微调 MoE 中遇到的这些挑战而感到沮丧,但最近的一篇论文 《MoEs Meets Instruction Tuning》 (opens new window) (2023 年 7 月) 带来了令人兴奋的发现。这篇论文进行了以下实验:
- 单任务微调
- 多任务指令微调
- 多任务指令微调后接单任务微调
当研究者们对 MoE 和对应性能相当的 T5 模型进行微调时,他们发现 T5 的对应模型表现更为出色。然而,当研究者们对 Flan T5 (一种 T5 的指令优化版本) 的 MoE 版本进行微调时,MoE 的性能显著提升。更值得注意的是,Flan-MoE 相比原始 MoE 的性能提升幅度超过了 Flan T5 相对于原始 T5 的提升,这意味着 MoE 模型可能从指令式微调中获益更多,甚至超过了稠密模型。此外,MoE 在多任务学习中表现更佳。与之前关闭 辅助损失 函数的做法相反,实际上这种损失函数可以帮助防止过拟合。

稀疏 VS 稠密,如何选择?
稀疏混合专家模型 (MoE) 适用于拥有多台机器且要求高吞吐量的场景。在固定的预训练计算资源下,稀疏模型往往能够实现更优的效果。相反,在显存较少且吞吐量要求不高的场景,稠密模型则是更合适的选择。
注意: 直接比较稀疏模型和稠密模型的参数数量是不恰当的,因为这两类模型基于的概念和参数量的计算方法完全不同。
让 MoE 起飞
最初的混合专家模型 (MoE) 设计采用了分支结构,这导致了计算效率低下。这种低效主要是因为 GPU 并不是为处理这种结构而设计的,而且由于设备间需要传递数据,网络带宽常常成为性能瓶颈。在接下来的讨论中,我们会讨论一些现有的研究成果,旨在使这些模型在预训练和推理阶段更加高效和实用。我们来看看如何优化 MoE 模型,让 MoE 起飞。
并行计算
让我们简要回顾一下并行计算的几种形式:
- 数据并行: 相同的权重在所有节点上复制,数据在节点之间分割。
- 模型并行: 模型在节点之间分割,相同的数据在所有节点上复制。
- 模型和数据并行: 我们可以在节点之间同时分割模型和数据。注意,不同的节点处理不同批次的数据。
- 专家并行: 专家被放置在不同的节点上。如果与数据并行结合,每个节点拥有不同的专家,数据在所有节点之间分割。
在专家并行中,专家被放置在不同的节点上,每个节点处理不同批次的训练样本。对于非 MoE 层,专家并行的行为与数据并行相同。对于 MoE 层,序列中的令牌被发送到拥有所需专家的节点。

容量因子和通信开销
提高容量因子 (Capacity Factor, CF) 可以增强模型的性能,但这也意味着更高的通信成本和对保存激活值的显存的需求。在设备通信带宽有限的情况下,选择较小的容量因子可能是更佳的策略。一个合理的初始设置是采用 Top-2 路由、1.25 的容量因子,同时每个节点配置一个专家。在评估性能时,应根据需要调整容量因子,以在设备间的通信成本和计算成本之间找到一个平衡点。
部署技术
您可以在
Inference Endpoints
部署 mistralai/Mixtral-8x7B-Instruct-v0.1 (opens new window)。
部署混合专家模型 (MoE) 的一个关键挑战是其庞大的参数规模。对于本地使用情况,我们可能希望使用更小的模型。为了使模型更适合部署,下面是几种有用的技术:
- 预先蒸馏实验: Switch Transformers 的研究者们进行了预先蒸馏的实验。他们通过将 MoE 模型蒸馏回其对应的稠密模型,成功保留了 30-40%的由稀疏性带来的性能提升。预先蒸馏不仅加快了预训练速度,还使得在推理中使用更小型的模型成为可能。
- 任务级别路由: 最新的方法中,路由器被修改为将整个句子或任务直接路由到一个专家。这样做可以提取出一个用于服务的子网络,有助于简化模型的结构。
- 专家网络聚合: 这项技术通过合并各个专家的权重,在推理时减少了所需的参数数量。这样可以在不显著牺牲性能的情况下降低模型的复杂度。
高效训练
FasterMoE (2022 年 3 月) 深入分析了 MoE 在不同并行策略下的理论性能极限,并且探索了一系列创新技术,包括用于专家权重调整的方法、减少延迟的细粒度通信调度技术,以及一个基于最低延迟进行专家选择的拓扑感知门控机制。这些技术的结合使得 MoE 运行速度提升高达 17 倍。
Megablocks (2022 年 11 月) 则专注于通过开发新的 GPU kernel 来处理 MoE 模型中的动态性,以实现更高效的稀疏预训练。其核心优势在于,它不会丢弃任何令牌,并能高效地适应现代硬件架构 (支持块稀疏矩阵乘),从而达到显著的加速效果。Megablocks 的创新之处在于,它不像传统 MoE 那样使用批量矩阵乘法 (这通常假设所有专家形状相同且处理相同数量的令牌),而是将 MoE 层表示为块稀疏操作,可以灵活适应不均衡的令牌分配。

开源混合专家模型
目前,下面这些开源项目可以用于训练混合专家模型 (MoE):
- Megablocks: https://github.com/stanford-futuredata/megablocks (opens new window)
- Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm (opens new window)
- OpenMoE: https://github.com/XueFuzhao/OpenMoE (opens new window)
对于开源的混合专家模型 (MoE),你可以关注下面这些:
- Switch Transformers (Google) (opens new window): 基于 T5 的 MoE 集合,专家数量从 8 名到 2048 名。最大的模型有 1.6 万亿个参数。
- NLLB MoE (Meta) (opens new window): NLLB 翻译模型的一个 MoE 变体。
- OpenMoE (opens new window): 社区对基于 Llama 的模型的 MoE 尝试。
- Mixtral 8x7B (Mistral) (opens new window): 一个性能超越了 Llama 2 70B 的高质量混合专家模型,并且具有更快的推理速度。此外,还发布了一个经过指令微调的模型。有关更多信息,可以在 Mistral 的 公告博客文章 (opens new window) 中了解。
一些有趣的研究方向
首先是尝试将稀疏混合专家模型 (SMoE) 蒸馏 回到具有更少实际参数但相似等价参数量的稠密模型。
MoE 的 量化 也是一个有趣的研究领域。例如,QMoE (opens new window) (2023 年 10 月) 通过将 MoE 量化到每个参数不到 1 位,将 1.6 万亿参数的 Switch Transformer 所需的存储从 3.2TB 压缩到仅 160GB。
简而言之,一些值得探索的有趣领域包括:
- 将 Mixtral 蒸馏成一个稠密模型。
- 探索合并专家模型的技术及其对推理时间的影响。
- 尝试对 Mixtral 进行极端量化的实验。
视觉模型中的专家混合模型
MoE 并不仅限于语言模型。视觉模型(如 ViT)使用 Transformer 架构,因此也有潜力使用 MoE。
快速回顾一下,ViT(Vision Transformer)是一种将图像分割为若干块,并将其作为 tokens 处理的架构。

这些图像块(或 tokens)会被投射到嵌入向量(加上额外的位置嵌入向量)中,然后输入到常规编码器中:

当这些图像块进入编码器时,它们会像 tokens 一样被处理,这使得这种架构非常适合 MoE。
Vision-MoE
Vision-MoE(V-MoE)是图像模型中第一个实现 MoE 的例子之一。它将 ViT 中的密集 FFNN 层替换为稀疏 MoE。

这种改进使得 ViT 模型(通常比语言模型小)能够通过增加专家的数量来大幅扩展。
为了降低硬件限制,每个专家都设置了一个较小的预定义容量,因为图像通常包含大量的图像块。然而,低容量往往会导致图像块被丢弃(类似于 token 溢出)。

为了保持容量较低,网络会为每个图像块分配重要性得分,并优先处理这些得分较高的图像块,从而避免溢出图像块的丢失。这种方法被称为批量优先路由(Batch Priority Routing)。

因此,即使 token 数量减少,我们仍然能够看到重要的图像块被成功路由。

优先路由使得在处理较少的图像块时,仍能聚焦于最重要的图像块。
从稀疏 MoE 到软 MoE
在 V-MoE 中,优先评分机制能够区分出重要和不重要的图像块。然而,图像块被分配给每个专家后,未被处理的图像块中的信息就会丢失。
软 MoE(Soft-MoE)旨在通过混合图像块,从离散的图像块(token)分配转变为软分配。
第一步,我们将输入 x(图像块嵌入)与一个可学习矩阵 Φ 相乘。这将生成路由信息,它告诉我们某个 token 与某个专家的相关程度。

然后对路由信息矩阵进行 softmax 操作(在列上),从而更新每个图像块的嵌入向量。

更新后的图像块嵌入本质上是所有图像块嵌入的加权平均。

从视觉上看,这就像是所有图像块被混合。这些组合后的图像块被发送到每个专家。生成输出后,它们再次与路由矩阵相乘。

路由矩阵在 token 层面影响输入,并在专家层面影响输出。
因此,我们获得了“软”图像块/token,这些 token 被处理,而不是离散输入。
Mixtral 8x7B 的激活与稀疏参数对比
MoE 的一个重要特点是其计算需求。由于在同一时刻只会使用部分专家,我们可以拥有比实际使用的更多的参数。
尽管给定的 MoE 拥有更多的参数(稀疏参数),但由于我们在推理时只使用部分专家(活跃参数),因此激活的参数较少。

换句话说,我们仍然需要将整个模型(包括所有专家)加载到设备中(稀疏参数),但在实际运行推理时,我们只需要使用部分参数(活跃参数)。MoE 模型需要更多的显存来加载所有专家,但推理时运行速度更快。
让我们以 Mixtral 8x7B 为例来探讨稀疏参数与活跃参数的数量差异。

在此例中,我们可以看到每个专家的参数量为 5.6B,而不是 7B(尽管一共有 8 个专家)。

我们需要加载 8x5.6B(46.7B)的参数(加上所有共享参数),但推理时只需要使用 2x5.6B(12.8B)的参数。
相关资源
- 混合专家模型 (MoE) 详解 (opens new window)
- 一份MoE 可视化指南 (opens new window)
- Adaptive Mixture of Local Experts (1991) (opens new window)
- Learning Factored Representations in a Deep Mixture of Experts (2013) (opens new window)
- Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (2017) (opens new window)
- GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (Jun 2020) (opens new window)
- GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (Dec 2021) (opens new window)
- Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (Jan 2022) (opens new window)
- ST-MoE: Designing Stable and Transferable Sparse Expert Models (Feb 2022) (opens new window)
- FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models(April 2022) (opens new window)
- MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (Nov 2022) (opens new window)
- Mixture-of-Experts Meets Instruction Tuning:A Winning Combination for Large Language Models (May 2023) (opens new window)
- Mixtral-8x7B-v0.1 (opens new window), Mixtral-8x7B-Instruct-v0.1 (opens new window).