ST-MoE: Designing Stable and Transferable Sparse Expert Models论文解读

本文是论文ST-MoE: Designing Stable and Transferable Sparse Expert Models的解读。

目录

• Abstract
• 1. Introduction
• 2. Background
• 3. 让稀疏模型的训练更稳定
  • 3.1 移除乘法交互时的稳定性和质量权衡
  • 3.2 增加噪音时的稳定性和质量权衡
  • 3.3 稳定性和质量权衡在约束激活和梯度方面
  • 3.4 选择精度格式：权衡效率和稳定性
• 4.稀疏模型的精细调整性能
  • 4.1 假设：泛化问题
  • 4.2 优化模型参数子集以提高泛化性能
  • 4.3 稀疏模型和密集模型需要不同的微调协议
  • 4.4 稀疏模型在微调期间丢弃的token鲁棒性
• 5. 设计稀疏模型
  • 5.1 设置专家数量
  • 5.2 选择容量因子和路由算法
• 6. 实验结果
  • 6.1 ST-MOE-L
  • 6.2 ST-MoE-32B
• 7. 追踪模型的token
  • 7.1 编码器专家表现出专业化
  • 7.2 解码器专家缺乏专业化
  • 7.3 多语言专家的专业化，但不是按语言划分
• 8.相关工作
• 9. 讨论

Abstract

规模化已经在自然语言处理中开辟了新的领域，但代价却很高。作为对此的回应，混合专家模型（Mixture-of-Experts，MoE）和开关Transformers（Switch Transformers）被提出作为一条节能高效的路径，以构建更大、更强大的语言模型。然而，在跨一系列自然语言任务上推动最新技术的发展受到了训练不稳定性和在微调过程中质量不确定性的阻碍。

我们的工作重点解决了这些问题，并充当了设计指南。最后，我们将一个稀疏模型扩展到2690亿参数，其计算成本与32亿参数的密集编码器-解码器Transformer（稳定且可迁移的混合专家模型，ST-MoE-32B）相当。首次，一个稀疏模型在跨多个任务的迁移学习中实现了最先进的性能，包括推理（SuperGLUE、ARC Easy、ARC Challenge）、摘要生成（XSum、CNN-DM）、闭卷问题回答（WebQA、自然问题）以及对抗构建的任务（Winogrande、ANLI R3）。

1. Introduction

稀疏专家神经网络展示了纯规模的优势，并提供了一种有效的替代方案，与今天通常使用的静态神经网络架构相比（Raffel等，2019；Brown等，2020；Rae等，2021）。与将相同的参数应用于所有输入不同，稀疏专家网络动态选择用于每个输入的参数（Shazeer等，2017）。这允许网络大大扩展其参数数量，同时保持每个token的浮点运算数（FLOPs）大致恒定。这些方法已经产生了最先进的翻译模型（Lepikhin等，2020）、4-7倍的预训练加速（Fedus等，2021；Artetxe等，2021），以及使用1/3的能量训练成本实现了与GPT-3相当的一次性性能（Du等，2021）。尽管参数数量惊人，稀疏模型将大型神经网络的训练碳足迹降低了一个数量级（Patterson等，2021）。然而，仍然存在一些困难。

Fedus等人（2021）观察到，一个稀疏的1.6T参数模型在先前的最新技术（Raffel等，2019）基础上实现了4倍的预训练加速，但在像SuperGLUE这样的常见基准数据集上进行微调时，它的性能落后于较小的模型。在Artetxe等人（2021）的研究中，当MoE语言模型在域外数据上进行微调时，也观察到了类似的差距。作为回应，提出了Switch-XXL，这是一个参数较少但计算量更大（FLOPs大约等于最大的T5模型的8倍）的模型，并提高了在自然语言理解任务中的质量。然而，在先前规模较小的研究中未能检测到的训练不稳定性影响了必要的预训练。这些不稳定性随后在其他稀疏模型中被发现（Du等，2021）。这些结果揭示了参数和计算之间的必要平衡，但在如何可靠地训练这类模型方面留下了一个开放的问题。

本文的目标是增加稀疏模型的实用性和可靠性。我们研究了这两个问题，并预训练了一个2690亿参数的稀疏模型，在微调过程中在许多竞争性NLP基准测试中取得了最先进的结果，包括SuperGLUE。我们还提出了对稀疏专家模型的额外分析和设计指南（或至少是我们的启发式方法）。此外，这项工作强调了共同优化上游预训练和下游微调指标，以避免差异（Tay等，2021）。

本文的主要贡献：

1. 对稳定性技术的质量和稳定性之间的权衡进行的大规模研究。
2. 引入了解决不稳定性问题的路由器z-loss，同时略微提高了模型质量。
3. 对稀疏和密集模型的微调分析，突显了对批量大小和学习率的不同超参数敏感性。我们展示了糟糕的超参数导致几乎没有比密集模型更好的微调效果，尽管进行了大量的预训练加速。
4. 在分布式环境中设计帕累托高效稀疏模型的架构、路由和模型设计原则。
5. 通过定性分析追踪专家层中的token路由决策。
6. 一个2690亿参数的稀疏模型（稳定可传递专家混合模型或ST-MoE-32B），在各种自然语言基准测试中取得了最先进的性能。

2. Background

稀疏专家模型通常通过一组专家替换神经网络层，每个专家具有独特的权重（Jacobs等，1991；Jordan和Jacobs，1994）。通常，层内的所有专家都是相同类型和形状的（同质的），然而，也可以有不同类型和形状的（异质的）专家。为了节省计算，输入只由专家的子集处理，因此必须添加一个机制来确定将每个输入发送到哪里。通常，路由器或门控网络确定要将输入（例如单词、句子、图像补丁等）发送到何处，但也提出了替代方案（Lewis等，2021；Roller等，2021；Zuo等，2021；Clark等，2022）。

具体而言，在自然语言处理中，Shazeer等人（2017）提出了一个混合专家（MoE）层，该层以token表示x作为输入，并将其路由到一组\(\{E_i(x)_{i=1}^N\}\)这N个专家中的最佳匹配的前k个专家。路由器参数$W_r$生成logits $h(x) = W_r \cdot x$，通过softmax分布对该层中可用的N个专家进行归一化。专家i的门值由以下公式给出：

并且token x被路由到具有最高前k个门值的专家（索引集合$\mathcal{T}$）。该层的输出是由门值加权和计算的每个专家的结果。

最初在LSTMs中提出（Hochreiter和Schmidhuber，1997），后来由Shazeer等人在Transformer中使用（Vaswani等人，2017；Shazeer等人，2018；Lepikhin等人，2020）。Fedus等人（2021）的后续工作进一步简化了MoE，使其将token路由到单个专家（top-1），并降低了其他成本以提高训练效率。

为了提高硬件利用率，大多数稀疏模型的实现对每个专家使用静态批处理大小（Shazeer等人，2017；2018；Lepikhin等人，2020；Fedus等人，2021）。专家容量指的是可以路由到每个专家的token数量。如果超出了这个容量（路由器向该专家发送了太多输入），则溢出的token不会应用计算，并通过残差连接传递到下一层。

批次 $\mathcal{B}$ 的输入token被分成 $\mathcal{G}$ 个在数据并行维度上唯一的组，每个组的大小为 $\mathcal{B}/\mathcal{G}$。专家容量等于 CF·tokens/experts，其中 CF 表示容量因子超参数，experts 是专家的数量，tokens 是组的大小。如果容量因子增加，它会创建额外的缓冲区，以便在负载不平衡的情况下减少丢弃的token数量。然而，增加容量因子还会增加内存和计算成本，因此存在权衡。

