ST-MOE: DESIGNING STABLE AND TRANSFERABLE SPARSE EXPERT MODELS

文章标题：ST-MoE：设计稳定且可迁移的稀疏专家模型
作者/机构：Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, William Fedus (Google Brain, Google Research)

A1 主要贡献

本文旨在提高稀疏专家模型（Sparse Expert Models）的实用性和可靠性，解决其在训练稳定性和微调（fine-tuning）质量方面存在的问题。作者们通过对这些问题进行研究，预训练了一个269B参数的稀疏模型，该模型在包括SuperGLUE在内的多个竞争性NLP基准测试中，经过微调后取得了最先进的结果。本文还为稀疏专家模型提供了额外的分析和设计指南（或启发式方法），并强调了联合优化上游预训练和下游微调指标以避免性能差异的重要性。

核心贡献如下：
1. 稳定性技术权衡研究：对多种稳定性技术进行了大规模研究，分析了它们在模型质量和训练稳定性之间的权衡。
2. 引入Router Z-loss：提出了一种新的辅助损失——Router Z-loss，它解决了训练不稳定问题，同时还能略微提升模型质量。
3. 微调分析：对稀疏模型和密集模型的微调过程进行了分析，强调了它们对批量大小和学习率等超参数的不同敏感性。研究表明，不当的超参数设置会导致稀疏模型相较于密集模型几乎没有微调增益，尽管其在预训练阶段有显著的加速。
4. 模型设计原则：提出了在分布式环境下设计帕累托效率（Pareto efficient）的稀疏模型的架构、路由和模型设计原则。
5. 定性分析：对令牌（token）在专家层之间的路由决策进行了定性分析和可视化。
6. ST-MoE-32B模型：提出了一个269B参数的稀疏模型（Stable Transferable Mixture-of-Experts，ST-MoE-32B），该模型在推理（SuperGLUE, ARC）、摘要（XSum, CNN-DM）、闭卷问答（WebQA, Natural Questions）和对抗性构建任务（Winogrande, ANLI R3）等多种自然语言基准测试中取得了最先进的性能。

A3 背景知识

稀疏专家模型基础
稀疏专家模型通常用一组具有独立权重的专家（experts）替换神经网络中的某个层。这种方法由【29, Adaptive mixtures of local experts, 1991】和【30, Hierarchical mixtures of experts and the em algorithm, 1994】等人提出。通常，一个层内的所有专家类型和形状都相同（同构），但也可以是变化的（异构）。为了节省计算量，输入只由一部分专家处理，因此需要一个机制来决定将每个输入发送到何处。通常由一个路由器（router）或门控网络（gating network）来决定输入（如单词、句子、图像块等）的去向，但也有其他方案被提出，如【42, BASE layers: Simplifying training of large, sparse models, 2021】, 【70, Hash layers for large sparse models, 2021】, 【94, Taming sparsely activated transformer with stochastic experts, 2021】, 【7, Unified scaling laws for routed language models, 2022】。

自然语言处理中的MoE层
在自然语言处理领域，【80, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017】等人提出了一种专家混合（Mixture-of-Experts, MoE）层。该层接收一个令牌表示$x$作为输入，并将其路由到从$N$个专家集合$\{E_i(x)\}_{i=1}^{N}$中选出的最佳匹配的top-k个专家。路由器变量$W_r$产生logits $h(x) = W_r \cdot x$，这些logits通过一个在$N$个可用专家上的softmax分布进行归一化。专家$i$的门控值（gate-value）由以下公式给出：


令牌$x$被路由到具有最高top-k门控值的专家（索引集合为$T$）。该层的输出是每个专家的计算结果与其门控值的加权和：

MoE层在Transformer中的演进
专家层最初在LSTM中被提出（【26, Long short-term memory, 1997】），后来被【82, Mesh-tensorflow: Deep learning for supercomputers, 2018】和【40, Gshard: Scaling giant models with conditional computation and automatic sharding, 2020】用于Transformer模型（【86, Attention is all you need, 2017】）。后续工作【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】进一步简化了MoE，将令牌路由到单个专家（top-1），并减少了其他开销以提高训练效率。

专家容量与令牌丢弃
为了提高硬件利用率，大多数稀疏模型的实现都为每个专家设置了静态的批量大小。专家容量（expert capacity）指的是可以路由到每个专家的令牌数量。如果这个容量被超出（即路由器向某个专家发送了过多的输入），溢出的令牌将不进行任何计算，而是通过残差连接（residual connection）直接传递到下一层。相关工作包括【80, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017】, 【82, Mesh-tensorflow: Deep learning for supercomputers, 2018】, 【40, Gshard: Scaling giant models with conditional computation and automatic sharding, 2020】, 【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】。

容量因子
一批输入令牌$B$在数据并行维度上被分成$G$个独立的组，每个组的大小为$B/G$。专家容量等于$CF \cdot \text{tokens}/\text{experts}$，其中$CF$是容量因子（capacity factor）超参数，$\text{experts}$是专家数量，$\text{tokens}$是组的大小。如果增加容量因子，会为负载不均衡的情况创建额外的缓冲空间，从而减少被丢弃的令牌数量。然而，增加容量因子也会增加内存和计算成本，因此存在一个权衡。

辅助负载均衡损失
最后，一个辅助的负载均衡损失（auxiliary load balancing loss）被用来鼓励令牌在专家之间大致均匀分布（【80, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017】）。这通过确保所有加速器并行处理大量数据块来提高硬件效率。该损失的细节在附录A中呈现。也存在其他替代方案：【42, BASE layers: Simplifying training of large, sparse models, 2021】和【7, Unified scaling laws for routed language models, 2022】将平衡的令牌分配视为一个分配问题，并完全移除了辅助损失。


表1：本文使用的术语。

A2 方法细节

3 稳定稀疏模型的训练

稀疏模型的训练不稳定性
稀疏模型经常遭受比标准密集激活Transformer更严重的训练不稳定性。图1展示了训练损失的发散情况。虽然很容易找到提高稳定性的方法（例如使用极小的学习率或严格的梯度裁剪），但这些方法通常会以牺牲模型质量为代价。本文将提高稳定性的方法分为三类并进行研究：（1）移除乘法交互作用；（2）注入模型噪声；（3）约束激活值和梯度。最终，本文推荐一种新的辅助损失——Router Z-loss，它能在不降低质量的情况下显著提高训练稳定性。这是对Mesh Tensorflow代码库（【82, Mesh-tensorflow: Deep learning for supercomputers, 2018】）中用于最终softmax logits的z-loss的改编。


