Small-scale proxies for large-scale Transformer training instabilities

https://arxiv.org/abs/2309.14322

要約

この論文は大規模 Transformer モデルの学習における不安定性を小規模モデルで再現し研究する方法を提案している

1. 学習率と損失の関係を異なるモデルサイズで調査することで不安定性を特定する手法を導入した

2. 先行研究で報告された 2 つの不安定性（アテンション層のロジット増大と出力ロジットの発散）を小規模モデルで高学習率を用いて再現し、大規模モデルで有効だった緩和策が小規模でも機能することを示した

3. ウォームアップ、重み減衰、μParam など他の既知の介入がモデルサイズに応じた学習率感度にどのような影響を与えるかを調査した

4. モデル特性（勾配ノルムやアクティベーションノルムなど）のスケーリング挙動を分析することで不安定性を事前に予測できる可能性を示した

5. デフォルトの AdamW の ε ハイパーパラメーターが大規模モデルでは大きすぎる可能性があることを発見し、より小さな値を使用することで性能が向上することを示した

この研究は大規模な計算リソースがなくても Transformer モデルの学習安定性に関する洞察を得られる可能性を示している。これにより大規模モデル学習の効率化や安定化に向けた新たな研究機会が開かれると考えられる

Abstract

大規模 Transformer モデルの学習で報告された不安定性を小規模モデルで再現・研究する方法を提案する

学習率と損失の関係を異なるスケールで測定することで小規模モデルでも高学習率下で不安定性が現れることを示した

具体的にはアテンション層のロジット増大と出力ロジットの発散という 2 つの不安定性を調査し、大規模モデルで有効だった緩和策が小規模でも機能することを確認した

さらにウォームアップや重み減衰などの介入が学習率感度に与える影響を調べ、モデル特性のスケーリング挙動を分析して不安定性を予測できる可能性を示した

この研究により大規模な計算リソースなしで Transformer モデルの学習安定性に関する知見を得られる新たな研究機会が開かれる

1 Introduction

Transformer のスケールアップは著しい進歩をもたらしたが全ての学習が成功するわけではない。大規模 Transformer の学習では学習を遅延または不安定化させる問題が報告されている。大規模実行に必要なリソースが増加し続けるため Transformer 学習の失敗モードを調査することが重要となっている

本論文では小規模 Transformer モデルで学習の不安定性を再現、研究、予測する方法を探る。学習率と損失の関係をスケール間で測定することが不安定性を特定する有用なツールであることを見出し、学習率 (LR) 感度という指標を導入した。これは学習率を 3 桁変化させたときの最適性能からの偏差を測定する

また先行研究で大規模モデルにおいて記述された 2 つの不安定性を小規模モデルで再現できることを示す。これにより大規模なリソースプールへのアクセスなしに研究が可能となる。さらに他の既知の最適化手法やモデル介入が学習率と損失の曲線形状にどのように影響するかを調査する

最後にモデル特性のスケーリング挙動を調べることで不安定性を予測できる可能性を示す。この手法を用いて現在十分に文書化されていない不安定性を探索する

2 Experimental methodology

1. 実験設定

   • GPT-2 に似た Flax 実装の小規模 Transformer モデルを使用
   • デコーダーのみのモデルで自己回帰損失で学習
   • AdamW オプティマイザー、線形ウォームアップ、コサイン減衰スケジュールを採用
   • qk-layernorm、z-loss 正則化、回転位置埋め込みなどの技術を使用
   • TPU 上で bfloat16 精度で学習

2. 学習率 (LR) vs 損失曲線と LR 感度

   • 異なるモデルサイズで LR と損失の関係をプロットし不安定性を調査
   • LR 感度 : 3 桁の LR 変動時の最適性能からの偏差を測定
   • 不安定性は高 LR での損失爆発として特徴付けられる

3. モデル特性のスケーリングトレンド

   • 勾配やアクティベーションノルムなどのモデル特性のスケーリング挙動を調査
   • 不安定性予測や新たな問題発見に有用

