[論文レビュー] 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) を提案する。入力に依存する残差のみを微調整することで、最小限の計算量で、フルMoE訓練と同等の性能を達成しつつ、3%未満の追加訓練コストに抑える。ADE20KではSwin-T、CvT-13、Swin-Lより1.1–1.0 mIoU高い性能を発揮し、MS-COCOでは0.6–1.6 APの向上を達成した。

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.

研究の動機と目的

• 大規模なMixture of Experts (MoE) 視覚変換器の微調整タスクにおける高い計算コストを軽減すること。
• 大規模な上流データセットにおけるフル再訓練を必要とせず、視覚変換器の効率的スケーリングを可能にすること。
• MoEモデルの訓練コストを低減しながら、上限性能に近い性能を維持すること。
• 公開利用可能な非MoEチェックポイントを、大規模なMoEモデルへの初期化として効果的に活用すること。
• エキスパート重みをコアと微調整可能な残差に分解することで、効率的な微調整を実現すること。

提案手法

• MoE重みを入力に依存しないコアと入力に依存する残差に分解し、コアは事前学習済みの非MoEモデルから固定する。
• 中間微調整段階で、上流データセットにおけるフルMoE再訓練を回避するため、下流データで残差成分のみを微調整する。
• 中間微調整中に勾配停止（stop gradient）技術を適用することで、性能の低下を防ぐ。
• エキスパートの特化を促進するために、コピーした重みにノイズ初期化スケールを適用する。
• 勾配スコアに基づく戦略を用いたMoE層選択が、Last-2 や Every-2 の手法を上回り、Every-Lastの性能に匹敵するが、より効率的である。
• 中間微調整と下流微調整の両方をサポートし、推論時のFLOPsと遅延にほとんど影響を与えない。

実験結果

リサーチクエスチョン

• RQ1 訓練済みのMoE変換器の重みをコアと残差成分に分解することで、効率的な微調整が可能になるか？
• RQ2 エキスパートの残差成分のみを微調整することで、フルMoE訓練と同等の性能が得られ、著しく低いコストで実現できるか？
• RQ3 既存の非MoE視覚変換器のチェックポイントを、大規模なMoEモデルへの初期化として効果的に活用できるか？
• RQ4 MoE層選択戦略（例：Every-Last、勾配スコア）の選択が、下流性能に与える影響は何か？
• RQ5 最適なパフォーマンスを達成するためのハイパーパrameter（例：top-k、エキスパート数、ノイズスケール）は何か？

主な発見

• ADE20KではSwin-Tに比べ+1.1 mIoU、MS-COCOでは+1.4 APの向上を達成し、3%未満の追加訓練コストで実現した。
• 非MoEモデルと比較すると、CvT-13では+0.9 mIoU、Swin-Lでは+1.0 mIoUの向上を示し、コスト増加は最小限に抑えられた。
• 勾配スコアに基づくMoE層選択は、Last-2 や Every-2 を上回り、Every-Lastの性能に匹敵するが、より効率的である。
• top-1ゲートと8エキスパートを用いることで最適な性能が得られ、16エキスパートに増やすと負荷バランス損失の干渉により結果に改善が見られない。
• 中間微調整時に勾配停止を適用することで性能低下を防ぎ、ADE20Kではノイズスケール0.01が最良の結果をもたらした。
• 推論時間とFLOPsはわずかに増加するのみで、RMoEがMoEモデルの効率的利点を維持していることが確認された。