当前位置:首页 > 新闻资讯 > FPGA之家动态 >

深度学习博士万字长文:MoE架构,条件计算的诞生和兴起

时间:2024-08-02      来源:网络搜集 关于我们 0

本文作者为Cameron R. Wolfe ,深度学习博士,Rebuy的 AI 总监。这是《深度(学习)焦点》时事通讯,通过从头开始概述相关主题来帮助读者了解 AI 研究。

//

大型语言模型 (LLM) 的现代进步大多是扩展定律的产物 [6]。随着底层模型的规模增加,我们会看到性能平稳提升(假设模型是在足够大的数据集上训练的)[7]。这样的扩展定律最终促使我们创建了 GPT-3,以及随后的其他(更强大的)LLM。在扩展这些密集的 LLM 时,由于内存成本增加以及模型计算占用空间对总参数数量的依赖,我们最终将达到硬件强加的限制。

因此,计算和数据已成为训练更好的 LLM 的主要瓶颈——如果能够获得更多的计算和数据,创建更好的模型相对简单。本文我们将研究如何通过训练稀疏激活语言模型来避免这一限制。

“由于巨型密集模型的训练达到了当今硬件资源可用性和能力的极限,与质量相当的密集模型相比,混合专家 (MoE) 模型由于其训练成本显著降低,已成为最有前途的模型架构之一。” - 摘自 [12]

混合专家 (MoE) 层很简单,它使我们能够增加语言模型的大小或容量,而无需相应增加计算量。我们只需将模型的某些层替换为具有自己参数的层的多个副本(称为“专家”)。然后,我们可以使用门控机制(稀疏地)选择用于处理每个输入的专家。

这个想法源于 20 世纪 90 年代早期对条件计算的研究 [15, 30],它使我们能够以易于处理的方式训练大量模型,这在语言建模等受益于具有额外容量的模型的领域中很有用。在这里,我们将研究 MoE、它的起源以及它在过去二十年中的发展历程。

TL;DR:什么是专家混合(MoE)?

仅解码器的变压器架构

上图显示了大多数生成式 LLM 使用的标准解码器专用变压器架构。在 LLM 的背景下,MoE 对此架构进行了简单的修改——我们用 MoE 层替换了前馈子层此 MoE 层由多位专家组成(即从几个专家 [13] 到数千位专家 [5]),其中每个专家都是自己的前馈子层,具有一组独立的参数。

Transformer 块中的 MoE 层(源自 [5])

我们可以类似地将编码器-解码器转换器(例如T5 [8])转换为 MoE 模型,方法是将编码器和解码器中的前馈子层替换为 MoE 层。但是,我们通常只将这些子层的一部分(例如,每隔一层)转换为 MoE 层。在本文中,我们将主要概述基于编码器-解码器转换器的 MoE 模型的工作,而基于 MoE 的自回归 LLM 的工作将在以后的文章中介绍。

“ST-MoE 模型有 32 位专家,专家层频率为 1/4(每四个 FFN 层被一个 MoE 层替换)。” - 摘自 [24]

稀疏专家。这种方法似乎有问题,因为它会给模型添加大量参数。MoE 模型在 Transformer 的每个前馈子层内都有多个独立的神经网络(即,而不是单个前馈神经网络)。但是,我们在前向传递中只使用每个 MoE 层专家的一小部分!给定一个 token 列表作为输入,我们使用路由机制来稀疏地选择一组将每个 token 发送到的专家。因此,MoE 模型前向传递的计算成本远低于具有相同数量参数的密集模型。

MoE 层的组件。当应用于 Transformer 模型时,MoE 层有两个主要组件:

稀疏 MoE 层:用几个结构类似的“专家”组成的稀疏层取代 Transformer 中的密集前馈层。 路由器(Router ):确定 MoE 层中的哪些令牌被发送给哪些专家。

稀疏 MoE 层中的每个专家都只是一个前馈神经网络,拥有自己独立的一组参数1。每个专家的架构都模仿标准 Transformer 架构中使用的前馈子层。路由器将每个 token 作为输入,并生成专家的概率分布,以确定每个 token 发送给哪个专家。

MoE 层的结构(来自 [5])

路由器有自己的一组参数,并与网络的其余部分联合训练。每个令牌可以发送给许多专家,但我们通过仅将令牌发送给其顶级K专家来施加稀疏性。例如,许多模型设置k=1或k=2,这意味着每个令牌分别由一个或两个专家处理。

容量更大,计算量固定。假设我们能够获得足够大的训练数据集,那么增加语言模型的大小和容量是提高性能的主要途径之一。然而,训练成本会随着底层模型的大小而增加,这使得较大的模型在实践中变得负担过重。MoE 模型通过在推理过程中仅使用模型参数的子集来避免这种成本。

“假设每个 token 只使用两个专家,推理速度就像使用 12B 模型(而不是 14B 模型),因为它计算 2x7B 矩阵乘法,但共享一些层” - 摘自HuggingFace MoE 博客