图1：稀疏模型的训练不稳定性。我们将训练不稳定性定义为训练损失的发散。上图是两个与T5-XL版本（【68, Exploring the limits of transfer learning with a unified text-to-text transformer, 2019】）FLOP匹配的稀疏模型的运行情况，每个模型都使用1M令牌的批量大小和Adafactor优化器（【79, Adafactor: Adaptive learning rates with sublinear memory cost, 2018】）进行训练。（左）一个不稳定的训练运行。（右）一个稳定的训练运行。

设计大规模稳定性研究
本文设计了一项大规模的稳定性研究，使用与T5-XL（【68, Exploring the limits of transfer learning with a unified text-to-text transformer, 2019】）FLOP匹配的稀疏模型，在多语言语料库mC4（【91, mt5: A massively multilingual pre-trained text-to-text transformer, 2020】）上进行预训练。每个稀疏模型有32个专家，每四个FFN层中插入一个稀疏MoE层。训练容量因子为1.25，评估容量因子为2.0。对于每种稳定性技术，记录其稳定运行的比例、平均质量（英语负对数困惑度）和跨种子的标准差。构建此研究的主要问题是，小模型很少不稳定，而大的不稳定模型运行成本太高。研究发现，与T5-XL FLOP匹配的稀疏模型是一个很好的研究对象，因为它大约有1/3的运行是不稳定的，但训练成本相对较低。此外，在多语言数据上进行不稳定性实验，因为这会加剧模型的不稳定性，从而允许在稍小的模型上进行实验。基线配置使用六个随机种子进行训练，而每个采用稳定性技术的配置使用三个随机种子。每个模型在mC4上使用掩码语言建模目标（【19, Maskgan: Better text generation via filling in the, 2018】；【14, Bert: Pre-training of deep bidirectional transformers for language understanding, 2018】）预训练20k步。

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

乘法交互作用的定义
一些架构改进涉及的乘法比加法更多，或者不同时对许多项求和。例如，矩阵乘法中每个加法对应一个乘法，因此我们不称其为“乘法”操作。本文在此呈现并分析Transformer中两种乘法交互作用实例的影响。

GELU门控线性单元（GEGLU）
第一个例子是门控线性单元（Gated Linear Unit, 【8, Language modeling with gated convolutional networks, 2017】），它是两个线性投影的逐元素乘积，其中一个首先通过sigmoid函数。【78, Glu variants improve transformer, 2020】将其扩展到其他变体，并提出了一个GELU-Linear（【24, Gaussian error linear units (gelus), 2016】）FFN层作为Transformer中常规ReLU（【55, Rectified linear units improve restricted boltzmann machines, 2010】）FFN的替代品。


这一质量增益在后来的工作中得到了证实（【56, Do transformer modifications transfer across implementations and applications?, 2021】）。

均方根缩放参数
第二个例子是均方根（RMS）归一化（【92, Root mean square layer normalization, 2019】）中的缩放参数。在Transformer内部，层与层之间不是直接调用，而是有一种内部结构（称为子层调用），以改善梯度传播和训练动态。本文的子层调用与【68, Exploring the limits of transfer learning with a unified text-to-text transformer, 2019】匹配，包括：（1）RMS归一化，（2）层调用（例如自注意力），（3）dropout（【84, Dropout: a simple way to prevent neural networks from overfitting, 2014】），（4）添加残差（【23, Deep residual learning for image recognition, 2015】）。RMS归一化根据均方根对输入向量$x \in R^d$进行逐元素缩放。然后，它通过与一个学习到的缩放参数$g$进行逐元素相乘来重新缩放输出。

移除乘法交互的实验结果
表2显示，移除GEGLU层或RMS缩放参数都能提高稳定性，但会显著降低模型质量。我们注意到，这些缩放参数（$g$）对模型质量的增益远大于其他地方的参数（例如FFN）。与我们的发现一致，【83, Normformer: Improved transformer pretraining with extra normalization, 2021】发现，在Transformer的残差连接中添加一个学习到的乘法标量会使它们变得更加不稳定。附录C进一步研究了在专家层中添加新的乘法交互对质量的影响，发现该操作在几乎不减慢模型步长时间的情况下带来了质量提升。


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

3.2 添加噪声时的稳定性与质量权衡

噪声注入对稳定性的影响
本文接下来探讨一个假设，即向模型中添加噪声可以提高训练稳定性（【58, Adding gradient noise improves learning for very deep networks, 2015】）。受到【85, Antifragile: Things that gain from disorder, 2012】提出的反脆弱性概念（系统通过噪声得到改善）以及微调过程（通过dropout注入噪声）很少不稳定的观察启发，我们研究了训练噪声是否能提高稀疏模型的稳定性。表3显示，与基线相比，稳定性有所提高，但代价是质量下降。我们还发现，由【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】引入的输入抖动（input-jitter）在XL规模下会降低质量，因此在我们的模型中将其剔除。输入抖动将输入到路由器的logits乘以一个在$[1 - 10^{-2}, 1 + 10^{-2}]$之间的均匀随机变量。在我们消融实验中的dropout应用于整个Transformer。正如之前所见，小规模设置下的改进在扩大规模时可能无法泛化，因此应始终监控和重新评估趋势（【32, Scaling laws for neural language models, 2020】）。


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

3.3 约束激活值和梯度时的稳定性与质量权衡

激活值和梯度的约束方法
稳定神经网络最成功的方法之一是约束激活值和梯度（【62, On the difficulty of training recurrent neural networks, 2013】；【28, Batch normalization: Accelerating deep network training by reducing internal covariate shift, 2015】；【73, Weight normalization: A simple reparameterization to accelerate training of deep neural networks, 2016】；【3, Layer normalization, 2016】）。一种流行的方法是裁剪梯度范数，以解决在深层网络中反向传播时的梯度爆炸问题（【62, On the difficulty of training recurrent neural networks, 2013】）。

Adafactor优化器中的更新裁剪
在本文中，我们因其内存效率而使用Adafactor优化器（尽管最近引入的8位优化器【13, 8-bit optimizers via block-wise quantization, 2021】可能提供更好的权衡）。Adafactor不使用梯度裁剪，而是使用更新裁剪（update clipping），即权重的变化被约束在某个范数以下。我们实验了将更新裁剪收紧到一个更小的值。

Router Z-loss
接下来，我们研究了对输入到路由器的logits的约束。路由器以float32精度计算专家上的概率分布（即选择性精度）（【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】）。然而，在最大规模下，我们发现这不足以保证可靠的训练。为了解决这个问题，我们引入了Router Z-loss：


其中$B$是令牌数，$N$是专家数，$x \in R^{B \times N}$是输入到路由器的logits。这会惩罚进入门控网络的大logits。

