Jeff Dean：我们写了一份「稀薄模型计划指南」，请查收

稀薄大师神经网络展示了纯规模的优势，并为当今常用的静态神经网络架构提供了一种有效的替代方案。稀薄大师网络不是对所有输入应用相同的参数，而是为每个输入动态选择运用哪些参数。这允许网络极大地扩展参数的数量，同时保持每个 token 的 FLOPs 大致不变。这些方式的采用已经带来了 SOTA 翻译模型、4-7 倍的预训练加速，以及仅运用 1/3 的训练成本就能达到 GPT-3 级的 one-shot 机能。尽管参数数量惊人，但稀薄模型将训练大型神经网络的碳足迹降低了一个数量级。然而，困难依然存在。Fedus et al. (2021)观察到，与之前的 SOTA 方式（Raffel et al., 2019）相比，稀薄 1.6T 参数模型实现了 4 倍的预训练加速，但在 SuperGLUE 等常用基准上进行微调时，却落后于较小的模型。在 Artetxe et al. (2021)中，研究者在域外数据上对 MoE 语言模型进行了微调，并观察到了相似的差距。为了处理这一问题，Switch-XXL 模型被提出，该模型参数较少，但计算占用空间增加到原来的 8 倍（FLOPs 大约等于最大的 T5 模型），在自然语言理解任务上的机能有所提高。然而，必要的预训练受到先前在小规模研究中未检测到的训练不稳定性的阻碍。这些不稳定性后来在其他稀薄模型中被识别出来。这些结果揭示了参数和计算的必要平衡，但如何可靠地训练这种模型依然是一个待处理的问题。这篇论文的目的就是提高稀薄模型的实用性和可靠性。他们研究了这两个问题，并给出了计划指南。最后，他们将稀薄模型的参数缩放到 269B，其计算成本与 32B 麇集编码器 – 解码器 Transformer（稳定、可迁移的 Mixture-of-Experts、ST-MoE-32B）相当。这是稀薄模型首次在迁移学习中实现 SOTA 机能，跨越了一系列不同的任务，包括推理（SuperGLUE、ARC Easy、ARC Challenge)、摘要（XSum、CNN-DM）、闭卷问答（WebQA、Natural Questions）和对抗式构造任务（Winogrande、ANLI R3）。

本文的贡献可以概括如下：1、开展了一项关于稳定性技术的质量 – 稳定性权衡（quality-stability trade-offs）大规模研究；2、引入了 router z-loss 来处理稳定性问题，同时略微提高了模型质量；3、给出了关于稀薄和麇集模型的微调分析，揭示了二者对批大小和学习率的不同超参数敏感性；他们发现，糟糕的超参数导致麇集模型上几乎没有微调增益，尽管预训练有很大的加速；4、给出了分布式环境下计划 Pareto 高效稀薄模型的架构、routing 和模型计划原则；5、给出了追踪跨大师层的 token routing 决策的定性分析；6、训练出了一个 269B 稀薄模型，在一组不同的自然语言基准上实现了 SOTA 机能。router z-loss稳定神经网络最成功的方式之一是对激活的约束和梯度。一种流行的方式是在通过深度网络反向传播时，裁剪梯度范数来弥补爆炸梯度。在这篇论文中，研究者运用 Adafactor 优化器是因为它的内存效率（尽管最近推出的 8 位优化器 (Dettmers et al., 2021) 可能会提供更好的 trade-off）。Adafactor 运用更新裁剪（update clipping），而不是梯度裁剪（gradient clipping），其中对权重的更改被限制在一定的范数以下。他们尝试将更新裁剪收紧到更小的值。接下来，他们研究了即将进入 router 的 logit 上的约束。router 以 float32 计算大师的概率分布。然而，研究者发现，在最大的规模下，这不足以带来可靠的训练结果。为了处理这个问题，他们引入了 router z-loss，其中，B 是 token 的数目，N 是大师数，x ∈ RB×N 是将要进入 router 的 logit。下表 4 显示，在三次运行中，update clipping 和 router z-loss 都稳定了模型，但是 update clipping 严重影响了模型的质量。因此，研究者运用 z-loss 方式来固定模型稳定性。

router z-loss 引入了另一个超参数 (c_z)，这是一个加权系数，作为优化的总丧失的一部分。总丧失是交叉熵丧失 (cross entropy loss, L_CE)、辅助负载平衡丧失 (auxiliary load balance loss, L_B) 和 router z-loss (L_Z) 的线性加权组合。

基于用超参数扫描进行预训练后的最佳模型质量，研究者选择 c_z = 0.001 的值。附录 B 记录了预训练过程中的丧失。稀薄模型的计划麇集模型的计划受到 Kaplan et al. (2020) 的基础工作的指导。但是到了稀薄模型这里，还有无数的额外问题需要处理，比如：(1)运用多少大师？(2)运用哪种 routing 算法？(3)容量因子（capacity factor）的值是多少？(4)硬件如何改变这些决策？在本文中，研究者给出的建议是：1、在他们的设置中，他们推荐容量因子为 1.25 的 top-2 routing，每个核心最多有一个大师；2、在评估期间，可以更改容量因子，以适应新的内存 / 计算要求；3、麇集层叠加和乘法偏置（multiplicative bias）可以提高质量。更多细节请参考原论文。