[论文解读] Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints

Aran Komatsuzaki, Joan Puigcerver+7 more|arXiv (Cornell University)|Dec 9, 2022

Domain Adaptation and Few-Shot Learning被引用 12

一句话总结

本文提出稀疏再利用（sparse upcycling），一种将已训练的密集Transformer模型转化为更强大的专家混合模型（Mixture-of-Experts, MoE）的方法，通过从密集检查点初始化MoE层。该方法在SuperGLUE和ImageNet上达到最先进性能，优于原始密集模型和从零开始训练的MoE模型，且仅使用原始预训练计算预算的约50%。

ABSTRACT

Training large, deep neural networks to convergence can be prohibitively expensive. As a result, often only a small selection of popular, dense models are reused across different contexts and tasks. Increasingly, sparsely activated models, which seek to decouple model size from computation costs, are becoming an attractive alternative to dense models. Although more efficient in terms of quality and computation cost, sparse models remain data-hungry and costly to train from scratch in the large scale regime. In this work, we propose sparse upcycling -- a simple way to reuse sunk training costs by initializing a sparsely activated Mixture-of-Experts model from a dense checkpoint. We show that sparsely upcycled T5 Base, Large, and XL language models and Vision Transformer Base and Large models, respectively, significantly outperform their dense counterparts on SuperGLUE and ImageNet, using only ~50% of the initial dense pretraining sunk cost. The upcycled models also outperform sparse models trained from scratch on 100% of the initial dense pretraining computation budget.

研究动机与目标

• 通过重用现有的密集模型检查点，解决从零开始训练大型稀疏模型的高计算成本问题。
• 在不引入新数据或完整重训练的前提下，实现大型模型的性能提升。
• 开发一种实用方法，以极小的额外计算成本，将密集模型升级为更高效、高容量的MoE模型。
• 证明在相同计算预算下，从密集检查点初始化的MoE模型能够超越原始密集模型和从零开始训练的MoE模型。

提出的方法

• 通过将预训练的密集Transformer检查点作为MoE层的热启动，初始化一个专家混合（MoE）模型。
• 将Transformer块中的MLP组件替换为一组专家前馈网络和一个路由器网络，以实现稀疏激活。
• 使用专家选择路由机制，动态为每个token选择最相关的专家，实现高效计算。
• 应用一种模型手术方案，包括对MoE组件进行仔细初始化和适应，以最小化结构变更带来的性能下降。
• 使用适度的额外计算预算（通常为原始密集模型预训练成本的10%至60%）训练再利用后的MoE模型。
• 使用与原始密集模型相同的训练数据和调度策略，确保不引入新数据。

实验结果

研究问题

• RQ1 能否以极小的额外计算成本，有效将预训练的密集模型再利用为性能更强的MoE模型？
• RQ2 从密集检查点初始化MoE层是否能带来优于继续训练原始密集模型的性能？
• RQ3 在总计算预算相同的情况下，再利用的MoE模型是否能超越从零开始训练的MoE模型？
• RQ4 与持续密集训练相比，实现稀疏再利用性能提升所需额外计算量是多少？
• RQ5 该再利用方法是否在不同模型架构和领域（如自然语言处理和视觉）中具有泛化能力？

主要发现

• 再利用的T5-Base、T5-Large和T5-XL模型在SuperGLUE上比其密集基线模型高出1.5–2个绝对点，仅使用原始预训练计算量的46–55%。
• 再利用的ViT-Base和ViT-Large模型在ImageNet 10-shot上实现1%以上的准确率提升，额外训练时间仅增加13%，而持续密集训练则需额外58%的训练时间。
• 在相同总计算预算下，稀疏再利用的性能优于从零开始训练的MoE模型，证明了重用预训练权重的价值。
• 对于ViT-Base，仅增加13%的额外计算量即可在ImageNet 10-shot上实现1%以上的准确率提升，而持续密集训练需额外58%的计算量才能达到相同效果。
• 该方法在多种模型尺寸和领域中均有效：在语言和视觉任务上，再利用模型在计算受限条件下均超越了密集模型和从零开始训练的MoE基线。
• 再利用方法具有鲁棒性和泛化能力，在不同模型架构和任务（包括SuperGLUE和ImageNet）中均表现出一致的性能提升。