约束方法的实验结果
表4显示，更新裁剪和Router Z-loss都在所有3次运行中稳定了模型，但更新裁剪严重损害了模型质量。因此，我们选择使用Z-loss方法来解决模型稳定性问题，因为它在提高质量和稳定性方面表现更优。


表4：约束权重更新和路由器logits。约束Adafactor中的更新裁剪提高了稳定性，但代价是质量的灾难性损失。更宽松的裁剪值不能可靠地稳定训练，因此我们在此排除它们。Router Z-loss稳定了模型，没有任何质量下降（在这种情况下，我们观察到质量略有提升）。

总损失函数
Router Z-loss引入了另一个超参数$c_z$，它是这个损失在总优化损失中的权重系数。总损失是交叉熵损失（$L_{CE}$）、辅助负载均衡损失（$L_B$）和Router Z-loss（$L_Z$）的线性加权组合，得到总损失：


我们根据超参数扫描后预训练得到的最佳模型质量，选择了$c_z = 0.001$。附录B记录了预训练过程中这些损失的变化情况。

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

混合精度训练
与大多数现代分布式Transformer一样，我们使用混合精度进行训练（【53, Mixed precision training, 2017】）。权重以float32格式存储用于梯度更新，然后在前向和后向传播中进行矩阵乘法时转换为bfloat16。此外，所有激活值都以bfloat16格式存储和操作，allreduce通信可以在bfloat16或float32数值精度下进行。对于本文中探索的最大模型（ST-MoE-32B），我们发现将allreduce的数值精度减半可以带来加速，但这也会使训练不稳定，因此本文始终将其保持为float32。

低精度格式的优缺点
较低的精度格式通过减少(a)处理器和内存之间的通信成本、(b)计算成本、(c)存储张量（如激活值）的内存，来使模型更高效。然而，较低的精度格式会带来更大的舍入误差，可能导致不可恢复的训练不稳定性。

精度格式与舍入误差
图2回顾了不同精度格式的属性及其在不同数值范围内的相应舍入误差。任何两个连续的2的幂次方范围内的数字（例如$[2,4)$和$[1024, 2048)$）都由固定数量的尾数位表示（bfloat16为7位，float32为23位）。因此，(1) bfloat16的舍入误差大约是float32的65536倍（即$2^{23-7} = 2^{16} = 65536$）；(2) 更大的数字有更大的舍入误差。由于8个指数位，数字可以变得非常大，达到约$3 \times 10^{38}$，这导致即使是float32在舍入误差方面也存在一些问题。


图2：数值精度格式和舍入误差。更大的数字有更大的舍入误差。bfloat16的舍入误差比float32差高达65536倍。Router Z-loss鼓励数字的绝对值变小，这不会妨碍模型性能，并能减少舍入误差。Router Z-loss在输入到那些较大误差会急剧改变相对输出的函数（如指数和正弦函数）时最为有效。

稀疏模型对舍入误差的敏感性
稀疏专家模型对舍入误差很敏感，因为它们由于路由器的存在而有更多的指数函数。这些额外的指数函数会加剧舍入误差，导致训练不稳定。虽然舍入误差不会改变softmax操作内概率的排序，但它会影响MoE中第二个令牌的路由，因为存在相对阈值（例如，一个令牌只有在第二个专家的门控概率是第一个专家的1/5大小时才会被路由到第二个专家）。此外，舍入误差会极大地改变用于缩放专家输出的概率——我们发现这很重要。最后，我们推测解码器-only模型（decoder-only models）具有更高的稳定性（此处未展示）是因为它们的指数函数较少。

关于Router Z-loss的补充说明
有人可能认为Router Z-loss是一种可以通过裁剪logits（【90, Google’s neural machine translation system: Bridging the gap between human and machine translation, 2016】）来替代的复杂方法。我们解释为什么不是这样。目标是最小化进入指数函数的大的舍入误差。裁剪logits发生在任何舍入误差之后——导致更大的不连续性。从一个角度看，裁剪本身就是一种舍入误差；相反，z-loss自然地鼓励模型产生值较小的logits，从而可以更精确地建模。由于这些动态，我们确保所有进行指数运算的张量都转换为float32。这暗示了可能存在更好的神经网络数值格式，因为当在网络中添加z-loss时，指数位并未被充分利用。

4 稀疏模型的微调性能

微调的背景
性能最好的语言模型通常通过两个步骤获得：（1）在大量数据（如互联网）上进行预训练，然后（2）在感兴趣的任务（如SuperGLUE）上进行微调。尽管出现了一些有前途的替代技术，如少样本推理（few-shot inference）（【6, Language models are few-shot learners, 2020】）、前缀微调（prefix tuning）（【43, Prefix-tuning: Optimizing continuous prompts for generation, 2021】）、提示微调（prompt tuning）（【41, The power of scale for parameter-efficient prompt tuning, 2021】）和适配器模块（adapter modules）（【27, Parameter-efficient transfer learning for nlp, 2019】），但与微调相比，质量上仍存在差距。因此，本文重点关注微调，但也强调了稀疏模型在少样本设置中的最新成功（【16, Glam: Efficient scaling of language models with mixture-of-experts, 2021】；【2, Efficient large scale language modeling with mixtures of experts, 2021】）。

4.1 假设：一个泛化问题

稀疏模型在微调中的表现
稀疏模型在大型数据集上表现出色，但在微调时有时表现不佳（【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】；【2, Efficient large scale language modeling with mixtures of experts, 2021】）。我们提出了一个假设：稀疏模型容易过拟合。我们通过SuperGLUE（【87, Superglue: A stickier benchmark for general-purpose language understanding systems, 2019】）中的两个任务来说明这个问题——Commitment Bank (CB)（【10, The commitmentbank: Investigating projection in naturally occurring discourse, 2019】）和ReCoRD（【93, Record: Bridging the gap between human and machine commonsense reading comprehension, 2018】）。CB有250个训练样本，而ReCoRD有超过10万个。这种显著的大小差异为在同一基准测试中选择的两个任务上研究过拟合提供了一个自然的研究环境。

过拟合现象的实验验证
在图3中，我们比较了Dense L和ST-MoE-L模型的微调特性。每个模型都在C4语料库（【68, Exploring the limits of transfer learning with a unified text-to-text transformer, 2019】）的500B令牌上进行了预训练。这些模型被设计为与T5-Large编码器-解码器模型（【68, Exploring the limits of transfer learning with a unified text-to-text transformer, 2019】）大致FLOP匹配的变体，拥有770M参数。ST-MoE模型有32个专家，专家层频率为1/4。预训练和微调的训练容量因子为1.25，评估容量因子为2.0。我们评估了在留出的验证集和训练集分区上的性能。


