1.6万亿参数的语言模型：谷歌大脑提出Switch Transformer，预训练速度可达T5的7倍

2021-01-14 03:18:08

机械之心报道

[本文来自：www.ii77.com]

编纂：魔王、杜伟、张倩 [原创文章：www.ii77.com]

方才，Google Brain 高级研究科学家 Barret Zoph 发帖透露，他们设计了一个名叫「Switch Transformer」的简化稀少架构，能够将说话模型的参数量扩展至 1.6 万亿（GPT-3 是 1750 亿）。在较量资源沟通的情形下，Switch Transformer 的练习速度能够达到 T5 模型的 4-7 倍。

在深度进修范畴，模型平日会对所有输入重用沟通的参数。但 Mixture of Experts (MoE，夹杂专家) 模型是个破例，它们会为每个输入的例子选择分歧的参数，究竟获得一个稀少激活模型——固然参数量惊人，但较量成本恒定。

今朝，MoE 模型已在机械翻译范畴取得了令人瞩目的成就，但因为模型复杂度高、通信成本高、练习不敷不乱，其普遍应用受到了必然的阻碍。

为认识决这些问题，Google Brain 的研究者提出了 Switch Transformer。在 Switch Transformer 的设计中，它们简化了 MoE 的路由算法（routing algorithm），设计了直观的改善模型，新模型的通信成本和较量成本都大大降低。此外，他们提出的练习手艺还提高了练习的不乱性，首次表明大型稀少模型也能够用低精度（bfloat16）进行练习。

论文链接：https://arxiv.org/pdf/2101.03961.pdf
代码链接：https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/moe.py

研究者还将新模型与 T5-Base 和 T5-Large 进行了对比，究竟表明，在沟通的较量资源下，新模型实现了最高 7 倍的预练习速度提拔。这一改善还能够扩展至多说话设置中，在所有的 101 种说话中都测到了新模型相对于 mT5-Base 版本的机能提拔。

最后，研究者在 Colossal Clean Crawled Corpus 长进行预练习，将说话模型的参数量提拔至上万亿，且比拟 T5-XXL 模型实现了 4 倍加快。

研究者还透露，固然这项工作着眼于规模，但它也表明，Switch Transformer 架构不光在具备超等较量机的情况下具有优势，在只有几个较量焦点的较量机上也是有效的。此外，研究者设计的大型稀少模型能够被蒸馏为一个小而浓密的版本，同时还能保留稀少模型质量提拔的 30%。

Switch Transformer 的设计道理

Switch Transformer 的首要设计原则是，以一种简洁且较量高效的体式最大化 Transformer 模型的参数量。Kaplan 等人（2020）已经对扩展的效益进行了详尽的研究，揭示了随模型、数据集巨细以及较量预算转变的幂定律缩放。主要的是，该研究首倡在相对较少数据上练习大型模型，将其作为较量最优方式。

基于这些，研究者在增加参数量的同时连结每个示例的 FLOP 不变。他们假设参数量与执行的总较量量无关，是能够零丁缩放的主要组件。所以，研究者经由设计一个稀少激活的模型来实现这一方针，该模型可以高效地行使 GPU 和 TPU 等为浓密矩阵乘法设计的硬件。

在分布式练习设置中，模型的稀少激活层在分歧设备上分派独一的权重。所以，模型权重随设备数量的增加而增加，同时在每个设备上连结可治理的内存和较量空间。

Switch Transformer 的编码器块如下图 2 所示：

简化稀少路由

Shazeer 等人（2017）提出了一个天然说话 MoE 层，它以 token 表征 x 为输入，然后将其发送给最果断的 top-k 专家（从 N 个专家构成的 {E_i(x)} ^N_i=1 鸠合中选出）。他们假设将 token 表征发送给 k>1 个专家是需要的，如许能够使 routing 函数具备有意义的梯度。他们认为若是没有对比至少两个专家的能力，则无法进修路由。

与这些设法分歧，谷歌大脑这项研究采用简化策略，只将 token 表征发送给单个专家。研究表明，这种简化策略连结了模型质量，降低了路由较量，而且机能更好。研究者将这种 k=1 的策略称为 Switch 层。

