人工智能

Mixture-of-Experts 在大型语言模型中的崛起

mm
Mixture of Experts Grok Mistral

在自然语言处理(NLP)领域,构建更大、更强大的语言模型的追求一直是推动许多最近进步的动力。然而,随着这些模型的增长,训练和推理的计算要求变得越来越苛刻,推动着可用硬件资源的极限。

这就是 Mixture-of-Experts(MoE)技术的用武之地,它承诺缓解计算负担,同时使得训练更大、更强大的语言模型成为可能。以下,我们将讨论 MoE、探索其起源、内部工作原理及其在基于 Transformer 的语言模型中的应用。

Mixture-of-Experts 的起源

Mixture-of-Experts(MoE)概念可以追溯到 20 世纪 90 年代初,当时研究人员探索了条件计算的想法,即神经网络的某些部分根据输入数据被选择性地激活。这个领域的开创性工作之一是 Jacobs 等人在 1991 年发表的《自适应局部专家混合》论文,该论文提出了一个用于神经网络集成的监督学习框架,每个神经网络专门处理输入空间的不同区域。

MoE 的核心思想是拥有多个“专家”网络,每个网络负责处理输入数据的一个子集。一个门控机制,通常是一个神经网络本身,决定哪个(些)专家应该处理给定的输入。这种方法允许模型更高效地分配其计算资源,只激活相关的专家来处理每个输入,而不是为每个输入使用整个模型容量。

多年来,研究人员探索和扩展了条件计算的思想,导致了分层 MoE、条件计算的低秩近似和通过随机神经元和硬门控激活函数估计梯度的技术等发展。

Mixture-of-Experts 在 Transformer 中

Mixture of Experts

Mixture of Experts

虽然 MoE 的想法已经存在了几十年,但其应用于基于 Transformer 的语言模型相对较新。Transformer 已经成为语言模型的事实标准,主要由多个层组成,每个层包含一个自注意力机制和一个前馈神经网络(FFN)。

将 MoE 应用于 Transformer 的关键创新是用稀疏的 MoE 层替换密集的 FFN 层,每个层由多个专家 FFN 和一个门控机制组成。门控机制决定哪个(些)专家应该处理每个输入令牌,使得模型可以选择性地激活每个输入序列的一个子集的专家。

早期展示 MoE 在 Transformer 中潜力的工作之一是 Shazeer 等人在 2017 年发表的《令人难以置信的庞大神经网络:稀疏门控的 Mixture-of-Experts 层》论文,该论文引入了稀疏门控 MoE 层的概念,该层使用门控机制添加了稀疏性和噪声到专家选择过程中,确保只有每个输入的一个子集的专家被激活。

此后,其他工作进一步推进了 MoE 在 Transformer 中的应用,解决了训练不稳定、负载平衡和高效推理等挑战。值得注意的例子包括 Switch Transformer(Fedus 等人,2021 年)、ST-MoE(Zoph 等人,2022 年)和 GLaM(Du 等人,2022 年)。

Mixture-of-Experts 对语言模型的益处

在语言模型中使用 MoE 的主要好处是能够在保持相对恒定的计算成本的同时扩大模型规模。在每个输入令牌上选择性地激活一个专家子集,MoE 模型可以实现大型密集模型的表达能力,同时需要的计算量大大减少。

例如,考虑一个具有 7 亿参数的密集 FFN 层的语言模型。如果我们用一个由 8 个专家组成的 MoE 层替换该层,每个专家具有 7 亿参数,则参数总数增加到 56 亿。然而,在推理过程中,如果我们每个令牌只激活两个专家,则计算成本等同于一个 14 亿参数的密集模型,因为它计算两个 7 亿参数的矩阵乘法。

这种推理过程中的计算效率在部署资源有限的场景中尤其有价值,例如移动设备或边缘计算环境。此外,训练过程中的计算需求减少可以带来显著的能量节省和更低的碳足迹,与可持续人工智能实践日益增长的重视相一致。

挑战和考虑

虽然 MoE 模型提供了令人信服的好处,但其采用和部署也带来了几个挑战和考虑因素:

  1. 训练不稳定性:MoE 模型比其密集对应物更容易出现训练不稳定性。这种问题源于专家激活的稀疏和条件性质,这可能导致梯度传播和收敛困难。Zoph 等人(2022 年)提出的路由器 z 损失等技术已被提出以缓解这些不稳定性,但仍需要进一步研究。
  2. 微调和过拟合:MoE 模型在微调过程中更容易过拟合,特别是当下游任务具有相对较小的数据集时。这一行为归因于 MoE 模型的增加容量和稀疏性,这可能导致对训练数据的过度专门化。需要仔细的正则化和微调策略来缓解这一问题。
  3. 内存要求:虽然 MoE 模型可以在推理过程中减少计算成本,但它们通常比类似大小的密集模型具有更高的内存要求。这是因为所有专家权重都需要加载到内存中,即使每个输入只激活一个子集。内存约束可能会限制 MoE 模型在资源受限设备上的可扩展性。
  4. 负载平衡:为了实现最佳的计算效率,平衡专家之间的负载至关重要,确保没有一个专家过载而其他专家闲置。这种负载平衡通常通过训练过程中的辅助损失和仔细调整容量因子来实现,容量因子决定了可以分配给每个专家的最大令牌数。
  5. 通信开销:在分布式训练和推理场景中,MoE 模型可能会引入额外的通信开销,这是由于需要在不同设备或加速器上驻留的专家之间交换激活和梯度信息。高效的通信策略和硬件感知模型设计对于减轻这一开销至关重要。

尽管存在这些挑战,MoE 模型在实现更大、更强大的语言模型方面的潜在好处已经激发了大量的研究工作来解决和缓解这些问题。