最后，辅助的负载平衡损失鼓励token在专家之间大致均匀分布（Shazeer 等人，2017）。这通过确保所有加速器并行处理大块数据来提高硬件效率，如上所述。损失的详细信息在附录 A 中呈现。然而，也存在替代方案：Lewis 等人（2021）和Clark 等人（2022）将平衡的token分配视为一个分配问题，并完全删除了辅助损失。

3. 让稀疏模型的训练更稳定

稀疏模型在训练过程中通常会遇到比标准密集激活的Transformer中观察到的更严重的训练不稳定性（图1）。

图1：稀疏模型的训练不稳定性。我们将训练不稳定性定义为训练损失的发散。上图显示了两个运行，这些运行的稀疏模型与T5-XL版本（Raffel等人，2019）匹配，每个运行使用Adafactor优化器（Shazeer和Stern，2018）以1M个token的批量进行训练。（左）一个不稳定的训练运行。 （右）一个稳定的训练运行。

找到改善稳定性的变化是直接的，然而，这些改变通常以无法接受的模型质量损失为代价（例如，使用任意小的学习率或使用紧密的梯度剪裁）。我们对改善稳定性的几种方法进行分类和研究。这些稳定性技术涵盖了对Transformer的通用修复以及对稀疏模型特定的修复：(1) 移除乘法交互 (2) 注入模型噪音 (3) 限制激活和梯度。我们最后提出了我们的建议：一种新的辅助损失，即路由器z-loss，它在不降低质量的情况下显著提高了训练稳定性。这是对Mesh Tensorflow代码库中用于最终softmax logits的z-loss的改编（Shazeer等人，2018）。

• 许多方法可以稳定稀疏模型，但代价是更差的质量。
• 路由器z-loss可以在不降低质量的情况下稳定模型。
• 具有更多乘法组件的Transformer修改（GEGLU、RMS normalization）会恶化稳定性，但提高质量。

设计大规模的稳定性研究。 我们设计了一个大规模的稳定性研究，使用了与T5-XL版本（Raffel等人，2019）FLOP匹配的稀疏模型，该模型在多语料库mC4（Xue等人，2020）上进行了预训练。每个稀疏模型有32个专家，我们为每个第四个FFN引入了一个稀疏MoE层。训练容量因子为1.25，评估容量因子为2.0。有关本文中使用的模型的更详细描述，请参见表11。对于每种稳定性技术，我们记录了稳定的比例、平均质量（对英语的负对数困惑度）以及在种子上的标准差。

构建这项研究的主要问题在于，小型模型很少不稳定，但大型不稳定模型的运行成本过高，无法运行足够的步骤和种子。我们发现，与T5-XL匹配的稀疏模型在大约1/3的运行中是不稳定的，但仍然相对便宜进行训练，因此它成为研究的良好对象。此外，我们在多语言数据上运行不稳定性实验，因为我们发现这会加剧模型的不稳定性，从而使我们能够在稍微较小的模型上进行实验。有关更多详细信息，请参见第9节。我们的基准配置使用六个随机种子进行训练，每个具有稳定性技术的配置使用三个随机种子。我们对基准使用六个种子以更好地表征不稳定性率，并对变体使用三个种子以节省计算资源。每个模型使用掩码语言建模目标（Fedus等人，2018；Devlin等人，2018）在mC4上进行了20,000步的预训练。

3.1 移除乘法交互时的稳定性和质量权衡

一些架构改进涉及更多的乘法而不是加法，或者一次不对许多项求和。例如，矩阵乘法每次加法都有一个乘法，因此我们不将其称为“乘法”操作。我们在这里介绍并分析了Transformer中两个实例的乘法交互的影响。

GELU门控线性单元（GEGLU）。 我们的第一个例子是门控线性单元（Dauphin等人，2017），它是两个线性投影的逐分量乘积，其中一个首先通过S形函数传递。Shazeer（2020）将其扩展到其他变体，并提出用GELU-Linear（Hendrycks和Gimpel，2016）FFN层替换Transformer中通常的ReLU（Nair和Hinton，2010）FFN。

\(FFN_{GEGLU}(x,W,V,b,c)=GELU(xW+b) \odot (xV+c)\) 这种质量的提升在后续工作中得到了证实（Narang等人，2021）。

均方根标度参数。 我们的第二个例子是均方根（RMS）归一化中的标度参数（Zhang和Sennrich，2019）。在Transformer内部，而不是依次调用层，有一种内部结构（称为子层调用），它改善了梯度传播和训练动态。我们的子层调用与Raffel等人（2019）相匹配，包括：（1）RMS归一化，（2）自注意力层/全连接层，（3）dropout（Srivastava等人，2014），（4）添加残差（He等人，2015）。RMS归一化按元素对输入向量 进行根均方归一化。然后，它通过与学习到的标度参数g逐元素相乘，对输出进行逐元素重新缩放。

表2显示，无论是移除GEGLU层还是RMS标度参数，都能提高稳定性，但对模型质量造成了显著损失。我们注意到这些标度参数g对模型质量的提升与其他地方的参数（例如FFN）相比具有不成比例的收益。与我们的发现一致，Shleifer等人（2021）发现在Transformer的残差连接中添加一个学习到的乘法标量会使其变得更加不稳定。

在附录C中，我们进一步研究了在专家层中添加新的乘法交互对质量的影响。我们发现，这个操作可以在几乎不减慢模型步骤时间的情况下提高质量。

表2：移除具有更多乘法交互的操作。乘法交互可以提高质量，但可能使训练不稳定。单独移除两个乘法组件的来源可以提高稳定性，但显著降低质量。当我们移除GEGLU层时，我们将其替换为一个等效的Dense-ReLU-Dense层，以匹配FLOPs和参数。

3.2 增加噪音时的稳定性和质量权衡

接下来，我们探讨一个假设，即向模型中添加噪音可以提高训练的稳定性（Neelakantan等人，2015）。Taleb（2012）认为某些系统具有反脆弱性的特性，即它们通过噪音变得更好。受到这一概念的启发，以及我们观察到微调（通过dropout注入噪音）很少不稳定，我们研究了训练噪音是否可以提高稀疏模型的稳定性。表3显示了相对基线的稳定性改善，但代价是较低的质量。我们还发现，由Fedus等人（2021）引入的输入抖动在XL规模下会降低质量，因此我们在模型中消除了它。输入抖动将路由器的输入logits乘以一个均匀随机变量，该变量位于$[1 − 10^{-2}，1 + 10^{-2}]$之间。在我们的消融中，dropout应用于整个Transformer。正如以前所见，小规模环境中的改进在扩大规模时可能无法推广，因此趋势应始终在不断增加的规模下进行监控和重新评估（Kaplan等人，2020）。

表3：训练期间注入噪音。输入抖动和dropout都提高了稳定性，但导致了模型质量的显著降低。大多数方法存在明显的权衡：当一种方法提高了稳定性时，通常会降低模型质量。我们的工作旨在找到在不损害质量的情况下提高稳定性的方法。。

3.3 稳定性和质量权衡在约束激活和梯度方面

