[论文解读] ST-MoE: Designing Stable and Transferable Sparse Expert Models

Barret Zoph, Irwan Bello+6 more|arXiv (Cornell University)|Feb 17, 2022

Topic Modeling被引用 46

一句话总结

本文提出 ST-MoE，一种大规模稀疏专家模型设计指南，引入路由器 z-loss 以稳定训练，并在多样化的 NLP 任务上展示了 269B 参数稀疏模型（ST-MoE-32B）的最先进的迁移性能。

ABSTRACT

Scale has opened new frontiers in natural language processing -- but at a high cost. In response, Mixture-of-Experts (MoE) and Switch Transformers have been proposed as an energy efficient path to even larger and more capable language models. But advancing the state-of-the-art across a broad set of natural language tasks has been hindered by training instabilities and uncertain quality during fine-tuning. Our work focuses on these issues and acts as a design guide. We conclude by scaling a sparse model to 269B parameters, with a computational cost comparable to a 32B dense encoder-decoder Transformer (Stable and Transferable Mixture-of-Experts or ST-MoE-32B). For the first time, a sparse model achieves state-of-the-art performance in transfer learning, across a diverse set of tasks including reasoning (SuperGLUE, ARC Easy, ARC Challenge), summarization (XSum, CNN-DM), closed book question answering (WebQA, Natural Questions), and adversarially constructed tasks (Winogrande, ANLI R3).

研究动机与目标

• 研究稀疏 MoE 模型训练中的不稳定性并识别稳定性-保持的设计选择。
• 开发一个稳定友好的路由机制（router z-loss），在不牺牲质量的前提下改善训练。
• 分析稀疏模型的微调协议，并与密集模型对比以理解泛化。
• 提供架构与训练指南，以设计适用于大规模分布式训练的Pareto高效稀疏模型。

提出的方法

• 在多尺度设置下回顾并量化稀疏 MoE 训练中的稳定性风险。
• 引入 router z-loss 以约束路由器 logits，稳定端到端训练且不损害模型质量。
• 评估精度格式与数值舍入对 MoE 路由的影响并提供缓解见解。
• 系统性研究微调方案，包括参数子集、 dropout 与哨兵标记，以提升泛化。
• 提出关于专家数量、容量因子与路由算法的设计选择，以获得 Pareto 高效的稀疏模型。
• 对一个 269B 的稀疏模型（ST-MoE-32B）进行预训练与评估，在多样化的 NLP 基准上显示最先进的迁移性能。

实验结果

研究问题

• RQ1 稳定性技术是否能在不牺牲模型质量的前提下稳定稀疏 MoE 训练？
• RQ2 面向路由器的损失（router z-loss）如何影响训练稳定性与最终性能？
• RQ3 与密集模型相比，哪些微调协议对 ST-MoE 模型在 NLP 任务中的泛化效果最好？
• RQ4 在大规模阶段，哪些架构与路由设计原则能够实现Pareto高效的稀疏模型？

主要发现

模型	训练 CF	评估 CF	辅助损失	丢弃的令牌百分比	SuperGLUE（↑）
稀疏	0.75	2.0	是	10.6%	86.5 ± 0.21
稀疏	1.25	2.0	是	0.3%	86.7
稀疏	2.0	3.0	是	0.0%	85.8
稀疏	4.0	5.0	是	0.0%	86.4
稀疏	0.75	2.0	否	15.6%	85.7
稀疏	1.25	2.0	否	2.9%	85.8
稀疏	2.0	3.0	否	0.4%	85.9
稀疏	4.0	5.0	否	0.0%	86.4

• 一项大规模稳定性研究显示，许多稳定化措施以稳定性换取质量，而 router z-loss 在不损失质量的情况下稳定了训练。
• router z-loss 提高稳定性并且可略微提升质量，与对更新幅度进行强力裁剪的做法不同，后者会损害表现。
• 稀疏模型需要与密集模型不同的微调协议，稀疏模型从较小的批量大小和较高的学习率中获益。
• 稀疏模型在微调过程中对令牌丢失显示鲁棒性，在某些容量设置下仍可在非零令牌丢失时保持性能。
• ST-MoE-32B 在包括 SuperGLUE、ARC Easy/Challenge、XSum、CNN-DM、WebQA、Natural Questions、Winogrande 以及 ANLI R3 等任务上实现了最先进的迁移性能。
• 本文提供了实践性设计指南，用于设定专家数量、容量因子与路由策略，以实现 Pareto 效率。