示例:Mixtral 8x7B 和 GLaM

为了说明 MoE 在语言模型中的实际应用,让我们考虑两个值得注意的示例:Mixtral 8x7B 和 GLaM。

Mixtral 8x7B 是由 Anthropic 开发的 Mistral 语言模型的 MoE 变体。它由 8 个专家组成,每个专家具有 7 亿参数,总共 56 亿参数。然而,在推理过程中,每个令牌只激活两个专家,有效地将计算成本降低到 14 亿参数密集模型的水平,因为它计算两个 7 亿参数的矩阵乘法。

Mixtral 8x7B 已经展示了令人印象深刻的性能,超越了 70 亿参数的 Llama 模型,同时提供了更快的推理速度。Mixtral-8x7B-Instruct-v0.1 是 Mixtral 8x7B 的指令微调版本,进一步增强了其遵循自然语言指令的能力。

另一个值得注意的示例是 GLaM(Google 语言模型),这是由 Google 开发的一个大规模 MoE 模型。GLaM 采用了解码器仅的 Transformer 架构,并在一个巨大的 1.6 万亿令牌数据集上进行了训练。该模型在少次和一次性评估中实现了令人印象深刻的性能,匹配了 GPT-3 的质量,同时只使用了训练 GPT-3 所需能量的三分之一。

GLaM 的成功可以归因于其高效的 MoE 架构,这使得训练一个具有大量参数的模型成为可能,同时保持了合理的计算要求。该模型还展示了 MoE 模型在能效和环境可持续性方面的潜力,相比其密集对应物具有优势。

Grok-1 架构

GROK MIXTURE OF EXPERT

GROK MIXTURE OF EXPERT

Grok-1 是一个基于 Transformer 的 MoE 模型,具有独特的架构,旨在最大限度地提高效率和性能。让我们深入了解其关键规格:

  1. 参数:具有 314 亿参数,Grok-1 是迄今为止最大的开源大型语言模型。然而,得益于 MoE 架构,只有 25% 的权重(约 86 亿参数)在任何时候是活跃的,从而增强了处理能力。
  2. 架构:Grok-1 采用了 8 个专家的混合架构,每个令牌在推理过程中由两个专家处理。
  3. :该模型由 64 个 Transformer 层组成,每个层包含多头注意力机制和密集块。
  4. 标记化:Grok-1 使用具有 131,072 个令牌的 SentencePiece 标记器。
  5. 嵌入和位置编码:该模型具有 6,144 维的嵌入,并采用旋转位置嵌入,使得与传统固定位置编码相比,数据的解释更加动态。
  6. 注意力:Grok-1 使用 48 个注意力头来查询和 8 个注意力头来处理键和值,每个头的大小为 128。
  7. 上下文长度:该模型可以处理最长 8,192 个令牌的序列,使用 bfloat16 精度进行高效计算。

性能和实现细节

Grok-1 已经展示了令人印象深刻的性能,超越了 LLaMa 2 70B 和 Mixtral 8x7B,MMLU 得分为 73%,展示了其在各种测试中的效率和准确性。

然而,值得注意的是,Grok-1 需要大量的 GPU 资源,因为其体积庞大。当前的开源版本实现重点是验证模型的正确性,并采用了一种低效的 MoE 层实现,以避免需要自定义内核。

尽管如此,该模型支持激活分片和 8 位量化,这可以优化性能并降低内存要求。

值得注意的是,xAI 已经在 Apache 2.0 许可下发布了 Grok-1,使其权重和架构对全球社区开放,以供使用和贡献。

开源版本包括一个 JAX 示例代码库,展示了如何加载和运行 Grok-1 模型。用户可以使用 torrent 客户端或直接通过 HuggingFace Hub 下载检查点权重,从而方便地访问这一开创性的模型。

Mixture-of-Experts 在语言模型中的未来

随着对更大、更强大的语言模型的需求持续增长,MoE 技术的采用预计将继续增强。正在进行的研究工作重点是解决剩余的挑战,例如改进训练稳定性、减轻微调过程中的过拟合以及优化内存和通信要求。

一个有前途的方向是探索分层 MoE 架构,其中每个专家本身由多个子专家组成。这种方法可能使得更大的可扩展性和计算效率成为可能,同时保持大型模型的表达能力。

此外,针对 MoE 模型的硬件和软件系统的开发是一个活跃的研究领域。专门为 MoE 模型的稀疏和条件计算模式设计的加速器和分布式训练框架可能会进一步增强其性能和可扩展性。

此外,将 MoE 技术与语言模型中的其他进步(如稀疏注意力机制、有效的标记化策略和多模态表示)相结合,可能会导致更强大、更通用的语言模型,能够处理广泛的任务。

结论

Mixture-of-Experts 技术已经成为构建更大、更强大的语言模型的有力工具。通过根据输入数据选择性地激活专家,MoE 模型为解决大型密集模型相关的计算挑战提供了一个有前途的解决方案。虽然仍然存在需要克服的挑战,例如训练不稳定性、过拟合和内存要求,但 MoE 模型在计算效率、可扩展性和环境可持续性方面的潜在好处使其成为一个令人兴奋的研究和开发领域。

随着自然语言处理领域的不断推进,MoE 技术的采用可能在实现下一代语言模型方面发挥至关重要的作用。通过将 MoE 与模型架构、训练技术和硬件优化的其他进步相结合,我们可以期待更强大、更通用的语言模型,它们能够以自然、无缝的方式真正理解和与人类进行交流。

我已经沉浸在了令人着迷的机器学习和深度学习世界中五年了。我的热情和专业知识让我为超过50个不同的软件工程项目做出了贡献,特别关注AI/ML。我的持续的好奇心也让我对自然语言处理产生了兴趣,这是一个我渴望进一步探索的领域。