稳定神经网络最成功的方法之一是对激活和梯度进行约束（Pascanu等，2013；Ioffe和Szegedy，2015；Salimans和Kingma，2016；Ba等，2016）。一种常见的方法是剪裁梯度范数，以解决在深度网络中反向传播时梯度爆炸的问题（Pascanu等，2013）。

在这项工作中，我们使用Adafactor优化器，因为它具有内存效率（尽管最近引入的8位优化器（Dettmers等，2021）可能提供更好的权衡）。Adafactor不使用梯度剪裁，而是使用更新剪裁，其中权重的更改被限制在某个范数以下。我们尝试将更新剪裁值缩小到一个较小的值。

接下来，我们研究传递到路由器的logits上的约束。路由器以float32精度（即选择性精度）计算专家之间的概率分布（Fedus等，2021）。然而，在最大规模上，我们发现这对于可靠的训练是不足够的。为了解决这个问题，我们引入了路由器z-loss：

其中B是token数量，N是专家数量，$x \in R^{B \times N}$是传递到路由器的logits。这对进入门控网络的大logits进行惩罚，第3.4节包含了为什么在路由器之前使用z-loss是有益的更详细的解释。

T表4 表明，无论在3次运行中使用更新剪裁还是路由器 z-loss，都可以稳定模型，但是更新剪裁显著损害了模型的质量。因此，我们选择使用 z-loss 方法来提高模型的质量和稳定性。

表格 4：约束权重更新和路由器 logits。在 Adafactor 中约束更新剪裁提高了稳定性，但以巨大的质量损失为代价。较宽松的剪裁值并不能可靠地稳定训练，因此我们在此将它们排除。路由器 z-loss 在不降低质量的情况下稳定了模型（在这种情况下，我们观察到了轻微的质量提升）。

路由器 z-loss 引入了另一个超参数（$c_z$），该超参数是用于加权作为优化的总损失的一部分。总损失是交叉熵损失（$L_{CE}$）、辅助负载平衡损失（$L_B$）和路由器 z-loss（$L_Z$）的线性加权组合，得到总损失。

我们选择 $c_z = 0.001 $ 的值，基于在超参数扫描后预训练后的最佳模型质量。附录 B 记录了预训练过程中得到的损失。

3.4 选择精度格式：权衡效率和稳定性

与大多数现代分布式Transformer一样，我们使用混合精度进行训练（Micikevicius等，2017）。权重在进行梯度更新时存储为float32，然后在前向和后向传递中进行矩阵乘法时转换为bfloat16。此外，所有激活都以bfloat16存储和操作，所有reduce通信可以在bfloat16或float32数值精度下完成。对于本文中探讨的最大模型（稍后介绍的ST-MoE-32B），我们发现将allreduce的数值精度减半可以提速，但这也可能导致训练不稳定，因此我们在本文中一直将其保持为float32。

较低精度格式通过减少（a）处理器和内存之间的通信成本，（b）计算成本，（c）存储张量（例如激活）的内存等方面实现了更高效的模型。然而，较低精度格式的代价是更大的舍入误差，可能导致无法恢复的训练不稳定性。

图2：数值精度格式和舍入误差。较大的数字具有较大的舍入误差。bfloat16的舍入误差最多比float32糟糕65,536倍。路由器z-loss鼓励数字的绝对幅度较小，这不会妨碍模型性能并减小舍入误差。路由器z-loss在对那些较大误差可能极大改变相对输出的函数中最为有效（例如指数和正弦函数）。

理解精度格式和舍入误差。 图2回顾了不同精度格式及其相应数值范围的舍入误差的特性。在两个连续的2的幂的任何范围内的数字（例如[2,4)和[1024, 2048)）由固定数量的尾数位表示（bfloat16为7，float32为23）。因此，（1）bfloat16的舍入误差大约为65,536倍（即23 - 7 = 16个额外的位和216 = 65536），如float32；（2）较大的数字具有较大的舍入误差。由于8位指数位，数字最大可达约$3e^{38}$，这使得即使是float32在舍入误差方面也存在一些问题。

稀疏专家模型对舍入误差非常敏感，因为它们由于路由器而具有更多的指数函数。 稀疏专家模型通过路由器引入了额外的指数函数，可能加剧舍入误差，导致训练不稳定性。虽然舍入误差不会改变softmax操作中概率的排序，但由于相对阈值处理（例如，仅当第二个专家的门控概率是第一个专家的1/5时，token才被路由到其第二个位置的专家）的原因，它确实会影响到MoE中的第二个token的路由。此外，舍入误差可能大幅改变缩放专家输出的概率，我们发现这一点非常重要。最后，我们推测我们观察到的解码器模型更稳定的原因（此处不显示）可能是因为它们具有较少的指数函数。第9节包含更详细的讨论。

关于路由器z-loss的说明。 有人可能认为路由器z-loss是一种可以用截断logits（Wu等，2016）替代的复杂方法。我们解释为什么这并非如此。目标是最小化输入指数函数的大舍入误差。截断logits发生在任何舍入误差之后，导致更大的不连续性。从某种程度上说，截断本身就是一种舍入误差；相反，z-loss自然地鼓励模型产生数值较小的logits，因此更准确地建模。由于这些动态特性，我们确保所有指数化张量都被转换为float32。这暗示了在整个网络中添加z-losses时神经网络可能具有更好的数值格式的可能性（见第9节）。

4.稀疏模型的精细调整性能

通常，获得性能最佳的语言模型是通过（1）在大量数据上进行预训练（例如互联网），然后（2）在感兴趣的任务上进行精细调整（例如SuperGLUE）而获得的。出现了一些有希望的新技术作为替代方案，包括少样本推断（Brown等人，2020），前缀调整（Li和Liang，2021），提示调整（Lester等人，2021）以及适配器模块（Houlsby等人，2019） - 然而，与精细调整相比，仍然存在一定的质量差距。因此，本文侧重于精细调整，但强调了Du等人（2021）；Artetxe等人（2021）在少样本设置中稀疏模型的最新成功案例。此外，我们将通过强化学习（Ouyang等人，2022）将大型语言模型调整的技术作为未来的研究方向。

4.1 假设：泛化问题

稀疏模型在大数据集领域表现出色，但在精细调整时有时表现不佳（Fedus等人，2021; Artetxe等人，2021）。我们提出一个（并不奇怪的）假设，即稀疏模型容易过拟合。我们通过SuperGLUE的两个任务 - Commitment Bank（De Marneffe等人，2019）和ReCORD（Zhang等人，2018）来说明这个问题。Commitment Bank（CB）有250个训练示例，而ReCORD有100,000多个。这种显著的大小差异为在相同基准的一部分选定的两个任务进行过拟合的自然研究提供了便利。

在图3中，我们比较了Dense L和ST-MoE-L模型的精细调整特性。每个模型都在C4语料库（Raffel等人，2019）的500B token上进行了预训练。这些模型被设计成与Raffel等人（2019）的T5-Large编码器-解码器模型的FLOP匹配变体，参数为770M。 ST-MoE模型有32个专家，专家层频率为1/4（每四个FFN层中的一个被MoE层替换）。预训练和精细调整的训练能力因子为1.25，评估能力因子为2.0。我们在留出的验证集和训练数据集分区上评估性能。

图3：稀疏模型容易过拟合。我们绘制了在CB任务（250个训练序列）和ReCoRD（138k个训练序列）上对ST-MoE-L和dense-L模型进行精细调整的训练和验证曲线。在两种情况下，稀疏模型在训练分区上学习速度更快（蓝线超过绿线）。然而，对于较小的CB任务，密集模型在验证集上表现优于稀疏模型（红色与橙色相比）。相反，在较大的ReCoRD任务中，稀疏模型表现优于密集模型数个百分点。

