谷歌推出首个万亿级语言模型Switch Transformers

2021-01-14 10:17:35来源：雷锋网

众所周知，参数量是机器学习算法的关键。在大规模参数量和数据集的支持下，简单的体系结构将远远超过复杂的算法。

在自然语言领域，被称为史上最强 NLP 的 GPT-3 拥有 1750 亿参数。近日，Google 将这一参数量直接拉高到了 1.6 万亿。

1 月 11 日，Google 在 arXiv 上发表论文《Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity》，提出了最新语言模型—Switch Transformer。

研究员介绍称，Switch Transformer 拥有 1.6 万亿参数，是迄今为止规模最大的 NLP 模型。

论文中指出，Switch Transformer 使用了稀疏激活(Sparsely Activated)技术，该技术只使用了神经网络权重的子集，或者是转换模型内输入数据的参数。在相同计算资源下，其训练速度上比 Google 之前研发的最大模型 T5-XXL 还要快 4 倍。

Switch Transformer 由混合专家(Mix of Expert，MoE)AI 模型范式发展而来。MoE 模型是指将多个专家或专门从事不同任务的模型放在一个较大的模型中，并有一个 “门控网络(Gating Network)”来选择为任何给定数据要咨询哪些专家。

自 90 年代初首次被提出以来，MoE 模型通在机器翻译方面取得了显著成果。但因复杂性高、通信成本高和训练不稳定，导致其广泛应用受到阻碍。

为了解决这些问题，研究人员简化了 MoE 路由算法，设计了直观的改进模型，使通信成本和计算成本大幅度降低。同时，他们提出的训练技术消除了训练的不稳定性，并证明了大型稀疏模型也可能以低精度(bfloat16)格式进行训练。

此外，与 T5-Base 和 T5-Large 模型相比，新模型在相同计算资源下将预训练速度提高了 7 倍之多。这些改进扩展到了多语言设置中，可以测量 101 种语言在 mT5-Base 版本上的收益。

最后，通过在 “Colossal Clean Crawled Corpus”上预先训练多达数万亿个参数模型来提高当前语言模型的规模，使速度比 T5-XXL 模型提高了 4 倍。

研究者还表示，Switch Transformer 架构不仅在具备超级计算机的环境下具有优势，在只有几个计算核心的计算机上也是有效的。此外，研究者设计的大型稀疏模型可以被蒸馏为一个小而稠密的版本，同时还能保留稀疏模型质量提升的 30%。

Switch Transformer 模型架构

Switch Transformer 的设计原则是以一种简单且高效计算的方式来最大化 Transformer 模型的参数数量。

基于此，模型需要在增加参数量同时保持每个示例的浮点运算(FLOPs)不变。研究人员假设，参数量与执行的总计算量无关，是可以单独缩放的重要组件，那么它可以通过设计稀疏激活模型来实现。

研究人员有效地利用了为密集矩阵乘法(Dense Matrix Multiplications, 一种广泛用于语言模型的数学运算)而设计的硬件，如 GPU、Google 的张量处理单元 TPU。在分布式训练设置中，模型将不同的权重分配到不同的设备上，可以保证权重随着设备数量的增加而增加，同时保证每个设备上均有可管理的内存和计算足迹。

Switch Transformer 编码器模块图示

如图，研究人员将 Transformer 的密集前馈网络(FFN)层替换为稀疏 Switch FFN 层(浅蓝色)。该层对序列中的令牌进行独立操作。他们两个标记令牌(X1=More 和 X2=Parameters )在四个 FFN 专家之间进行匹配(实线)，其中路由器独立地传输每个令牌。Switch FFN 层返回输出的 FFN 并乘以路由器门值(虚线)

简化稀疏通道

研究人员提出了 Switch Routing 简化策略：针对不同的输入，匹配最适合的一个专家。实验证明，这种简化可以保留模型质量，减少路由计算并提高性能。具体优势如下：

1、通道计算量大大降低，因每个样本仅需要一个专家通道参与计算。

2、每个专家通道所计算样本的 batchsize 被大大缩小 (至少减半)。

3、每个专家通道的实现复杂度减小了，专家间的通信开销降低。

高效稀疏通道

研究人员使用 Mesh-Tensorflow(MTF)库—与 Tensorflow 类似的语义和 API，能够能够高效支持分布式的数据与模型结构。它将物理内核集抽象为处理器的逻辑网格。然后利用张量和计算按指定的维度进行切分，便于跨维度轻松划分模型。

与 MoE Transformer 的比较

第一个测试是在 “Colossal Clean Crawled Corpus”数据集上进行的。该数据集也被成为 C4，含 750GB 的文本语料，全部来自于 Reddit、Wikipedia 和其他网络资源中。研究人员使用了 32 个 TPU 核在 C4 数据集上预先训练了几种不同的 Switch Transformer 模型。

模型的任务是预测段落中 15% 被遮住的单词，以及通过检索文本回答一些有难度的问题。

从以上对比结果中，可以发现：

1、在速度 - 质量的指标上，Switch-Transformer 超过了精密微调的稠密模型，与 MoE Transformer。在有限的计算量与工作时间下，Switch Transformer 都得到了最优的结果。