下图 3 展示了具有分歧专家容量因子（expert capacity factor）的路由示例：

高效稀少路由

研究者使用了 Mesh-Tensorflow 库 (MTF)，它具有雷同于 TensorFlow 的语义和 API，可促进高效分布式数据和模型并行架构。研究者在设计模型时考虑到了 TPU，它需要静态巨细。

分布式 Switch Transformer 实现：所有张量外形在编译时均获得静态确定，但因为练习和揣摩过程中的路由决议，较量是动态的。鉴于此，一个主要的手艺难题显现了：若何设置专家容量？

专家容量（每个专家较量的 token 数量）的较量体式为：每个批次的 token 数量除以专家数量，再乘以容量因子。如公式（3）所示：

若是将太多 token 发送给一个专家（下文称为「丢弃的 token」），则会跳过较量，token 表征经由残差保持直接传递到基层。但增加专家容量也不是没出缺点，数值太高将导致较量和内存虚耗。这当着的衡量如上图 3 所示。

实证研究发现，将丢弃的 token 比例连结在较低水平对于稀少专家模型的扩展很主要。设计决议对模型质量和速度的影响拜见下表 1。

Switch Transformer

研究者首先在 Colossal Clean Crawled Corpus (C4) 数据集上对 Switch Transformer 进行了预练习测试，使用了遮掩说话建模义务。在预练习设置中，他们遵循 Raffel 等人（2019）确定的最优方案，去掉了 15% 的 token，然后使用单个 sentinel token 来替代遮掩序列。为了对照模型机能，研究者供应了负对数疑心度的究竟。

Switch Transformer 与 MoE Transformer 的对照究竟如下表 1 所示。究竟表明，Switch Transformer 在速度 - 质量（speed-quality）根蒂上优于精心调整的浓密模型和 MoE Transformer，并在固定较量量和挂钟时间情形下取得了最佳究竟；Switch Transformer 的较量占用空间比 MoE Transformer 小；Switch Transformer 在低容量因子（1.0, 1.25）下示意更好。

提拔练习和微调的技能

与原版 Transformer 模型比拟，稀少专家模型在练习时或者加倍难题。所有这些层中的 hard-swithing（路由）决议都或者导致模型的不不乱。此外，像 bfloat16 如许的低精度花样或者加剧 router 的 softmax 较量问题。研究者接纳了以下几种技能来战胜练习难题，并实现不乱和可扩展的练习。

对大型稀少模型使用可选择行精度（Selective precision with large sparse models）
为实现不乱性使用更小的参数初始化（Smaller parameter initialization for stability）
正则化大型稀少模型（Regularizing large sparse models）

预练习可扩展性

在预练习时代，研究者对 Switch Transformer 的可扩展性进行了研究。在此过程中，他们考虑了一个算力和数据都不受限制的机制。为了避免数据受限，研究者使用了大型 C4 数据库，里面包含 180B 的方针 token。在视察到收益递减之前，他们一向进行练习。

专家的数量是扩展模型最有效的维度。增加专家的数量几乎不会改变较量成本，因为模型只为每个 token 选择一个专家，这与专家的总体数量无关。router 必需基于更多的专家较量概率分布，但这是一个轻量级的较量成本 O(d_model × num experts)。个中，d_model 是层与层之间所传递的 token 的嵌入维度。在这一部门，研究者以固定的较量成本考虑基于步数和时间的可伸缩性。

基于步数的可扩展性

下图 4 展示了多个模型在练习步数恒定、专家数量增加时示意出的可扩展性提拔情形。从中能够视察到一个趋势：在连结每个 token 的 FLOPS 不变时，拥有更多的参数（专家）能够提高练习速度。

基于时间的可扩展性

如上图 4 所示，跟着专家数量的增加，模型的机能会络续提拔。固然模型的每个 token 拥有与基线近乎沟通的 FLOPS，但 Switch Transformers 会发生额外的跨设备通信成本，路由机制也会带来额外的较量开销。是以，在基于步数的设置中视察到的样本效率提拔未必能转化为时间上的模型质量提拔。这就引出了一个问题：

