GoogleBrain提出简化的稀疏架构，预训练速度比T5快7倍，参数数量扩展到1.6万亿（GPT-3是1750亿）。在相同的计算资源下，SwitchTransformer的训练速度可以达到T5模型的4-7倍。在深度学习中，模型通常对所有输入重复使用相同的参数。专家混合(MoE)模型是个例外，它为每个输入示例选择不同的参数，从而导致稀疏激活模型——参数数量惊人但计算成本恒定。目前，MoE模型在机器翻译领域取得了骄人的成绩，但由于模型复杂度高、通信成本高、训练不稳定等问题，在一定程度上阻碍了其广泛应用。为了解决这些问题，GoogleBrain研究人员提出了SwitchTransformer。在SwitchTransformer的设计中，他们简化了MoE的路由算法（routingalgorithm），设计了直观的改进模型，大大降低了新模型的通信成本和计算成本。此外，他们提出的训练技术还提高了训练稳定性，首次表明大型稀疏模型也可以以低精度进行训练(bfloat16)。论文链接：https://arxiv.org/pdf/2101.03961.pdf代码链接：https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/moe.py研究人员还将新模型与对比了T5-Base和T5-Large，结果表明，在相同的计算资源下，新模型实现了高达7倍的预训练速度提升。这种改进也可以扩展到多语言设置，在所有101种语言中测量新模型相对于mT5-Base版本的性能改进。最后，研究人员在ColossalCleanCrawledCorpus上进行预训练，将语言模型的参数数量增加到万亿级，相比T5-XXL模型实现了4倍的加速。研究人员还表示，虽然这项工作着眼于规模，但它也表明SwitchTransformer架构不仅在拥有超级计算机的环境中具有优势，而且在只有几个计算核心的计算机上也有效。此外，研究人员设计的大型稀疏模型可以提炼成小型密集版本，同时保留稀疏模型30%的质量提升。SwitchTransformer的设计原则SwitchTransformer的主要设计原则是以简单且计算高效的方式最大化Transformer模型的参数数量。Kaplan等人详尽地研究了缩放的好处。(2020)，揭示了随模型、数据集大小和计算预算而变化的幂律缩放。重要的是，该研究提倡在相对较少的数据上训练大型模型作为计算上的最佳方法。基于这些，研究人员在增加参数数量的同时保持每个示例的FLOPs不变。他们假设参数的数量与执行的计算总量无关，并且是可以独立缩放的重要组成部分。因此，研究人员通过设计一个具有稀疏激活的模型来实现这一目标，该模型可以有效地利用专为密集矩阵乘法设计的硬件，例如GPU和TPU。在分布式训练设置中，模型的稀疏激活层在不同设备上分配了唯一的权重。因此，模型权重随设备数量而变化，同时在每个设备上保持可管理的内存和计算空间。SwitchTransformer的编码器块如下图2所示：SimplifiedSparseRoutingShazeeretal.(2017)提出了一个自然语言MoE层，它将标记表示x作为输入并将其发送给最忠诚的top-k专家（从^N_i=1组N专家中选出）。他们假设将令牌表示发送给k>1个专家对于启用路由功能的有意义的梯度是必要的。他们认为，如果不比较至少两名专家的能力，就无法学习路由。与这些想法不同，GoogleBrain研究采用了一种简化的策略，即只向单个专家发送标记表示。研究表明，这种简化策略可以保持模型质量、减少路由计算并获得更好的性能。研究人员将这种k=1策略称为Switch层。下面的图3显示了具有不同专家容量因子的路由示例：高效稀疏路由研究人员使用Mesh-Tensorflow库(MTF)，它具有类似TensorFlow的语义和API，以促进高效的分布式数据和模型并行架构。研究人员在设计模型时考虑到了TPU，这需要一个静态尺寸。DistributedSwitchTransformer实现：所有张量形状都是在编译时静态确定的，但由于训练和推理期间的路由决策，计算是动态的。鉴于此，一个重要的技术难点就出现了：如何设置专家容量？专家容量（每个专家计算出的代币数量）的计算方式为：每批次的代币数量除以专家数量，再乘以容量因子。如等式（3）所示：如果发送给专家的令牌太多（以下称为“丢弃的令牌”），则跳过计算，并通过残差连接将令牌表示直接传递给下层。但是，增加专家容量也不是没有缺点，过高的值会导致计算和内存的浪费。手头的权衡如上图3所示。实证研究发现，保持较低的丢弃令牌比例对于稀疏专家模型的扩展很重要。设计决策对模型质量和速度的影响如下表1所示。SwitchTransformer研究人员首先使用屏蔽语言建模任务在ColossalCleanCrawledCorpus(C4)数据集上测试了SwitchTransformer预训练。在预训练设置中，他们遵循Raffel等人确定的最佳方案。(2019)，丢弃了15%的标记，然后用单个哨兵标记替换了掩码序列。为了比较模型性能，研究人员提供了负对数困惑度结果。SwitchTransformer和MoETransformer的对比结果如下表1所示。结果表明，SwitchTransformer在速度-质量基础上优于经过良好调优的密集模型和MoETransformer，并在恒定计算和挂钟时间下取得最佳结果；SwitchTransformer的计算足迹比MoETransformer小；开关变压器在低容量因数（1.0、1.25）下表现更好。改进的训练和微调技术稀疏专家模型可能比香草Transformer模型更难训练。所有这些层中的硬切换（路由）决策可能导致模型不稳定。此外，像bfloat16这样的低精度格式会加剧路由器的softmax计算问题。研究人员采用了以下技术来克服训练困难并实现稳定和可扩展的训练。大型稀疏模型的选择性精度较小的参数初始化以确保稳定性正则化大型稀疏模型预训练可扩展性在预训练期间，研究人员研究了SwitchTransformer的可扩展性。在这个过程中，他们考虑了一种计算能力和数据都不受限制的机制。为了避免数据限制，研究人员使用了一个大型C4数据库，其中包含180B个目标令牌。他们继续训练，直到观察到收益递减。专家的数量是扩展模型的最有效维度。增加专家数量几乎不会改变计算成本，因为该模型只为每个标记选择一名专家，而与专家总数无关。路由器必须计算基于更多专家的概率分布，但这是一个轻量级的计算，成本为O(d_model×numexperts)。其中，d_model是层间传递的token的embedding维数。在这一部分中，我们考虑以固定的计算成本进行基于步骤和时间的可扩展性。基于步骤的可扩展性下面的图4显示了在使用恒定步数进行训练并增加专家数量时，几个模型的可扩展性改进。可以观察到一个趋势：拥有更多参数（专家）可以提高训练速度，同时保持每个令牌的FLOPS不变。基于时间的可扩展性如上图4所示，随着专家数量的增加，模型的性能不断提高。虽然模型的每个token都有和baseline几乎一样的FLOPS，但是SwitchTransformers会产生额外的跨设备通信成本，路由机制也会带来额外的计算开销。因此，在基于步骤的设置中观察到的样本效率增益可能不一定转化为时间模型质量增益。这就引出了一个问题：在给定固定训练时间和计算成本的情况下，我们应该训练密集模型还是稀疏模型？下面的图5和图6解决了这个问题。图5显示了随着训练时间的增加，预训练模型质量的变化。在训练时间和计算代价一定的情况下，SwitchTransformer的速度优势非常明显。在此设置下，Switch-Base64专家模型的训练时间仅为T5-Base模型的1/7，即可达到类似的困惑度。可扩展性VS。一个更大的密集模型上面的实验表明，一个计算无限的密集模型已经被Switch超越了。图6考虑了一种不同的情况：如果将计算资源分配给更大的密集模型会怎样？为了测试这个问题，研究人员将Switch-Base与更强的基线T5-Large进行了比较。实验结果表明，虽然T5-Large每token使用的FLOPs是Switch-Base的3.5倍，但后者的采样效率仍然更高，速度是前者的2.5倍。另外，如果设计一个和T5-Large要求的FLOPs相同的Switch模型（Switch-Large），上述的提升会更加明显。下游任务中结果的微调这里使用的基线方法是高度调整的T5-Base和739M参数的T5-Base和T5-Large。针对这两个模型，作者设计了一个参数更多的FLOP-matchedSwitchTransformer。在多项自然语言任务中，SwitchTransformer带来了显着的性能提升。最明显的是SuperGLUE，其上与FLOP匹配的SwitchTransformer分别优于T5-Base和T5-Large4.4%和2%，类似的情况也出现在Winogrande，闭书TriviaQA和XSum上。唯一没有观察到性能改进的基准是AI2推理挑战(ARC)数据集：T5-Base在ARC挑战数据集上的表现优于Switch-Base；T5-Large在ARCeasy数据集Large上的表现优于Switch-Base。总体而言，SwitchTransformer模型在多项推理和知识任务中带来了显着的性能提升。这表明模型架构不仅对预训练有用，还可以通过微调将质量改进传递给下游任务。蒸馏部署具有数十亿和数万亿个参数的大型神经网络并非易事。为此，本文研究了如何将大型稀疏模型提炼为小型密集模型。下表7显示了本研究中使用的蒸馏技术：在使用表7中的最佳蒸馏技术后，研究人员将多个稀疏模型蒸馏成密集模型。他们提炼出参数大小从11亿到147亿不等的Switch-Base模型，具体取决于专家的数量。该研究能够将具有11亿个参数的模型压缩82%，同时保持37%的性能增益。在最极端的情况下，模型被压缩了99%，同时保持了28%的性能增益。最后，研究人员将微调的稀疏模型提炼成密集模型。下面的表9显示了为SuperGLUE任务-223MT5-Base微调的74亿参数Switch-Base模型的蒸馏结果。与预训练结果类似，蒸馏后的模型仍然保留了30%的性能增益。这可能有助于识别和提取特定专家以进行微调任务，从而实现更好的模型压缩。多语言学习在下游任务的实验中，研究人员权衡了模型质量和速度之间的权衡，并在101种不同的语言上对模型进行了预训练。下面的图7显示了SwitchT5Base模型和mT5-Base在所有语言上的质量改进（负对数困惑度）。在对两个模型进行100万步预训练后，SwitchTransformer的最终负对数困惑度与基线相比有所提高。下面的图8显示了SwitchTransformer与mT5-Base相比每一步的加速比。前者实现了平均5倍的加速，其中91%的语言至少实现了4倍的加速。这表明SwitchTransformer是一个高效的多任务多语言学习器。使用数据、模型和专家并行化来设计模型随意增加专家数量存在收益递减的问题（参见上图4），本研究介绍了一些互补的扩展策略，涉及组合数据的权衡，模型和专家并行化。结合数据、模型和专家并行化构建万亿参数模型SwitchTransformer在设计过程中，研究人员试图平衡FLOPspertoken和参数量。当专家数量增加时，参数数量增加而每个令牌的FLOPs不变。要增加FLOPs，需要增加d_ff维度（这样也会增加参数的数量，但是比较少）。这是一个权衡：增加d_ff维度会导致每核内存耗尽，因此必须增加m。但由于核心数N是固定的N=n×m，n必须减少，这意味着需要使用更小的批处理大小。在将模型并行化和专家并行化相结合后，将令牌发送给正确的专家和模型并行化引起的内部all-reduce通信会产生all-to-all通信成本。当结合这三种方法时，如何平衡FLOPs、通信成本和每核内存变得非常复杂。该研究结合数据、模型和专家并行化，设计了两个分别具有3950亿个参数和1.6万亿个参数的大型SwitchTransformer模型，并研究了这些模型在上游预训练语言模型和下游微调任务中的表现。不同模型的参数量、FLOPs和超参数量如下表10所示：关于SwitchTransformer还有很多问题在论文的最后部分，谷歌大脑研究人员探索了一些关于SwitchTransformer和稀疏专家模型的问题（这里是sparse指的是权重，而不是注意力模式）。问题一：开关变压器性能更好吗？是参数量太大的原因吗？性能确实更好，但原因不在于参数数量，而在于设计。参数有助于扩展神经语言模型，而更大的模型确实表现更好。但本研究提出的模型在使用相同的计算资源的情况下具有更高的样本效率。问题2：没有超级计算机可以使用这个方法吗？尽管本文侧重于非常大的模型，但研究人员发现了一个由两名专家组成的模型，该模型既可以提高性能，又可以轻松适应常用GPU或TPU的内存限制。因此，研究人员认为该技术可用于小规模环境。问题3：在速度-精度帕累托曲线上，稀疏模型是否优于密集模型？是的。在几种不同的模型大小下，稀疏模型在每步和挂钟时间上都优于密集模型。对照实验表明，对于固定的计算量和时间，稀疏模型优于密集模型。问题四：万亿参数模型部署不了，能不能缩容？虽然万亿参数模型的质量无法完全保持，但通过将稀疏模型提炼成密集模型，可以实现10-100倍的压缩比，同时实现比专家模型约30%的质量提升。问题5：为什么使用SwitchTransformer而不是modelparalleldensemodel？在时间基础上，SwitchTransformer比使用分片参数的密集模型更有效。同时，这种选择并不相互排斥。模型并行化也可以用在SwitchTransformer中，可以提高FLOPspertoken，但是也会导致传统模型并行化的速度变慢。问题六：为什么稀疏模型没有得到广泛应用？尝试稀疏模型的想法因密集模型的压倒性成功而受挫。此外，稀疏模型面临多个问题，包括模型复杂性、训练难度和通信成本。开关变压器缓解了这些问题。
