首个兆级模型！Google 重量级推出语言模型 Switch Transformers，1.6 兆参数秒杀 GPT-3

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

自然语言领域，称为史上最强 NLP 的 GPT-3 拥有 1,750 亿参数。近日，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）选择为任何给定数据要咨询哪些专家。

自 1990 年代初首次提出，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 层（浅蓝色）。该层对序列的令牌独立操作。两个标记令牌（X ₁ =More和X ₂ =Parameters ）在 4 个 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），可提高 4 个较小下游任务的性能。

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

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

步态的可扩展性

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

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

时间的可扩展性

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

图展示预训练模型品质随训练时间增加产生的变化。训练时间和计算成本都固定的情况下，Switch Transformer 的速度优势非常明显。如果要达到相似困惑度，Switch-Base 64 专家模型的训练时间仅 T5-Base 模型七分之一。

微调

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

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

蒸馏技术

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

使用最优的蒸馏技术，可将具备 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. 进一步提高大型模型训练的稳定性。目前正进行的方法有：正则化函数、适度梯度裁剪等。
2. 一般来说预训练越好，下游任务的效果也就越好，但一些任务发现，1.6T 参数的 Switch-C 会较小模型结果更低。
3. 目前方法使用的同类专家，未来可尝试透过更灵活的基础结构支持异构专家。
4. 调查 FFN 层之外的专家层，初步证据表明，这同样可以改善模型品质。
5. 目前工作仅考虑语言任务，未来可尝试将模型稀疏性类似应用到其他模态（如图像）或多模态。

• SWITCH TRANSFORMERS: SCALING TO TRILLION PARAMETER MODELS WITH SIMPLE AND EFFICIENT SPARSITY
• Google trained a trillion-parameter AI language model

（本文由 雷锋网 授权转载；首图来源：Unsplash）