在两个任务中，稀疏模型更快地收敛到100％的训练集准确率，证明稀疏模型在数据分布变化下能够有效优化。在更大的任务ReCORD上，稀疏模型的验证质量跟随训练的提升显著超过了密集模型。然而，在较小的任务CB上，稀疏模型在留出数据上落后于其密集对应物。根据Fedus等人（2021）的建议，我们考虑增加专家隐藏状态内的辍学（即专家辍学），但发现在这个规模上，更高的值只能适度提高质量（图4）。我们将在4.2节研究精细调整的进一步改进，并在4.3节研究超参数的敏感性。

Figure 4: 稀疏模型精细调整的正则化研究。对于每种设置，我们在SuperGLUE上使用三个随机种子进行训练，直至收敛。我们发现通过增加dropout来提供适度的提升。 （左侧）展示了在全局dropout为0.1时的最高SuperGLUE精细调整质量。较高的值过度正则化，并严重损害质量。 （右侧）从已知的最佳全局dropout0.1开始，我们有选择地增加专家dropout（专家隐藏激活的独立dropout）。这带来了进一步的泛化效益，并符合Fedus等人（2021）的发现。

4.2 优化模型参数子集以提高泛化性能

为了对抗过拟合，我们尝试在精细调整期间仅更新模型参数的子集。图5测量了更新5个不同参数子集的质量：所有参数（All），仅非MoE参数（Non MoE），仅MoE参数（MoE），仅自注意力和编码-解码注意力参数（Attention）以及仅非MoE的那些FFN参数（FFN）。

图5：在精细调整期间仅更新模型参数的子集。为了提高稀疏模型的泛化性并对抗过拟合，我们对模型参数的子集进行精细调整。所有结果均基于ST-MoE-L模型，是5个不同随机种子的平均值。我们观察到，更新3/5的参数子集似乎效果差不多，而仅微调MoE参数会导致显著的质量降低。

我们观察到，仅更新非MoE参数的效果大致与更新所有参数相当，而仅更新FFN参数的效果略好一些。仅更新MoE参数显著降低了精细调整性能，其中约80%的模型参数存在。

仅更新非MoE参数可以是一种有效的加速和减少精细调整内存的方式。

我们假设仅微调MoE参数导致性能不佳，因为专家层仅每4层出现一次，每层token最多只会看到两个专家。因此，更新MoE参数将影响的层数和FLOP数比我们尝试的任何其他参数子集都要少。仅更新MoE参数导致训练损失较大，尽管参数数量明显更多。我们进一步观察到，更新所有非MoE参数的训练损失高于更新所有参数，但不幸的是，这种正则化效果并没有转化为更好的验证性能。

此外，我们尝试了一个正则化器，即在训练期间随机屏蔽整个专家的dropout变体。然而，在我们的初步研究中，这并未改善泛化效果。附录J对这个实验进行了扩展并包含其他负面结果。

4.3 稀疏模型和密集模型需要不同的微调协议

稀疏模型和密集模型对微调协议的敏感性如何？我们研究了两个超参数：批大小和学习率。我们预先在C4的500Btoken上对Dense-L和ST-MoE-L进行预训练，然后在SuperGLUE上进行微调。图6总结了我们的实验，完整的数据见表20（附录F）。在所有的超参数设置中，稀疏模型（橙色）表现优于密集模型（蓝色） - 然而，对于每个模型，最佳设置可能会显著改变结果。在不同的批大小和学习率设置下，稀疏模型和密集模型表现迥异。稀疏模型受益于较小的批大小和较高的学习率。与过拟合假设一致（见第4.1节），这两个变化可能通过在微调过程中引入更多噪音来改善泛化效果。最后，我们强调在微调过程中正确调整批大小和学习率的重要性。简单地使用对密集模型有效的相同微调超参数可能掩盖了稀疏模型获得的任何预训练改进。

图6：批大小和学习率敏感性。我们衡量了密集模型（蓝色）和稀疏模型（橙色）之间微调协议的差异和敏感性。每个条形图是使用不同超参数的6个不同运行的平均值。在SuperGLUE上，稀疏模型受益于包括小批大小和高学习率在内的更多噪音超参数。密集模型表现几乎相反。详细数据请参见附录F。

4.4 稀疏模型在微调期间丢弃的token鲁棒性

稀疏模型在每个层面将token路由到一个或多个专家。为了使这些模型在具有现代硬件的SPMD范例中变得高效，专家容量（每个专家处理的token数）需要提前确定（有关详细信息，请参见第2节）。当专家接收的token超过其容量时，额外的token将被丢弃，不对这些token应用任何计算。我们再次尝试通过以下两种方式来防止这种情况：（1）使用促使将相等数量的token发送到每个专家的辅助损失进行预训练，以及（2）使用容量因子（一个超参数），该因子在每个专家处增加了额外的token容量。我们尝试在微调期间关闭辅助损失并使用不同的容量因子。表5显示了一个令人惊讶的结果，即微调质量在丢弃多达10-15%的token时并没有受到实质性影响。对ST-MoE-32B的研究证实，高容量因子并不提高微调质量。这与Yang等人（2021年）发现的负载不平衡可能不会显着影响模型质量的结果一致。

表5：在微调期间，稀疏模型对于丢弃的token表现出鲁棒性。我们发现，在探索的值范围内，对SuperGLUE的微调质量没有显著影响。有趣的是，丢弃10-15%的token的模型性能几乎与丢弃<1%的模型相当。我们还观察到负载平衡损失（Aux Loss）可以提高微调质量。丢弃的token百分比对应于在达到最佳验证准确性时在所有专家层中丢弃的token的分数。

哨兵token表示在跨度破坏目标（Fedus等人，2018年；Devlin等人，2018年）中的屏蔽序列。这与我们可能会遇到的任何微调任务都不同，导致了预训练和微调之间的域不匹配。表6说明了这种差异。我们研究了修改微调任务以使其更像预训练任务是否会影响结果。

表6：在微调过程中插入哨兵token模仿了预训练跨度目标。我们突显了跨度破坏和微调之间的典型差异。我们建议通过插入哨兵token修改微调任务，使其类似于预训练。

在表7中，我们发现在微调过程中添加哨兵token只能改善语法错误纠正（Grammar Error Correction，GEC）（Rothe等人，2021），而对于SuperGLUE则没有改善。我们尝试通过插入多个哨兵token（与模型在预训练过程中遇到的情况相同）来进一步减少数据分布差异，但同样没有普遍的好处。然而，尽管在验证数据上没有一致的好处，我们发现对于密集模型和稀疏模型来说，训练收敛速度加快了。

表7：哨兵token对微调的影响。在微调过程中添加哨兵token（与Lester等人（2021）中使用的类似概念）在我们考虑的两个任务上表现出了不同的性能。SuperGLUE记录了平均分数，而GEC记录了精确匹配情况。尽管我们发现它并不能改善泛化效果，但哨兵token可以加速训练的收敛。

5. 设计稀疏模型

