人工智能
Mixture-of-Experts 在大型语言模型中的崛起

作者
Aayush Mittal Mittal
在自然语言处理(NLP)领域,构建更大、更强大的语言模型的追求一直是推动许多最近进步的动力。然而,随着这些模型的增长,训练和推理的计算要求变得越来越苛刻,推动着可用硬件资源的极限。
这就是 Mixture-of-Experts(MoE)技术的用武之地,它承诺缓解计算负担,同时使得训练更大、更强大的语言模型成为可能。以下,我们将讨论 MoE、探索其起源、内部工作原理及其在基于 Transformer 的语言模型中的应用。
Mixture-of-Experts(MoE)概念可以追溯到 20 世纪 90 年代初,当时研究人员探索了条件计算的想法,即神经网络的某些部分根据输入数据被选择性地激活。这个领域的开创性工作之一是 Jacobs 等人在 1991 年发表的《自适应局部专家混合》论文,该论文提出了一个用于神经网络集成的监督学习框架,每个神经网络专门处理输入空间的不同区域。
MoE 的核心思想是拥有多个“专家”网络,每个网络负责处理输入数据的一个子集。一个门控机制,通常是一个神经网络本身,决定哪个(些)专家应该处理给定的输入。这种方法允许模型更高效地分配其计算资源,只激活相关的专家来处理每个输入,而不是为每个输入使用整个模型容量。
多年来,研究人员探索和扩展了条件计算的思想,导致了分层 MoE、条件计算的低秩近似和通过随机神经元和硬门控激活函数估计梯度的技术等发展。
虽然 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 年)。
在语言模型中使用 MoE 的主要好处是能够在保持相对恒定的计算成本的同时扩大模型规模。在每个输入令牌上选择性地激活一个专家子集,MoE 模型可以实现大型密集模型的表达能力,同时需要的计算量大大减少。
例如,考虑一个具有 7 亿参数的密集 FFN 层的语言模型。如果我们用一个由 8 个专家组成的 MoE 层替换该层,每个专家具有 7 亿参数,则参数总数增加到 56 亿。然而,在推理过程中,如果我们每个令牌只激活两个专家,则计算成本等同于一个 14 亿参数的密集模型,因为它计算两个 7 亿参数的矩阵乘法。
这种推理过程中的计算效率在部署资源有限的场景中尤其有价值,例如移动设备或边缘计算环境。此外,训练过程中的计算需求减少可以带来显著的能量节省和更低的碳足迹,与可持续人工智能实践日益增长的重视相一致。
虽然 MoE 模型提供了令人信服的好处,但其采用和部署也带来了几个挑战和考虑因素:
尽管存在这些挑战,MoE 模型在实现更大、更强大的语言模型方面的潜在好处已经激发了大量的研究工作来解决和缓解这些问题。
为了说明 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 是一个基于 Transformer 的 MoE 模型,具有独特的架构,旨在最大限度地提高效率和性能。让我们深入了解其关键规格:
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 下载检查点权重,从而方便地访问这一开创性的模型。
随着对更大、更强大的语言模型的需求持续增长,MoE 技术的采用预计将继续增强。正在进行的研究工作重点是解决剩余的挑战,例如改进训练稳定性、减轻微调过程中的过拟合以及优化内存和通信要求。
一个有前途的方向是探索分层 MoE 架构,其中每个专家本身由多个子专家组成。这种方法可能使得更大的可扩展性和计算效率成为可能,同时保持大型模型的表达能力。
此外,针对 MoE 模型的硬件和软件系统的开发是一个活跃的研究领域。专门为 MoE 模型的稀疏和条件计算模式设计的加速器和分布式训练框架可能会进一步增强其性能和可扩展性。
此外,将 MoE 技术与语言模型中的其他进步(如稀疏注意力机制、有效的标记化策略和多模态表示)相结合,可能会导致更强大、更通用的语言模型,能够处理广泛的任务。
Mixture-of-Experts 技术已经成为构建更大、更强大的语言模型的有力工具。通过根据输入数据选择性地激活专家,MoE 模型为解决大型密集模型相关的计算挑战提供了一个有前途的解决方案。虽然仍然存在需要克服的挑战,例如训练不稳定性、过拟合和内存要求,但 MoE 模型在计算效率、可扩展性和环境可持续性方面的潜在好处使其成为一个令人兴奋的研究和开发领域。
随着自然语言处理领域的不断推进,MoE 技术的采用可能在实现下一代语言模型方面发挥至关重要的作用。通过将 MoE 与模型架构、训练技术和硬件优化的其他进步相结合,我们可以期待更强大、更通用的语言模型,它们能够以自然、无缝的方式真正理解和与人类进行交流。
我已经沉浸在了令人着迷的机器学习和深度学习世界中五年了。我的热情和专业知识让我为超过50个不同的软件工程项目做出了贡献,特别关注AI/ML。我的持续的好奇心也让我对自然语言处理产生了兴趣,这是一个我渴望进一步探索的领域。