Jeff Dean署名《深度学习稀疏专家模型》综述论文

2022 年 10 月 4 日 专知

机器之心报道

机器之心编辑部

30年时间，稀疏专家模型已逐渐成为一种很有前途的解决方案。

稀疏专家模型是一个已有 30 年历史的概念，至今依然被广泛使用，是深度学习中的流行架构。此类架构包括混合专家系统（MoE）、Switch Transformer、路由网络、BASE 层等。稀疏专家模型已经在自然语言处理、计算机视觉和语音识别等多个领域展示出良好的性能。

近日，谷歌 AI 负责人 Jeff Dean 等人撰写了一篇稀疏专家模型的综述，回顾了稀疏专家模型的概念，提供了通用算法的基本描述，最后展望了未来的研究方向。

论文地址：https://arxiv.org/pdf/2209.01667.pdf

通过增加计算预算、训练数据和模型大小，机器学习（尤其是自然语言）取得了显著进步。著名的里程碑语言模型包括 GPT-2（Radford et al., 2018）、BERT（Devlin et al., 2018）、T5（Raffel et al., 2019）、GPT-3（Brown et al., 2020）、Gopher (Rae et al., 2021)、Chinchilla (Hoffmann et al., 2022) 和 PaLM (Chowdhery et al., 2022)。

然而，最先进的模型现在需要数以千计的专用互连加速器，训练数周或数月的时间，因此这些模型的生产成本很高（Patterson et al., 2021）。随着机器学习系统规模的扩大，该领域寻求更有效的训练和服务范式。稀疏专家模型已成为一种很有前途的解决方案。

稀疏专家模型（其中，混合专家系统（MoE）是最流行的变体）是一种特殊的神经网络，其中一组参数被划分为 “专家”，每个“专家” 具有唯一的权重。

在训练和推理期间，模型将输入样本给到特定的专家权重，让每个样本仅与网络参数的子集交互，这与将整个网络用于每个输入的常规方法不同。由于每个样本仅使用一小部分专家，因此相对于总模型来说计算量大幅减小了。

许多现代稀疏专家模型从 Shazeer et al. (2017)中汲取了灵感。该研究训练了当时最大的模型，并取得了最先进的语言建模和翻译结果。与 Transformer 语言模型结合使用时，稀疏专家模型的受欢迎程度进一步飙升（Lepikhin et al., 2020; Fedus et al., 2021）。虽然大部分工作都在自然语言处理方面，但稀疏专家模型也已成功用于各种领域，包括计算机视觉（Puigcerver et al., 2020）、语音识别（You et al., 2021）和多模态学习（Mustafa et al., 2022）。Clark et al. (2022) 研究了稀疏专家模型在不同模型大小和专家数量下的扩展特性。此外，许多基准测试的最新结果目前由 ST-MoE 等稀疏专家模型持有（Zoph et al., 2022）。随着研究和工程的进步，该领域正在迅速发展。

这篇综述论文将调查范围缩小到狭义深度学习时代（从 2012 年开始）的稀疏专家模型，回顾了最近的进展并讨论了有希望的未来途径。

稀疏专家模型

机器学习中的 MoE 概念可以追溯到至少 30 年前，在早期的概念中，专家们定义了一个完整的神经网络，MoE 类似于一种集成方法。

Eigen et al. (2013) 提出在抖动（jittered）的 MNIST 上使用堆叠层专家混合架构，这项工作为后来模型的高效实现奠定了基础。

Shazeer et al. (2017)提出在两个 LSTM 层之间插入一个 MoE 层，由此产生的稀疏模型在机器翻译中实现 SOTA 性能。然而，此方法虽然取得了成功，但后续研究就像冬眠了一样，停滞不前，大部分研究都转向了 Transformer。

2020 到 2021 年间，GShard 和 Switch Transformer 发布，两者都用专家层取代了 Transformer 中的前馈层。

尽管使用一层专家层的方法已经成为主导范式，但近两年的研究重新审视了专家模型作为完全独立模型的概念，以实现模块化和可组合性。

下图 2 是 Shazeer et al. (2017) 提出的原始 top-k 路由机制，这是许多后续工作的基础。该综述论文在第四章中详细讲解了路由算法的新进展。

