[论文解读] Task-Specific Expert Pruning for Sparse Mixture-of-Experts

Tianyu Chen, Shaohan Huang+6 more|arXiv (Cornell University)|Jun 1, 2022

Indoor and Outdoor Localization Technologies被引用 9

一句话总结

本文提出了一种针对特定任务的专家剪枝方法，将稀疏的专家混合模型（MoE）转换为高效的单专家稠密模型，以用于下游任务。通过在微调过程中基于贡献度分数逐步剪枝非专业专家，该方法在保持99.3%的MoE性能的同时，实现了2倍的推理加速，并消除了设备间通信开销。

ABSTRACT

The sparse Mixture-of-Experts (MoE) model is powerful for large-scale pre-training and has achieved promising results due to its model capacity. However, with trillions of parameters, MoE is hard to be deployed on cloud or mobile environment. The inference of MoE requires expert parallelism, which is not hardware-friendly and communication expensive. Especially for resource-limited downstream tasks, such sparse structure has to sacrifice a lot of computing efficiency for limited performance gains. In this work, we observe most experts contribute scarcely little to the MoE fine-tuning and inference. We further propose a general method to progressively drop the non-professional experts for the target downstream task, which preserves the benefits of MoE while reducing the MoE model into one single-expert dense model. Our experiments reveal that the fine-tuned single-expert model could preserve 99.3% benefits from MoE across six different types of tasks while enjoying 2x inference speed with free communication cost.

研究动机与目标

• 为解决稀疏MoE模型在资源受限的下游部署中因专家并行和通信开销导致的效率低下问题。
• 探究在预训练MoE模型中，是否仅最擅长的专家对特定下游任务有显著贡献。
• 开发一种剪枝策略，以识别并保留最擅长该任务的专家，同时剔除其他专家，从而实现高效推理。
• 证明剪枝后的单专家模型优于稠密基线模型，并可匹配或超越完整MoE微调后的性能。
• 优化剪枝的时间点和标准，以最大化最终模型性能和推理效率。

提出的方法

• 基于固定长度滑动训练窗口内计算的熟练度分数对专家进行剪枝。
• 在微调过程中定期应用一个下钻阈值；低于该阈值的专家将被移除。
• 剪枝过程持续进行，直到每个MoE层仅剩一个专家，或直到微调进行到一半。
• 提出一种“激进”剪枝策略，即在训练中点时移除非专业专家，以使第二阶段能集中优化选定的专家。
• 在训练剩余时间中，对最终选定的专家进行独立微调，以最大化性能。
• 该方法避免知识蒸馏，直接将预训练MoE模型的知识迁移至单个稠密专家。

实验结果

研究问题

• RQ1 能否从预训练的MoE模型中识别出一个专家，使其在特定下游任务上保留完整MoE架构的大部分性能优势？
• RQ2 在微调过程中，剪枝非专业专家的最佳时机和阈值是什么，以最大化最终模型的准确率？
• RQ3 训练窗口长度如何影响所选专家的性能？
• RQ4 剪枝非专业专家是否能在不损失性能的前提下，实现更快的推理速度和更低的通信开销？
• RQ5 剪枝后的单专家模型是否能超越从零开始微调的稠密模型？

主要发现

• 剪枝后的单专家模型在六个多样化的下游任务上，保留了完整MoE模型99.3%的性能增益。
• 在MNLI任务上，剪枝模型的推理速度达到完整MoE模型的2倍；在SQuAD任务上，推理速度提升了1.28倍。
• 尽管所选专家中包含子-MoE层，剪枝模型的推理效率仍达到稠密模型的80%。
• ‘激进’剪枝策略（在训练中点时剪枝非专业专家）相比分阶段或两阶段方法，能获得更优的推理性能。
• 该方法完全消除了设备间通信开销，使其在云环境和移动设备部署中极具适用性。
• 所选专家在训练初期即主导贡献，这为早期剪枝提供了依据，使第二阶段可集中优化，从而提升性能。