タスク特化型エキスパートプルーニングによる非効率性を評価してSMoE言語モデルを再考する

Sparse Mixture of Expert (SMoE) モデルは、言語モデリングにおいて密なモデルのスケーラブルな代替手段として登場しました。これらのモデルはトランスフォーマーブロック内で条件付きに活性化されるフィードフォワードサブネットワークを利用し、モデル全体のパラメータ数と例ごとの計算量の分離を可能にします。しかし、大規模なトークンルーティング型SMoEモデルは重要な課題に直面しています。それは推論中に、シーケンスまたはバッチ全体でモデル全体を使用しなければならず、分散環境でのレイテンシが高くなり、トークン単位のスパース活性化の利点が相殺されてしまうことです。本研究では、SMoEアーキテクチャ設計の決定に役立てるためにタスク特化型モデルプルーニングを探求し、主に事前学習時のエキスパート数の選択を調整します。プルーニングされたモデルが、タスクごとに個別に評価・比較した際、ゼロからトレーニングされたより小規模なSMoEモデルより利点を持つか否かを検討します。そのために、トレーニング後のオフラインでMoE層ごとのエキスパート数を削減する適応的タスク認識プルーニング技術UNCURLを導入しました。研究結果は、事前学習時に使用したエキスパート数に依存するスレッショルド値を示しており、それを超える削減はモデル性能の低下を招きます。これらの知見は、SMoEアーキテクチャによる事前学習時のモデル設計の理解に貢献し、特に後の段階でのタスク特化型推論最適化を検討する際に有用です。