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。在 SwITch TRansfoRMeR 的设计中，他们简化了 MoE 的路由算法，设计了改进模型，新模型的通信成本和计算成本都大大降低。他们提出的训练技术提高了训练的稳定性，首次表明大型稀疏模型也可以用低精度（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 版本的性能提升。
研究者在 ColoSSal Clean CRawled CoRpUS 上进行预训练，将语言模型的参数量提升至上万亿，且相比 T5-XXL 模型实现了 4 倍加速。
研究者表示，虽然这项工作着眼于规模，但它也表明，SwITch TRansfoRMeR 架构不仅在具备超级计算机的环境下具有优势，在只有几个计算核心的计算机上也是有效的。此外，研究者设计的大型稀疏模型可以被蒸馏为一个小而稠密的版本，同时还能保留稀疏模型质量提升的 30%。
SwITch TRansfoRMeR 的设计原理
SwITch TRansfoRMeR 的主要设计原则是，以一种简单且计算高效的方式最大化 TRansfoRMeR 模型的参数量。Kaplan 等人已经对扩展的效益进行了详尽的研究，揭示了随模型、数据集大小以及计算预算变化的幂定律缩放。重要的是，该研究提倡在相对较少数据上训练大型模型，将其作为计算最优方法。
在分布式训练设置中，模型的稀疏激活层在不同设备上分配唯一的权重。所以，模型权重随设备数量的增加而增加，同时在每个设备上保持可管理的内存和计算空间。
SwITch TRansfoRMeR 的编码器块如下图 2 所示：
简化稀疏路由
与其他模型不同，谷歌大脑的研究采用了简化策略，只将 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 的比较结果如下表 1 所示。结果表明，SwITch TRansfoRMeR 在速度 – 质量基础上优于精心调整的稠密模型和 MoE TRansfoRMeR，并在固定计算量和挂钟时间情况下取得了最佳结果；SwITch TRansfoRMeR 的计算占用空间比 MoE TRansfoRMeR 小；SwITch TRansfoRMeR 在低容量因子下表现更好。
提升训练和微调的技巧
与原版 TRansfoRMeR 模型相比，稀疏专家模型在训练时可能更加困难。所有这些层中的 haRd-swITHing决策都可能导致模型的不稳定。此外，像 bfloat16 这样的低精度格式可能加剧 RouteR 的 softMax 计算问题。研究者采取了以下几种技巧来克服训练困难，并实现稳定和可扩展的训练。
对大型稀疏模型使用可选择行精度
为实现稳定性使用更小的参数初始化
正则化大型稀疏模型
预训练可扩展性
在预训练期间，研究者对 SwITch TRansfoRMeR 的可扩展性进行了研究。在此过程中，他们考虑了一个算力和数据都不受限制的机制。为了避免数据受限，研究者使用了大型 C4 数据库，里面包含 180B 的目标 Token。在观察到收益递减之前，他们一直进行训练。
专家的数量是扩展模型最有效的维度。增加专家的数量几乎不会改变计算成本，因为模型只为每个 Token 选择一个专家，这与专家的总体数量无关。RouteR 必须基于更多的专家计算概率分布，但这是一个轻量级的计算成本。
基于步数的可扩展性
下图 4 展示了多个模型在训练步数恒定、专家数量增加时表现出的可扩展性提升情况。从中可以观察到一个趋势：在保持每个 Token 的 FLOPS 不变时，拥有更多的参数（专家）可以提高训练速度。
基于时间的可扩展性
如上图 4 所示，随着专家数量的增加，模型的性能会不断提升。虽然模型的每个 Token 拥有与基线近乎相同的 FLOPS，但 SwITch TRansfoRMeRs 会产生额外的跨设备通信成本，路由机制也会带来额外的计算开销。因此，在基于步数的设置中观察到的样本效率提升未必能转化为时间上的模型质量提升。
在训练时间和计算成本都固定的前提下，我们应该训练一个稠密模型还是稀疏模型？
下图 5 和 6 解决了这个问题。图 5 展示了预训练模型质量随训练时间增加所产生的变化。在训练时间和计算成本都固定的情况下，SwITch TRansfoRMeR 的速度优势非常明显。在这种设置下，如果要达到相似的困惑度，SwITch-Base 64 专家模型的训练时间仅为 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），上述提升还会更加明显。
下游任务中的结果
微调
这里使用的基线方法是具备 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）数据集：在 ARC challenge 数据集上 T5-Base 的性能超过 SwITch-Base；在 ARC easy 数据集上，T5-LaRge 的性能超过 SwITch