2.1 Experimental set-up

実験では GPT-2 に類似した設定の小規模 Transformer モデルを使用した

1. デコーダーのみのモデル、自己回帰損失で学習

2. AdamW オプティマイザーを使用

   • $\beta_1 = 0.9$, $\beta_2 = 0.95$, $\epsilon = 1e-8$
   • 勾配クリッピング : グローバルノルム 1

3. 学習スケジュール

   • 5e3 ステップのウォームアップ
   • 総ステップ数 : 1e5
   • 線形ウォームアップ後にコサイン減衰
   • 最小学習率 : 1e-5

4. 重み減衰: 1e-4（独立）

5. z-loss 正則化 : 係数 1e-4

6. アーキテクチャー

   • 事前正規化 Transformer
   • qk-layernorm 使用
   • バイアスなし
   • レイヤーノーム $\epsilon = 1e-6$
   • 埋め込みサイズ、深さ、ヘッド数を同時にスケール
   • 回転位置埋め込み

7. データセット : C4

8. その他

   • バッチサイズ : 256（各要素 512 トークン）
   • TPU 上で bfloat16 精度で学習
   • Flax と JAX を使用

2.2 LR vs. loss curves and learning rate sensitivity

異なるモデルサイズで学習率 (LR) と損失の関係をプロットすることが不安定性の調査に有用であることを示している

1. LR vs 損失曲線

   • 不安定性は高学習率での損失爆発として現れる
   • モデルサイズに応じて最小の不安定学習率がどう変化するかを示す

2. LR 感度

   • LR vs 損失曲線を要約する指標
   • 3 桁の LR 変動時の最適性能からの偏差を測定
   • 定義 : $E_{\eta \in [a,b]} [\min(\ell(A(\eta)), \ell_0) - \ell^*]$
     ここで
     • $\theta = A(\eta)$ は学習率 $\eta$ で学習したモデルの重み
     • $\ell(\theta)$ は重み $\theta$ での検証損失
     • $\ell^*$ は $[a,b]$ 範囲内の最良の損失
     • $\ell_0$ は初期化時の損失

3. 使用上の注意

   • LR の意味を変える介入では比較に適さない
   • 最適 LR のシフトを考慮する必要がある
   • 損失のスケールに不変なため LR vs 損失曲線と併せて解釈すべき

2.3 Scaling trends for model characteristics

この節ではモデル特性のスケーリングトレンドを調査する方法について説明している

1. 勾配やアクティベーションノルムなどのモデル特性のスケーリング挙動を調べることが不安定性の研究に有用であることを指摘

2. この手法は不安定性を予測するのに役立つ可能性がある

3. 従来のスケーリング研究がモデルスケールと損失の関係に主に焦点を当てていたのに対しこの手法は異なるアプローチを取る

4. アテンションロジット成長の不安定性を例として、この手法が不安定性を予測する能力を示す

5. さらに、この方法を用いて現在十分に文書化されていない不安定性を探索する可能性を示唆

3 Results

1. 既知の不安定性を小規模再現
   a) アテンションロジット成長 : 小規模モデルでも高学習率で発生、qk-layernorm が効果的
   b) 出力ロジット発散 : 小規模モデルでも再現、z-loss が効果的

2. 他の介入効果
   a) ウォームアップ : 長いウォームアップが LR 感度を低減
   b) 独立重み減衰 : LR 感度を低減
   c) 幅 vs 深さのスケーリング : 深さのスケーリングが LR 感度を急増させる
   d) μParam: 最適 LR を安定化するが、LR 感度は改善せず

3. アテンションロジット成長不安定性予測

   • モデル特性のスケーリング挙動を利用して大規模モデルでの不安定性を予測

4. 新たな不安定性を探索

   • 勾配 RMS が減少傾向を示し、デフォルトの AdamW $\epsilon$ に近づく問題を発見
   • $\epsilon$ を 1e-15 に下げることで性能が向上

