大規模Transformerトレーニングの不安定性に対する小規模プロキシ

Small-scale proxies for large-scale Transformer training instabilities

September 25, 2023

著者: Mitchell Wortsman, Peter J. Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D. Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-dickstein, Kelvin Xu, Jaehoon Lee, Justin Gilmer, Simon Kornblith

cs.AI

要旨

大規模なTransformerベースのモデルを訓練したチームは、小規模での訓練では現れなかった訓練の不安定性が大規模で現れると報告しています。このような不安定性の原因は科学的に興味深いものですが、それらを再現するために必要なリソースの量が調査を困難にしています。本研究では、小規模で訓練の安定性と不安定性を再現し、研究する方法を探ります。まず、先行研究で説明されている訓練不安定性の2つの原因に焦点を当てます：アテンション層におけるロジットの増大（Dehghani et al., 2023）と、出力ロジットが対数確率から乖離すること（Chowdhery et al., 2022）。学習率と損失の関係をスケール間で測定することで、これらの不安定性は高学習率で訓練する場合に小規模モデルでも現れ、大規模で用いられた緩和策がこの領域でも同様に有効であることを示します。これにより、他の既知のオプティマイザやモデルの介入が最終的な損失の学習率変化に対する感度にどの程度影響を与えるかを調査することを促します。この目的のために、ウォームアップ、重み減衰、muParam（Yang et al., 2022）などの方法を研究し、学習率の変動幅が数桁にわたっても同様の損失を達成する小規模モデルを訓練するために技術を組み合わせます。最後に、探索を締めくくるために、モデルの活性化と勾配ノルムのスケーリング挙動を調べることで、不安定性が現れる前に予測できる2つのケースを研究します。

English

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 muParam (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.