稠密模型的设计受到Kaplan等人（2020年）基础工作的指导。但稀疏模型提出了许多额外的问题：（1）使用多少专家？（2）使用哪种路由算法？（3）容量因子选择什么值？（4）硬件如何影响这些决策？在本节中，我们对这些问题进行评论并提出了建立帕累托有效稀疏模型的建议。与此同时，Clark等人（2022年）提供了额外的设计建议，包括更高层频率和根据Fedus等人（2021年）的建议进行的top-1路由。

设计稀疏模型

• 在我们的设置中，我们建议使用top-2路由，容量因子为1.25，每个核最多使用一个专家。
• 可以在评估过程中更改容量因子，以适应新的内存/计算要求。
• 密集层叠加和乘法偏差可以提高质量（附录C）。

5.1 设置专家数量

关于要使用多少专家是最初的问题之一。Fedus等人（2021年）介绍了Switch Transformer的缩放属性，该属性在C4上呈现出单调的预训练效益（基于步骤），他们最多使用了512个专家，Kim等人（2021年）使用了64个专家，Clark等人（2022年）用了512个专家。但随着专家数量的增加（>256），或者等效地说，随着模型的稀疏度增加（<1%的激活专家），额外的收益迅速减小。

然而，对特定硬件系统的考虑可以进一步指导此选择。计算内存比（操作强度）可以用作不同操作效率的估计（Williams等人，2009年；Shazeer，2019年）。如果将张量加载到计算核心（例如ALU/MMU）所需的时间远远超过对张量进行计算所需的时间，则模型受到内存限制。在现代GPU和TPU上，增加此计算到内存比可以提高效率。

回到稀疏专家模型，每个核心使用多个专家会增加内存传输，可能损害效率。增加专家数量不会改变所执行的计算（稀疏模型对每个输入应用固定数量的计算），但会增加内存传输要求（必须从设备内存加载附加专家变量）。这会降低计算到内存比率。

在我们的TPU系统上，我们建议每个核心使用一个专家（或更少）。我们最大的模型同时使用数据并行性和模型并行性，其中数据并行性是在“行”上进行的，模型并行性是在逻辑网格的“列”上进行的。我们在每个数据并行性行上使用≤ 1个专家，以确保计算到内存比率较高，并减少评估和推理所需的核心数量。此外，使用较少的专家使我们能够将更多的核心分配给模型并行性的“列”，以在我们的模型中获得更多的FLOP。附录H解释了在专家数量少于数据并行性行数时我们的网格布局。

5.2 选择容量因子和路由算法

我们将top-1路由（Fedus等人，2021年；Roller等人，2021年）和top-2路由（Shazeer等人，2017年；Lepikhin等人，2020年）推广为研究top-n路由，其中每个token最多由n个专家处理。在这项研究中，所有模型都进行了100k步的预训练，每批处理1Mtoken，稀疏模型有32个专家，并且与T5-Large Raffel等人（2019年）的FLOP匹配。我们得出了两个关键结论。

首先，增加训练和评估容量因子（CF）都会提高质量，通过比较表8的分段块可以看到。例如，从1.0 → 1.25训练CF时，top-1路由的neg. log perp.提高了+0.011，而当从1.25 → 2.0训练CF时，top-2路由提高了+0.009。为了提供这些数字的背景：将密集模型的大小增加三倍（从Dense-L到Dense-XL）会提高+0.090的neg. log perp.。因此，这些CF提升约为该幅度的1/10。但这是有代价的。线性增加容量因子会导致einsums成本、激活内存、all2all通信成本以及专家层的模型并行通信成本的增加。

其次，对于给定的固定容量因子，top-(n+1)路由相对于top-n路由有小的增益（表8）。例如，在训练CF为1.25时，top-2路由相对于top-1提高了+0.004，约为密集模型三倍提升的1/20。这修改了Fedus等人（2021年）早期的建议。这些实验设置之间的主要区别在于计算规模。Fedus等人（2021年）为50Btoken训练了220M-FLOP匹配的模型。我们发现在训练规模8倍大（100Btoken的1B-FLOP匹配模型）的情况下，与一个专家相比，选择多个专家有小的增益。此外，在更大的实验规模下，top-n与top-(n + 1)路由的速度差异可以忽略不计。在Fedus等人（2021年）中观察到速度差异是因为路由器计算占总模型计算的较大部分。

表8：比较容量因子（CF）和路由算法。增加训练和评估CF都会提高性能。如果在评估时有更多或更少的计算资源，增加或减少评估CF可以提供额外的控制杆。其次，在不同的容量因子下，top-(n + 1)路由相对于top-n路由有较小的增益。由于质量提高，但随着CF的增加，速度减慢，Pareto有效的CF必须由具体的硬件系统确定。

具体的硬件-软件系统将决定最佳的n和容量因子。例如，如果系统支持快速的all2all和allreduce通信，那么在top-n路由中使用更大的容量因子和更大的n可能是最优的。然而，如果all2all和/或allreduce通信速度较慢，那么较小的容量因子可能更为主导。在我们的情况下，硬件-软件堆栈是TPU和Mesh TensorFlow。我们在表9中记录了我们的ST-MoE-L和ST-MoE-32B模型随着增加训练容量因子的训练速度。随着模型规模的增加，较高的容量因子使模型变得越来越慢。ST-MoE-L不需要模型并行性（它适应于加速器内存，这意味着没有额外的allreduce通信），使其更适合于比我们的ST-MoE-32B模型更高的容量因子。因此，对于我们最大的模型，我们继续使用Fedus等人（2021年）推荐的更小的训练容量因子1.25，以实现帕累托效率，与其他使用更大且更昂贵的2.0容量因子的研究有所不同（Lepikhin等人，2020年；Du等人，2021年）。

表9：在TPU上对稀疏模型进行性能分析。将训练容量因子从1.25增加到2.0，大型（1B）模型的步骤时间增加了+7%，但我们的32B模型增加了+14%。随着模型规模的增加，我们发现从表8中获得的更高训练容量因子的小幅质量提升被显著的14%减速所抵消。注意：ST-MoE-L和ST-MoE-32B之间的步骤时间不可比较，因为它们使用了不同数量的核心。

本节中，我们的结果主要关注top-n路由，但我们还在附录J中尝试了各种其他路由技术。我们发现，与top-n路由相比，大多数表现相似或更差。然而，我们发现Batch Prioritized Routing（BPR），在Riquelme等人（2021年）中引入，对于容量因子小于一的情况显著提高了性能（附录D）。我们建议在all2all和allreduce更昂贵、较小容量因子更为优越的大型模型中使用BPR。

6. 实验结果

鉴于我们在训练稳定性、微调和模型设计方面的改进，我们首先验证了一个与T5-Large（Raffel等人，2019）近似FLOP匹配的稀疏模型。我们通过设计和训练一个269B的稀疏参数模型（与32B的密集模型匹配的FLOP），在一系列广泛的自然语言处理任务中实现了最先进的质量，并以此来结束这一部分。