5. 出力ノルム成長との関連

   • パラメーターノルム成長がアテンションロジット成長と AdamW $\epsilon$ 不安定性に関連

3.1 Reproducing two known instabilities at small scale

1. アテンションロジット成長不安定性

   • 大規模モデルで報告された問題を小規模モデルで高学習率を用いて再現
   • qk-layernorm が効果的な緩和策として機能
   • アテンションロジット $z_{ij} = \langle q_i, k_j \rangle / \sqrt{d_h}$ が大きくなり、注意重みが one-hot ベクトルに収束
   • クエリーとキーの成長が原因で、その整列度合いは変化しない

2. 出力ロジット発散不安定性

   • 小規模モデルでも高学習率で再現可能
   • 出力ロジット $y$ が負の値に発散
   • z-loss（補助損失 $\log^2 Z$, $Z = \sum_j e^{y_j}$）が効果的な緩和策
   • 重み減衰も大きなモデルではこの不安定性を緩和

3. 共通点

   • 両不安定性とも小規模モデルで高学習率を用いることで再現可能
   • 大規模モデルで有効だった緩和策が小規模モデルでも機能
   • LR vs 損失曲線が不安定性と緩和策の効果を示すのに有用

3.1.1 Attention logit growth

1. 問題の概要

   • Transformer 学習中にアテンションロジットが大きくなる問題
   • 式 : $z_{ij} = \langle q_i, k_j \rangle / \sqrt{d_h}$（ここで $d_h$ はヘッド次元）
   • アテンション重みが one-hot ベクトルに収束 (attention entropy collapse)

2. 小規模モデルでの再現

   • 高学習率を用いることで小規模モデルでも発生
   • モデルサイズが大きくなるほど発散する学習率が低下

3. qk-layernorm の効果

   • クエリーとキーに LayerNorm を適用する手法
   • LR 感度を大幅に低減
   • 1.2B パラメーターのモデルを学習率 0.3 で学習可能に

4. 原因の分析

   • クエリーとキーのノルム増大が主な原因
   • クエリーとキーの整列度合いは変化しない
   • パラメーターノルムへの 2 次依存性が要因の可能性

5. 実験結果

   • 小さなネットワークで異なる最大アテンションロジット値を強制
   • $\kappa \approx 10^3$ で損失が悪化し始め $\kappa \approx 10^4$ でバイグラムモデルより悪化

3.1.2 Output logit divergence

1. 問題の概要

   • 出力ロジットが対数確率から発散する不安定性
   • モデルの出力ロジット $y$ を用いてクラス確率 $p_i = e^{y_i} / Z$ を計算（ここで $Z = \sum_j e^{y_j}$）

2. 小規模モデルでの再現

   • 高学習率を使用することで小規模モデルでも発生
   • 学習後半でロジットが負の値に発散

3. z-loss による緩和

   • 補助損失 $\log^2 Z$ を係数 1e-4 で追加
   • $\log Z$ を 0 に近づけることを奨励

4. 実験結果

   • 重み減衰なしのモデルでスケールに関わらず不安定性が発生
   • z-loss がこの不安定性を解決
   • 大きなモデルでは重み減衰もこの不安定性を緩和

5. 重み減衰との関係

   • 重み減衰なしの場合、スケールに関わらず不安定性が発生
   • 大きなモデルでは重み減衰が効果的

3.2 Measuring the effect of other known interventions

1. 目的

   • 既知のモデルおよび最適化手法の介入が LR vs 損失曲線の形状にどう影響するかを調査
   • これらの介入が LR 感度にどのような影響を与えるかを検討

2. ウォームアップ

   • 長いウォームアップ期間が LR 感度を低減
   • 特に大規模モデルで効果が顕著

3. 独立重み減衰

   • Loshchilov と Hutter が推奨した手法
   • LR 感度を低減
   • 更新式: $\theta \leftarrow \theta - s_t(η\Delta - λ\theta)$（$s_t$ はスケジュール、$η$ は学習率、$λ$ は重み減衰係数）

