稀疏模型在深度学习领域发挥着越来越重要的作用。对于给定的 token 或样本，它可以只激活模型的一小部分，从而在拥有很大的参数量的同时也能做到计算友好。但是，如何可靠地训练这类模型依然是一个需要解决的问题。在这篇文章中，来自谷歌的 Barret Zoph、Irwan Bello、William Fedus 、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）可以提高质量。
更多细节请参考原论文。