硬件

现代稀疏专家模型已与用于训练最大神经网络的分布式系统共同设计。

大型神经网络（Brown et al., 2020; Rae et al., 2021; Chowdhery et al., 2022）的研究表明神经网络已经远远超过单个加速器的内存容量，因此权重，激活函数，优化器变量等张量需要使用各种并行策略分片。

三种常见的并行方法包括：数据并行（复制模型权重，对数据进行分片）、张量模型并行（数据和权重张量跨设备分片）和 pipeline 并行（整个层或层组被跨设备分片），混合专家模型通常能够适应这些并行方案。

在MoE 模型的训练和部署方面，Jaszczur et al. (2021)将 Transformer 模型的所有层进行稀疏化，从而实现了 37 倍的推理加速；Kossmann et al. (2022)通过 RECOMPILE 库解决了静态专家批处理大小的约束问题。

除了数据并行、模型并行和专家并行外，Rajbhandari et al.(2022) 提出 DeepSpeed-MoE 库，以支持 ZeRO partitioning 和 ZeRO-Offload，实现了 10 倍推理提升和 SOTA 翻译性能，从而增加了模型在生产服务中的实用性。

稀疏专家模型的扩展特性

密集神经语言模型的交叉熵损失表现为一种关于模型参数计数、数据量和计算预算（Kaplan et al., 2020）的幂律。幂律系数后来在 Hoffmann et al. (2022)中得到纠正，这表明计算最优模型需要数据和参数扩展更紧密的平衡。相比之下，稀疏专家模型的早期研究以启发式方式扩展获得了强有力的实证结果，但没有仔细描述扩展定律。此外，一些工作强调了上游（例如预训练）和下游（例如微调）行为之间的差异（Fedus et al., 2021; Artetxe et al., 2021），进一步复杂化了对稀疏专家模型的理解。

上游扩展

稀疏专家模型在大型数据集上训练时表现出色。自然语言处理中的常见模式是先执行上游训练（例如预训练），然后对特定感兴趣的数据分布进行下游训练（例如微调）。在上游阶段，稀疏专家模型一直比密集对应模型产生更高的收益。Shazeer et al. (2017) 在 10 亿字语言建模基准（Chelba et al., 2013）上提出了关于模型参数和计算预算的扩展曲线，与密集版本相比取得了显著收益。Lepikhin et al. (2020)提出了模型扩展函数的改进版本，并在其最大的 600B 参数稀疏模型上获得了 13.5 的 BLEU 分数增益。Switch Transformer (Fedus et al., 2021) 在 T5 模型上使用相同的计算资源测量到 4-7 倍的 wall-time 加速。该工作还研究了作为参数计数函数的交叉熵损失扩展，但观察到超过 256 位专家时，收益减少了。

下游扩展

然而，可靠的上游扩展并没有立即在下游任务上产生一致的收益。在一项强调迁移挑战的工作中，Fedus et al. (2021)使用低计算、高参数编码器 - 解码器 Transformer（1.6T 参数，每个稀疏层有 2048 个专家）观察到 4 倍的预训练改进，但它在 SuperGLUE 等密集型推理任务上的微调效果不佳。这一发现暗示了进一步的必要研究以及计算和参数之间可能需要的平衡。

Du et al.(2021) 展示了稀疏 GLaM 模型的扩展，范围从 1B-64B FLOP，每个稀疏层使用 64 位专家。GLaM 取得了 SOTA 结果，在 zero-shot 和 one-shot 性能方面优于 175B 参数的 GPT-3 模型（Brown et al., 2020），同时在推理时每个token使用的 FLOP 减少了 49%，功耗降低了 65% （如下图4(左)所示）。下图 4 (右) 是稀疏模型在少样本推理上表现良好的另一个例子。

Srivastava et al. (2022) 研究了稀疏模型在多项选择 BIG-Bench 任务上的校准，即度量预测概率与正确概率匹配的程度。研究结果如下图 5 所示，虽然较大的密集模型和稀疏模型的校准都得到了改善，但稀疏模型与使用 10 倍以上 FLOP 的密集模型的校准相当。

扩展专家层的数量、大小和频率