例如,假设我们有一个 7B 参数的 LLM,并将其每个前馈子层替换为由八个专家组成的 MoE 层,其中每个 token 激活两个专家。这正是Mixtral-8x7B [13] 使用的架构,它是Mistral-7B [14]的 MoE 变体。完整模型有大约 47B 个参数2,所有这些都必须加载到内存中。但是,该模型的推理成本与 14B 参数模型相当。仅使用两个专家来处理每个标记,这会产生 ~2 x 7B 矩阵乘法3这样,我们无需花费任何成本,就实现了~50B参数模型的容量!

优点和缺点。MoE 模型被广泛使用,因为它们能够以固定的计算预算训练更大的模型,但使用 MoE 风格的 LLM 既有优点也有缺点!与密集模型相比,MoE 模型的预训练速度更快,并且与具有相同数量参数的密集模型相比,其推理速度更快。MoE使我们能够增加模型容量,同时保持我们使用的计算量(相对)较低。然而,MoE 风格的 LLM 也消耗更多的 VRAM,因为所有专家都必须加载到内存中。此外,与密集模型相比,MoE 模型容易过度拟合,并且微调难度极大,这使得它们在实际应用中的使用更加复杂。

专家组合的起源

尽管 MoE 在最近的 AI 研究中非常流行,但这个想法已经存在了很长时间。事实上,条件计算的概念(或动态打开/关闭神经网络的某些部分)源于 20 世纪 90 年代早期的研究!在本节中,我们将探索这项关于条件计算的早期研究,并研究它是如何演变成 MoE 的想法的。最终,这些想法被应用于(并取得了一些成功)早期基于LSTM的 [4] 语言模型的训练。

条件计算的早期工作

“条件计算,即网络的各部分根据每个示例处于活动状态,在理论上已被提出,可以大幅提高模型容量,而无需按比例增加计算量” - 摘自 [1]

MoE 的概念起源于 Geoffrey Hinton 在 20 世纪 90 年代初所做的工作 [15],他提出了一个由多个网络组成的监督学习框架。每个网络都应处理不同的训练数据子集,而使用哪个网络则由门控机制决定。换句话说,该系统中的每个网络都是一个专家,专门处理输入数据的某个域(或区域)。

(摘自[15])

自 [15] 提出以来,已有多项研究探索并扩展了条件计算的这一理念。例如,[30](也是在 20 世纪 90 年代初撰写的!)中的作者提出了一种分层(树结构)MoE,可以使用期望最大化算法以监督方式进行训练。有许多此类研究深入研究了这一主题。

(摘自[30])

估计或通过随机神经元传播梯度以进行条件计算 [16]。这项工作探索了四种可能的技术来估计随机神经元和硬阈值激活函数的梯度;例如,通过使用REINFORCE 、直通估计器或加法/乘法噪声。虽然这些术语可能看起来陌生,但条件计算是硬阈值激活函数的一个例子——某些神经元的激活被从计算中完全消除(或设置为零)。[16] 中提出的估计器对于使用条件计算的训练网络是必需的。[16] 中探讨的神经网络结构由许多计算单元组成,这些计算单元通过分布式门网络连接,可用于消除大量计算。通过关闭该网络的某些部分,我们可以大大降低大型神经网络的计算成本。

(摘自[17])

深度神经网络中条件前馈计算的低秩近似 [17]。这项工作探索了一种神经网络设置,其中网络节点补充了门控单元,以确定是否应计算节点。由于ReLU将神经网络中的所有负激活值都设置为零,因此可以将任何具有负预激活值的节点从计算中完全移除。基于这一想法,作者在 [17] 中表明,可以通过计算权重矩阵的低秩近似来估计 ReLU 之前的隐藏激活的符号,当网络稀疏时,可显著提高效率。

(摘自[18])

在深度混合专家中学习因子表示 [18]。这项工作考虑了由多个专家网络组成的 MoE 层,这些专家网络专门处理输入空间的不同区域。使用学习门控机制将输入映射到这些专家,从而允许稀疏地计算更大的网络。更进一步,[18] 中的作者考虑了具有多个连续 MoE 层(即“堆叠”MoE)的更深层网络,如上图所示。在 MNIST 数据集上测试这些网络时,我们发现深度 MoE 在每一层都学习独特的特征;例如,第一层学习位置相关特征,而第二层学习类别特定特征。

“条件计算是一种在不增加所需计算量的情况下增加深度神经网络容量的方法,通过根据每个示例按需激活一些参数。” ——摘自[19]

深度学习中条件计算的容量与计算比率呈指数级增长 [19]。众所周知,模型越大,性能越好4。如果计算量不是问题,我们可以使用更大的数据集和模型来提高泛化能力,但深度网络使用的计算量会随着参数数量的增加而增加。为了避免这个问题并更好地利用条件计算,[19] 中的作者提出了一种新的权重矩阵参数化方法,当观察到特定的隐藏激活模式时,它会关闭参数组5。这种方法可以成倍地增加参数与计算的比例。

