[論文レビュー] ST-MoE: Designing Stable and Transferable Sparse Expert Models

Barret Zoph, Irwan Bello+6 more|arXiv (Cornell University)|Feb 17, 2022

Topic Modeling被引用数 46

ひとこと要約

本論文は大規模スパースエキスパートモデル設計ガイドである ST-MoE を提案し、学習を安定化させる router z-loss を導入し、269B パラメータのスパース MoE モデル（ST-MoE-32B）による多様な NLP タスクで最新の転移性能を示す。

ABSTRACT

Scale has opened new frontiers in natural language processing -- but at a high cost. In response, Mixture-of-Experts (MoE) and Switch Transformers have been proposed as an energy efficient path to even larger and more capable language models. But advancing the state-of-the-art across a broad set of natural language tasks has been hindered by training instabilities and uncertain quality during fine-tuning. Our work focuses on these issues and acts as a design guide. We conclude by scaling a sparse model to 269B parameters, with a computational cost comparable to a 32B dense encoder-decoder Transformer (Stable and Transferable Mixture-of-Experts or ST-MoE-32B). For the first time, a sparse model achieves state-of-the-art performance in transfer learning, across a diverse set of tasks including reasoning (SuperGLUE, ARC Easy, ARC Challenge), summarization (XSum, CNN-DM), closed book question answering (WebQA, Natural Questions), and adversarially constructed tasks (Winogrande, ANLI R3).

研究の動機と目的

• スパース MoE モデルのトレーニング不安定性を調査し、安定性を保つ設計選択を特定する。
• 品質を損なうことなくトレーニングを改善するための安定性重視のルーティング機構（router z-loss）を開発する。
• スパースモデルのファインチューニングプロトコルを分析し、密モデルと比較して一般化を理解する。
• 大規模分散トレーニングに適した Pareto効率の高いスパースモデルを設計するためのアーキテクチャとトレーニングの指針を提供する。

提案手法

• 複数のスケール設定に渡るスパース MoE トレーニングの安定性リスクを検討・定量化する。
• ルータのロジットを制約し、モデル品質を損なうことなくエンドツーエンドのトレーニングを安定化させるために router z-loss を導入する。
• MoE ルーティングに対する数値精度形式と丸め誤差の影響を評価し、緩和策を示す。
• パラメータサブセット、ドロップアウト、センチネル・トークンを含むファインチューニングレジームを体系的に研究し、一般化を改善する。
• Pareto効率の高いスパースモデルを実現するための専門家数、キャパシティファクター、ルーティングアルゴリズムの設計選択を提示する。
• 269B のスパースモデル（ST-MoE-32B）を事前学習し、多様な NLP ベンチマークで最先端の転移性能を示す評価を行う。

実験結果

リサーチクエスチョン

• RQ1 安定化技術はモデル品質を犠牲にすることなく、スパース MoE トレーニングを安定化できるか？
• RQ2 ルータを重視したロス（router z-loss）はトレーニングの安定性と最終性能にどう影響するか？
• RQ3 NLPタスク全体で、ST-MoEモデルを密モデルと比較して最も良く一般化させるファインチューニングプロトコルは何か？
• RQ4 スケール時にPareto効率を満たすスパースモデルを生むアーキテクチャおよびルーティング設計原理は何か？

主な発見

モデル	訓練CF	評価CF	補助損失	ドロップされたトークンの割合	SuperGLUE (↑)
Sparse	0.75	2.0	Yes	10.6%	86.5 ± 0.21
Sparse	1.25	2.0	Yes	0.3%	86.7
Sparse	2.0	3.0	Yes	0.0%	85.8
Sparse	4.0	5.0	Yes	0.0%	86.4
Sparse	0.75	2.0	No	15.6%	85.7
Sparse	1.25	2.0	No	2.9%	85.8
Sparse	2.0	3.0	No	0.4%	85.9
Sparse	4.0	5.0	No	0.0%	86.4

• 大規模な安定性研究は、多くの安定化手法が安定性のために品質を犠牲にする一方、router z-loss は品質の低下なしにトレーニングを安定化させることを示している。
• Router z-loss は安定性を改善し、品質をわずかに向上させる一方、過度なアップデートクリッピングは性能を損なう。
• スパースモデルは密モデルとは異なるファインチューニングプロトコルを必要とし、スパースモデルは小さめのバッチサイズと高い学習率の恩恵を受ける。
• スパースモデルはファインチューニング時のトークンドロップに対して頑健で、特定の容量設定下で非ゼロのトークンドロップを維持して性能を保つことができる。
• ST-MoE-32B は SuperGLUE, ARC Easy/Challenge, XSum, CNN-DM, WebQA, Natural Questions, Winogrande, ANLI R3 などのタスクで最先端の転移性能を達成する。
• 本論文は、Pareto効率を達成するための専門家数、キャパシティファクター、ルーティング戦略の実用的な設計指針を提供する。