我们在整个工作中研究了SuperGLUE（Wang等人，2019）基准，该基准包括情感分析（SST-2）、词义消歧（WIC）、句子相似性（MRPC、STS-B、QQP）、自然语言推理（MNLI、QNLI、RTE、CB）、问答（MultiRC、RECORD、BoolQ）、共指解析（WNLI、WSC）和句子完成（COPA）以及句子可接受性（CoLA）等任务。我们经常观察到在SuperGLUE上的良好性能与（但不能保证）在许多自然语言处理任务上的性能相关。我们还包括一系列其他基准。CNN-DM（Hermann等人，2015）和BBC XSum（Narayan等人，2018）数据集用于评估总结文章的能力。我们使用SQuAD数据集（Rajpurkar等人，2016）以及ARC Easy和ARC Reasoning Challenge（Clark等人，2018）上的小学科学问题来探讨问题回答。与Roberts等人（2020年）一样，我们通过在三个封闭书籍问答数据集上进行微调来评估模型的知识：自然问题（Kwiatkowski等人，2019），Web问题（Berant等人，2013）和Trivia QA（Joshi等人，2017）。封闭书籍仅指在没有额外参考或上下文材料的情况下提出的问题。为了衡量模型的常识推理能力，我们在Winogrande Schema Challenge（Sakaguchi等人，2020）上进行评估。最后，我们在Adversarial NLI Benchmark（Nie等人，2019）上测试我们模型的自然语言推理能力。

6.1 ST-MOE-L

为了简化操作并轻松涵盖数十个任务，我们训练了混合任务的模型，而不是在每个任务上分别微调模型。然而，由于任务的大小差异较大，按照样本数量进行等量采样会过度采样大型任务并欠采样小型任务。因此，我们按照每个任务“train” split中的样本数量的比例混合每个任务（最大样本数为65536，如Raffel等人，2019年所述）。这意味着包含超过65536个训练样本的任务被加权，就好像它们只包含最大样本数。

表10总结了一个密集的T5-Large（L）模型和一个稀疏模型的质量，它们具有大致相同数量的FLOP，在C4数据集（Raffel等人，2019）上进行了500k步的预训练，批量大小为1M（524B个token）。编码器的序列长度为512，解码器为114。我们观察到在验证（dev）集上在广泛的任务中（包括自然语言理解、问题回答和摘要）取得了改进。正如在Fedus等人（2021年）中所见，封闭书籍问答（Roberts等人，2020年）中观察到了显著的增益。此外，在支持第4.1节中提出的过拟合假设的情况下，我们观察到两个最小的任务CB和WSC（分别有250和259个训练示例）是稀疏模型没有在这两个任务中超过密集模型的唯一任务。这再次表明，为稀疏模型提供改进的正则化形式可能会释放更大的性能。

表10：FLOP匹配的密集和稀疏模型的微调性能比较。密集-L基准与稀疏FLOP匹配版本的对比（数字越高越好）。我们观察到在各种任务中都取得了一致的收益，使用大致相同的计算量。唯一两个没有从稀疏模型中获得改进的任务是两个最小的任务：CB有250个训练示例，WSC有259个。

6.2 ST-MoE-32B

在对T5-Large规模的质量进行验证后，我们着手通过ST-MoE-32B推动稀疏模型的能力。在设计这个模型时，我们追求FLOP和参数之间的平衡。在Fedus等人的工作中，高FLOP的稀疏模型在我们的设置中（即编码器-解码器模型，Adafactor优化器）以前是不稳定的，但路由器的z-loss使我们能够继续进行。为了提高计算效率，我们扩展了专家的隐藏大小（表11中的$d_{ff}$）。最后，我们将dkv增加到128，以提高在我们的硬件上的性能。相对于Switch-C和Switch-XXL，最显著的变化是总参数减少，每个token的FLOP增加。我们的ST-MoE-32B仅有269B参数，大致与具有32B参数的密集Transformer匹配的FLOP相匹配。与Switch-C和Switch-XXL相比，减少的参数数量减轻了在服务和微调中的负担。最后，我们使用附录C中描述的稀疏-密集堆叠。

表11：模型比较。对Dense-L和T5-XXL，两个最大的Switch Transformer变体（Switch-XXL和Switch-C），以及ST-MoE-L和ST-MoE-32B进行比较。$d_{model}$指的是模型隐藏状态的大小，$d_{ff}$是FFN层的内部大小。$d_{kv}$是每个注意力头的维度。Expert Layer Freq. 是用稀疏层替换的FFN层的比例。Sparse-Dense指的是附录C中描述的架构变体。

我们在仅含英语的C4数据集（Raffel等人，2019）和GLaM（Du等人，2021）数据集的混合物上进行了1.5Ttoken的预训练，附录E对此进行了总结。我们使用每批1Mtoken，Adafactor优化器的默认超参数，以及10k步的学习率预热，然后是倒数平方根衰减。我们的模型遵循Fedus等人（2021）提出的初始化方案。

表12：ST-MoE-32B与先前最佳推理技术和经过微调的模型进行比较。 “dev/test”拆分是指零射击和一射的dev拆分以及精调质量的测试拆分。数据不可用时填入“–”。上标字母表示结果：$^a$：Raffel et al.（2019）$^b$：Roberts et al.（2020）$^c$：Karpukhin et al.（2020）$^d$：Brown et al.（2020）$^e$：Du et al.（2021）$^f$：Wang et al.（2021）$^g$：UnifiedQA + ARC MC/DA + IR，$^h$：Zhang et al.（2020）。

表12评估了我们的ST-MoE-32B模型与以前的最先进方法，包括仅推理（零样本、一样本）以及微调。在SuperGLUE上，我们的模型改进了先前的最先进模型，在测试服务器上获得了91.2的平均分数（93.2的验证准确性），超过了估计的人类能力一百分之一以上。对于摘要数据集XSum和CNN-DM，我们的模型在没有对训练或微调进行其他更改的情况下实现了最先进水平（Raffel等人，2019；Liang等人，2021）。ST-MoE-32B改进了ARC Easy（92.7 → 94.8）和ARC Challenge（81.4 → 86.5）的测试服务器提交的最先进水平。在三个闭卷QA任务中的两个上，我们超过了先前的最先进水平。闭卷WebQA达到了47.4的准确性（先前最佳为来自Roberts等人（2020）的42.8），并超过了ERNIE 3.0 Titan 260B密集参数模型的零样本性能结果（Wang等人，2021）。闭卷NatQA的准确性提高到41.9（先前最佳为Karpukhin等人（2020）的41.5）。我们在对抗性构建的数据集（ANLI R3和WinoGrande XL）上取得了显著的改进。ANLI R3（Nie等人，2019）将最先进水平提高到74.7（先前最佳为53.4）。

我们注意到我们模型的一些弱点。ST-MoE-32B在小的SQuAD数据集上表现平平，精确匹配得分为90.8，不及T5-XXL设定的91.3的旧基准。此外，虽然在总体上创造了SuperGLUE的新最先进水平，但某些任务，包括CB、WSC等小任务，未能取得改进。最后，在闭卷Trivia QA上，我们的模型在SSM的微调基线上取得了改进，但未能产生对GPT-3和GLAM的双重收益。虽然这不是本文的重点，但我们展示了最近在零样本学习和在这些任务上进行微调的推理技术方面的质量差异（GPT-3（Brown等人，2020），GLAM（Du等人，2021）和Gopher（Rae等人，2021））。正如预期和以前观察到的那样，微调胜过零/一样本学习，但它的缺点是需要额外的训练和每个任务不同的模型。

7. 追踪模型的token

到目前为止，我们已经呈现了定量的度量和性能指标。我们改变方法，通过可视化token在专家之间如何路由来探索定性特征。我们通过将一批token传递给模型并手动检查每个层次上的token分配来实现这一点。我们考虑我们的ST-MoE-L模型在单语C4语料库（Raffel et al.，2019）或多语种mC4语料库（Xue et al.，2020）上进行的预训练。在编码器和解码器上，该模型具有六个稀疏层，每个层都有32个专家。

