Connect with us

人工智能

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、条件计算的低秩近似和通过随机神经元和硬阈值激活函数估计梯度等技术的发展。

Transformer中的Mixture-of-Experts

Mixture of Experts

Mixture of Experts

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

将MoE应用于Transformer的关键创新是用稀疏的MoE层替换密集的FFN层,每个MoE层由多个专家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层的语言模型。如果我们用一个由八个专家组成的MoE层替换这个层,每个专家都有7亿参数,总参数数量增加到56亿。然而,在推理过程中,如果我们每个令牌只激活两个专家,计算成本等同于一个14亿参数的密集模型,因为它计算两个7亿参数的矩阵乘法。

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

挑战和考虑

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

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

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

示例:Mixtral 8x7B和GLaM

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

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

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

另一个值得注意的例子是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. 参数:拥有惊人的3140亿参数,Grok-1是迄今为止最大的开放LLM。然而,感谢MoE架构,只有25%的权重(大约860亿参数)在任何时候都是活跃的,从而增强了处理能力。
  2. 架构:Grok-1采用8个专家的Mixture-of-Experts架构,每个令牌在推理过程中由两个专家处理。
  3. :该模型由64个Transformer层组成,每个层都包含多头注意力机制和密集块。
  4. 标记化:Grok-1使用SentencePiece标记器,词汇表大小为131,072个令牌。
  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个不同的软件工程项目做出贡献,特别注重人工智能/机器学习。我的持续好奇心也使我对自然语言处理产生了兴趣,这是一个我渴望进一步探索的领域。