《Switch Transformers：用简单高效的稀疏性扩展到万亿参数模型》

不把自己全部用上的模型

大多数语言模型是稠密的：输入的每个词元都要经过模型的全部参数，因此模型变大，处理每一条输入的开销也跟着变大。混合专家模型（Mixture of Experts，MoE）走的是另一条路。它为每一条进来的样本挑选不同的参数，从而得到一个稀疏激活（sparsely activated）的模型：参数量可以非常大，但计算成本保持恒定^[1]。对任何一条输入来说，网络里只有一小部分被启用，而启用哪一部分是按样本决定的，不是事先固定死的。

这个思路有吸引力，因为它把稠密模型绑在一起的两件事拆开了：一是模型「知道」多少东西，大致由参数量衡量；二是处理一条样本要做多少运算，用浮点运算次数衡量。不过，MoE 模型此前并没有被广泛采用，有三个问题挡在路上：路由算法复杂、通信代价高、训练不稳定^[1]。

论文做了什么

William Fedus、Barret Zoph 与 Noam Shazeer 于 2021 年 1 月 11 日把《Switch Transformers》提交到 arXiv^[1]。论文保留了稀疏激活的思路，同时直接针对上述障碍：它把 MoE 的路由算法简化，使此前阻碍这一路线推广的复杂性、通信成本和训练不稳定变得可控^[1]。

第二个贡献在数值格式上。论文提出的训练技术，让大型稀疏模型第一次可以用更低精度的格式训练，也就是 bfloat16^[1]。bfloat16 是一种 16 位浮点格式，它保留了训练中常用的 32 位格式的指数范围，只存很少的尾数位。由于它把激活值占用的内存和需要在设备之间搬运的数据量都减半，对稀疏模型尤其重要——这类模型的成本主要不在算术上，而在参数和激活值的搬运上。

结果

论文给出的对比对象是 T5 系列，即作为稠密基线的文本到文本模型。基于 T5-Base 和 T5-Large 的模型在相同计算资源下，预训练速度最高提升到 7 倍^[1]。这一优势不限于英语：在与 mT5-Base 对比时所测的全部 101 种语言上，增益都成立^[1]。

作者还把语言模型的规模继续往上推，在 Colossal Clean Crawled Corpus（C4，一个大型公开网页文本语料）上预训练了参数量最高达万亿的模型，相对 T5-XXL 取得 4 倍加速^[1]。论文中最大的模型 Switch-C 有 15,710 亿（1,571 billion）个参数，而每个序列只用 890 GFLOPs^[1]。1 GFLOP 是 10 亿次浮点运算，所以这个数字是把一条序列送过模型要付的算术账单——它由网络被点亮多少决定，而不由网络存了多少参数决定。

为什么重要

把这些数字放在一起，能看出一笔交换：参数量涨了几个数量级，而每条序列的成本仍留在同一量级，因为对任何一条样本来说只用到了其中一部分参数^[1]。这正是万亿参数模型能够训练的前提，也是为什么论文对路由的简化是核心动作，而不是实现细节。更简单的路由决策算起来更便宜，也更容易分摊到大量设备上，从而降低了让早期 MoE 系统难以训练的通信开销^[1]。

Switch-C 把这个练习的终点摆了出来：一个拥有 15,710 亿参数的模型，每条序列仍然只花 890 GFLOPs^[1]。这类模型是否在每个方面都更好，论文并没有这样断言；它展示的是一套可行的做法——参数量增长而每条样本的计算量不增长，以及一组让大型稀疏模型可以用 bfloat16 训练的技术^[1]。简化后的路由与低精度训练这两块，是后来稀疏模型与混合专家语言模型继续沿用的部分。