前提条件 跨度破坏目标是恢复输入中被屏蔽的可变长度连续段的跨度。格式如下：

输入：I went to <extra id 0> to buy <extra id 1> 目标：<extra id 0> the store <extra id 1> milk

在我们的编码器-解码器架构中，输入将传递给编码器，而目标将传递给解码器。

每组token与专家进行联合路由，通过辅助损失进行专家之间的负载平衡，这是Shazeer et al.（2017）提出的（有关详细信息，请参见附录A）。token将与其组中的其他token竞争专家分配，而不是整个批次，专家的专业化受到每个组中token分布的重大影响。引入组的概念是为了限制将正确的token分派给正确的专家的成本。

7.1 编码器专家表现出专业化

我们的第一个观察是，在每一层，至少有一个专家专门处理token（代表要填充的空白的token）。此外，一些编码器专家表现出明显的专业化，其中一些专家主要处理标点、动词、专有名称、计数等。表13展示了编码器专家之间的一些值得注意的专业化示例。虽然我们发现许多专业化的实例，但这些实例是从许多示例中特别提取出来的，而没有清晰的语义或句法专业化。

**表13：编码器专家专业化的显著示例。我们发现有些专家专门处理标点、连词和冠词、动词、视觉描述、专有名称、计数和数字。在所有层中（未显示），我们观察到主要处理token的专家（token为）。请注意，如果一个token在词汇表中不存在，SentencePiece模型（Kudo和Richardson，2018）将拆分该token，例如Kenneth可能会变成Ken、ne、th。**

7.2 解码器专家缺乏专业化

相比之下，解码器中专业化要少得多。不仅是token相对均匀地路由到解码器的专家中（参见表14），而且我们还没有观察到解码器专家中有意义的专业化（语义或语法）。

我们假设这种缺乏有意义的专业化是由跨度破坏目标引起的目标token的分布。特别是，（a）由于编码器中的较长序列长度（例如，在我们的设置中，编码器中的组大小为2048，而解码器中的组大小为456），在解码器中，较少数量的token是联合路由的；和（b）在解码器中，token的更高比例是sentinel token。因此，每个组中的目标token通常涵盖较小的语义空间（与编码器相比），这可能解释了解码器中缺乏专家专业化的现象。架构和训练目标之间的这种错综复杂的相互作用促使进一步研究如何更好地利用解码器中的稀疏性和专家专业化。或者，未来的工作可以研究简单地删除解码器层中的专家，这在自回归解码期间也会带来好处（Kudugunta et al.，2021a）。

表14：在编码器和解码器层之间路由的token的熵。我们通过计算路由的token的熵来支持我们的定性观察，即编码器专家专业化，而解码器专家没有专业化。编码器路由的熵较低，但解码器路由的熵很高，几乎等于均匀路由。由于每个层有32个专家，完全均匀的分布的熵为3.5。

7.3 多语言专家的专业化，但不是按语言划分

接下来，我们考虑一个在不同语言混合中预训练的多语言稀疏模型，并检查编码器中的专业化情况。与单语情况一样，我们发现有专家专门处理token、数字、连词和冠词以及专有名称等方面的强烈证据。表15呈现了专家专门处理token、数字、连词和冠词以及专有名称等方面的一些示例。

表15：多语言专家（编码器）专业化的示例。多语言专家也表现出专业化，有时跨越不同语言（例如，“for”和“pour”）。在多语言混合中训练的专家不表现出语言专业化。

人们可能期望专家在语言方面有所专门化，这似乎是在专家之间分配数据批次的自然标准。然而，我们没有发现语言专业化的证据（见表15）。相反，路由器会毫不区分地传递来自英语、日语、法语和中文的token，而专家似乎是多语言的。但考虑到token路由和负载平衡的机制，这种缺乏语言专业化并不令人惊讶。由于每组token可能只包含一个或最多几种语言（在我们的设置中，一组通常包含2-4个序列），因此鼓励所有专家处理来自所有语言的token。我们尝试了全局负载平衡损失，但通常会导致负载平衡和模型性能较差，因此我们将进一步改进多语言专家模型留作开放性工作的领域（第9节）。

我们的可视化揭示了我们模型中学到的明显的专业化（表13、表15），这是编码器层的情况。在Shazeer等人的附录中还观察到了其他专业化。然而，这引发了一个有趣的问题，即 Roller等人（2021）；Zuo等人（2021）所消除的学习路由的架构似乎表现良好。对学习与随机路由的扩展研究可能会在未来的工作中提供帮助，并有助于我们更好地理解路由行为。

8.相关工作

混合专家（MoE）可以追溯到至少三十年前的Jacobs等人的工作（1991年）；Jordan和Jacobs（1994年）。在最初的概念中，MoE定义了整个类似集成方法的神经网络。但后来Eigen等人（2013年）将将MoE作为更深层次网络的组成部分进行了扩展。Shazeer等人（2017年）然后将这个想法扩展到了一个含有1370亿参数的模型中，以在机器翻译方面取得了最先进的成果。后来的大多数工作（包括我们自己的工作）都遵循了这种将MoE作为组成部分的方法。

自然语言处理中的规模。自然语言处理中规模的显著成功（Kaplan等人，2020年；Brown等人，2020年）通过最近的研究激发了MoE的研究，如Lepikhin等人（2020年）、Fedus等人（2021年）、Yang等人（2021年）、Kim等人（2021年）、Du等人（2021年）、Artetxe等人（2021年）、Zuo等人（2021年）、Clark等人（2022年）。稀疏专家模型被提出作为以更高效的方式实现大规模密集模型结果的方法。Fedus等人（2021年）展示了与T5-XXL（Raffel等人，2019年）相比4倍的预训练加速，并且Du等人（2021年）在仅使用1/3的能量的情况下匹配了GPT-3（Brown等人，2020年）的质量。在过去的十二个月中，多个团队已经成功地高效训练了万亿参数的深度神经网络（Fedus等人，2021年；Yang等人，2021年；Du等人，2021年），最近，Lin等人（2021年）引入了一种训练10T参数模型的技术。一个副注是，稀疏专家模型最近的显著成功往往出现在具有大量数据且没有分布偏移的环境中 - 两个例子是语言建模/跨度破坏和机器翻译（Shazeer等人，2017年；Lepikhin等人，2020年；Kim等人，2021年；Fedus等人，2021年）。相比之下，Fedus等人（2021年）；Narang等人（2021年）；Artetxe等人（2021年）观察到了稀疏模型在强大的预训练质量和弱的微调质量之间的差异，但我们期望对正则化技术的进步将继续提高下游质量。

朝着更好的路由算法。BASE层（Lewis等人，2021年）将token路由重新构建为线性分配问题 - 消除了对负载平衡辅助损失的需要。该工作还展示了单专家层的功效。Clark等人（2022年）深入研究了几种不同路由算法的规模特性，并提出了他们自己的基于最优传输公式的BASE层变体。Yang等人（2021年）引入了M6-T架构和专家原型化，将专家分成不同的组，并应用k top-1路由过程（与其他地方普遍使用的top-k路由形成对比）。Hazimeh等人（2021年）提出了一种连续可微的稀疏门，展示了相对于普通的top-k门的改进。其他工作（Bengio等人，2016年）考虑将路由选择视为强化学习问题。更激进的版本则完全消除了学习路由。Hash层（Roller等人，2021年）表明随机固定路由（每个哈希函数一个）导致了与学习路由竞争的竞争性性能。Zuo等人（2021年）还提出了一种在训练和推断过程中随机选择专家的算法，并发现在Switch Transformers上取得了2 BLEU分的增益，并与Kim等人（2021年）的更大模型的竞争性分数相媲美。最后，Fan等人（2021年）设计了一种具有显式语言特定子层的体系结构（而不是像Lepikhin等人（2020年）所做的那样允许任意路由），以获得+1 BLEU的增益。