2、Switch Transformer 有更小的计算足迹(computational footprint)。

3、在更小的 capactiy factor 下，Switch Transformer 工作的效果更佳。

提升训练与微调技术

稀疏专家模型相比普通 Transformer 更难训练，这是由于切换通道带来的不稳定引起的，同时，低精度计算会导致恶化 softmax 计算中的问题。

研究人员通过在局部引入高精度的训练，而非全局的方法，达到了理想的稳定性，测试效果如图：

同时作为不稳定性的另一种补救措施，研究人员减小了默认的 Transformer 初始化比例，将 S=1.0 降低了 10 倍。这既提高了质量，又降低了不稳定训练的可能性。如图：

“预训练 + 微调”的 NLP 方法，在处理下游任务时容易发生过拟合现象。对此，研究人员提出了 “expert dropout”的解决方案——通过在微调过程中，增加在每个 expert 中的 dropout 比例来缓解过拟合。

我们观察到，简单地增加所有层的 dropout 率会导致性能下降。但是，在非专家层设置较小的 dropout 率(0.1)和在专家层设置较大的 dropout 率(0.4)，可以提高四个较小下游任务的性能。

性能测试：预训练、微调和多任务训练

论文中，研究人员对预训练阶段 Switch Transformer 的预训练可扩展性进行了研究。为了避免数据受限，他们采用了大型 C4 数据库，并以固定的计算成本测试了基于步数和时间的可伸缩性。

步态的可扩展性

下图为多个模型在训练步数恒定、专家数量增加时表现出的可扩展性提升情况。可以看到：在保持每个 token 的 FLOPS 不变时，拥有更多的参数(专家)可以提高训练速度。

此外还可以发现，专家数量的增加会导致更有效的样本模型。Switch-Base 64 专家模型在 60k 和 450k 步态达到了与 T5-Base 模型相同的性能，这是 7.5 倍的加速。

时间的可扩展性

从基于步数的设置中，可以观察到：样本效率提升未必能转化为时间上的模型质量提升。那么，在训练时间和计算成本都固定的前提下，是应该训练一个稠密模型还是稀疏模型?下图解决了这个问题。

图中展示了预训练模型质量随训练时间增加所产生的变化。在训练时间和计算成本都固定的情况下，Switch Transformer 的速度优势非常明显。在这种设置下，如果要达到相似的困惑度，Switch-Base 64 专家模型的训练时间仅为 T5-Base 模型的 1/7。

微调

针对 T5-Base 和 T5-Large 模型，研究人员设计了具备更多参数的 FLOP-matched Switch Transformer。在 SuperGLUE 基准上，FLOP-matched Switch Transformer 相比 T5-Base 和 T5-Large 的性能分别提升了 4.4% 和 2%。同时，该模型在多项推理和知识任务中也带来了显著提升。

这说明该模型架构不只对预训练有用，还可以通过微调将质量改进迁移至下游任务中。

蒸馏技术

为了解决超大规模神经网络的部署问题，研究人员提出了多种蒸馏技术：将大型稀疏模型蒸馏到小型密集模型中。

使用表 7 中最优的蒸馏技术，可以将具备 11 亿参数量的模型压缩 82%，同时保留 37% 的性能提升。最极端的情况下，将模型压缩了 99%，且维持了 28% 的性能提升。

多语言学习

在最后一组下游实验中，研究人员利用模型对 101 种不同语言的混合进行了预训练。

如图展示了 Switch T5 Base 模型与 mT5-Base 在所有语言上的质量提升情况(负对数困惑度)。对两个模型经过 100 万步预训练后，Switch Transformer 的最终负对数困惑度相较基线有所提升。

此外，透过 mT5-Base 使用 Switch Transformer 模型的加速直方图，可以发现，mT5-Base 的平均速度提高了 5 倍，同时，91%的语言速度至少提升了 4 倍。这表明 Switch Transformer 能够有效的多任务和多语言。

总结及未来研究方向

研究人员表示，Switch Transformer 模型是具有可扩展性的高效自然语言处理模型，它通过简化 MoE，形成了快速理解、训练稳定且比同等大小模型更有效的样本体系结构。实验证明，这些模型在不同的语言任务和训练模式中都表现了出色的性能，包括预训练、微调和多任务训练。同时，这些结果也让训练具有上千亿、上万亿参数量的模型成为可能。

最后研究人员指出，虽然 Switch Transformer 改进了训练程序，并对稀疏模型进行了扩展的研究。但未来在此方向还有许多未完成的工作，如

进一步提高大型模型训练的稳定性。目前正在进行的方法有：正则化函数、适度的梯度裁剪等。

一般来说预训练越好，下游任务的效果也就越好，但在一些任务上发现，1.6T 参数的 Switch-C 会比较小模型的结果更低。

目前方法使用的同类专家，未来可尝试通过更灵活的基础结构支持异构专家。

调查 FFN 层之外的专家层，初步的证据表明，这同样可以改善模型质量。

目前的工作仅考虑了语言任务，未来可尝试将模型稀疏性类似地应用到其他模态(例如图像)或多模态中。