[论文解读] Efficient Large Scale Language Modeling with Mixtures of Experts

Mikel Artetxe, Shruti Bhosale+22 more|arXiv (Cornell University)|Dec 20, 2021

Topic Modeling被引用 13

一句话总结

本文提出并实证评估了用于高效大规模语言建模的专家混合（MoE）语言模型，利用条件计算仅激活每 token 的部分参数。MoE 模型在使用高达四分之一计算量的情况下，性能与密集模型相当，其中 1.1 万亿参数的 MoE 模型优于计算量相当的 67 亿参数密集模型，在多种自然语言处理任务中展现出显著的计算效率。

ABSTRACT

Mixture of Experts layers (MoEs) enable efficient scaling of language models through conditional computation. This paper presents a detailed empirical study of how autoregressive MoE language models scale in comparison with dense models in a wide range of settings: in- and out-of-domain language modeling, zero- and few-shot priming, and full-shot fine-tuning. With the exception of fine-tuning, we find MoEs to be substantially more compute efficient. At more modest training budgets, MoEs can match the performance of dense models using $\sim$4 times less compute. This gap narrows at scale, but our largest MoE model (1.1T parameters) consistently outperforms a compute-equivalent dense model (6.7B parameters). Overall, this performance gap varies greatly across tasks and domains, suggesting that MoE and dense models generalize differently in ways that are worthy of future study. We make our code and models publicly available for research use.

研究动机与目标

• 研究专家混合（MoE）模型与密集模型在大规模语言建模中的效率与性能表现。
• 在多种设置下评估 MoE 模型的表现：领域内与领域外语言建模、零样本提示与少样本提示、全量微调。
• 确定 MoE 模型在规模扩大时是否表现出与密集模型不同的泛化特性，以及在计算资源受限条件下其性能优势是否得以保持。
• 证明 MoE 模型可在显著降低计算成本的前提下实现最先进性能，尤其在中等训练预算下优势更为突出。

提出的方法

• 作者训练了从数亿到超过一万亿参数的自回归 MoE 和密集 Transformer 基础语言模型。
• MoE 层使用路由网络为每个 token 选择前 2 个专家，实现条件计算，即仅激活部分参数。
• 训练采用混合精度训练（FP16 优化器状态），并使用激活检查点技术以减少训练期间的内存占用。
• 采用完全分片数据并行（FSDP）技术，通过在多个 GPU 上分片模型参数、梯度和优化器状态，实现对大模型规模的训练扩展。
• 通过分析方法计算 FLOP，考虑了激活检查点和 MoE 前馈网络带来的额外 FLOP，MoE 模型的 FLOP 相比密集模型增加了 32Tlh²。
• 在多个基准测试中比较了模型性能，包括领域内与领域外语言建模（Pile）、零样本提示和微调，在等效计算预算下进行评估。

实验结果

研究问题

• RQ1 MoE 模型是否能在中等训练预算下实现与密集模型相当的性能，同时显著降低计算量？
• RQ2 在多种自然语言处理任务中，MoE 模型在零样本和少样本提示设置下的表现相对于密集模型如何？
• RQ3 随着规模扩大，MoE 与密集模型之间的性能差距是持续存在还是逐渐缩小？MoE 架构是否表现出不同的泛化特性？
• RQ4 MoE 模型在不同模型规模和训练预算下的计算效率增益如何？
• RQ5 MoE 模型与密集模型在任务和领域之间是否表现出不同的泛化行为？

主要发现

• 在中等训练预算下，MoE 模型使用约四分之一的计算量即可达到与密集模型相当的性能，展现出强大的计算效率。
• 即使在大规模训练（约 5000 GPU 天）下，最大的 MoE 模型（1.1 万亿参数）在下游任务中仍优于计算量相当的密集模型（67 亿参数）。
• MoE 与密集模型之间的性能差距在不同任务和领域间存在显著差异，表明其泛化行为不同。
• MoE 模型在零样本和少样本设置中保持了优异性能，表明条件计算并未损害其上下文学习能力。
• FSDP 与激活检查点的结合使得在标准硬件上训练超过一万亿参数的 MoE 模型成为可能，使大规模 MoE 训练具备可行性。
• 分析性 FLOP 估算表明，MoE 训练带来的计算增加是可预测且可控的，主要来自额外的专家前馈网络。