图3：稀疏模型容易过拟合。我们绘制了ST-MoE-L和dense-L模型在CB任务（250个训练序列）和ReCoRD（138k个训练序列）上微调的训练和验证曲线。在这两种情况下，稀疏模型在训练集上学习得更快（蓝线超过绿线）。然而，对于较小的CB任务，密集模型在留出的验证集上优于稀疏模型（红线 vs. 橙线）。相比之下，在较大的ReCoRD任务上，稀疏模型比密集模型高出几个百分点。

过拟合分析
在这两个任务中，稀疏模型都更快地收敛到100%的训练集准确率，这支持了稀疏模型在数据分布偏移下能有效优化的观点。在较大的任务ReCoRD上，稀疏模型的验证质量随着训练的提升而提升，并显著超过密集模型。然而，在较小的任务CB上，稀疏模型在留出数据上的表现落后于其密集模型对应物。根据【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】的建议，我们考虑增加专家隐藏状态内的dropout（即专家dropout），但发现在这个规模下，更高的值只能适度提高质量（图4）。


图4：稀疏模型微调的正则化研究。对于每种设置，我们在SuperGLUE上训练三个随机种子直到收敛。我们发现通过dropout增加正则化提供了适度的提升。（左）展示了在全局dropout率为0.1时SuperGLUE微调质量达到峰值。更高的值会过度正则化并严重损害质量。（右）从已知的最佳全局dropout率0.1开始，我们选择性地增加专家dropout（专家隐藏激活上的一个独立dropout率）。这带来了进一步的泛化好处，并与【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】的发现一致。

4.2 微调模型参数的子集以改善泛化

微调参数子集的实验
为了对抗过拟合，我们实验了在微调期间只更新模型参数的一个子集。图5测量了更新5个不同参数子集的质量：所有参数（All）、仅非MoE参数（Non MoE）、仅MoE参数（MoE）、仅自注意力和编码器-解码器注意力参数（Attention）以及仅非MoD FFN参数（FFN）。


图5：在微调期间仅更新模型参数的一个子集。为了改善稀疏模型的泛化并对抗过拟合，我们微调了模型参数的一个子集。所有结果都使用ST-MoE-L模型，并且是5个不同随机种子的平均值。我们观察到，更新3/5的参数子集表现大致相同，而仅微调MoE参数会导致质量急剧下降。

参数子集微调的观察结果
我们观察到，更新非MoE参数的效果与更新所有参数差不多，而只更新FFN参数的效果稍好一些。只更新MoE参数会显著降低微调性能，而这部分参数占了模型参数的约80%。只更新非MoE参数可以成为加速和减少微调内存的有效方法。

性能差异的假设
我们假设只微调MoE参数导致性能不佳，因为专家层只在每1/4层出现一次，且一个令牌在每层最多只会看到两个专家。因此，更新MoE参数影响的层数和FLOPs远少于我们尝试的任何其他参数子集。只更新MoE参数导致的训练损失远大于更新非MoE参数，即使前者参数更多。我们进一步观察到，更新所有非MoE参数导致的训练损失高于更新所有参数，但遗憾的是，这种正则化效果并未转化为更好的验证性能。

其他正则化尝试
此外，我们尝试的一种正则化器是dropout的变体，即在训练期间随机地将整个专家掩码掉。然而，在我们的初步研究中，这未能改善泛化。附录J详细阐述了此实验和其他负面结果。

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

微调超参数敏感性研究
稀疏模型和密集模型对微调协议的敏感度如何？我们研究了两个超参数：批量大小和学习率。我们预训练了一个Dense-L和ST-MoE-L模型，在C4上训练了500B个令牌，然后在SuperGLUE上进行微调。图6总结了我们的实验，完整数据见表20（附录F）。在所有超参数设置中，稀疏模型（橙色）都优于其密集模型（蓝色）——然而，每种模型的最佳设置可能会显著改变结果。

超参数敏感性的主要发现
稀疏模型和密集模型在不同批量大小和学习率下的性能差异巨大。稀疏模型受益于较小的批量大小和较高的学习率。这与过拟合假设（4.1节）一致，这两个变化可能通过在微调过程中增加噪声来改善泛化。最后，我们指出在微调期间正确调整批量大小和学习率的重要性。简单地使用对密集模型有效的微调超参数，可能会掩盖稀疏模型在预训练中获得的任何改进。


图6：批量大小和学习率敏感性。我们测量了密集（蓝色）和稀疏（橙色）模型之间微调协议的差异和敏感性。每个条形是6次不同超参数运行的平均值。在SuperGLUE上，稀疏模型受益于更嘈杂的超参数，包括小批量大小和高学习率。密集模型的行为几乎相反。所有数据见附录F。

4.4 稀疏模型在微调期间对丢弃的令牌具有鲁棒性

令牌丢弃机制与鲁棒性实验
稀疏模型在每一层将令牌路由到一个或多个专家。为了使这些模型在现代硬件的SPMD范式下高效，专家容量（每个专家处理的令牌数）需要提前固定。当一个专家接收到的令牌超过其容量时，多余的令牌会被丢弃——不对这些令牌应用任何计算。我们通过（1）使用辅助损失进行预训练以促进令牌均匀分配和（2）使用容量因子（一个超参数）为每个专家增加额外空间来防止这种情况。我们实验了在微调期间关闭辅助损失并使用不同的容量因子。

对令牌丢弃的鲁棒性结果
表5揭示了一个令人惊讶的结果：微调质量并未因丢弃高达10-15%的令牌而受到实质性影响。对ST-MoE-32B的研究证实，高容量因子并不能提高微调质量。这与【92, M6-t: Exploring sparse expert models and beyond, 2021】的发现一致，即不均衡的负载可能不会显著影响模型质量。


表5：稀疏模型在微调时对丢弃的令牌具有鲁棒性。我们发现SuperGLUE上的微调质量在所探索的值范围内没有显著影响。有趣的是，丢弃10-15%的令牌可以表现得和丢弃<1%的模型差不多。我们还观察到负载均衡损失（Aux Loss）可以改善微调。丢弃的令牌百分比对应于在验证准确率峰值时所有专家层丢弃的令牌比例。

4.5 在微调期间插入哨兵令牌

哨兵令牌与领域不匹配问题
哨兵令牌（Sentinel tokens）在跨度损坏目标（span-corruption objective）（【19, Maskgan: Better text generation via filling in the, 2018】；【14, Bert: Pre-training of deep bidirectional transformers for language understanding, 2018】）中表示被掩码的序列。这与我们可能遇到的任何微调任务都不同，导致预训练和微调之间的领域不匹配。表6说明了这种差异。我们研究了修改微调任务使其更像预训练任务是否会影响结果。