在练习时间和较量成本都固定的前提下，我们应该练习一个浓密模型照样稀少模型？

下图 5 和 6 解决了这个问题。图 5 展示了预练习模型质量随练习时间增加所发生的转变。在练习时间和较量成本都固定的情形下，Switch Transformer 的速度优势非常显着。在这种设置下，若是要达到相似的疑心度，Switch-Base 64 专家模型的练习时间仅为 T5-Base 模型的 1/7。

可扩展性 VS. 一个更大的浓密模型

上面的实验表明，一个较量不受限制的浓密模型已经被 Switch 超越。图 6 考虑了一种分歧的情形：若是把较量资源分给一个更大的浓密模型会怎么样？

为了验证这个问题，研究者将 Switch-Base 与更强的基线 T5-Large 进行了对比。实验究竟表明，尽管 T5-Large 每个 token 所用的 FLOPs 是 Switch-Base 的 3.5 倍，但后者的样本效率依然更高，并且速度是前者的 2.5 倍。此外，若是设计一个与 T5-Large 所需 FLOPs 沟通的 Switch 模型（Switch-Large），上述提拔还会加倍显着。

粗俗义务中的究竟

微调

这里使用的基线方式是经由高度调参、具备 223M 参数的 T5-Base 和具备 739M 参数的 T5-Large 模型。针对这两个模型，该研究作者设计了具备更多参数的 FLOP-matched Switch Transformer。

在多项天然说话义务中，Switch Transformer 带来了显著机能提拔。最显着的是 SuperGLUE，在该基准上 FLOP-matched Switch Transformer 比拟 T5-Base 和 T5-Large 的机能离别提拔了 4.4% 和 2%，在 Winogrande、closed book Trivia QA 和 XSum 上也显现了雷同情形。独一没有视察到机能提拔的基准是 AI2 推理挑战赛（ARC）数据集：在 ARC challenge 数据集上 T5-Base 的机能跨越 Switch-Base；在 ARC easy 数据集上，T5-Large 的机能跨越 Switch-Large。

整体而言，Switch Transformer 模型在多项推理和常识义务中带来了显著机能提拔。这解说该模型架构不只对预练习有效，还能够经由微调将质量改善迁徙至粗俗义务中。

蒸馏

布置具备十亿、万亿参数量的大型神经收集并非易事。为此，该论文研究了若何将大型稀少模型蒸馏为小型浓密模型。下表 7 展示了该研究所用的蒸馏手艺：

使用表 7 中最优的蒸馏手艺后，研究者将多个稀少模型蒸馏为浓密模型。他们对 Switch-Base 模型进行蒸馏，因为专家数量的分歧，其参数量在 11 亿至 147 亿之间。该研究能够将具备 11 亿参数量的模型压缩 82%，同时保留 37% 的机能提拔。最极端的情形下，将模型压缩了 99%，且维持了 28% 的机能提拔。

最后，研究者将微调稀少模型蒸馏为浓密模型。下表 9 展示了对 74 亿参数 Switch-Base 模型（该模型针对 SuperGLUE 义务进行了微调）的蒸馏究竟——223M T5-Base。与预练习究竟雷同，蒸馏后的模型仍保留 30% 的机能提拔。这或者有助于确定用于微调义务的特定专家并进行提取，从而获得更好的模型压缩。

多说话进修

鄙人游义务实验中，研究者权衡了模型质量和速度的衡量，模型在 101 种分歧说话长进行了预练习。下图 7 展示了 Switch T5 Base 模型与 mT5-Base 在所有说话上的质量提拔情形（负对数疑心度）。对两个模型经由 100 万步预练习后，Switch Transformer 的最终负对数疑心度相较基线有所提拔。

下图 8 展示了 Switch Transformer 相较 mT5-Base 的每一步加快情形，前者实现了平均 5 倍的加快，个中在 91% 的说话上实现了至少 4 倍加快。这表明 Switch Transformer 是高效的多义务和多说话进修器。

使用数据、模型和专家并行化来设计模型

随意地增加专家数量会显现收益递减问题（拜见上图 4），该研究介绍了一些增补性的扩展策略，涉及连系数据、模型与专家并行化的衡量。

