[论文解读] Residual Mixture of Experts

Lemeng Wu, Mengchen Liu+4 more|arXiv (Cornell University)|Apr 20, 2022

Domain Adaptation and Few-Shot Learning被引用 8

一句话总结

该论文提出残差专家混合模型（Residual Mixture of Experts, RMoE），一种通过将专家权重分解为与输入无关的主干和与输入相关的残差部分，从而高效扩展视觉变换器的训练流程。通过仅微调残差部分并仅增加极少计算量，RMoE在训练成本增加不足3%的情况下，实现了与完整MoE训练相当的性能，优于Swin-T、CvT-13和Swin-L等非MoE模型，在ADE20K上mIoU提升1.1–1.0，在MS-COCO上AP提升0.6–1.6。

ABSTRACT

Mixture of Experts (MoE) is able to scale up vision transformers effectively. However, it requires prohibiting computation resources to train a large MoE transformer. In this paper, we propose Residual Mixture of Experts (RMoE), an efficient training pipeline for MoE vision transformers on downstream tasks, such as segmentation and detection. RMoE achieves comparable results with the upper-bound MoE training, while only introducing minor additional training cost than the lower-bound non-MoE training pipelines. The efficiency is supported by our key observation: the weights of an MoE transformer can be factored into an input-independent core and an input-dependent residual. Compared with the weight core, the weight residual can be efficiently trained with much less computation resource, e.g., finetuning on the downstream data. We show that, compared with the current MoE training pipeline, we get comparable results while saving over 30% training cost. When compared with state-of-the-art non- MoE transformers, such as Swin-T / CvT-13 / Swin-L, we get +1.1 / 0.9 / 1.0 mIoU gain on ADE20K segmentation and +1.4 / 1.6 / 0.6 AP gain on MS-COCO object detection task with less than 3% additional training cost.

研究动机与目标

• 解决在下游任务中训练大规模专家混合（MoE）视觉变换器带来的高计算成本问题。
• 在无需在大规模上游数据集上重新训练的前提下，实现视觉变换器的高效扩展。
• 在保持接近MoE训练上限性能的同时，降低MoE模型的训练成本。
• 利用公开的非MoE检查点作为大规模MoE适配的初始化。
• 证明专家权重可被分解为核心部分与可训练残差部分，以实现高效微调。

提出的方法

• 该方法将MoE权重分解为与输入无关的核心部分和与输入相关的残差部分，其中核心部分从预训练的非MoE模型中固定获取。
• 通过在下游数据上使用低成本更新方式对残差部分进行微调，避免在上游数据上重新训练完整MoE模型。
• 在中间微调阶段使用梯度截断技术，防止性能下降。
• 对复制的权重应用噪声初始化尺度，以促进专家专业化。
• 采用基于梯度得分的MoE层选择策略，优于Last-2和Every-2方法。
• 该流程支持中间微调和下游微调，对推理FLOPs和延迟的影响极小。

实验结果

研究问题

• RQ1 是否可以将已训练MoE变换器的权重分解为核心部分与残差部分，以实现高效微调？
• RQ2 仅微调MoE专家的残差部分是否能实现与完整MoE训练相当的性能，同时显著降低计算成本？
• RQ3 现有的非MoE视觉变换器检查点是否可有效用作大规模MoE模型的初始化？
• RQ4 MoE层选择策略（如Every-Last、梯度得分）的选择如何影响下游性能？
• RQ5 哪些超参数（如top-k、专家数量、噪声尺度）能在最小额外成本下实现最优性能？

主要发现

• 与Swin-T相比，RMoE在ADE20K上实现+1.1 mIoU提升，在MS-COCO上实现+1.4 AP提升，训练成本增加不足3%。
• 与非MoE模型相比，RMoE在ADE20K上分别实现+0.9 mIoU（CvT-13）和+1.0 mIoU（Swin-L）的性能提升，成本增加可忽略。
• 基于梯度得分的MoE层选择策略优于Last-2和Every-2，性能接近Every-Last，同时更具效率。
• 使用top-1门控机制搭配8个专家可获得最优性能，而增加至16个专家无法提升结果，因负载均衡损失产生干扰。
• 在中间微调阶段应用梯度截断可防止性能下降，且在ADE20K上采用0.01的噪声尺度可获得最佳结果。
• 推理时间和FLOPs仅略有增加，证实RMoE保持了MoE模型的效率优势。