有几个重要的超参数控制着稀疏专家模型的扩展，包括：1）专家数量，2）每个专家的大小，以及 3）专家层的频率。这些决策可能对上游和下游扩展产生重大影响。

许多早期的工作扩展到每层数千个相对较小的专家，从而产生了出色的预训练和翻译质量（Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2021）。然而，稀疏模型的质量在域偏移（Artetxe et al., 2021）或对不同任务分布进行微调（Fedus et al., 2021）时不成比例地降低。为了使计算和参数获得更好的平衡，用于小样本推理 (GLaM (Du et al., 2021)) 和微调 (ST-MoE (Zoph et al., 2022)) 的 SOTA 稀疏模型最多仅能使用 64 个更大的专家。由于专家维度的增加，这些模型需要有针对加速器的特定系统级分片策略才能有效运行（Du et al., 2021; Rajbhandari et al., 2022）。

路由算法

路由算法是所有稀疏专家架构的关键特性，它决定了将样本发送到哪里。该领域已被广泛研究，包括使用固定的、非学习路由模式的反直觉方法（Roller et al., 2021）。由于要对选择哪些专家做出离散决策，因此路由决策通常是不可微的。

专家选择问题后来被重新定义为 Bandit 问题，并且已有一些工作使用强化学习来学习专家选择（Bengio et al., 2016; Rosenbaum et al., 2017; 2019; Clark et al., 2022）。Shazeer et al. (2017) 提出了一种可微的启发式算法，以避开强化学习的难点。

该论文详细阐述了路由算法的分类，并进一步说明了该领域的关键问题——负载均衡。

稀疏专家模型的迅速发展

稀疏专家模型的影响正在迅速蔓延到NLP以外的其他领域，包括计算机视觉、语音识别和多模态应用。尽管领域不同，但模型的架构和算法大致相同，下表 1 显示了不同域的稀疏层输入表征。

稀疏专家模型近几年的发展非常迅速。以NLP领域为例，Shazeer et al. (2017) 为 LSTM 语言建模和机器翻译引入了混合专家层，这些层被插入到 LSTM 模型的标准层之间。

2020 年，Lepikhin et al. (2020)首次将 MoE 层引入到 Transformer，当研究人员将每个专家层扩展到 2048 名专家时，模型在 100 种不同语言上实现了 SOTA 翻译结果。

Fedus et al. (2021) 创建了一个稀疏的 1.6T 参数语言模型，实现了 SOTA 预训练质量。

新研究正在推动小样本学习推理和微调基准的发展。Du et al. (2021) 训练了一个纯 MoE 解码器语言模型，取得了小样本上的 SOTA 结果，并且只需要训练 GPT-3 所需的 1/3 计算量。Zoph et al. (2022) 提出了 ST-MoE，这是一种稀疏编码器 - 解码器模型，在大量推理和生成任务中都实现了 SOTA。在 SuperGLUE 上进行微调时，ST-MoE 的性能优于 PaLM-540B ，同时仅使用约 1/20 的预训练 FLOP 和 1/40 的推理 FLOP。

什么时候使用稀疏模型

一个常见的问题是，如果给你一个固定的计算或 FLOP 预算（例如 100 个 GPU 20 小时），你应该训练什么类型的模型来获得最佳性能？

从根本上说，稀疏模型允许通过增加专家的数量来大幅增加模型中的参数数量，同时保持每个样本的 FLOP 大致恒定。这种做法有好有坏，具体取决于模型的用途。

当你有许多加速器（例如 GPU/TPU）来承载使用稀疏性带来的所有附加参数时，稀疏性是很有利的。

使用稀疏性还需要仔细考虑下游任务。假如你有许多机器来进行预训练，但是用来微调或者服务的机器却很少，那么稀疏度（例如专家的数量）应该根据下游用例中可用的内存量进行调整。

在有些情况下，稀疏模型看起来总是比密集模型差。例如当所有参数都保存在加速器内存中，这种情况稀疏模型就不如密集模型了。当你有能力在多台机器上并行训练或服务，以便托管来自专家的附加模型参数时，稀疏模型就非常合适。

此外，该综述论文还介绍了对稀疏模型训练的改进、可解释性和未来的研究方向等，感兴趣的小伙伴可以查看原论文，了解更多研究内容。

专知便捷查看