[論文レビュー] Staged Training for Transformer Language Models

Sheng Shen, Pete Walsh+4 more|arXiv (Cornell University)|Mar 11, 2022

Topic Modeling被引用数 4

ひとこと要約

この論文では、Transformer言語モデルの段階的訓練を提案しており、まず小さなモデルを訓練した後、損失と学習ダイナミクスを保持するように学習された演算子を用いて、段階的に深さや幅を拡大する。過去の計算リソースを再利用し、スケーリング法則を用いて段階的移行を最適化することで、初期から訓練するのと比較して最大22%の計算コスト削減を達成しながら、性能を同等または上回る。

ABSTRACT

The current standard approach to scaling transformer language models trains each model size from a different random initialization. As an alternative, we consider a staged training setup that begins with a small model and incrementally increases the amount of compute used for training by applying a "growth operator" to increase the model depth and width. By initializing each stage with the output of the previous one, the training process effectively re-uses the compute from prior stages and becomes more efficient. Our growth operators each take as input the entire training state (including model parameters, optimizer state, learning rate schedule, etc.) and output a new training state from which training continues. We identify two important properties of these growth operators, namely that they preserve both the loss and the "training dynamics" after applying the operator. While the loss-preserving property has been discussed previously, to the best of our knowledge this work is the first to identify the importance of preserving the training dynamics (the rate of decrease of the loss during training). To find the optimal schedule for stages, we use the scaling laws from (Kaplan et al., 2020) to find a precise schedule that gives the most compute saving by starting a new stage when training efficiency starts decreasing. We empirically validate our growth operators and staged training for autoregressive language models, showing up to 22% compute savings compared to a strong baseline trained from scratch. Our code is available at https://github.com/allenai/staged-training.

研究の動機と目的

• 大規模なTransformer言語モデルを訓練する際の高い計算コストを低減すること。
• 標準的な訓練における非効率性に対処し、より大きなモデルでは初期段階では計算効率が低く、後にのみ効率が向上することを改善すること。
• モデルサイズの増加に伴い、損失と学習ダイナミクスの両方を保持する成長演算子を設計すること。
• スケーリング法則を用いて、計算効率を最大化するための原則的で整合性のある段階的移行スケジュールを開発すること。
• GPT-2スタイルの自己回帰的モデルを用いて、標準的なベンチマーク（Wikitext-103 や LAMBADA など）で、この方法の実証的妥当性を検証すること。

提案手法

• 本手法は、全訓練状態（重み、最適化子状態、学習率スケジュール）を入力とし、新しいより大きなモデル状態を出力する深さおよび幅の成長演算子を用いる。
• 成長演算子は損失を保持し、学習ダイナミクスを保持するように設計されており、成長後の損失と損失減少率が、初期から訓練したモデルと一貫していることを保証する。
• 段階的スケジュールは、Kaplanら（2020）のスケーリング法則に基づいて決定され、訓練効率が低下し始める時期と、新しい段階を開始する時期を予測する。
• 正確なスケーリング法則が未知である現実世界の設定において、最適スケジュールの実用的近似が導出される。
• 標準的なベンチマーク（Wikitext-103 や LAMBADA など）を用いて、GPT-2スタイルの自己回帰的言語モデル上で、本手法の有効性が検証される。
• 過去の訓練段階を再利用し、冗長な計算を回避することで、大幅に少ない計算リソースで大規模モデルを訓練可能となる。

実験結果

リサーチクエスチョン

• RQ1 モデル成長を伴う段階的訓練は、大規模なTransformer言語モデルの訓練に要する総合計算コストを削減できるか？
• RQ2 安定的かつ効率的な段階的訓練を確保するために、成長演算子が満たすべき性質は何か？
• RQ3 計算コスト削減を最大化するために、最適な段階的移行スケジュールをどのように決定できるか？
• RQ4 最適化子と学習率を含む、完全な訓練状態をモデル成長中にどの程度保持できるか？
• RQ5 学習ダイナミクスを保持することで、単純な成長手法と比較して、収束性と性能が向上するか？

主な発見

• 提案された成長演算子は、損失と学習ダイナミクスの両方を保持しており、成長後の損失減少率が、初期から訓練したモデルと一致することを保証する。
• スケーリング法則を用いて段階的移行を決定することで、初期から訓練するのと比較して最大22%の計算コスト削減が達成される。
• GPT-2スタイルのモデルにおける実証的結果から、成長させたモデルは、大幅に少ない計算リソースで同等またはより良い検証損失とゼロショットパープレキシティを達成することが示された。
• 本手法により、GPT-2 baseモデルの15%の計算リソースで、Wikitext-103で37.5 PPLを達成し、同等のLAMBADA精度を得るのに33%の計算リソースで十分であった。
• 成長を早すぎるか遅すぎる段階でするとうまくいかず、計算コストの削減が失われるため、正確な段階的スケジューリングの重要性が確認された。
• 本手法は、異なるモデルサイズや学習目的においても性能を維持するという点で、頑健で汎用的であることが示された。