神经网络中的条件计算可加快模型速度 [20]。这项工作探索了使用条件计算的神经网络的强化学习训练策略(即 REINFORCE)。我们发现,这种使用策略梯度算法的方法可用于以保持准确性同时最大化计算速度的方式训练网络。此外,我们可以使用正则化来多样化网络中观察到的条件计算模式。实验表明,这种学习策略可以产生在计算成本和性能之间实现良好平衡的网络。

(摘自[21])

动态容量网络 (DCN) [21]。通过定义低容量和高容量子网络,自适应地为输入数据的不同部分分配容量。对于大多数数据,应用低容量子网络。但是,可以使用基于梯度的注意机制来选择输入的部分(基于网络对数据这一区域的敏感度),将高容量子网络应用于这些部分。我们在实验中看到,DCN 可以减少计算量,同时实现与传统卷积神经网络 (CNN) 相当(或更好)的性能。

(摘自[22])

深度顺序神经网络 [22]摒弃了神经网络架构中使用的传统“层”概念,而是选择构建一整套候选转换,这些转换可以有选择地应用于每一层。在处理输入时,每一层都会从候选集中选择一个元素,形成基于转换序列的DAG类架构。根据输入数据的属性,所选的转换序列可能会有所不同,与经典多层网络相比,这增加了模型的表达能力。与之前的工作类似,这种顺序架构是使用策略梯度方法进行训练的。

超大型神经网络:稀疏门控专家混合层[1]

“我们在模型容量上获得了超过 1000 倍的提升,而计算效率仅略有损失,并且显著提高了公共语言建模和翻译数据集上的最新成果。” - 摘自 [1]

条件计算的想法具有令人难以置信的前景。特别是在语言建模领域,训练数据集往往非常庞大6,能够在不导致计算量相应增加的情况下增加模型底层容量的能力至少可以说是诱人的。尽管已经研究了二十多年,但由于各种技术挑战,MoE 未能实现其承诺。在 [1] 中,作者尝试通过稀疏门控混合专家层的提议来克服其中的一些挑战,表明 MoE 可以应用于语言建模和翻译领域

MoE 的先前问题。尽管 [1] 中的工作是最早将 MoE 应用于语言建模的工作之一,但 MoE 的想法可以追溯到 20 世纪 90 年代初。考虑到这一点,我们可能会想:是什么阻碍了 MoE 得到更广泛的采用?几个技术障碍导致了这种缺乏采用:

GPU 擅长高效执行算术,但它们在分支(即条件计算的主要组成部分)方面并不那么擅长。 需要较大的批量大小才能有效地训练神经网络,但 MoE 会减小批量大小(即每个专家只接收输入批量的一部分)。 在研究具有较大训练数据集的领域时,增加模型容量最有影响力,而 MoE 在计算机视觉领域主要是针对训练数据集不够大的问题进行研究。

(摘自[1])

什么是稀疏 MoE?大多数现代 LLM 使用的 MoE 层风格类似于 [1] 中提出的稀疏门控专家混合层(如上所示),它是一个具有两个组件的通用神经网络模块:

专家:每一层都有几个“专家”,它们是独立的神经网络模块或具有独立参数集的层。 路由器:一种参数化(且可学习)的门控机制,选择一组(稀疏)专家来处理每个输入。

在 [1] 中,MoE 中的每个专家都是具有相同架构的前馈神经网络。但是,可以使用更复杂的架构;例如,我们甚至可以通过将每个专家实现为另一个 MoE 来创建“分层”MoE 模块!专家和门控机制都通过梯度下降与其他网络参数一起进行联合训练。

计算 MoE 模块的输出

为了计算 MoE 模块的输出,我们采用专家输出的加权组合,其中权重由路由器提供。路由器输出一个N- 维权重向量(即N专家数量)。虽然这种方法最初看起来可能没什么用,但当路由器的输出稀疏时,奇迹就会发生——在计算 MoE 的输出时,不再考虑权重为零的专家。因此,MoE 允许我们创建、训练和使用非常大的神经网络,而无需大量计算,因为在任何给定时间只使用一部分模型参数。

简单的softmax门控机制(来自[1])

门控机制。已经提出了许多不同的策略用于 MoE 内的路由。最简单的方法是将我们的输入乘以权重矩阵并应用softmax。但是,这种方法不能保证专家的选择是稀疏的7为了解决这个问题,[1] 中的作者提出了一种改进的门控机制,为这种简单的 softmax 门控机制增加了稀疏性和噪声。

稀疏 MoE 的路由机制实现(来自 [1])

上述门控机制执行路由的方式与 softmax 门控机制类似,但我们增加了两个步骤:

在应用 softmax 之前,将可调节量的高斯噪声添加到路由器的输出。 除了顶级专家的输出之外的所有内容K都被屏蔽(即设置为 -∞)以确保专家的选择是稀疏的。

平衡专家。MoE的一个问题是,网络倾向于在训练期间重复使用相同的少数专家。门控机制不会学习统一使用所有专家,而是会收敛到始终为每个输入选择同一组专家的状态。这是一个自我实现的循环——如果一位专家被选中的次数最多,那么他的训练速度就会更快,因此会继续被选中,而不是其他专家!先前的研究提出了几种解决这个问题的方法 [2, 3],但我们在 [1] 中看到,可以通过在训练损失中添加一个简单的“软”约束来平衡专家。