表6：在微调期间插入哨兵令牌以模仿预训练的跨度目标。我们突出了跨度损坏和微调之间的典型差异。我们建议通过插入哨兵令牌来修改微调任务，使其类似于预训练。

插入哨兵令牌的实验结果
在表7中，我们发现在微调时添加哨兵令牌仅对语法错误纠正（GEC）（【71, A simple recipe for multilingual grammatical error correction, 2021】）有改善，而对SuperGLUE没有。我们尝试通过插入多个哨兵令牌（就像模型在预训练时会遇到的那样）来进一步减少数据分布的偏移，但同样没有发现普遍的好处。然而，尽管在留出数据上没有一致的好处，我们发现对于密集和稀疏模型，训练收敛都加速了。


表7：哨兵令牌对微调的影响。在微调期间添加哨兵令牌（与【41, The power of scale for parameter-efficient prompt tuning, 2021】中使用的概念相似）在我们考虑的两个任务上表现不一。SuperGLUE记录平均分，GEC记录精确匹配。虽然我们发现它不能改善泛化，但哨兵令牌可以加速训练收敛。

5 设计稀疏模型

稀疏模型的设计挑战
密集模型的设计已由【32, Scaling laws for neural language models, 2020】的基础工作指导。但稀疏模型带来了一系列额外的问题：（1）使用多少个专家？（2）选择哪种路由算法？（3）容量因子的值设为多少？（4）硬件如何改变这些决策？本节将评论这些问题，并为构建帕累托效率的稀疏模型提供建议。

5.1 设置专家数量

专家数量的扩展性
一个首要问题是使用多少专家。【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】展示了Switch Transformer的扩展属性，在C4上（按步数计算）预训练效益单调增加，最多可达512个专家；【34, Scalable and efficient moe training for multitask multilingual models, 2021】最多到64个专家，【7, Unified scaling laws for routed language models, 2022】最多到512个专家。但当专家数量很多（>256）或模型非常稀疏（<1%的专家被激活）时，增量效益会迅速减弱。

硬件考量：计算与内存比
考虑特定的硬件系统可以进一步指导这一选择。计算与内存的比率（操作强度）可以作为评估不同操作效率的指标（【89, Roofline: an insightful visual performance model for multicore architectures, 2009】；【77, Fast transformer decoding: One write-head is all you need, 2019】）。如果将张量加载到计算核心（如ALU/MMU）的时间远超在张量上进行计算所需的时间，则模型是内存受限的。在现代GPU和TPU上，增加这个计算与内存的比率可以提高效率。

专家数量与效率的关系
回到稀疏专家模型，每个核心使用超过一个专家会增加内存传输，可能损害效率。增加专家数量不会改变所做的计算（稀疏模型对每个输入应用固定量的计算），但会增加内存传输需求（必须从设备内存加载额外的专家变量）。这降低了计算与内存的比率。

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

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

Top-n路由研究
我们泛化了top-1路由（【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】；【70, Hash layers for large sparse models, 2021】）和top-2路由（【80, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017】；【40, Gshard: Scaling giant models with conditional computation and automatic sharding, 2020】）来研究top-n路由，其中每个令牌最多由n个专家处理。在这项研究中，所有模型都预训练了100k步，每批1M个令牌，稀疏模型有32个专家，并且与T5-Large（【68, Exploring the limits of transfer learning with a unified text-to-text transformer, 2019】）的FLOPs相匹配。我们得出两个关键结论。

容量因子的影响
首先，增加训练和评估的容量因子（CF）可以提高质量，如表8中分块比较所示。例如，当训练CF从1.0增加到1.25时，top-1路由的负对数困惑度提高了+0.011；当训练CF从1.25增加到2.0时，top-2路由提高了+0.009。将密集模型的大小增加三倍（Dense-L到Dense-XL）会带来+0.090的提升。因此，这些CF的提升大约是其1/10。但这有代价：增加容量因子会线性增加einsums成本、激活值的内存、all2all通信成本和专家层的模型并行allreduce通信成本。

路由算法的影响
其次，在给定固定容量因子的前提下，top-(n+1)相对于top-n路由的增益很小（表8）。例如，在训练CF为1.25时，top-2路由比top-1路由提升了+0.004，大约是密集模型扩大三倍增益的1/20。这修正了【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】早期的建议。这些实验设置的主要区别在于计算规模。我们在8倍大的训练规模下发现，路由到多个专家有小的增益。此外，在更大的实验规模下，top-n与top-(n+1)路由的速度差异可以忽略不计。


表8：比较容量因子（CF）和路由算法。增加训练和评估的CF可以提高性能。在评估时增加或减少评估CF，为不同计算资源提供了额外的调节手段。其次，在不同容量因子下，top-(n+1)相对于top-n路由的增益较小。因为随着CF增加，质量提高但速度变慢，帕累托效率的CF必须由具体的硬件系统决定。

硬件系统与最优选择
具体的硬件-软件系统将决定最优的n和容量因子。在我们的TPU和Mesh Tensorflow硬件-软件栈中，我们记录了ST-MoE-L和ST-MoE-32B模型在增加训练容量因子时的训练速度（表9）。随着模型规模的扩大，更高的容量因子会使模型越来越慢。ST-MoE-L不需要模型并行，使其更适合高容量因子。因此，对于我们最大的模型，我们继续使用【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】倡导的较小的训练容量因子1.25，以实现帕累托效率。


表9：在TPU上分析稀疏模型。将训练容量因子从1.25增加到2.0，大型（1B）模型的步长时间增加了+7%，而我们的32B模型则增加了+14%。随着模型尺寸的增加，我们发现表8中更高训练容量因子带来的微小质量增益被显著的14%的减速所抵消。注意：ST-MoE-L和ST-MoE-32B的步长时间不可比，因为它们使用了不同数量的核心。

其他路由技术
本节的结果主要关注top-n路由，但我们也在附录J中实验了多种其他路由技术。我们发现大多数技术与top-n路由相比，性能相似或更差。然而，我们发现由【69, Scaling vision with sparse mixture of experts, 2021】引入的批量优先路由（Batch Prioritized Routing, BPR）在容量因子小于1时显著提升性能（附录D）。我们推荐BPR用于all2all和allreduce通信更昂贵且较低容量因子为优的大型模型。

7 追踪模型中的令牌

