[論文レビュー] 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

ひとこと要約

この論文では、下流タスク向けにスパースなMixture-of-Experts（MoE）モデルを、効率的な単一エキスパートの密度モデルに変換する、タスク固有のエキスパートプルーニングを提案する。貢献度スコアに基づき、微調整中に非専門的エキスパートを段階的にプルーニングすることで、MoE性能の99.3％を保持しながら、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を微調整した性能と同等またはそれを上回ることを示す。
• 最終的なモデル性能と推論効率を最大化するための、エキスパートプルーニングのタイミングと基準を最適化する。

提案手法

• 専門性スコアは、固定長のスライディングトレーニングウィンドウ上で計算される。
• 定期的な間隔でドロップしきい値が適用され、しきい値未満のエキスパートは削除される。
• 1つのエキスパートが各MoEレイヤーに残るか、または微調整の半分までプルーニングが継続される。
• 「エージャー（即時）」プルーニング戦略が提案され、非専門的エキスパートが中間点で削除され、第二段階で選択されたエキスパートの最適化に集中できるようにする。
• 最終的に選択されたエキスパートは、残りのトレーニングスケジュール中、独立して微調整される。
• 知識蒸留を回避し、MoE事前学習モデルから直接単一の密度エキスパートへ知識を転送する。

実験結果

リサーチクエスチョン

• RQ1 事前学習済みMoEモデルから、特定の下流タスクに対して、完全なMoEアーキテクチャの性能向上の大部分を保持する単一エキスパートを特定できるか？
• RQ2 微調整中に非専門的エキスパートをプルーニングする最適なタイミングとしきい値は何か？ これは最終的なモデル精度を最大化するためのものである。
• RQ3 トレーニングウィンドウの長さが、選択されたエキスパートの性能にどのように影響するか？
• RQ4 非専門的エキスパートをプルーニングすることで、性能を損なわず、推論速度が向上し、通信コストが削減されるか？
• RQ5 プルーニング後の単一エキスパートモデルは、スクラッチから微調整された密度モデルを上回る性能を示すか？

主な発見

• プルーニング後の単一エキスパートモデルは、6つの多様な下流タスクにおいて、完全なMoEモデルの性能の99.3％を保持する。
• MNLIタスクでは、完全なMoEモデルと比較して、2倍の高速な推論速度を達成し、SQuADでは1.28倍の高速化を実現する。
• 選択されたエキスパートにサブ-MoEレイヤーが含まれても、プルーニングモデルの推論効率は、密度モデルの80％に達する。
• 「エージャー」プルーニング戦略（中間点で非専門的エキスパートを削除）は、段階的または二段階法よりも優れた推論性能を示す。
• この手法により、デバイス間通信コストは完全に排除され、クラウドおよびモバイルデプロイメントに非常に適している。
• 選択されたエキスパートは、訓練の初期段階で貢献を支配するため、早期にプルーニングすることで、第二段階での集中最適化が正当化される。