4. 幅 vs 深さのスケーリング

   • 深さのスケーリングが幅のスケーリングより急速に LR 感度を増加
   • 最大スケールでは深さのスケーリングが低い検証損失を達成

5. μParam

   • 最適 LR を安定化するが 損失や LR 感度は改善しない
   • qk-layernorm の必要性は軽減されない

6. その他の介入

   • 学習ステップ数の変更 : LR 感度に大きな影響なし
   • バッチサイズの増加 : LR 感度に大きな影響なし
   • 重み減衰の変更 : 最適 LR を若干右にシフト

3.2.1 Warm-up

1. 主な発見

   • より長いウォームアップ期間が LR 感度を低減する
   • この効果は特に大規模モデルで顕著

2. 実験設定

   • 総ステップ数を 1e5 に固定
   • すべてのモデルで qk-layernorm を使用

3. 結果の詳細

   • 大規模モデルは長いウォームアップなしでは学習率 3e-1 で安定しない
   • ウォームアップの長さを増やすことで、より高い学習率での学習が可能に

4. 先行研究との関連

   • ウォームアップの重要性は以前から指摘されていたがスケーリング挙動は測定されていなかった

5. 解釈

   • ウォームアップは初期の学習段階でモデルパラメーターを徐々に調整し、急激な変化を防ぐ
   • 大規模モデルほどこの緩やかな調整の恩恵を受ける可能性がある

6. 実践的示唆

   • 大規模モデルの学習時には、より長いウォームアップ期間を検討すべき
   • ウォームアップ長さはモデルサイズに応じて調整する必要がある可能性

3.2.2 Independent weight decay

1. 主な発見

   • 学習率から独立して重み減衰をパラメーター化することで LR 感度が低減する

2. 手法の詳細:

   • Loshchilov と Hutter が推奨した方法
   • 更新式 : $\theta \leftarrow \theta - s_t(η\Delta - λ\theta)$
     ここで $s_t$ はスケジュール、$η$ は学習率、$λ$ は重み減衰係数

3. 従来の実装との違い

   • PyTorch や Optax のデフォルト実装 : $\theta \leftarrow \theta - s_tη(\Delta - λ\theta)$
   • 独立減衰では $η$ が $\Delta$ にのみ適用される

4. 実験設定

   • 独立減衰使用時 : $λ = 1e-4$
   • 非独立減衰時 : $λ = 0.1$

5. 結果

   • 独立重み減衰が LR 感度を低減
   • 特に大規模モデルで効果が顕著

6. 実践的示唆

   • 大規模モデル学習時には独立重み減衰の使用を検討すべき
   • デフォルトの実装を使用する場合、この違いに注意が必要

7. 追加の観察

   • 重み減衰値を増やすと最適な学習率が若干右にシフトする傾向

3.2.3 Scaling width vs. depth

1. 主な発見

   • 深さのスケーリングは幅のスケーリングよりも急速に LR 感度を増加させる
   • 最大スケールでは深さのスケーリングがより低い検証損失を達成

2. 実験設定

   • 幅のスケーリング : $d = 512$, レイヤー数は 6 で固定
   • 深さのスケーリング : レイヤー数を増加, $d = 512$ で固定
   • ヘッド数は幅に比例してスケール

3. 結果の詳細

   • 深さのスケーリングは LR 感度を急増させるが最終的な性能は優れている
   • 幅のスケーリングは LR 感度の増加が緩やかだが最大スケールでの性能は劣る

4. 追加の観察

   • qk-layernorm なしの場合、スケールに関わらずアテンションロジット成長不安定性が頻繁に発生

5. 解釈

   • 深いモデルはより複雑な機能を学習できるが学習が不安定になりやすい
   • 幅広いモデルは学習が安定しやすいが表現力に制限がある可能性

6. 実践的示唆

   • モデルスケーリング時は深さと幅のバランスが重要
   • 大規模モデルでは深さの増加が有効だが安定性に注意が必要