定性分析设置
到目前为止，我们呈现了定量测量和性能指标。现在我们转向探索定性特征，通过可视化令牌如何在专家之间路由。我们通过向模型传递一批令牌并手动检查每一层的令牌分配来做到这一点。我们考虑的是在单语C4语料库（【68, Exploring the limits of transfer learning with a unified text-to-text transformer, 2019】）或多语mC4语料库（【91, mt5: A massively multilingual pre-trained text-to-text transformer, 2020】）上预训练的ST-MoE-L模型。在编码器和解码器上，该模型都有六个稀疏层，每个层有32个专家。

预备知识：跨度损坏与分组路由
跨度损坏（span corruption）目标是恢复输入中被掩码掉的可变长度连续片段。每个令牌组是联合路由的，通过辅助损失激励专家间的负载均衡。令牌是与同组内的其他令牌竞争专家分配，而不是整个批次，因此专家专业化受每组中令牌分布的严重影响。引入“组”的概念是为了限制向正确专家分派和收集正确令牌的成本。

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

编码器专家的专业化现象
我们的第一个观察是，在每一层，至少有一个专家专门处理哨兵令牌（代表待填补空白的掩码令牌）。此外，一些编码器专家表现出明显的专业化，一些专家主要处理标点符号、动词、专有名词、计数等。表13展示了编码器专家专业化的一些显著例子。虽然我们发现了许多专业化的实例，但这些是从许多没有明确语义或句法专业化的例子中特别提取出来的。


表13：编码器专家专业化的显著例子。我们发现专家专门处理标点符号、连词和冠词、动词、视觉描述、专有名词、计数和数字。在所有层中（未显示），我们观察到主要处理哨兵令牌（标记为）的专家。注意，SentencePiece模型（【35, Sentencepiece: A simple and language independent subword, 2018】）会分割词汇表中不存在的令牌，例如Kenneth可能变成Ken, ne, th。

7.2 解码器专家缺乏专业化

解码器专家缺乏专业化的现象
相比之下，解码器中的专家专业化远不明显。不仅哨兵令牌在解码器专家中路由得相当均匀（见表14），我们也没有观察到解码器专家有意义的专业化（语义或句法）。


表14：跨编码器和解码器层路由的哨兵令牌的熵。我们通过计算哨兵令牌路由的熵来支持我们的定性观察，即编码器专家专业化，而解码器专家不专业化。编码器路由熵很低，但解码器路由是高熵的，几乎等于均匀路由。因为每层有32个专家，一个完全均匀的分布的熵是3.5。

缺乏专业化的原因假设
我们假设这种缺乏有意义的专家专业化是由跨度损坏目标引起的目标令牌分布造成的。特别是，(a) 由于编码器中序列长度较长，解码器中联合路由的令牌数量较少（例如，在我们的设置中，编码器中的组大小为2048，而解码器中为456），以及 (b) 解码器中哨兵令牌的比例较高。因此，每组中的目标令牌通常覆盖较小的语义空间（与编码器相比），这可能解释了解码器中缺乏专家专业化的原因。这种架构和训练目标之间复杂的相互作用，为进一步研究如何更好地利用解码器中的稀疏性和专家专业化提供了方向。

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

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


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

缺乏语言专业化的原因
人们可能期望专家会按语言进行专业化，这似乎是在专家之间分配数据批次的自然标准。然而，我们没有发现语言专业化的证据（见表15）。路由器不加区分地传递来自英语、日语、法语和汉语的令牌，专家似乎是多语言的。但当我们考虑到令牌路由和负载均衡的机制时，这种缺乏语言专业化就不那么令人惊讶了。由于每组令牌可能只包含一到几种语言（在我们的设置中，一个组通常由2-4个序列组成），因此所有专家都被鼓励处理来自所有语言的令牌。我们实验了全局负载均衡损失，但这通常会导致更差的负载均衡和更差的模型性能。

A4 实验环境

  • 数据集：
    • 预训练：使用单语C4语料库（【68, Exploring the limits of transfer learning with a unified text-to-text transformer, 2019】）、多语言mC4语料库（【91, mt5: A massively multilingual pre-trained text-to-text transformer, 2020】）以及GLaM数据集（【16, Glam: Efficient scaling of language models with mixture-of-experts, 2021】）的混合。GLaM数据集包括英语C4、网页、书籍、维基百科、代码和新闻。ST-MoE-32B在1.5T令牌上进行预训练。
    • 微调：SuperGLUE（包含SST-2, WIC, MRPC, STS-B, QQP, MNLI, QNLI, RTE, CB, MultiRC, RECORD, BoolQ, WNLI, WSC, COPA, CoLA等任务）、CNN-DM（【25, Teaching machines to read and comprehend, 2015】）、BBC XSum（【57, Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization, 2018】）、SQuAD（【69, Squad: 100,000+ questions for machine comprehension of text, 2016】）、ARC Easy/Challenge（【8, Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018】）、闭卷问答（Web Questions（【5, Semantic parsing on freebase from question-answer pairs, 2013】）、Natural Questions（【39, Natural questions: a benchmark for question answering research, 2019】）、TriviaQA（【31, Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension, 2017】））、Winogrande（【72, Winogrande: An adversarial winograd schema challenge at scale, 2020】）和ANLI（【59, Adversarial nli: A new benchmark for natural language understanding, 2019】）。
  • 模型架构：
    • 基于T5编码器-解码器Transformer架构。
    • ST-MoE-L：约770M参数，与T5-Large在FLOPs上匹配。包含32个专家，专家层频率为1/4（每4个FFN层替换为1个MoE层）。
    • ST-MoE-32B：269B参数，与32B参数的密集模型在FLOPs上匹配。包含2048个专家。具体参数：d_model=2048, num_heads=32, d_kv=128, d_ff=65536。采用了附录C中描述的“稀疏-密集”堆叠架构。
  • 硬件配置：在TPU系统上进行训练和评估。
  • 软件配置：
    • 代码实现基于Mesh Tensorflow（【82, Mesh-tensorflow: Deep learning for supercomputers, 2018】）。
    • 优化器使用Adafactor（【79, Adafactor: Adaptive learning rates with sublinear memory cost, 2018】），学习率采用10k步的预热，然后是逆平方根衰减。
    • 初始化方案遵循【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】。

A4 实验结果