用于平衡专家利用率的辅助损失(来自[1])

我们首先为输入数据批次中的每个专家定义一个“重要性”分数,该分数可以通过对批次中每个专家的门值求和来计算。简而言之,在批次中经常被选中的专家将具有较高的重要性分数。然后,我们可以通过取专家重要性分数的平方变异系数(CV) 来计算辅助损失函数。如果所有专家的重要性分数都非常相似,则 CV 会很小,反之亦然。可以将这个损失项添加到模型的训练目标中,以鼓励专家在每个批次中获得同等的重要性。

实证评估。为了测试稀疏 MoE 层的性能,[1] 中的作者训练了一个语言模型,该模型在LSTM的堆叠层之间插入了一个 MoE 模块8。MoE 模块以卷积方式应用,这意味着多个 LSTM 时间步骤可同时(批量)处理。此外,还提出了一种新的分布式训练策略,该策略结合了模型和数据并行训练,以消除 MoE 中批量缩小的问题。这种策略与其他一些系统优化相结合,允许以合理的计算要求训练具有多达 137B 个参数的语言模型!

(摘自[1])

这些模型在不同规模的语言建模数据集上进行评估(即 10 亿、1000 亿和 1000 亿个单词)。平面和分层 MoE 都使用不同数量的专家进行测试,并且所有模型都经过训练,使得它们的总计算成本大致相等(即活跃专家的数量k相同)9。即使在研究最小的数据集时,我们也会看到,最大的 MoE 模型在困惑度方面有明显的改善(越低越好)。

(摘自[1])

如果训练数据集的大小足够小,那么通过增加专家数量来增加容量的收益就会递减,如上所示。但是,我们发现在更大的数据集上,性能会继续提高,直到参数大小达到 68B(即 65,536 位专家)!这一发现强调了 MoE 层与语言建模领域之间的协同作用——如果训练语料库足够大,增加模型容量会很有帮助。[1] 中的工作首次表明,这些好处可以在实践中实现,而不会明显增加计算成本,为 MoE 成为 LLM 的流行工具奠定了基础。

“两条线之间的差距越来越大,这(毫不奇怪)表明,增加模型容量对更大的训练集更有帮助。” ——摘自[1]

将混合专家技术应用于 Transformer

现在我们已经研究了条件计算的早期工作,我们可以看看 MoE 在 Transformer 架构中的一些应用。

如今,基于 MoE 的 LLM 架构(例如Mixtral [13] 或Grok )已经变得流行,但 MoE 在语言模型中的探索是分几个阶段的。在这里,我们将研究将 MoE 应用于编码器-解码器 Transformer 架构。这项工作与现代 LLM 应用相关,并为在实践中有效使用 MoE 提供了无数经验教训。这样的研究为后续探索基于 MoE 的生成式 LLM 的工作奠定了基础,这将在未来的概述中介绍。

Switch Transformers:通过简单高效的稀疏性扩展到万亿参数模型[5]

在 [1] 中提出稀疏门控 MoE 之后,Transformer 和语言模型方面的工作尚未开始使用这些想法——MoE的普遍复杂性以及高通信成本和训练不稳定等问题阻碍了其应用。为了解决这些问题,[5] 中的作者提出了一种基于 MoE 的编码器-解码器 Transformer 架构,称为 Switch Transformer,它使用简化的门控机制使训练更加稳定,从而使 MoE 成为语言建模应用更现实、更实用的选择。

“简单的架构——在充足的计算预算、数据集大小和参数数量的支持下——胜过更复杂的算法” - 摘自[5]

当 Switch Transformer 被提出时,研究人员才刚刚开始研究神经缩放定律[6],该定律表明,随着语言模型规模的增加,其性能会平稳提升10考虑到这一点,[5] 中的工作是基于这样一个事实:MoE 模型允许我们为神经缩放定律的分析增加另一个维度——我们可以增加模型大小,同时保持模型前向传递的计算复杂度不变

(摘自[5])

