众所周知,参数量是机器学习算法的关键。在大规模参数量和数据集的支持下,简单的体系结构将远远超过复杂的算法。
在自然语言领域,被称为史上最强 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、每个专家通道的实现复杂度减小了,专家间的通信开销降低。