6.1 ST-MoE-L

  • 实验内容：将一个在C4数据集上预训练了524B令牌的密集T5-Large (L)模型与一个FLOPs近似匹配的稀疏模型ST-MoE-L进行比较。模型在多个任务的混合数据上进行微调。
  • 实验结果（表10）：ST-MoE-L在自然语言理解、问答和摘要等多种任务的验证集上均取得了提升。尤其在闭卷问答任务上，增益非常显著。
  • 分析结论：唯一没有提升的两个任务是CB和WSC，它们分别是训练样本最少的任务（250和259个样本），这再次支持了稀疏模型容易过拟合的假设（4.1节）。

    表10：FLOP匹配的密集和稀疏模型的微调性能比较。比较了dense-L基线和稀疏FLOP匹配版本（数字越大越好）。我们观察到，在不同任务上，使用大致相同的计算量，稀疏模型都有一致的增益。仅有的两个没有从稀疏模型中获得改善的任务是最小的两个：CB（250个训练样本）和WSC（259个）。

6.2 ST-MoE-32B

  • 实验内容：设计并训练了一个269B参数的稀疏模型ST-MoE-32B，它与一个32B参数的密集模型在FLOPs上匹配。该模型在C4和GLaM数据集的混合数据上预训练了1.5T令牌，并在多个NLP基准上进行微调。模型设计中平衡了FLOPs和参数，并使用Router Z-loss来保证训练稳定性。
  • 实验结果（表12）：ST-MoE-32B在多个任务上取得了最先进的（SOTA）性能。
    • SuperGLUE：测试集平均分达到91.2，超过了先前SOTA和人类水平的估计。
    • 摘要：在XSum和CNN-DM上取得了SOTA。
    • 推理：在ARC Easy (94.8) 和 ARC Challenge (86.5) 的测试服务器提交中刷新了SOTA。
    • 闭卷问答：在WebQA (47.4) 和NatQA (41.9) 上刷新SOTA，WebQA结果超过了ERNIE 3.0 Titan（260B密集模型）的零样本性能。
    • 对抗性任务：在ANLI R3 (74.7) 和WinoGrande XL上取得显著提升。
  • 分析结论：
    • 模型弱点：在小规模的SQuAD数据集上表现平平（90.8），未超过T5-XXL。尽管在SuperGLUE总分上创下新高，但在一些小任务（如CB, WSC）上未能改进。在闭卷TriviaQA上，虽优于基线但未超过GPT-3和GLaM。
    • 与少样本学习的比较：结果再次证实，微调的性能优于零/单样本学习，但需要为每个任务进行额外的训练并维护不同的模型。

      表11：模型比较。比较了Dense-L、T5-XXL、两种最大的Switch Transformer变体（Switch-XXL和Switch-C）以及ST-MoE-L和ST-MoE-32B。dmodel指模型隐藏状态大小，dff是FFN层内部大小。dkv是每个注意力头的维度。Expert Layer Freq.是被稀疏层替换的FFN层的比例。Sparse-Dense指附录C中描述的架构变体。

      表12：ST-MoE-32B与之前最佳的仅推理技术和微调模型的比较。“dev/test”的分割指的是Zero-Shot和One-Shot的dev集，以及Fine-Tune的test集质量。不可用的数据用“–”填充。

A5 结论

本文通过展示一个参数量为先前工作1/5，但在计算量（FLOPs）与参数之间取得更好平衡的模型，是一个更有效的稀疏学习器，从而缓和了【20, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021】中对规模的过度热情。此外，这种设计降低了部署时的内存开销，从而提高了稀疏模型的可用性。利用本文提出的稀疏模型变体，在多个最具竞争力的公开基准测试中取得了最先进的成果。作者希望这项工作能展示模型稀疏化的力量，并加速这类模型的普及。

A6 附录

A 令牌负载均衡描述

辅助负载均衡损失
本文同样使用了【80, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017】提出的辅助负载均衡损失来平衡专家间的令牌分配。假设有$N$个专家（索引从$i=1$到$N$）和一批包含$T$个令牌的$B$。辅助损失计算为向量$f$和$P$之间的缩放点积：


其中，$f_i$是分派给专家$i$的令牌比例：

$P_i$是为专家$i$分配的路由器概率的比例：

损失函数的目标
由于我们期望批次中的令牌在$N$个专家之间均匀路由，我们希望这两个向量的值都为$1/N$。公式7的辅助损失鼓励均匀路由，因为它在均匀分布下最小化。该目标函数可以被微分，因为$P$向量是可微的，而$f$向量不是。最终的损失乘以专家数$N$以在专家数变化时保持损失的恒定。最后，一个超参数$\alpha$是这些辅助损失的乘法系数；在本文中，我们使用$\alpha = 10^{-2}$。

B Router Z-loss训练动态

Router Z-loss系数扫描
图7绘制了对公式5中的Router Z-loss进行系数扫描的结果，其中最佳值$c_z=0.001$在编码器和解码器中均以绿色标出。当系数$c_z$的值为$1e-2$或更大时，z-loss会缩小到接近零。


图7：对Router Z-Loss的损失系数（cz）进行扫描。我们绘制了在没有Router Z-loss（蓝色）和随着cz值增加（我们为所有后续实验选择了与绿色曲线相关的系数）的情况下，预训练过程中的Router Z-loss。当值为1e-2或更大时，z-loss收缩到接近零。左图显示编码器层，右图显示解码器层。

C 改进的架构修改

稀疏-密集层堆叠
第一个修改是在每个MoE层之前或之后立即添加额外的FFN层（称为Sparse-Dense）。表16揭示了在每个稀疏层之前或之后立即放置一个FFN层的有效性，并且这些额外的FFN层在网络中其他地方添加时帮助较小。保证所有令牌在每个注意力层之间至少被一个FFN应用似乎是有用的。


表16：在每个稀疏层之前或之后立即放置一个密集的FFN层可以提高质量。在每个稀疏层之前或之后立即插入一个额外的密集FFN层，其质量提升效果是（随机）将密集层放置在网络其他地方的两倍。为了公平比较，所有非基线模型都添加了相同数量的FFN层。

专家层中的乘法偏置
其次，我们在专家层中引入了一个额外的偏置。我们所有的模型都使用GELU-Linear FFN（【78, Glu variants improve transformer, 2020】），而不是ReLU FFN。加法偏置（Additive bias）是在FFN层的第一个矩阵乘法后添加的一个学习权重（$B$），形状为[batch, d_ff]。乘法偏置（Multiplicative bias）（也称为缩放参数）是形状相同但进行逐元素乘法的学习权重。我们将加法偏置初始化为零，乘法偏置初始化为一。


表17显示了我们不同方法的结果。加法和乘法偏置基本上是免费的：计算成本低，增加的新参数少，并且在模型和专家并行时不会产生额外的通信成本。当使用我们的Router Z-loss时，我们没有观察到由乘法偏置引起的不稳定性。我们确实看到乘法交互提高了性能，与我们的强稀疏基线相比，收敛时间加快了4%。这暗示了未来架构研究的一个有前途的途径是寻找在网络中添加更多乘法交互的新方法。