将 MoE 应用于转换器。要创建编码器-解码器转换器的 MoE 变体,我们可以简单地将模型的前馈子层转换为 MoE 层。前馈变换以逐点方式应用,这意味着每个 token 都单独通过前馈网络。因此,序列中的每个 token 都单独路由到其对应的专家集。例如,序列中的每个 token都通过路由函数,形成专家的概率分布。然后,我们为每个单独的 token[“I”, “love”, “LLM”, “#s”]选择顶级专家——同一序列中的 token 并不总是发送给相同的专家

更好的路由。在之前的工作中,任何 MoE 层中活跃专家的最小数量为 2。使用至少两个专家被认为是在路由函数中实现非平凡梯度的必要条件11在 [5] 中,作者建议将每个 token 路由到单个专家——这称为交换层。通过路由到单个专家,我们简化了路由功能,减少了计算开销,减少了通信成本,同时提高了模型的性能。

开关变压器中的 Softmax 路由(来自 [5])

Switch Transformer 使用的路由函数(如上所示)只是一种 softmax 门控机制。我们将每个 token 向量传递到一个线性层,该层产生一个大小为N(即专家数量)的输出,然后应用 softmax 变换将此输出转换为专家的概率分布。

计算 Switch Transformer 的 MoE 输出(摘自 [5])

从这里开始,我们通过以下方式计算切换层的输出:i)选择单个专家;ii)根据路由函数分配给该专家的概率缩放该专家的输出。值得注意的是,这种方法不同于 [1],因为它计算的是所有专家的分布,而不是仅考虑顶级K专家的分布。换句话说,我们仍然N在路由函数中计算所有专家的概率分布,然后通过根据所选专家的概率缩放其输出来计算层的输出。这种方法允许我们在 时训练 MoE 模型K=1。

简单的负载平衡。在 [1] 中,作者使用多个辅助损失函数来平衡重要性得分并在专家之间执行负载平衡(即,每个专家都会从批次中收到大致相等数量的标记)。我们在 [5] 中看到,这两个目标都可以通过应用于模型中每个交换层的单个辅助损失函数来实现。

简化的负载平衡辅助损耗(来自[5])

这种损失对于来说是可微的,P并且可以轻松地纳入训练中,它鼓励分配给每个专家的令牌分数和分配给每个专家的路由器概率分数都为1/N,这意味着专家同等重要并获得均衡数量的令牌

容量因子。在 Switch Transformer 中,我们设置了一个全局“专家容量”变量,该变量决定了可以路由到任何 MoE 层中每个专家的最大令牌数量。专家容量公式如下所示。

(摘自[5])

在交换层中,每个令牌都会被路由到路由机制分配到概率最高的专家。如果向单个专家发送了过多令牌(即超出专家容量),则将跳过这些令牌的计算。这些“丢弃”的令牌将通过残差连接直接传递到下一层。将容量因子设置为大于 1 可让 MoE 处理令牌在专家之间不完美平衡的情况。

(摘自[5])

必须调整专家容量以避免丢弃过多令牌(负载平衡损失也可以通过将令牌均匀地路由给专家来帮助解决此问题),同时不会浪费计算或内存。除了 Switch Transformer 之外,容量因子是任何 MoE 模型的重要超参数,我们可以调整它,通过在专家之间平衡计算来提高硬件利用率。

实际评估。[5] 中的作者评估了基于T5 [8] 的 Switch Transformer 模型。该模型在C4上使用掩码语言建模 (MLM) 目标和使用格式进行某些操作的混合精度策略进行预训练。通过仅在选定操作上使用低精度(例如,选择以全精度执行路由函数),我们可以提高效率而不会导致训练不稳定。有趣的是,作者指出,他们增加了 MoE 层的 dropout 率以避免过度拟合。bfloat16

“Switch Transformers 的指导原则是以简单且计算高效的方式最大化 Transformer 模型的参数数量。” ——摘自[5]

使用这种方法,可以训练具有高达 1.6T 参数的 Switch Transformer 模型。为了创建这些模型,我们从 T5 架构开始,将每个前馈子层转换为 MoE,并增加每个 MoE 层中的专家数量!尽管这些模型规模很大,但每个 token 上花费的计算量并没有太大变化。事实上,在某些情况下,训练Switch Transformer 的速度比训练基本 T5 模型快 7 倍。MoE 模型具有更好的统计学习效率 12与 T5 相比,在给定固定计算量或训练时间的情况下表现更佳。

(摘自[5])

与标准 MoE 模型 [1] 相比,Switch Transformer 占用的计算资源更少(因为活跃专家更少),并且在速度质量方面表现更佳。换句话说,Switch Transformer 的表现优于使用相同计算量或相同时间进行训练的同类 MoE 模型。我们还在上面看到,与 MoE 相比,Switch Transformer 在较低容量因子下表现更好,这使我们能够通过更好地平衡专家之间的负载来提高训练效率。

(摘自[5])

扩展属性。[5] 中的结果表明,MoE 模型(尤其是 Switch Transformer 变体)对语言建模应用大有裨益。为什么 MoE 在语言建模中如此受欢迎?原因很简单,MoE 模型允许我们以稳定的计算成本扩大模型规模。如上图所示,随着专家数量的增加,LLM 性能持续提高,产生的扩展趋势与密集模型相当。MoE 模型使我们能够获得更大模型的性能,同时保持合理的训练和推理成本。

ST-MoE:设计稳定且可迁移的稀疏专家模型[24]

(摘自[24])

尽管 Switch Transformer 让 MoE 模型更加实用/可用,但它仍然存在一些问题。该模型存在训练不稳定的问题,[5] 中的作者表明,尽管 Switch Transformer 在训练前比密集模型实现了加速,但在SuperGLUE等流行基准上进行微调时,其性能仍落后于之前最先进的模型。简而言之,这些基于 MoE 的语言模型只有在数据集非常大时才能表现良好,否则它们的性能会因过度拟合而下降。在 [24] 中,作者深入研究了 MoE 模型的不稳定性及微调困难,为稀疏模型的实用可靠训练提供了设计指南(如上所述)。使用此指南,作者训练了一个稀疏的 269B 参数编码器-解码器模型,这是第一个实现最先进迁移学习的基于 MoE 的模型十三结果。

