众所皆知，参数量是机器学习算法的关键。在大规模参数量和数据库支持下，简单的体绑定构将远远超过复杂的算法。
自然语言领域，称为史上最强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层（浅蓝色）。该层对串行的令牌独立操作。两个标记令牌（X1=More和X2=Parameters）在4个FFN专家之间配对（实线），其中路由器独立传输每个令牌。Switch FFN层返回输出的FFN并乘以路由器门值（虚线）。
简化稀疏信道
研究人员提出Switch Routing简化策略：针对不同输入，匹配一个最适合的专家。实验证明，这简化可保留模型品质，减少路由计算并提高性能。具体优势如下：
- 信道计算量大大降低，因每个样本仅需一个专家信道参与计算。
- 每个专家信道所计算样本的batchsize大大缩小（至少减半）。
- 每个专家信道的完成复杂度减少，专家通信开销降低。
高效稀疏信道
研究人员使用Mesh-Tensorflow（MTF）数据库──与Tensorflow类似的语义和API，高效支持分布式数据与模型结构。将物理内核集抽象为处理器的逻辑网格，然后利用张量和计算按指定的维度切分，便于跨维度轻松划分模型。
与MoE Transformer比较
第一个测试是在“Colossal Clean Crawled Corpus”数据库进行，也称为C4，含750GB文本语料，全部来自Reddit、Wikipedia和其他网络资源。研究人员使用32个TPU核在C4数据库预先训练几种不同Switch Transformer模型。
模型任务是预测段落15%被遮住的单词，以及通过检索文本回答有难度的问题。
从对比结果可发现：
- 速度─品质指标，Switch-Transformer超越精密微调的稠密模型与MoE Transformer。在有限计算量与工作时间下，Switch Transformer都得到最优结果。
- Switch Transformer有更小的计算足迹（computational footprint）。
- 更小的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.6T参数的Switch-C会较小模型结果更低。
- 目前方法使用的同类专家，未来可尝试通过更灵活的基础结构支持异构专家。
- 调查FFN层之外的专家层，初步证据表明，这同样可以改善模型品质。
- 目前工作仅考虑语言任务，未来可尝试将模型稀疏性类似应用到其他模态（如图像）或多模态。