単純なSGDではANDREAを訓練できない
確率的勾配降下法、始まりの点
バックプロパゲーションはすべてのパラメータに対して勾配gを計算します。単純な確率的勾配降下法(SGD)は、各パラメータをp -= lr * gで更新します。一つの学習率、ステップごとの一方向、過去の勾配の記憶なし。
単純なSGDはスケールで2つの理由で壊れます:
1. 勾配の大きさがパラメータ間で極端に異なる。まれなトークンの埋め込みはほとんどのステップで小さな勾配を受け取り、layernormのスケールは大きな勾配を受け取る。一つの学習率では両方に適さない。
2. 勾配が振動する。16ソースのコーパスからのノイジーなミニバッチがパラメータを左に押し、次に右に、次に左に押しやる。Plain SGDは自身と戦ってステップを無駄にする。
Adam (Kingma & Ba, 2015) は、パラメータごとに2つの移動平均で両方を修正する。
第一モーメント & 第二モーメント
m: 平滑化された方向
第一モーメント m は最近の勾配を指数的に平均する:
m = beta1 m + (1 - beta1) g
beta1 = 0.9 で。何度かステップを進めた後、m は平滑化された方向性を保持します。一つの悪いバッチではほとんどずれない。
v: 平滑化された大きさ
2番目のモーメント v は最近の二乗勾配の平均を取ります:
v = beta2 v + (1 - beta2) g^2
beta2 = 0.999 で、v は各パラメータの勾配が通常どれくらいの大きさになるかを追跡します。勾配の大きいパラメータは大きな v を得、小さな勾配のパラメータは小さな v を得ます。
パラメータごとの適応的学習率
平滑化された方向を平滑化された大きさの平方根で割ることで、すべてのパラメータを同等の基盤に再スケーリングします:
adam_step = m / sqrt(v + eps)
小さな勾配の埋め込みはスケールアップされ、大きな勾配のlayernormはスケールダウンされます。これで1つのグローバルなlrがすべてのパラメータに適します。
モーメントの読み取り
初期ステップでバイアス補正が必要な理由
コールドスタート・バイアス
m & v はゼロから開始します。ステップ1後、m = 0.1 g_1 & v = 0.001 g_1^2。両方の推定値は長期平均を大幅に下回ります。補正なしでは、オプティマイザは慎重に開始し、ゆっくりと加速し、表現が形成される貴重な初期ステップを無駄にします。
補正
Adam は各推定値を 1 / (1 - beta^t) でスケーリングします。ここで t はステップ番号です:
m_hat = m / (1 - beta1^t)
v_hat = v / (1 - beta2^t)
ステップ1で beta1 = 0.9 の場合、分母 (1 - 0.9) = 0.1 となり、m_hat = m / 0.1 = 10 * m となります。バイアス補正された推定値は、長期的平均が予測するものと一致します。t が増加するにつれ、beta^t は0に近づき、補正値は1に近づき、補正された値と補正されていない値が収束します。
デカップルド重み減衰 (AdamWの革新)
L2正則化 vs 重み減衰
古典的なL2正則化は、損失にペナルティを追加します:L_total = L_data + (lambda / 2) sum(p^2)。バックプロパゲーションはそのペナルティを勾配の一部として見なし、g_total = g_data + lambda p となります。L2項はAdamの m および v 更新を通過し、パラメータごとの大きさによって平滑化され、再スケーリングされます。
Loshchilov & Hutter (2019) は、Adam を通じて正則化項を平滑化すると両方が破綻することを証明しました。Adam の適応的スケーリングは、大勾配パラメータ(ここで減衰がオーバーフィッティングと最も強く戦うべき場所)での weight decay を縮小し、小勾配パラメータではそれを増幅します。
AdamW: 減衰を直接適用
AdamW は weight decay を勾配から分離します。減衰はパラメータ更新中に各パラメータに直接適用され、m や v には一切触れません:
p -= lr (m_hat / (sqrt(v_hat) + eps) + weight_decay p)
これで各ステップを駆動する項が 2 つになります:
1. Adam項: m_hat / (sqrt(v_hat) + eps) は、勾配の方向をパラメータごとの大きさの履歴で再スケーリングします。
2. 減衰項: weight_decay * p は、Adamのスムージングを経由せずに、すべてのパラメータを一様にゼロに向かって縮小します。
ANDREA-120M v2 は weight_decay = 0.01 を設定します。毎ステップ、すべてのパラメータがAdam項が行うことに加えて、1%ずつゼロに向かって縮小します。
なぜDecoupledが重要か
実証的証拠
v1 の崩壊(weight decay なし)
ANDREA-120M v1 は vanilla Adam で 165K ステップ訓練。サンプル出力:
- ステップ 80K: region region region region region region region
- ステップ 110K: ''''' ''''' '' ''' '' ''' '''?' ''' ' '' '' '
- ステップ 140K: games, games, games, games, games, games, games
- ステップ 165K: Budy Budy Budy Budy Budy Budy Budy Budy Budy
損失数は妥当な範囲を維持(ステップ 110K で EMA 最小値 3.23、ランダム推測の 9.04 に比べて)。損失だけでは繰り返し崩壊が隠れてしまう:1つのトークンを永遠に記憶するモデルは、そのトークンが現れるすべてのステップで低いクロスエントロピーを達成する。
v2 の安定性 (weight_decay = 0.01)
v2 では AdamW (プラス勾配クリッピング、LR ウォームアップ、サンプル監視) を追加しました。ステップ ~112K で生成されたサンプル:
- Carolina parakeet was declared extinct in 1939 (事実的に正しい)
- The Fourier transform decomposes signals into frequency components (教科書的な定義)
- Rain's rhythmic refrain, Rivulets on the window, Respite from life's pain (俳句の制約を満たしている)
外部評価では v2 のサンプルを 9.5/10 と評価し、「この規模での印象的な一貫性と知識保持」と称賛しました。
AdamWなしで12Mが生存した理由は?
ANDREA-12MはバニラAdamで訓練され、崩壊しませんでした。12Mパラメータでは、重み行列が十分に小さく保たれるため、Adamの適応スケーリングが個々の重みを繰り返しを引き起こす暴走した大きさに押し込むことができません。120Mスケールでは、重みの大きさがステップごとにさらにドリフトし蓄積します。均一な減衰はゼロに向けた一定の復元力を適用します。モデルがスケールするにつれて、デカップルされた重み減衰の重要性が増します。
weight_decay = 0.01 の選択
関連アクティビティ
AdamWはこのコースの3つの姉妹アクティビティと相互接続しています:
- アクティビティ11: LR warmup + cosine decay. AdamW単独では、新規初期化された重みでの即時ピーク学習率からモデルを救うことはできません。Warmupは2000ステップにわたってlrを徐々に増加させ、AdamWのバイアス補正と重み減衰が表現を安定させる時間を確保します。
- アクティビティ12: 勾配クリッピング. AdamWは勾配の大きさが制限されていることを前提としています。ANDREAのbanditで発生するソース遷移(7〜42ステップごと)は時折勾配スパイクを生み出します;クリッピングはAdamWがm、v、またはpに触れる前に、それらをL2ノルム1.0で制限します。
- アクティビティ13: FP32 / FP16 / FP8 精度. AdamWはパラメータごとにmとvを保存するため、重み単独のメモリ使用量を2倍にします。FP16はそれを半分に、FP8はさらに半分に削減します。精度の選択はオプティマイザの安定性と相互作用します。
AdamW、warmup、clipping、precision は四つ葉のクローバーを形成します。一枚の葉を落とせば、ANDREA が崩壊するのを観察せよ。