其他模态中的稀疏专家模型。MoE和稀疏专家模型在除语言以外的其他领域取得了进展。Riquelme等人（2021年）设计了一个具有150亿参数的V-MoE，以用更少的计算资源匹配ImageNet（Deng等人，2009年）模型的最新技术。Lou等人（2021年）同样通过在图像块和通道维度上使用MoE层来显示相对于密集视觉模型的好处。此外，语音识别也通过SpeechMoE变体（You等人，2021a;b）得到了改进。Kumatani等人（2021年）在Sequence-to-Sequence Transformer和Transformer Transducer中使用MoE模型减少了单词错误率。

改善稀疏模型的部署。最初的专家设计（包括本文）将每个token单独路由到该层的专家。其中一个问题是，这类架构可能很难提供，因为它需要足够的内存来存储参数。Fedus等人（2021年）已经证明蒸馏是相当有效的，但最近的方法修改了路由，将完整的句子或任务（Kudugunta等人，2021b；Zuo等人，2021年）路由到专家，这样就可以在服务时提取子网络（例如，仅部署与新任务相关的网络）。作为蒸馏的一种替代方案，Kim等人（2021年）考虑直接剪切对任务不重要的专家。

具有MoE的多任务学习。我们通过多任务环境中的成功来结束我们对最近MoE研究的介绍。Ma等人（2018年）建议为每个任务使用单独的门控或路由器网络，这个想法可能很快会被用于Transformer架构。最后，Gururangan等人（2021年）建议在基于领域/任务标签或通过推断标签条件激活专家时更大程度地模块化语言模型。

9. 讨论

尽管这项工作主要关注稀疏模型，但这些模型与机器学习中许多其他有趣的主题相交，如自适应计算、低精度训练、扩展原则和神经网络架构的进步。因此，我们的讨论涵盖了在这项研究期间出现的更广泛的一系列主题。

在多语言数据上进行预训练时的不可预测动态。我们经常观察到，在相同的模型在多语言数据上进行预训练时，将产生较小的预训练加速并且更不稳定。一个假设是这是由于跨批次的组内序列方差。作为提醒，我们鼓励组内的token保持负载平衡。通常每个组只有2-8个序列（较高会变得昂贵），其中每个序列都是用单一语言编写的。因此，即使在训练时使用超过100种语言，最多也只能在专家之间平衡2-8种语言。这导致组和批次之间的方差很大，导致路由混乱和不可预测。在一个后续实验中（仅出于简洁而突出显示），我们在英语C4的混合物加上微调任务的一小部分进行了预训练，同样导致了一个不稳定的模型。

稀疏模型的鲁棒性。尽管一篇论文专注于稀疏模型的细节，但总体上我们发现它们对一系列超参数和架构变化具有稳健性。稀疏模型在各种路由算法、放弃高比例token以及不同超参数下均取得了出色的性能。虽然我们确实指出了在微调中调整批处理大小和学习速率的重要性，与Kaplan等人（2020年）的观点一致，我们的直觉是，真正的赢家是规模。例如，表8显示通过简单增加容量因子（即FLOPs）而不是通过更复杂的路由（即算法）可以获得更大的收益。

自适应计算。稀疏模型是自适应计算模型的一个子类，因为每个输入都被应用不同的计算。在稀疏模型中，一个token被路由到其选择的专家。当容量因子小于1时，模型学会不对某些token应用计算。这在计算机视觉（Riquelme等人，2021年）和我们的语言实验中（附录D）显示出了希望。我们设想未来的模型将通过异构专家（例如，每个专家应用不同的计算）来扩展这一点。直观地说，不同的输入示例可能需要根据难度的不同而进行不同量的处理。未来的这方向上的模型将通过新兴的计算基础设施（Dean，2021年）得到高效的支持。

从小规模推广到大规模的发现。我们在整个工作中面临的一个关键问题是识别能够反映更大规模实验的小规模模型和训练设置。这在我们在第3节的稳定性研究中是显而易见的，其中必须使用XL大小的模型来展示相关动态。对于我们的架构和路由算法实验，我们经常发现在训练时间更长或模型更大时，改进会消失，甚至逆转。例如，Fedus等人（2021年）的top-n发现在我们这里呈现的8倍更大规模的实验中被颠倒，这揭示了top-(n + 1)路由相对于top-n路由的小增益（请参见表8）。

使用更低精度训练模型。我们发现稳定我们的模型而不损害（有时甚至改进）质量的最佳方法是路由器z-loss。这是一种辅助损失，鼓励模型的logits具有更小的绝对值。鉴于float32和bfloat16可以支持的数字的最大范围（约为$3e^{38}$），这使我们相信这个范围的大部分都是不需要的，并且压缩它实际上可能会改善模型训练动态。因此，未来的精度格式可能会考虑通过更紧凑的指数范围来训练某些类别的模型。

设计具有更多乘法交互的新操作。第3.1节显示，比起加法更具有乘法交互的操作，或者那些不在许多数字上累积的操作，可以提高模型性能。通过在专家层中注入更多的乘法交互，我们进一步测试了这一点，其在不改变步骤时间的情况下提高了4%的预训练速度（附录C）。我们认为这暗示了模型的有希望的架构改进，并且可能是一个良好的设计原则。最近，只累积3-5个元素的深度卷积也被证明极大地提高了Transformer的性能（So等人，2021年）。由于逐元素乘法通常在使用模型并行性时不引入任何通信开销（这使得深度卷积和我们的乘法交互等操作非常高效），因此这些操作尤其令人兴奋。尽管我们在第3.1节中注意到这些方法会增加模型的不稳定性，但在我们的模型中使用路由器z-loss防止了进一步的不稳定性。

限制激活以缓解其他不良的模型扩展动态。我们观察到训练不稳定性的另外两个来源。 （1）相对于固定量的FLOPs，编码器-解码器模型比仅解码器模型更不稳定。编码器-解码器模型具有更高的注意力层比（例如，由于在解码器上的每个FFN都有自注意力和enc-dec注意力层，所以它有更多的指数函数）。 （2）相对于固定量的FLOPs，较深的模型比较浅的模型更不稳定。较深的模型还通过额外的注意力层引入更多的指数函数。我们假设这两个观察的一个 contributing 因素只是网络中发现的指数函数数量的增加。未来的工作可以通过在非稀疏模型的注意力 softmax 上添加z-loss罚款来解决这些训练动态，特别是因为我们观察到添加这些罚款并不改变模型的质量。

稠密和稀疏模型对超参数的依赖不同。我们在第4.3节的微调分析中展示了稠密模型和稀疏模型之间最佳微调超参数之间存在显著差异。在某些情况下，对于稠密模型起作用的微调超参数掩盖了对稀疏模型的任何改进（尽管有很大的预训练加速）。对于新的模型类别，我们建议研究人员和从业者在过早放弃一种方法之前对关键超参数进行广泛测试。

• ← Previous Post
• Next Post →