(摘自[24])

解决训练不稳定性问题。尽管稀疏 MoE 模型具有诸多优势,但它们在训练过程中更容易出现不稳定现象。在 [24] 中,作者考虑了几种缓解这种训练不稳定性的技术,并根据i)模型的性能和ii)稳定的训练运行比例对它们进行了比较。发现有几种技术可以提高稳定性,但代价是模型性能下降,例如:

通过丢失或输入抖动向模型注入噪声(即将路由器的 logit 乘以范围内的随机值[0.99, 1.01])。 根据固定阈值剪切梯度或更新权重,以补救反向传播期间的梯度爆炸。 从架构中移除具有乘法相互作用的组件(例如,GeGLU 激活 [27] 或 RMNSorm 中的比例因子 [28])。

我们在 [24] 中了解到,限制路由器的 logits 值(即在应用 softmax 之前路由器的输出值)对训练稳定性的影响最大。在先前的工作 [5] 中,我们已经看到路由机制无法以低精度计算,因为 MoE 模型对舍入误差很敏感。由于路由函数中的指数(如下所示),此类舍入误差可能会出现在路由器中并导致训练不稳定。

MoE 路由器中的 Softmax 变换

因此,过大的路由器逻辑会对训练稳定性产生负面影响。但是,简单地裁剪(即应用硬阈值来删除较大的值)路由器逻辑可能会损害模型性能。相反,[25] 提出了路由器 z-loss(公式如下),作为一种辅助损失,鼓励在路由机制中使用较小的值(和更少的舍入误差)。

路由器 z-loss(源自 [24])

此损失只会在 MoE 的门控机制中惩罚较大的逻辑值。要将路由器 z-loss 纳入 MoE 模型的训练过程,我们可以简单地将其与正常训练损失和辅助负载平衡损失相结合,其中我们根据这些损失的相对重要性对其进行加权。路由器 z-loss 显著提高了训练稳定性,而不会导致模型质量明显下降。

结合 MoE 模型的损失(源自 [24])

调试微调问题。尽管稀疏模型在大型数据集上表现良好,但我们已经看到微调这些模型可能很困难。在 [24] 中,作者通过研究预训练的稀疏和密集模型在小型 ( CommitmentBank ) 和大型 ( ReCoRD ) 训练数据集的下游任务上的表现,证实了这些微调问题是由于过度拟合造成的14在这些实验中,我们看到稀疏模型在两种情况下都比密集模型收敛得更快。然而,在数据较少的下游任务上,稀疏模型的验证质量比密集模型差得多。换句话说,由于缺乏足够的数据,稀疏模型会过拟合

(摘自[24])

为了缓解过度拟合,作者首先尝试仅微调 MoE 模型参数的某些部分。有趣的是,仅微调 MoE 层会导致性能急剧下降,而选择性微调其他参数组的性能与微调所有参数相当。

(摘自[24])

进一步说,我们发现稀疏模型和密集模型往往需要不同的超参数设置才能有效微调。特别是,MoE 模型受益于高学习率和较小的批处理大小,这两者都通过在训练过程中引入更多噪音来提高泛化能力。此外,我们发现容量因子不会显著影响微调结果。事实上,我们可以在微调过程中消除辅助负载平衡损失而不会出现任何问题——稀疏模型对于丢弃多达 10-15% 的标记具有很强的鲁棒性

(摘自[24])

然而,[24] 中没有提供明确的解决方案来解决微调 MoE 模型的困难。作者确认这些问题是由于过度拟合而产生的,并分析了有助于对抗过度拟合的不同选择。然而,即使是 [24] 实验部分中分析的优化设计模型,在较小的下游任务上仍然会过度拟合!改进的正则化形式以微调 MoE 仍然是一个活跃的研究领域。

(摘自[24])

最佳 MoE 设计。基于广泛的分析,[24] 中的作者概述了有效设计和使用 MoE 模型的高级要点(如上所示)15。此分析主要侧重于选择正确的容量因子和 MoE 中的专家数量。简而言之,我们在 [24] 中看到,MoE 模型的最佳设置取决于硬件。例如,作者建议每个 TPU/GPU 核心最多使用一名专家,并注意到容量因子和专家数量的最佳值取决于硬件。由于向每个专家发送的最大令牌数量增加而导致的容量因子变化会产生几个下游影响:

激活对内存的要求更高。 沟通成本增加。 增加了每个专家前向传递的计算量。

另一方面,更高的容量因子可以提高模型质量!对于具有快速互连和强大 GPU 的系统,计算和通信成本的激增可能是值得的。然而,MoE 参数的选择是一个权衡,它在很大程度上取决于可用的硬件。