7. 追加の比較

   • 幅と深さを同時にスケーリングする標準的な方法が最大スケールで最も良い性能を示す
   • この方法はスケーリング予測の外挿もより信頼性が高い

この結果はモデルアーキテクチャーの設計において深さと幅のトレードオフを慎重に考慮することの重要性を強調している

3.2.4 μParam

1. μParam 概要

   • Yang と Hu が導入したニューラルネットワークのパラメーター化手法
   • モデル幅をスケーリングする際に最適な学習率を一定に保つことが可能

2. 実験結果

   • μParam は試験したスケールで最適な学習率を安定化させることに成功
   • しかし損失の改善や LR 感度の低減は見られなかった
   • qk-layernorm なしの高学習率では μParam は不安定性を軽減しない

3. 実装詳細

   • μParam (simple) : 線形レイヤーの学習率を base-fan-in/fan-in でスケール
   • μParam (full) : 追加の特徴として (i) ヘッドの初期化、(ii) アテンションレイヤーのスケーリング因子変更、(iii) クエリー射影重みのゼロ初期化を含む

4. 追加の観察

   • アテンションレイヤーのスケーリング因子を $1/\sqrt{d_h}$ から $1/d_h$ に変更しても qk-layernorm の必要性は軽減されない
   • クエリー射影重みのゼロ初期化は顕著な効果を示さなかった

5. 解釈

   • μParam は最適学習率の安定化に効果的だが学習の全体的な安定性や性能向上には寄与しない
   • 実用的な観点では最適学習率が変化しない場合 LR 感度の低減は重要ではない可能性がある

6. 実践的示唆

   • μParam は学習率選択を簡略化する可能性があるが、他の安定性問題には別のアプローチが必要
   • 大規模モデルでの使用時には他の安定化技術（例 : qk-layernorm）との併用を検討すべき

3.2.5 Additional interventions

1. 学習ステップ数

   • 1e5 から 5e4 または 2e5 に変更しても LR 感度に大きな影響はない

2. qk-layernorm 適用方法

   • モデル次元全体に適用するよりもヘッドごとに個別に適用する方が性能が向上
   • ヘッドごとの適用では学習可能パラメーター数が少なくて済む

3. バッチサイズ

   • 256 から 512 または 1024 に増加しても LR 感度に大きな影響はない
   • 各バッチ要素は 512 トークンを含む

4. 重み減衰

   • 重み減衰を増やすと最適な学習率が若干右にシフトする傾向

5. ソフトマックスがない注意機構

   • ポイントワイズ注意機構（squared-relu を使用）でもロジット成長不安定性は発生
   • これは不安定性がソフトマックスに起因するものではないことを示唆

3.3 Predicting attention logit growth instability from scaling behavior of model characteristics

1. 目的

   • モデル特性のスケーリング挙動からアテンションロジット成長不安定性を予測する方法の探索

2. 手法

   • モデルスケールと最大アテンションロジット値の関係を追跡
   • データに二次曲線をフィットさせ外挿予測を行う

3. 主な発見

   • 最大アテンションロジットが $1e4$ を超えるとモデルが発散
   • 学習率 $1e-2$ で次のモデルスケールがこの閾値を超えると予測

4. 予測の検証

   • 予測に基づき 4.8B パラメーターのモデルを学習率 $1e-2$ で訓練
   • 予測通り、このモデルは発散
   • フィットした曲線が最大アテンションロジット値を正確に予測

5. 閾値の妥当性検証

   • 10M パラメーターの小規模ネットワークで異なる最大アテンションロジット値を強制
   • クエリーとキーに $g(z) = \sqrt{\kappa} \cdot z / \sqrt{E_i[z^2_i]}$ を適用
   • $\kappa \approx 1e3$ で損失が悪化し始め $\kappa \approx 1e4$ でゼロレイヤーのバイグラムモデルより悪化

6. 実践的示唆

   • この手法により大規模モデルでの不安定性を事前に予測可能
   • モデル設計や学習戦略の最適化に活用できる

