訓練の両端で生じる2つの問題
初期ステップの問題:新鮮な重みは大きなステップを取れない
ステップ 0 では、すべての重みがほぼゼロの分布からサンプリングされた小さなランダム値として初期化されます。活性化はほぼゼロのままです。勾配は最終解に関するほとんど情報を含んでいません。このような勾配にピーク学習率を適用すると、モデルは意味のある構造をエンコードしない方向へ初期化点から大きく飛び出してしまいます。
ANDREA-120M v1 はこのミスを犯しました。 ウォームアップなし。ステップ 1 で新しく初期化された重みに対して lr = 0.0003 を使用しました。結果:モデルは数百ステップ以内に悪いパラメータ盆地に陥りました。損失値は妥当に見えましたが、ステップ 80K までにサンプルが繰り返しのループを生成し、回復することはありませんでした。
後期ステップの問題:大きなステップでは解を洗練できない
ステップ100K時点で、モデルは大まかな構造を学習し終えている。勾配は今や細かな信号を運んでいる:どのトークンの重みをわずかに調整すべきか、どのアテンションヘッドをわずかに再調整すべきか。この段階で最大学習率を適用すると、すべての細かな調整を飛び越えてしまい、最適値の周辺で振動し続け、収束できなくなる。
2つの問題は、訓練の両極端に位置している。1つのスケジュールで2つの領域を扱う:穏やかに立ち上げ、滑らかに減衰させる。
線形ウォームアップ:最初の2000ステップ
数式
ANDREA-120M v2 は2000ステップにわたる線形ウォームアップを使用する:
lr(t) = lr_scheduled(t) * min(1, (t + 1) / warmup_steps)
ここで、t はステップ番号(0から始まる)、warmup_steps = 2000、および lr_scheduled(t) はウォームアップを無視した場合のコサインスケジュールが定める値です。
式の読み方:
- t = 0 のとき:lr = lr_scheduled(0) min(1, 1/2000) = lr_scheduled(0) 0.0005。最初のステップは非常に小さな値になります。
- t = 1000 のとき: lr = lr_scheduled(1000) min(1, 1001/2000) = lr_scheduled(1000) 0.5。半分の強さ。
- t = 2000 のとき: lr = lr_scheduled(2000) min(1, 2001/2000) = lr_scheduled(2000) 1.0。フル強度。
- t > 2000 のとき: クランプにより乗数は 1 に固定され、ウォームアップは影響しなくなり、コサイン減衰のみが作用する。
ゼロからの線形ランプにより、モデルは 2000 ステップかけて粗い表現を形成し、その後 AdamW と勾配クリッピングがフル強度の更新を受け取る。ステップ 2000 時点で重みは十分に移動しており、ピーク lr が悪い盆地へ押しやることはなくなる。
ウォームアップ中の LR の計算
ウォームアップ後のコサイン減衰
曲線
ステップ2000でウォームアップが終了した後、学習率は残りのステップでピークからゼロまでコサイン曲線に従って減少します:
lr(t) = lr_min + (lr_peak - lr_min) 0.5 (1 + cos(pi * progress))
ここで progress = (t - warmup_steps) / (total_steps - warmup_steps)。progress = 0(ウォームアップ直後)では cos(0) = 1 となり、lr = peak。progress = 1(最終ステップ)では cos(pi) = -1 となり、lr = lr_min(通常は 0 または極小のフロア値)。
なぜ Cosine なのか?Linear や Exponential ではなく?
Cosine decay は、開始時は緩やか(ピーク付近でほぼ平坦)、中盤で加速し、終盤で再び減速する。3 つの利点がある:
1. ピーク付近のプラトー。ウォームアップ直後のステップでもほぼ最大の lr が維持され、モデルが長い高学習率期間を利用して表現を構築できる。
2. 中盤での滑らかな遷移。AdamW が吸収しなければならない急激な変化がない。
3. ゼロ近傍のプラトー。 最終ステップでは微小な lr により微調整が行われ、シミュレーテッドアニーリングに似た挙動を示します。
ANDREA-120M は合計 200K ステップで学習します。そのうち 198K ステップは、2000 ステップのウォームアップ後に続くコサイン減衰領域です。
ANDREA-12M のステップ 25K におけるウォームリスタート
プラトー
ANDREA-12M は lr = 0.0004 のピークからコサイン減衰を用いて 60K ステップ学習しました。ステップ 22K 付近で損失が EMA ≈ 2.4 でプラトー状態になりました。コサイン減衰により lr は約 0.00015 まで低下していました。バンディットは多様なデータを供給し続けていましたが、モデルの改善は止まりました。
診断: lr が現在の盆地から抜け出すには減衰しすぎていた。Hermes データがカリキュラムに入る直前(ステップ 25K)で、590K の新しい会話が追加される予定だった。モデルはそのデータショックを吸収するためのエネルギーを必要としていた。
再起動
ステップ 25K で、スケジュールはウォーム再起動を実行した: lr を減衰後の 0.00015 から元のピーク値 0.0004 まで急上昇させ、その後残りのステップでコサイン減衰を再開した。
Loshchilov & Hutter (2017) はこの手法を「SGDR」(stochastic gradient descent with warm restarts)と名付けた。直観的には、高い lr が十分な運動エネルギーを与えて局所的な盆地から脱出し、隣接する盆地を探索できるようにする。その後のコサイン減衰により、より良い盆地へと再アニーリングされる。
Outcome. Loss EMA dropped from 2.40 to 2.10 over the next 10K steps after the restart. Model shipped at step 43.6K with SMMA loss 2.0, demonstrating coherent Q&A turn structure.
ANDREA-120M v2 chose NOT to use warm restarts: with 200K steps available & a much larger parameter count, smooth monotonic decay produced steadier convergence. Restart works best when training is short & a plateau coincides with a known data shift.
v1の失敗を診断する
実践におけるスケジュール選択
関連アクティビティ
LRスケジュールに関連する3つの兄弟リンク:
- Activity 10: AdamW. ウォームアップにより、AdamWのバイアス補正が安定する時間を確保できます。ウォームアップがない場合、ステップ1での10倍増幅がノイズを含む勾配を増幅しますが、ウォームアップがあることで乗数は実際の信号に作用します。
- Activity 12: Gradient clipping. 勾配クリッピングは、AdamWの前に勾配のL2ノルムを1.0に制限します。ウォームアップはlrを抑制し、クリッピングはgを抑制します。これらを組み合わせることで、衝撃を受けやすいカリキュラムでも初期ステップを安全に保てます。
- Activity 22: Checkpointing. ウォームリスタートでは、チェックポイントからオプティマイザ状態(m、v、ステップカウンター)を読み込み、実行中にスケジュールを変更する必要があります。ANDREA-12Mのステップ25Kでのリスタートはこの処理を示しており、状態読み込みロジックを正しく動作させるのに2回の試行が必要でした。
スケジュール、オプティマイザ、クリッピングは安定性の三角形を形成します。頂点を1つ落とすと、ANDREAはv1の崩壊を繰り返します。