[論文レビュー] Small-scale proxies for large-scale Transformer training instabilities

Mitchell Wortsman, Peter J. Liu+14 more|arXiv (Cornell University)|Sep 25, 2023

Explainable Artificial Intelligence (XAI)被引用数 4

ひとこと要約

この論文は、大規模なトランスフォーマーにおける学習不安定性を研究するための小スケールプロキシ実験を導入し、注意層におけるログットの増大や確率の対数からの逸脱といった主要な不安定性が、qk-layernorm や z-loss といった手法を用いて小スケールモデルでも再現・軽減可能であることを示している。著者らは、学習率感受性が信頼できる診断指標であることを示し、活性化および勾配ノルムのスケーリング行動が、不安定性が現れる前に対応できる予測可能性を示している。

ABSTRACT

Teams that have trained large Transformer-based models have reported training instabilities at large scale that did not appear when training with the same hyperparameters at smaller scales. Although the causes of such instabilities are of scientific interest, the amount of resources required to reproduce them has made investigation difficult. In this work, we seek ways to reproduce and study training stability and instability at smaller scales. First, we focus on two sources of training instability described in previous work: the growth of logits in attention layers (Dehghani et al., 2023) and divergence of the output logits from the log probabilities (Chowdhery et al., 2022). By measuring the relationship between learning rate and loss across scales, we show that these instabilities also appear in small models when training at high learning rates, and that mitigations previously employed at large scales are equally effective in this regime. This prompts us to investigate the extent to which other known optimizer and model interventions influence the sensitivity of the final loss to changes in the learning rate. To this end, we study methods such as warm-up, weight decay, and the $μ$Param (Yang et al., 2022), and combine techniques to train small models that achieve similar losses across orders of magnitude of learning rate variation. Finally, to conclude our exploration we study two cases where instabilities can be predicted before they emerge by examining the scaling behavior of model activation and gradient norms.

研究の動機と目的

• 同じハイパーパrameterを用いても小スケールでは現れないが、大規模トランスフォーマーで観察される学習不安定性を調査すること。
• 大規模な計算リソースを要しない小スケールプロキシモデルを開発し、それらを用いて不安定性を研究すること。
• qk-layernorm や z-loss などの有効な干渉手法が、学習率の変動にわたって安定した学習を可能にするかを同定および検証すること。
• 最適化手法およびモデルの変更（例：ウォームアップ、重み減衰、μ-Param）が、小スケールモデルにおける学習率感受性に与える影響を分析すること。
• 活性化および勾配ノルムのスケーリング傾向が、訓練中に不安定性が現れる前に対応できるかを調査すること。

提案手法

• Flax を用い、AdamW 最適化手法とデフォルトのハイパーパrameterを用いて、GPT-2 スタイルのアーキテクチャの小スケールデコーダー専用トランスフォーマーを訓練する。
• 学習率を3桁のスケールで系統的に変化させ、最終評価損失を測定し、学習率（LR）感受性を要約統計量として算出する。
• qk-layernorm および z-loss 正則化を導入・評価し、広い範囲の学習率変動にわたって安定した学習を実現する。
• ウォームアップ期間、重み減衰、μ-Param が、小スケールモデルにおける学習率感受性および学習安定性に与える影響を検討する。
• モデルスケールおよび学習率にわたる活性化ノルムおよび勾配ノルムのスケーリング行動を分析し、不安定性の兆候を早期に検出する。
• 大規模な訓練においてデフォルトの AdamW epsilon が大きすぎるため、更新が不適切に小さくなることにより不安定性が生じることに着目し、不安定性のメカニズムと関連付ける。

実験結果

リサーチクエスチョン

• RQ1 大規模トランスフォーマーで観察される学習不安定性（例：注意のログット増大、出力ログットの逸脱）は、小スケールモデルでも再現可能か？
• RQ2 既知の大規模な緩和手法（例：qk-layernorm、z-loss）は、小スケールプロキシ環境でも同様に有効か？
• RQ3 標準的な最適化手法およびモデルの干渉（例：ウォームアップ、重み減衰、μ-Param）は、小スケールモデルにおける学習率感受性にどのように影響するか？
• RQ4 活性化および勾配ノルムのスケーリング傾向は、訓練中に不安定性が現れる前に対応できるか？
• RQ5 AdamW の epsilon ハイパーパrameterは、特に勾配ノルムスケーリングと関連して、不安定性にどのような役割を果たすか？

主な発見

• 高学習率で訓練した場合、小スケールモデルでも注意のログット増大や出力ログットの逸脱といった学習不安定性が再現可能である。
• qk-layernorm および z-loss 正則化により、学習率感受性が低減され、小スケールモデルにおいて学習率の3桁の変動にわたって安定した学習が可能になる。
• スケーリングの深さが幅よりも学習率感受性を速く増大させるため、より深いモデルは学習率の選択に対してより感受性が高くなる。
• モデルスケールおよび学習率にわたって勾配ノルムが減少するため、デフォルトの AdamW epsilon（1e-8）が大きすぎることが判明し、パラメータ更新が不適切に小さくなる原因となっている。
• 進行的鋭さ（progressive sharpening）と適応的最適化手法（例：β₂ を有する Adam）のダイナミクスの相互作用が、周期的な損失スパイクを引き起こし、β₂ を低くすることで前処理の適応速度を上げることでこれを防げる。
• 活性化および勾配ノルムのスケーリング行動は、不安定性が現れる前に対応できる予測が可能であり、大規模な実行を伴わずに学習安定性の診断ツールとして利用可能である。