AdamWトレーニングダイナミクスにおけるワイブル重みスケールパラメータの進化
この研究は、2パラメータワイブル枠組みに基づき、AdamWトレーニング中にワイブル重みスケールパラメータλが成長、オーバーシュート、緩和する理由を分析する。AdamW更新から3つの力（アライメント力、注入力、減衰力）の分解を導出し、アライメント力が上昇段階を支配し（88-94%）、飽和時にアライメントと減衰がバランスすることを示す。スプライン変位法により、スパースチェックポイントからのアライメント力を約92-94%の精度で復元する。ピークλはトレーニングデータのコヒーレンスによって変化し、データ依存の重みスケール成長成分を示唆する。
最近、arXivに投稿された論文「Weibull Weight-Scale Parameter Evolution under AdamW Training Dynamics」は、AdamW最適化器を用いたTransformerのトレーニングにおける重み分布の進化の力学を深く掘り下げています。この研究は、2パラメータワイブル枠組みに基づき、重みスケールパラメータλがトレーニング中に成長、オーバーシュート、緩和する現象を理論的に説明しています。
著者のTiexin Ding氏は、AdamW更新則から二乗重みノルムの主要な3力分解を導出しました。このモデルには、重みと適応更新方向の相関を測るアライメント力、適応ステップサイズに由来する注入力、そして分離された重み減衰による減衰力が含まれます。自己トレーニングしたPythia-70Mモデルを用いた実験では、アライメント力がλの上昇段階の支配的要因であり、4つのランダムシードにおいて絶対力予算の88%から94%を占め、スーパーウェイトの除去に対しても堅牢であることが示されました。
トレーニングが飽和に近づくと、アライメント力と減衰力はバランスを迎え、これがλの成長から緩和への移行を説明します。これらの力のダイナミクスは、λ(t)の基盤となる二乗ノルム成分を直接制御します。残りのRMSからワイブルへの再構成オフセットは測定可能で、ブリッジ成分と統合成分に分解され、密にサンプリングされた領域では合計で約5-6%になります。
オプティマイザのモーメントが利用できない実モデルに分析を拡張するため、研究者らはスプライン変位法を導入しました。この手法はスパースチェックポイントからアライメント力を約92-94%の精度で復元し、単純な2点ベースラインの約2倍の精度を達成します。さらに、実験においてλ(t)のピーク値がトレーニングデータのコヒーレンスによって変化することが観察され、データに依存する重みスケール成長の成分が示唆されました。著者はこの点を今後の制御実験でさらに調査する予定です。
この研究のコードとデータはGitHubで公開されており、機械学習コミュニティにTransformerトレーニングの力学を診断・理解するための強力なツールを提供しています。