モデル特性のスケーリング挙動を分析することで大規模モデルでの不安定性を小規模実験から予測できる。これにより計算資源を効率的に使用しながら安定して大規模モデルを開発する事が可能になる可能性がある

3.4 Searching for new instabilities via scaling trends of model characteristics

1. 目的

   • モデル特性のスケーリング挙動を用いてデフォルトのモデルとハイパーパラメーター設定における新たな問題を予測

2. 主な発見

   • 勾配の二乗平均平方根（RMS）が $\text{RMS}(g) = \sqrt{E_i[g^2_i]}$ モデルサイズと学習率の増加に伴い減少
   • 大規模モデルと高学習率では勾配 RMS が AdamW の $\epsilon$ ハイパーパラメーター（1e-8）に近づく

3. 問題点

   • 勾配 RMS が $\epsilon$ と同程度になると AdamW の更新 $\Delta = v / (\sqrt{u} + \epsilon)$ が減少
   • これによりパラメーターが適切に学習信号を受け取れなくなる可能性がある

4. 緩和策

   • AdamW の $\epsilon$ を 1e-15 に下げることで性能が向上
   • 4.8B パラメーターのモデルで学習率 0.3 での実験で確認

5. メカニズム

   • 大きなネットワークと高学習率では、Transformer 出力の RMS が成長
   • 最終レイヤーノームの勾配が入力 RMS の逆数でスケールされるため、Transformer への勾配が縮小

6. 追加観察

   • $\epsilon$ を 1e-6 に増やすと不安定性が発生

7. 実践的示唆

   • 大規模モデルでは AdamW の $\epsilon$ をデフォルト値より小さくすることを検討すべき
   • この調整はモデルのスケールが大きくなるほど重要になる可能性がある

4 Related work

1. 主な関連研究

   • Dehghani et al. と Zhai et al. によるアテンション成長不安定性の観察
   • Chowdhery et al. と Thilak et al. による最終ロジット発散問題の研究

2. 類似の実験手法

   • Yang et al. : スケールに応じた学習率と損失の関係を測定（最適値の中心化に焦点）
   • Zhai et al. : 学習率を 2 倍にしてベースモデルの不安定性を誘発
   • Dettmers et al. : スケールに応じた外れ値特徴の存在を測定

3. Edge of Stability (EoS) と高速スパイク

   • Cohen et al. : 大規模バッチ学習での進行性シャープニングと自己安定化の相互作用を研究
   • Adaptive EoS (AEoS) : Adam の $\beta_2$ パラメーターが損失スパイクの発生に影響

4. パラメーターフリー手法

   • Ivgi et al. : 学習率を指定する必要のない手法を提案

5. その他のパラメーター化

   • 学習率 vs 損失の関係を探る他のパラメーター化手法の存在を示唆

6. 損失スパイクに関する研究

   • Shazeer and Stern, Chen et al., Zhang et al. : Transformer の適応的最適化器での損失スパイクを研究
   • $\beta_2$ の低減が損失スパイクを解決する観察

7. 本研究の位置づけ

   • 既存の不安定性と介入効果を小規模モデルで再現・研究する手法を提案
   • モデル特性のスケーリング挙動を分析し新たな不安定性を予測・特定する手法を導入

The Edge of Stability and fast spikes

1. 従来の理解

   • 勾配降下の不安定性は学習率が $2/\lambda_{\max}(H)$ を超えた時のみ発生（$H$ はヘッシアン）

2. Edge of Stability (EoS) の発見

   • 大規模バッチ SGD での学習動態に関する新たな洞察
   • 進行性シャープニングと自己安定化の相互作用により曲率が常に進化

3. 進行性シャープニング

   • 学習率 < $2/\lambda_{\max}(H)$ の時、曲率が徐々に増加
   • 安定性の閾値を超えると高速損失スパイクが発生

