[論文レビュー] Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints

Aran Komatsuzaki, Joan Puigcerver+7 more|arXiv (Cornell University)|Dec 9, 2022

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

ひとこと要約

この論文では、事前に訓練された密度型Transformerモデルを、MoE（Mixture-of-Experts）モデルに変換するためのスパースアップサイクリングという手法を提案する。MoEレイヤーを密度型チェックポイントから初期化することで、元の事前学習計算コストの約50%の追加計算量で、SuperGLUEおよびImageNetで最先端の性能を達成し、元の密度型モデルやからくりから訓練されたMoEモデルを上回る。

ABSTRACT

Training large, deep neural networks to convergence can be prohibitively expensive. As a result, often only a small selection of popular, dense models are reused across different contexts and tasks. Increasingly, sparsely activated models, which seek to decouple model size from computation costs, are becoming an attractive alternative to dense models. Although more efficient in terms of quality and computation cost, sparse models remain data-hungry and costly to train from scratch in the large scale regime. In this work, we propose sparse upcycling -- a simple way to reuse sunk training costs by initializing a sparsely activated Mixture-of-Experts model from a dense checkpoint. We show that sparsely upcycled T5 Base, Large, and XL language models and Vision Transformer Base and Large models, respectively, significantly outperform their dense counterparts on SuperGLUE and ImageNet, using only ~50% of the initial dense pretraining sunk cost. The upcycled models also outperform sparse models trained from scratch on 100% of the initial dense pretraining computation budget.

研究の動機と目的

• 大規模なスパースモデルをからくりから訓練する際の高い計算コストを軽減するため、既存の密度型モデルのチェックポイントを再利用すること。
• 新しいデータや完全な再訓練を必要とせずに、大規模モデルにおける性能向上を実現すること。
• 最小限の追加計算量で、より効率的かつ高容量なMoEモデルに密度型モデルをアップグレードする実用的な手法を開発すること。
• 同じ計算予算内で、密度型チェックポイントから初期化されたMoEモデルが、元の密度型モデルおよびからくりから訓練されたMoEモデルを上回ることを示すこと。

提案手法

• 事前学習済みの密度型Transformerチェックポイントを、MoEレイヤーのウェルコンディショニングとして再利用することで、Mixture-of-Experts（MoE）モデルを初期化する。
• Transformerブロック内のMLP部品を、多数のエキスパートフィードフォワードネットワークとルーティングネットワークに置き換えることで、スパースな活性化を可能にする。
• エキスパート選択ルーティングを用い、各トークンに対して最も関連性の高いエキスパートを動的に選択することで、計算を効率化する。
• 構造的変更に伴う性能低下を最小限に抑えるために、MoE部品の慎重な初期化と適応を含むモデルサurgeryレシピを適用する。
• 通常、元の密度型モデルの事前学習コストの10%～60%程度の追加計算量で、アップサイクルドされたMoEモデルを訓練する。
• 元の密度型モデルと同一のデータおよびトレーニングスケジュールを用いるため、新しいデータは一切導入しない。

実験結果

リサーチクエスチョン

• RQ1 事前学習済みの密度型モデルを、最小限の追加計算量で、より強力なMoEモデルに効果的にアップサイクルできるか？
• RQ2 密度型チェックポイントからMoEレイヤーを初期化することで、元の密度型モデルの継続的訓練よりも優れた性能が得られるか？
• RQ3 同じ合計計算量が与えられた場合、アップサイクルドされたMoEモデルは、からくりから訓練されたMoEモデルを上回れるか？
• RQ4 継続的密度型学習と比較して、スパースアップサイクリングによる性能向上を達成するために、どの程度の追加計算量が必要か？
• RQ5 このアップサイクリング手法は、自然言語処理やビジョンなど、異なるモデルアーキテクチャとドメインに一般化可能か？

主な発見

• T5-Base、T5-Large、T5-XLのアップサイクルドモデルは、SuperGLUEで元の密度型モデルを1.5～2.0絶対ポイント上回り、元の事前学習計算量の46～55%で実現した。
• ViT-BaseおよびViT-Largeのアップサイクルドモデルは、ImageNet 10-shotで1%以上の精度向上を達成し、継続的密度型学習の58%に対してわずか13%の追加トレーニング時間で実現した。
• 同じ合計計算量が与えられた場合、スパースアップサイクリングは、からくりから訓練されたMoEモデルよりも優れた性能を発揮し、事前学習済み重みを再利用する価値を示した。
• ViT-Baseでは、13%の追加計算量でImageNet 10-shotで1%以上の向上が達成されたが、継続的密度型学習では同じ向上を得るには58%の追加計算量が必要だった。
• この手法は、複数のモデルサイズおよびドメインで有効であり、制限付きの計算環境下でも、言語およびビジョンの両タスクで、密度型モデルおよびからくりから訓練されたMoEベースラインを上回る性能を発揮した。
• アップサイクリング手法は、ロバストで汎用的であり、SuperGLUEおよびImageNetを含むさまざまなモデルアーキテクチャとタスクで、一貫した性能向上を示した。