连系数据、模型与专家并行化，构建万亿参数模型

Switch Transformer 设计过程中，研究者试图均衡 FLOPs per token 和参数量。当专家数量增加时，则参数量增加，但不改变 FLOPs per token。要想增加 FLOPs，则需增加 d_ff 维度（这也会带来参数量的增加，但相对较少）。这就是一种衡量：增加 d_ff 维度会导致每个焦点内存的耗尽，因而必需增加 m。但因为焦点 N 的数量是固定的 N = n × m，是以必需降低 n，也就是说需要使用更小的批巨细。

在连系模型并行化和专家并行化之后，发送 token 到准确的专家以及模型并行化导致的内部 all-reduce 通信会带来 all-to-all 通信成本。在连系这三种方式时，若何均衡 FLOPs、通信成本和每个焦点的内存变得非常复杂。

该研究连系数据、模型与专家并行化，设计了两个大型 Switch Transformer 模型，离别具备 3950 亿参数和 1.6 万亿参数，并研究了这些模型在上游预练习说话模型和粗俗微调义务中的机能。参数量、FLOPs 和分歧模型的超参数拜见下表 10：

关于 Switch Transformer 还有好多问题

在论文最后部门，谷歌大脑研究者商量了一些关于 Switch Transformer 和稀少专家模型的问题（这里稀少指的是权重，而不是注重力模式）。

问题 1：Switch Transformer 的机能更好吗？原因是否在于巨量参数？

机能切实更好，但原因不在参数量，而在于设计。参数有助于扩展神经说话模型，大模型的机能的确会好一些。然则该研究提出的模型在使用沟通较量资源的情形下具备更高的样本效率。

问题 2：没有超等较量机的情形下，我能使用该方式吗？

尽管这篇论文聚焦非常大型的模型，但研究者仍找到了具备两个专家的模型，既能提拔机能又能够轻松适应常用 GPU 或 TPU 的内存限制。是以，研究者认为该手艺可用于小规模设置中。

问题 3：在速度 - 正确率帕累托曲线上，稀少模型的示意优于浓密模型吗？

是的。在多种分歧模型规模情形下，稀少模型在每一步和墙上时钟时间方面都优于浓密模型。受控实验表明，对于固定的较量量和时间而言，稀少模型的示意跨越浓密模型。

问题 4：我无法布置万亿参数模型，能够将模型缩小吗？

固然无法完整维持万亿参数模型的质量，但经由将稀少模型蒸馏为浓密模型，可实现 10-100 倍的压缩率，同时获得专家模型约 30% 的质量改善。

问题 5：为什么要使用 Switch Transformer 取代模型并行浓密模型？

以时间为基准，Switch Transformer 要比使用分片参数（sharded parameter）的浓密模型高效得多。同时，这一选择并非互斥，Switch Transformer 中也能够使用模型并行化，这能够提高 FLOPs per token，但也会导致传统模型并行化的减速。

问题 6：为什么稀少模型未获得普遍使用？

测验稀少模型的设法被浓密模型的伟大成功所阻挠。而且，稀少模型面临着多个问题，包罗模型复杂度、练习难度、通信成本等。而 Switch Transformer 缓解了这些问题。

AAAI 2021线上分享 | 行使多轮阅读懂得的框架解决实体链接问题

在华为云被 AAAI 2021 会议领受的论文中，研究者提出行使多轮阅读懂得的框架来解决短文本的实体链接问题，并在多个中英文数据集上均取得了当前 SOTA 的实体链接结果。

1月14日20:00，华为云语音语义立异 Lab 算法工程师Xavier为人人具体解读此前沿研究。

添加机械之心小助手（syncedai5），备注「AAAI」，进群一路看直播。

转载请关联本公家号获得授权

投稿或追求报道：content@jiqizhixin.com

上一篇：VR+行业|疫情之下变身“工具人”的VR，并非所有行业都适合
下一篇：常州电信用户速领，至高888元话费

1.6万亿参数的语言模型：谷歌大脑提出Switch Transformer，预训练速度可达T5的7倍

热门文章

小编推荐

推荐文章