新年伊始,大规模预训练模型军备竞赛进入万亿参数时代。
文:梦佳、周寅张皓、贾伟
近日,Google Brain的研究人员William Fedus、Barret Zoph、Noam Shazeer等在arxiv上提交了一篇新论文,“Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity”,提出了稀疏激活专家模型Switch Transformer。
链接:arxiv.org/abs/2101.03961
研究人员表示,这个1.6万亿参数模型似乎是迄今为止最大的模型,其速度是Google之前开发的最大语言模型(T5-XXL)的4倍,参数规模几乎是1750亿参数的GPT-3的十倍!
这应该是人类历史上发布的第一个万亿级人工智能模型。
研究人员在论文中指出,大规模训练是通向强大模型的有效途径,具有大量数据集和参数计数的简单架构可以远远超越复杂的算法,但目前有效的大规模训练主要使用稠密模型。
作为对比,William等人提出的 Switch Transformer 采用了“稀疏激活”技术。所谓稀疏,指的是对于不同的输入,只激活神经网络权重的子集。
根据作者介绍,Switch Transformer是在MoE的基础上发展而来的,而MoE则是90年代初首次提出的AI模型。MoE 将多个“专家”或专门从事不同任务的模型放在一个较大的模型中,并有一个“门控网络”来选择对于任何给定数据要咨询哪些/个“专家”。
尽管MoE取得了一些显著成功,但复杂性、通信成本和训练不稳定阻碍了其广泛采用。
Switch Transformer的新颖之处在于,它有效地利用了为稠密矩阵乘法(广泛用于语言模型的数学运算)而设计的硬件——例如GPU和Google TPU。研究人员为不同设备上的模型分配了唯一的权重,因此权重会随着设备的增多而增加,但每个设备上仅有一份内存管理和计算脚本。
Switch Transformer 在许多下游任务上有所提升。研究人员表示,它可以在使用相同计算资源的情况下使预训练速度提高7倍以上。他们证明,大型稀疏模型同样可以用于创建较小的、稠密的模型,通过微调,这些模型相对大型模型会有30%的质量提升。
论文一作William Fedus 也在twitter上表示,「我们的模型采样更加高效,相比于流行的模型,T5-Base,T5-Large、T5-XXL等能实现4到7倍的增速。」
在一项测试中,Switch Transformer模型以在100多种不同语言之间的翻译测试中,研究人员观察到“普遍改进”,与基准模型相比,91%的语言翻译有4倍以上的提速。
研究人员认为,在未来的工作中,Switch Transformer可以应用到其他模态或者跨模态的研究当中。模型稀疏性可以多模态模型中发挥出更大的优势。
01
模型框架
模型的设计原则是,在保持计算开销与效率的同时,尽可能的提升模型的参数量。该工作通过固定每个数据样本所进行的浮点数操作,对比不同参数量可能带来的影响。实际运行中&#