4. 自己安定化

   • 不安定性により $\lambda_{\max}(H)$ が小さい領域にパラメーターが振動
   • 安定した学習と進行性シャープニングが再開

5. Adaptive Edge of Stability (AEoS)

   • Cohen et al. による適応的最適化器での EoS の分析
   • 事前条件付き鋭さ $\lambda_{\max}(P^{-1}H)$ が最適化器固有の閾値周辺で振動

6. $\beta_2$ パラメーターの影響

   • $\beta_2$ が大きい場合 : $H$ が縮小し高速損失スパイクが発生
   • $\beta_2$ が小さい場合 : $P^{-1}$ が縮小し損失スパイクは観察されない

7. 関連研究との関連

   • Shazeer and Stern, Chen et al., Zhang et al. による損失スパイクの研究を AEoS の観点から説明可能
   • $\beta_2$ の低減が損失スパイクを解決する観察と一致

5 Conclusion

1. 研究の背景

   • 大規模 Transformer モデルの学習に必要な計算リソースが増加し続けている
   • 学習の安定性を理解することが重要になっている

2. 主な成果

   • 小規模 Transformer モデルを用いて学習の安定性に関する有用な洞察を得られることを示した
   • 学習率と損失の関係を異なるスケールで測定することが不安定性を特定する有効なツールであることを発見

3. 研究手法

   • 大規模モデルで報告された不安定性を小規模モデルで再現
   • モデル特性（勾配ノルムやアクティベーションノルムなど）のスケーリング挙動を分析

4. 主な発見

   • アテンション層のロジット増大と出力ロジットの発散という 2 つの不安定性を小規模モデルで再現
   • 大規模モデルで有効だった緩和策が小規模モデルでも機能することを確認
   • モデル特性のスケーリング挙動を分析することで、不安定性を予測できる可能性を示した

5. 研究の意義

   • 大規模な計算リソースがなくても Transformer モデルの学習安定性に関する研究が可能になる
   • 大規模モデル学習の効率化や安定化に向けた新たな研究機会を提供

6. 今後の展望

   • この研究アプローチが大規模モデル学習に有益な洞察をもたらす可能性がある

不安定性を低減する方法まとめ

LLM の学習不安定性を低減する方法をまとめると以下のようになる

1. qk-layernorm の使用

   • アテンション層のクエリーとキーに LayerNorm を適用
   • アテンションロジット成長不安定性を緩和

2. z-loss 正則化

   • 出力ロジットの発散を防ぐ補助損失を追加
   • $\log^2 Z$ を損失関数に加える（$Z$ は出力ロジットの和）

3. 適切なウォームアップ

   • より長いウォームアップ期間を使用
   • 特に大規模モデルで効果的

4. 独立重み減衰

   • 学習率とは独立して重み減衰を適用
   • LR 感度を低減

5. AdamW の $\epsilon$ パラメーター調整

   • 大規模モデルでは $\epsilon$ をデフォルト値（1e-8）より小さく設定
   • 例 : 1e-15 に設定して性能向上

6. 深さと幅のバランス

   • モデルスケーリング時に深さと幅を適切にバランス
   • 深さの増加は性能向上につながるが不安定性も増加

7. μParam の適用

   • 最適学習率の安定化に効果的
   • ただし他の不安定性には直接影響しない

8. 適応的最適化器のハイパーパラメーター調整

   • Adam の $\beta_2$ パラメーターを適切に設定
   • 小さな $\beta_2$ 値で損失スパイクを抑制

9. モデル特性のモニタリング

   • 勾配 RMS やアテンションロジットの最大値など重要な特性を監視
   • スケーリング挙動に基づいて不安定性を予測し事前に対策

10. バッチサイズと学習率の適切な設定

    • モデルサイズに応じて適切に調整

これらの方法を組み合わせることで LLM の学習不安定性を大幅に低減し、より効率的で安定した学習が可能になる。ただしモデルのアーキテクチャーやタスクに応じて、これらの方法の効果は異なる可能性があるため実験的な検証が重要である