表17：更多的乘法交互提高了稀疏模型的质量。加法偏置和乘法偏置几乎不增加任何参数或计算。

D 容量因子较低时的批量优先路由

默认丢弃顺序与BPR
令人惊讶的是，即使令牌路由是按序列从左到右的顺序进行的，top-1和top-2路由在CF小于1.0时也表现良好。默认的丢弃顺序很重要：我们从左到右丢弃令牌，以避免模型作弊。批量优先路由（Batch Prioritized Routing, BPR），由【69, Scaling vision with sparse mixture of experts, 2021】在视觉Transformer中引入，旨在拥有所有令牌的全局视图来决定哪些令牌应该被丢弃，而不是从左到右的顺序。该算法通过查看所有被发送到专家$i$的$N$个令牌，然后只路由其中具有最高路由器概率的$M$个。

BPR实验结果
表18显示，BPR top-1路由的性能优于top-2路由，尤其是在容量因子小于1.0时。一个重要的限制是，BPR路由只能在编码器-解码器模型的编码器端进行。在解码器上使用BPR会导致模型通过使用未来的令牌信息来作弊，从而提高当前令牌的预测。

E 预训练数据集详情

数据集构成
用于训练我们的稀疏32B模型的预训练数据集是C4（【68, Exploring the limits of transfer learning with a unified text-to-text transformer, 2019】）和GLaM（【16, Glam: Efficient scaling of language models with mixture-of-experts, 2021】）中引入的数据集的混合。


表19：训练集中的数据和混合权重。我们根据“混合中的权重”按比例从不同的数据源中采样。列出的令牌数量以十亿（B）为单位。

F 完整的微调敏感性数据

敏感性原始数据
表20包含了图6的原始数据，测量了微调协议的敏感性。Dense和Sparse是与T5-Large FLOP匹配的编码器-解码器模型，它们在C4语料库上用1M令牌的批量大小预训练了500k步。结果显示，在微调期间重置优化器状态会损害性能。稀疏模型和密集模型的最佳批量大小和学习率存在差异。


表20：微调协议敏感性。我们改变了批量大小、学习率以及是否为密集和稀疏模型重置优化器槽变量。在微调期间重置优化器状态会损害性能。我们观察到稀疏模型与密集模型在最佳批量大小和学习率上存在差异。某些超参数微调设置使稀疏和密集模型的性能几乎完全相同，显示了正确调整超参数的重要性。

G 优化设置路由阈值

Top-n路由算法描述
MoE top-n路由算法如下：首先，路由器找到分配了最高路由器分数（gate1）的专家，并始终将令牌发送给该专家。然后，令牌以一定概率被发送到其他n-1个专家。这个概率基于一个预定义的阈值threshold（通常为0.2）。

阈值调整实验
对于top-3路由与top-2路由，归一化专家分数的和更大，因此我们实验了降低阈值。我们的实验结果如表21所示。有趣的是，我们确实观察到top-3路由从较低的阈值中略微受益，而top-2路由则相反。我们还实验了绝对阈值策略，发现如果阈值被调整得当，它可以达到同样好的性能。


表21：不同阈值下top-2和top-3路由的性能。Top-3路由在较低阈值下比top-2路由表现稍好。

H 专家数较少时的数据、模型和专家并行网格布局

Mesh-Tensorflow并行策略
我们使用Mesh-Tensorflow（【82, Mesh-tensorflow: Deep learning for supercomputers, 2018】）进行数据和模型并行分区。该策略首先形成一个大小为$d \times m$的逻辑2D网格，其中行对应数据维度（$d$），列对应模型维度（$m$），乘积等于总核心数$n = d \times m$。

专家数少于数据并行维度的问题与解决方案
如果我们的专家数少于$d$，这个布局将无法工作。为了允许专家数少于网格中的数据并行行数，我们将数据维度分解为两个新维度：内部（$i$）和外部（$o$），其中$i \times o = d$，并且专家数等于$i$。这将逻辑2D网格$d \times m$转换为3D网格$o \times i \times m$。图8展示了这两种网格的可视化。


图8：用于分布模型的据和模型并行网格。在这个例子中，总共有32个处理器（例如n=32）。（左）如果专家数大于或等于数据并行维度，这是一个有效的2D网格。（右）当我们专家数少于数据并行维度时，这是一个有效的3D网格。批次维度被分解为两个新维度：内部数据和外部数据维度。现在我们每个内部数据维度（i）有1个专家。左图中的8个数据行变成了外部批次（o）中的4个和内部批次（i）中的2个，专家数为2而不是8。

I 分布式模型的通信成本说明

通信操作及其影响
通信操作（allreduce和all2all）可以显著影响稀疏模型训练的吞吐量。allreduce调用沿着模型和批次维度执行，通常由模型维度allreduce调用主导，这些调用对来自工作节点的局部矩阵乘法操作结果进行求和。当矩阵乘法跨多个核心分区时（即模型并行），需要这些调用。

优化通信
为了增加allreduce的吞吐量，可能需要将更多的工作节点分配给模型维度（而不是批次维度）。然而，增加工作节点数量可能会减少每个工作节点的计算量，导致更高的通信开销，从而抵消了allreduce更高通信吞吐量带来的一些增益。对于本文中的结果，我们首先探索了各种模型分区策略。接下来，根据性能基准测试分配了预训练作业的形状，这些基准测试显示了allreduce和all2all中最低的累积通信开销。

J 负面结果

未成功的实验
我们总结了一些在我们的设置中产生负面结果的想法。
- 向路由器添加令牌是否被丢弃的信息：我们实验了让专家层了解令牌在之前的专家层中是被路由还是被丢弃的信息。这没有带来性能差异。
- 添加显式的专家位置信息：我们实验了在专家层的输出中添加显式的位置信息，但这没有提高性能。
- 添加预训练噪声以修复预训练和微调的差异：为了帮助解决预训练困惑度和微调之间的差距，我们尝试在预训练中添加各种类型的噪声（dropout、丢弃整个专家、路由器熵最大化辅助损失）。不幸的是，所有方法要么过多地损害了预训练质量，要么最终没有帮助微调。
- 在top-n路由中对较低的n-1个专家进行负载均衡：我们实验了在top-n路由中为其他n-1个专家添加辅助负载均衡项，但发现这只提供了微不足道的益处。
- 混合预训练和微调数据以防止过拟合：我们尝试在微调期间混合不同比例的预训练跨度损坏数据，但这没有帮助微调性能，反而增加了训练损失。