[論文レビュー] LEMON: Lossless model expansion

Yite Wang, Jiahao Su+7 more|arXiv (Cornell University)|Oct 12, 2023

Advanced Neural Network ApplicationsComputer Science被引用数 3

ひとこと要約

LEMONは、事前学習済みの小さなモデルを用いてより大きなTransformerモデルを初期化する損失なしのモデル拡張フレームワークであり、性能の低下を伴わずに高速な学習を可能にする。対象となるVision TransformerとBERTに対して、それぞれ56.7%および33.2%の計算コスト削減を達成しており、対称性を破る重み初期化と最適化された学習率スケジューラーを用いることで、学習から始めることと比較して実現している。

ABSTRACT

Scaling of deep neural networks, especially Transformers, is pivotal for their surging performance and has further led to the emergence of sophisticated reasoning capabilities in foundation models. Such scaling generally requires training large models from scratch with random initialization, failing to leverage the knowledge acquired by their smaller counterparts, which are already resource-intensive to obtain. To tackle this inefficiency, we present $ extbf{L}$ossl$ extbf{E}$ss $ extbf{MO}$del Expansio$ extbf{N}$ (LEMON), a recipe to initialize scaled models using the weights of their smaller but pre-trained counterparts. This is followed by model training with an optimized learning rate scheduler tailored explicitly for the scaled models, substantially reducing the training time compared to training from scratch. Notably, LEMON is versatile, ensuring compatibility with various network structures, including models like Vision Transformers and BERT. Our empirical results demonstrate that LEMON reduces computational costs by 56.7% for Vision Transformers and 33.2% for BERT when compared to training from scratch.

研究の動機と目的

• 大規模モデルを学習から始めることによる非効率を、小さな事前学習済みモデルからの知識を活用することで是正すること。
• 幅や深さの増分が非可除な場合でも対応可能な、さまざまなTransformerアーキテクチャに適合する損失なしのモデル拡張手法を開発すること。
• モデル性能を維持したまま、学習時間と計算コストを削減すること。
• 特に、拡張済みモデルに最適な学習レシピ（特に学習率スケジューリング）を最適化すること。

提案手法

• LEMONは、拡張されたレイヤー内の重複するニューロンに、等しくないファンアウト重みを割り当てることで、冗長性を回避し、真の容量向上を実現する対称性を破る初期化戦略を採用する。
• 幅拡張時のLayerNormレイヤーに対応するための平均化拡張を導入し、幅増分が非可除であっても（例：512 → 768）互換性を保証する。
• アーキテクチャの違いに適応することで、標準的およびPre-LN Transformerを含むさまざまなTransformerバリアントをサポートする。
• 学習から始めることで使用された最大学習率を同じままに保ちつつ、より速い減衰を適用する独自の学習率スケジューラーを採用し、拡張済みモデルの収束を改善する。
• 小さな事前学習済みモデルから、より大きなターゲットモデルへの直接的な重み転送を可能にし、機能的同等性を維持する。
• 実証的検証により、拡張済みモデルが学習から始めたモデルと同等の性能を達成し、著しく短縮された学習時間で実現していることが確認された。

実験結果

リサーチクエスチョン

• RQ1 幅や深さの増加が非可除であっても、性能の低下なしに、事前学習済みの小さなTransformerモデルをより大きなモデルに拡張できるか？
• RQ2 モデル拡張時にどのように対称性を破れば、拡張済みモデルが元のモデルを上回る真の容量を獲得できるか？
• RQ3 拡張済みモデルの高速収束を実現するにあたり、特に学習率スケジューリングを含めた最適な学習レシピは何か？
• RQ4 提案された拡張手法は、ViT や BERT を含む多様なTransformerアーキテクチャに一般化可能か？
• RQ5 学習から始めることと比較して、損失なしのモデル拡張は、計算コストをどの程度削減できるか？

主な発見

• LEMONは、学習から始めることと比較して、Vision Transformerでは56.7%、BERTでは33.2%の計算コスト削減を達成した。
• 拡張済みのViT（12,768）は、学習から始めたモデルと同等の性能を、たった85エポックで達成しており、学習時間は28.3%削減された。
• この方法により、元のモデルと完全に機能的同等性が保たれ、拡張後も性能の低下がないことが確認された。
• 最大学習率は同じだが、より速い減衰を適用する最適化された学習率スケジューラーは、拡張済みモデルの学習効率を顕著に向上させた。
• 非可除な幅増分に対しても、ViT、BERT、Pre-LNバリアントを含む多様なTransformerアーキテクチャと互換性がある。
• 実証的結果により、拡張済みモデルが元のモデルの知識を損なわずすべて継承していることが確認され、拡張の損失なし性が裏付けられた。