[論文レビュー] Sparse MoE as the New Dropout: Scaling Dense and Self-Slimmable Transformers

Tianlong Chen, Zhenyu Zhang+3 more|arXiv (Cornell University)|Mar 2, 2023

Domain Adaptation and Few-Shot Learning被引用数 6

ひとこと要約

本稿では、固定されたランダム初期化ルーターを用いて訓練中に段階的に有効化されるエキスパートの数を増やしていく、プラグアンドプレイな訓練フレームワークであるSMoE-Dropoutを提案する。これにより、崩壊を伴わずにTransformerがより高い精度を達成できる。この手法は「自己スリミング可能（self-slimmable）」な性質を誘発し、利用可能なリソースに応じて推論および微調整時に滑らかに性能向上が得られる。BERTモデルは推論タスクで、密度型モデルを最大1.09%上回る性能を発揮する。

ABSTRACT

Despite their remarkable achievement, gigantic transformers encounter significant drawbacks, including exorbitant computational and memory footprints during training, as well as severe collapse evidenced by a high degree of parameter redundancy. Sparsely-activated Mixture-of-Experts (SMoEs) have shown promise to mitigate the issue of training efficiency, yet they are prone to (1) redundant experts due to representational collapse; and (2) poor expert scalability for inference and downstream fine-tuning, primarily due to overfitting of the learned routing policy to the number of activated experts during training. As recent research efforts are predominantly focused on improving routing policies to encourage expert specializations, this work focuses on exploring the overlooked scalability bottleneck of SMoEs and leveraging it to effectively scale dense transformers. To this end, we propose a new plug-and-play training framework, SMoE-Dropout, to enable scaling transformers to better accuracy in their full capacity without collapse. Specifically, SMoE-Dropout consists of a randomly initialized and fixed router network to activate experts and gradually increases the activated expert number as training progresses over time. Transformers trained by SMoE-Dropout naturally exhibit a self-slimmable property subject to resource availability, offering smooth and consistent performance boosts with an increase in activated experts during inference or fine-tuning. Our extensive experiments demonstrate the superior performance and substantial computation savings of SMoE-Dropout, compared to dense training baselines with equivalent parameter counts. In particular, our trained BERT outperforms its densely trained counterpart with consistent improvements of {1.03%, 0.78%, 1.09%} on challenging reasoning tasks {ASDiv-A, MAWPS, SVAMP}, respectively.

研究の動機と目的

• 推論および微調整時のスパースMixture-of-Experts（SMoE）モデルにおける表現崩壊とスケーラビリティの低さを解消すること。
• 特に固定されたエキスパート数にルーティングポリシーが過学習するための、SMoE訓練における見過ごされたスケーラビリティのボトルネックを同定すること。
• 崩壊を伴わずに全容量のモデルスケーリングを可能にするプラグアンドプレイな手法を開発し、エキスパート使用量の増加に伴い一貫した性能向上を保証すること。
• 「自己スリミング可能（self-slimmable）」なTransformerアーキテクチャを通じて、リソースに応じた動的な推論および微調整能力を提供すること。

提案手法

• 学習中に更新されない固定されたランダム初期化ルーターを導入し、学習可能なルーティング機構に置き換える。
• 訓練の過程で段階的に有効化されるエキスパート数（k）を増加させることで、段階的な容量スケーリングを実現する。
• SMoEモジュールをTransformerアーキテクチャの後段の層に挿入することで、初期層への挿入よりも実験的に優れた性能が得られることを確認した。
• 特定のTransformer層にSMoEブロックをモジュラーに挿入する手法を採用し、アブレーションにより後段の層への配置が最も効果的であることが示された。
• 過学習を回避し、暗黙の正則化によってエキスパートの特化を促進するため、ランダムルーティングポリシーを採用する。
• リソースの可用性に応じて有効化されるエキスパート数を調整可能にすることで、推論および微調整時の滑らかな性能スケーリングを実現する。

実験結果

リサーチクエスチョン

• RQ1 固定されたランダムルーティングポリシーは、スパースMixture-of-Experts（SMoE）モデルにおける表現崩壊の緩和とスケーラビリティの向上に寄与するか？
• RQ2 訓練中に段階的に有効化されるエキスパート数を増加させることで、推論および微調整時の一般化性能と性能向上が達成されるか？
• RQ3 学習可能なルーティングと密度型訓練と比較して、SMoE-Dropoutは過学習とモデル容量の活用度において優れているか？
• RQ4 TransformerにSMoEモジュールを挿入する最適な位置はどこか？性能とスケーラビリティを最大化するには？
• RQ5 SMoE-Dropoutは、リソース制約のあるデプロイメントに適した、より蒸留可能なモデルを生成できるか？

主な発見

• SMoE-Dropoutで訓練されたBERTモデルは、それぞれの密度型モデルに対して、ASDiv-A、MAWPS、SVAMPの推論タスクで1.03%、0.78%、1.09%の一貫した性能向上を達成した。
• この手法は「自己スリミング可能（self-slimmable）」な性質を誘発し、推論および微調整時に有効化されるエキスパート数に応じて滑らかに性能が向上する。これにより、リソースに応じた動的な効率-性能トレードオフが可能になる。
• enwik8テストセットでは、16エキスパートを使用したSMoE-Dropoutは、ビット/文字（BPC）が2.96×10⁻²にまで低下し、完全にランダムおよび決定的ハッシュルーティングのベースラインを著しく上回った。
• SMoE-Dropoutモデルは、12層または16エキスパートでも過学習の兆候を示さず、一方で学習可能なSMoEモデルは過学習による性能低下を示した。
• 蒸留実験の結果、SMoE-DropoutモデルはSST-2タスクにおいて、密度型モデルや標準的なSMoEモデルに比べて、1エキスパートモデルに蒸留された際に0.76%〜1.89%高い精度を達成した。
• アブレーションスタディにより、段階的にkを増加させることの重要性が確認され、このスケジューリングを無効化すると、すべてのパラメータを使用しても性能が悪化した。