实践经验。[24] 中的作者总结了他们的工作,利用上面概述的经验教训,基于 T5-Large [8] 训练了一个 269B 参数编码器-解码器转换器模型,其计算复杂度与 32B 参数密集模型相同。与之前的 MoE 模型相比,ST-MoE-32B 模型的参数更少,计算复杂度更高。每个 MoE 层中的专家更少,但规模更大。以前,这种高FLOP MoE 在训练过程中被发现不稳定,但增加路由器 z-loss 使 ST-MoE-32B 模型能够成功训练。由于与之前的 MoE 模型相比,ST-MoE-32B 的参数更少,因此更易于微调和部署。

“我们在 FLOP 和参数之间寻求平衡。高 FLOP 稀疏模型以前在我们的设置中不稳定(即编码器-解码器模型、Adafactor 优化器),但路由器 z-loss 使我们能够继续前进。” - 来自 [24]

该模型使用来自 C4 和用于训练 GLaM 的数据集的 1.5 万亿个 token 进行预训练 [25]。之后,对所有下游任务进行微调 —其中每个任务根据其训练集大小按比例采样16— 评估之前。该模型的结果如下所示,我们可以看到它在广泛的 NLP 任务中设定了新的最先进性能。值得注意的是,ST-MoE-32B 是第一个在迁移学习领域取得最先进结果的稀疏模型,因为之前的模型在微调过程中严重受到过度拟合的影响。

(摘自[24])

其他值得注意的 MoE 论文

在 MoE 风格的 LLM 流行之前,除了我们在 [5] 和 [24] 中探讨的研究之外,还有很多关于这个主题的论文。下面概述了该领域的一些著名出版物,以更全面地介绍语言建模领域中对 MoE 进行的研究。

(摘自[9])

GShard [9]提出了一个框架和一套工具,用于轻松实现神经网络中的并行计算模式,例如稀疏门控 MoE 层。然后,该框架用于将基于 Transformer 的神经机器翻译模型扩展到超过 600B 个参数的大小。使用 GShard,该模型可以在 4 天内在超过 2048 个 TPU 上进行高效训练,并且在将 100 多种不同语言翻译成英语方面超越了之前的模型的质量。在 GShard 中有效训练 MoE 模型的一些有用技巧包括:

编码器-解码器转换器中的每个其他前馈子层都被 MoE 层取代。 编码器和解码器均使用Top-2路由。 路由机制总是选择顶级专家,但第二位专家则是按照与其权重成比例的概率随机选择的。

“数据并行可以看作是沿批处理维度拆分张量和操作,而在 Mesh-TensorFlow 中,用户可以指定任意张量维度,将其拆分到多维处理器网格的任意维度上。” - 摘自 [31]

Mesh-Tensorflow [31]极大地简化了对高效、模型并行分布式训练策略的探索,该策略将模型本身分散到多个设备(即 GPU 或 TPUS)上17用于神经网络。以前,数据并行策略(将输入数据批次拆分到多个设备)是神经网络分布式训练的主要方法,但数据并行训练有几个缺点;例如,训练大型模型更加困难,因为每个设备都必须在内存中保留模型的完整副本,并且批量较小时效率会下降。简而言之,Mesh-Tensorflow 只是一种可用于指定分布式计算的通用语言。在此工具发布后,MoE 层被重新引入到 Transformer 架构中,因为编程数据并行和模型并行训练策略变得非常简单。在 [31] 中,作者使用 Mesh-Tensorflow 训练了一个 5B 参数 Transformer 模型,该模型在机器翻译任务上的结果超越了最先进的结果。

(摘自[10])

M6 [10]提出了一个新的(仅限中文)数据集,其中包含 1.9TB 的图像和 292GB 的文本,用于预训练多模态语言模型。从这个数据集中,使用一种新颖的跨模态预训练策略对多个语言模型进行预训练。[10] 中训练的模型之一是 100B 参数 MoE 模型,该模型用由 1024 位专家组成的 MoE 层替换了上文所述 Transformer 中的所有前馈子层。

(摘自[11])

使用稀疏混合专家扩展视觉 [11]。尽管 MoE 模型显然有利于语言建模,但目前尚不清楚计算机视觉是否会有类似的好处。在 [11] 中,作者探索了一种视觉 MoE (V-MoE) 模型,该模型将视觉转换器架构与稀疏 MoE 层相结合。简而言之,我们只是用 MoE 层替换视觉转换器中前馈子层的子集(例如每隔一层或最后n一层)。这种方法在性能和计算之间产生了理想的权衡。更进一步,[11] 中的作者提出了一种新的路由算法,可以在推理时自适应地为某些输入利用更多的计算。

“V-MoE 的性能可与最先进的网络相媲美,同时推理时只需要一半的计算量。” ——摘自 [11]

