[论文解读] Sparse MoE as the New Dropout: Scaling Dense and Self-Slimmable Transformers

Tianlong Chen, Zhenyu Zhang+3 more|arXiv (Cornell University)|Mar 2, 2023

Domain Adaptation and Few-Shot Learning被引用 6

一句话总结

本文提出 SMoE-Dropout，一种即插即用的训练框架，通过使用固定且随机初始化的路由器，在训练过程中逐步增加激活专家的数量，使 Transformer 模型在不发生崩溃的情况下实现更高的准确率。该方法诱导出一种‘自瘦身’特性，使模型在推理和微调过程中可根据可用资源实现平滑的性能提升，其中 BERT 模型在推理任务上的表现相比稠密模型最高提升 1.09%。

ABSTRACT

Despite their remarkable achievement, gigantic transformers encounter significant drawbacks, including exorbitant computational and memory footprints during training, as well as severe collapse evidenced by a high degree of parameter redundancy. Sparsely-activated Mixture-of-Experts (SMoEs) have shown promise to mitigate the issue of training efficiency, yet they are prone to (1) redundant experts due to representational collapse; and (2) poor expert scalability for inference and downstream fine-tuning, primarily due to overfitting of the learned routing policy to the number of activated experts during training. As recent research efforts are predominantly focused on improving routing policies to encourage expert specializations, this work focuses on exploring the overlooked scalability bottleneck of SMoEs and leveraging it to effectively scale dense transformers. To this end, we propose a new plug-and-play training framework, SMoE-Dropout, to enable scaling transformers to better accuracy in their full capacity without collapse. Specifically, SMoE-Dropout consists of a randomly initialized and fixed router network to activate experts and gradually increases the activated expert number as training progresses over time. Transformers trained by SMoE-Dropout naturally exhibit a self-slimmable property subject to resource availability, offering smooth and consistent performance boosts with an increase in activated experts during inference or fine-tuning. Our extensive experiments demonstrate the superior performance and substantial computation savings of SMoE-Dropout, compared to dense training baselines with equivalent parameter counts. In particular, our trained BERT outperforms its densely trained counterpart with consistent improvements of {1.03%, 0.78%, 1.09%} on challenging reasoning tasks {ASDiv-A, MAWPS, SVAMP}, respectively.

研究动机与目标

• 解决稀疏专家混合模型（SMoE）在推理和微调过程中存在的表征崩溃与可扩展性差的问题。
• 探索 SMoE 训练中被忽视的可扩展性瓶颈，特别是路由策略对固定专家数量的过拟合问题。
• 开发一种即插即用的方法，实现全容量模型扩展而不会崩溃，确保随着专家使用量的增加，性能持续提升。
• 通过‘自瘦身’Transformer 架构，提供动态、资源感知的推理与微调能力。

提出的方法

• 引入一个固定且随机初始化的路由器网络，其在训练过程中不更新，替代可学习的路由机制。
• 在训练过程中逐步增加激活专家的数量（k），实现容量的渐进式扩展。
• 将 SMoE 模块插入 Transformer 架构的深层，实证表明其性能优于早期层插入。
• 通过模块化方式将 SMoE 模块插入特定的 Transformer 层，消融实验表明深层插入最为有效。
• 采用随机路由策略，避免过拟合，并通过隐式正则化促进专家专业化。
• 通过根据资源可用性调整激活专家数量，实现在推理和微调过程中平滑的性能扩展。

实验结果

研究问题

• RQ1 固定且随机的路由策略是否能缓解稀疏专家混合（SMoE）模型中的表征崩溃问题，并提升其可扩展性？
• RQ2 在训练过程中逐步增加激活专家数量，是否能带来更好的泛化能力，并在推理和微调中实现性能提升？
• RQ3 SMoE-Dropout 与可学习路由及稠密训练相比，在过拟合和模型容量利用率方面表现如何？
• RQ4 在 Transformer 中插入 SMoE 模块的最优位置在哪里，以最大化性能与可扩展性？
• RQ5 SMoE-Dropout 是否能生成更适合资源受限环境部署的蒸馏模型，相比标准 SMoE 或稠密模型？

主要发现

• SMoE-Dropout 训练的 BERT 模型在 ASDiv-A、MAWPS 和 SVAMP 推理任务上，分别相比其稠密基线模型实现了 1.03%、0.78% 和 1.09% 的一致性能提升。
• 该方法诱导出一种‘自瘦身’特性，即在推理和微调过程中，随着激活专家数量的增加，性能可实现平滑提升，从而支持动态的效率-性能权衡。
• 在 enwik8 测试集上，SMoE-Dropout 在使用 16 个专家时达到 2.96×10⁻² 的比特每字符（BPC），显著优于完全随机和确定性哈希路由基线。
• 即使在 12 层或 16 个专家的情况下，SMoE-Dropout 模型也未表现出过拟合迹象，而可学习 SMoE 模型则因过拟合导致性能下降。
• 蒸馏实验表明，SMoE-Dropout 模型在将知识蒸馏到单专家模型时，在 SST-2 任务上相比稠密模型或标准 SMoE 模型，准确率高出 0.76% 至 1.89%。
• 消融研究证实，逐步增加 k 是实现自瘦身特性的关键，若禁用此调度策略，即使使用全部参数，性能仍会下降。