Google Brain的高级研究科学家BaRRet Zoph指出，他们开发了一种名为“SwITch TRansfoRMeR”的简化稀疏架构，该架构可以将语言模型的参数量扩展至1.6万亿（相比之下，GPT-3为1750亿）。在相同的计算资源下，SwITch TRansfoRMeR的训练速度是T5模型的4到7倍。
通常情况下，模型会对所有输入使用相同的参数，而MixtuRe of ExpeRts（MoE，混合专家）模型则有所不同，它为每个输入选择不同的参数，从而形成一个稀疏激活模型，保证了计算成本的恒定。
MoE模型在机器翻译领域取得了显著成就，但由于其复杂性、通信成本高和训练不稳定，导致其广泛应用受到限制。
为了应对这些挑战，Google Brain的研究人员提出了SwITch TRansfoRMeR。在该模型的设计中，他们简化了MoE的路由算法，并进行了模型改进，大幅降低了通信成本和计算成本。提出的训练技术提升了训练的稳定性，首次证明大型稀疏模型也可以使用低精度（bfloat16）进行训练。
研究者将新模型与T5-Base和T5-LaRge进行了比较，结果显示，在相同的计算资源条件下，新模型实现了最高7倍的预训练速度提升。此外，这一改进在多语言设置中同样有效，在101种语言中均显示出相较于MT5-Base版本的性能提升。
在ColoSSal Clean CRawled CoRpUS上进行的预训练中，研究者将语言模型的参数量提升至上万亿，并在与T5-XXL模型相比时实现了4倍的加速。
研究者表示，尽管该项工作聚焦于规模，但SwITch TRansfoRMeR架构的优势不仅在超级计算机环境下显现，在拥有几个计算核心的计算机上同样有效。此外，设计的大型稀疏模型可以被蒸馏为一个小而稠密的版本，同时保留稀疏模型质量提升的30%。
SwITch TRansfoRMeR的设计理念是以一种简单且计算高效的方式实现TRansfoRMeR模型参数的最大化。Kaplan等人已对扩展效益进行了深入研究，揭示了模型、数据集和计算预算变化下的幂定律缩放。该研究强调了在相对较少数据上训练大型模型作为计算最优的方法。
在分布式训练环境中，模型的稀疏激活层在不同设备上分配独特的权重。因此，随着设备数量的增加，模型权重也随之增加，同时在每个设备上保持可管理的内存和计算空间。
SwITch TRansfoRMeR的编码器块如下图所示：
简化稀疏路由
与其他模型不同，Google Brain的研究采用了简化策略，仅将Token表征发送至单一专家。研究表明，这种简化策略不仅保持了模型质量，还降低了路由计算，并提升了性能。研究者将这一k=1的策略称为SwITch层。
高效稀疏路由
研究者使用了Mesh-TensorFlow库（MTF），它与TensorFlow具有类似的语义和API，旨在促进高效的分布式数据和模型并行架构。在设计模型时，研究者考虑到了TPU的需求，确保其静态大小。
分布式SwITch TRansfoRMeR实现中，所有张量形状在编译时静态确定，但由于训练和推断过程中路由决策的影响，计算是动态的。因此，一个重要的技术难题是如何设定专家容量。
专家容量（每个专家处理的Token数量）的计算方式为：每批Token数量除以专家数量，乘以容量因子。
若将过多Token发送给某个专家，则会跳过计算，Token表征通过残差连接直接传递至下层。然而，增加专家容量也有其缺点，数值过高将导致计算和内存资源的浪费。
实证研究表明，保持较低的丢弃Token比例对稀疏专家模型的扩展至关重要。
SwITch TRansfoRMeR
研究者首先在ColoSSal Clean CRawled CoRpUS（C4）数据集上对SwITch TRansfoRMeR进行了预训练测试，采用了掩蔽语言建模任务。在预训练设置中，他们遵循RaFFel等人确立的最优方案，去除了15%的Token，并使用单个sentinel Token替代掩蔽序列。
SwITch TRansfoRMeR与MoE TRansfoRMeR的比较结果如下表所示。结果显示，SwITch TRansfoRMeR在速度与质量的基础上优于经过精心调整的密集模型和MoE TRansfoRMeR，并在固定计算量和挂钟时间的情况下获得最佳结果；此外，SwITch TRansfoRMeR的计算占用空间低于MoE TRansfoRMeR，并在低容量因子下表现更为出色。
提升训练与微调技巧
与原版TRansfoRMeR模型相比，稀疏专家模型在训练时可能面临更多挑战。这些层中的haRd-swITHing决策可能导致模型不稳定。此外，像bfloat16这样的低精度格式可能加剧RouteR的softMax计算问题。研究者采取了以下几种技巧以克服训练困难，实现稳定且可扩展的训练。
对大型稀疏模型使用可选行精度
使用更小的参数初始化以实现稳定性
对大型稀疏模型进行正则化
预训练可扩展性
在预训练期间，研究者对SwITch TRansfoRMeR的可扩展性进行了研究，考虑了算力和数据均不受限制的机制。为避免数据受限，研究者使用了包含180B目标Token的大型C4数据库，直到观察到收益递减为止。
专家数量是扩展模型最有效的维度。增加专家数量几乎不会改变计算成本，因为模型仅为每个Token选择一个专家，这与专家总数无关。RouteR必须基于更多专家计算概率分布，但这是一项轻量级的计算任务。
基于步数的可扩展性
下图展示了多个模型在训练步数恒定、专家数量增加时的可扩展性提升情况。可观察到一个趋势：在保持每个Token的FLOPS不变时，更多的参数（专家）可以提升训练速度。
基于时间的可扩展性
如上图所示，随着专家数量的增加，模型性能持续提升。尽管模型每个Token的FLOPS与基线相近，但SwITch TRansfoRMeRs会产生额外的跨设备通信成本，路由机制也带来额外的计算开销。因此，在基于步数的设置中观察到的样本效率提升未必能转化为时间上的模型质量提升。
在训练时间和计算成本均固定的情况下，究竟是训练稠密模型还是稀疏模型更好？
下图解决了这个问题。图展示了预训练模型质量随着训练时间的变化。在训练时间和计算成本固定的情况下，SwITch TRansfoRMeR表现出明显的速度优势。在这种情况下，若要达到相似的困惑度，SwITch-Base 64专家模型的训练时间仅为T5-Base模型的1/7。
可扩展性与更大稠密模型的对比
上述实验表明，计算不受限制的稠密模型已被SwITch超越。考虑不同情况：如果将计算资源分配给更大的稠密模型会如何？
为验证此问题，研究者将SwITch-Base与更强的基线T5-LaRge进行了对比。实验结果表明，尽管T5-LaRge每个Token所用的FLOPs是SwITch-Base的3.5倍，但后者的样本效率更高，速度是前者的2.5倍。此外，如果设计一个与T5-LaRge所需FLOPs相当的SwITch模型（SwITch-LaRge），上述提升将更加显著。
下游任务中的结果
微调
本研究使用的基线方法是具备223M参数的T5-Base和具备739M参数的T5-LaRge模型。针对这两个模型，研究者设计了具备更多参数的FLOP-Matched SwITch TRansfoRMeR。
在多项自然语言任务中，SwITch TRansfoRMeR显著提升了性能。最明显的提升体现在SupeRGLUE基准上，FLOP-Matched SwITch TRansfoRMeR相比T5-Base和T5-LaRge分别提升了4.4%和2%；在WinogRande、closed book TRivia QA和XSuM上同样表现出类似的提升。而在AI2推理挑战赛（ARC）数据集中，T5-Base在ARC challenge上的表现超越SwITch-Base；在ARC easy数据集上，T5-LaRge的表现也优于SwITch。