DeepSpeed-MoE [12]观察到使用 MoE 模型进行高效推理很困难。扩展密集的 LLM 最终会遇到硬件限制,但 MoE 模型可以 (几乎) 任意扩展。因此,MoE 模型是一种很有前途的架构进步,但它们的采用通常受到其较大尺寸和独特架构的限制。为了使条件计算在实践中更易于使用,[12] 中的作者创建了 DeepSpeed-MoE 框架,用于 MoE 模型的端到端训练和推理。该框架定义了几种 MoE 架构,提出了可将 MoE 大小减少 3.7 倍的模型压缩技术,并实现了可将延迟提高 7.3 倍的优化推理引擎。

(摘自[23])

超越以英语为中心的多语言机器翻译 [23]。这项工作创建了一个多对多机器翻译模型,可以直接在 100 种语言中的任意一对之间进行翻译。虽然之前的模型都是以英语为中心的,但 [23] 中的作者旨在创建一个能更好地反映全球翻译需求的模型。首先,创建一个开源训练数据集,其中包含数千个语言对的监督语言翻译示例。要从这些数据中学习,我们需要一个具有足够容量的模型。[23] 中的作者探索了将密集扩展(或仅增加模型的大小)与稀疏的、特定于语言的参数组相结合的高容量模型。

(摘自[25])

GLaM [25]将 Switch Transformers(即应用于编码器-解码器变压器架构 [26])的工作扩展到仅解码器的语言模型。这项工作的目标是更高效地训练GPT-3质量模型——使用少一个数量级的计算量。GLaM 是 Generalist Language Model 的缩写,它提供了一套 LLM,这些模型使用 MoE 层来扩展模型容量,同时保持合理的训练成本。该套件中最大的模型总共有 1.2 万亿个参数(即比 GPT-3 大 7 倍),在近 30 个 NLP 任务中的表现优于 GPT-3,训练时仅使用 GPT-3 所用能量的 1/3,并且在推理时所需的计算量是 GPT-3 的一半。

(摘自[25])

GLaM 模型使用前两个路由(即每层有两个活跃专家)。[25] 中的作者还探索了更大的容量因子,发现该指标对于控制训练或推理期间使用的计算量很有用。

最后的想法

“在深度学习中,模型通常会对所有输入重复使用相同的参数。混合专家 (MoE) 模型则打破了这一传统,而是为每个传入示例选择不同的参数。结果是稀疏激活模型(具有大量参数)但计算成本却保持不变。” - 摘自 [5]

在研究了二十多年的 MoE 研究后,我们应该对这个想法、它为什么有效以及如何应用它来训练更好的语言模型有了很好的了解。关于这个主题的研究经历了几个阶段,但所有这些工作的关键要点可以总结如下。

核心思想。MoE 背后的基本思想是将模型的参数数量与其使用的计算量分离开来。为此,我们可以简单地用几个稀疏激活的专家替换模型中的层。然后,门控机制可以确定应该为给定的输入使用哪个专家。通过这种方式,我们增加了模型的容量,同时在前向传递中仅使用模型总参数的一部分。

MoE 的组成部分。MoE 有两个主要组成部分:i)专家和ii)路由机制。每一层都有几个结构相似的专家,可用于处理输入数据。给定一组 token 作为输入,我们使用路由机制(稀疏地)选择每个 token 应发送到的专家。然后,我们计算每个 token 对应的专家输出,并结合他们的结果来得出 MoE 层的输出。许多 MoE 模型使用某种 softmax 门控机制变体作为路由器。

将 MoE 应用于 Transformer。在语言建模领域,我们通常基于 Transformer 架构创建 MoE 模型。为此,我们考虑 Transformer 中的每个前馈子层(或这些层中的一些较小子集;例如,每隔一层),并将其替换为 MoE 层。MoE 中的每个专家都是一个前馈神经网络,具有一组独立的参数,这些参数与初始前馈子层的架构相匹配。在大多数情况下,Transformer MoE 模型使用前一或前二路由。

优点和缺点。正如我们在这篇文章中看到的,与密集模型相比,MoE 模型有很多优点,但也有很多缺点!给定足够大的预训练数据集,MoE 模型往往比计算匹配的密集模型学习得更快。我们还可以显著提高 MoE 模型的容量,从而以较小的计算负担实现更大模型的性能。MoE 的这一特性在具有大量训练数据的环境中非常有用,因此 MoE 在语言建模领域很受欢迎。但是,MoE 模型也:

消耗更多内存(即我们必须将所有专家存储在内存中)。 训练稳定性困难(尽管[24]中的工作有助于解决这个问题)。 如果没有足够的训练数据,微调期间往往会过度拟合。

综上所述,是否使用 MoE 在很大程度上取决于具体应用。如果我们拥有几乎无限量的训练数据(大多数现代语言模型都是如此),那么 MoE 就大有裨益,让我们能够大幅扩展模型容量,甚至超越硬件限制。然而,在规模较小的应用中,与密集模型相比,MoE 的应用并不那么简单,甚至可能表现更差!返回搜狐,查看更多

责任编辑:


注明:本内容来源网络,不用于商业使用,禁止转载,如有侵权,请来信到邮箱:429562386ⓐqq.com 或联系本站客服处理,感谢配合!

用户登陆

    未注册用户登录后会自动为您创建账号

提交留言