学習率のwarmupは本当に訓練を安定させるのか。崩壊寸前のLRで20シード×3条件、実際に数えてみた
「学習率を高くして訓練するときは、最初にwarmup期間を置かないと訓練が不安定になる・序盤で発散する」——深層学習をやっていると、ほぼ常識のようによく聞く話だ。以前の記事でSGD・Momentum・RMSprop・Adamを競走させたときも、Adamの適応的な学習率について色々調べる中で、この「warmup神話」に何度も行き当たった。
ただ、自分の手で実際に確かめたことは一度もなかった。なので今回は、崩壊するかしないか際どいLRを実測で探し、そのLRでwarmupあり/なしを何十シードも回して、本当に効くのかを数字で見てみることにした。
先に結論を書いておく。このセットアップでは、warmupは訓練前半の暴れを確かに抑えたが、最終的に「崩壊」したシードの割合はwarmupなし・linear warmup・cosine warmupでほとんど同じだった(20シード中10/12/11)。むしろ後半のピーク損失はwarmupありの方がわずかに高かった。よく聞く教えほど単純ではなかった。
実験セットアップ: あえて正規化なしの深いMLP
わざと不安定になりやすい構成にした。
- モデル: 入力64(8×8のグレースケール画像)→128→128→128→128→10の全結合ReLU MLP(隠れ層4層)。BatchNormもLayerNormも残差接続も入れていない、素の多層パーセプトロン
- 重み初期化: He初期化(
N(0, 2/fan_in)) - データ: sklearn付属の
load_digits(手書き数字、1797サンプル、8×8=64次元、10クラス)。各特徴量を平均0・分散1に標準化し、8:2でtrain/testに分割(train 1437件、test 360件)。MNISTほど大きくない代わりに数十秒で何十回も訓練を回せるので、多シード実験向きの「ちょうどいい」実データにした - 最適化: Adam(β1=0.9, β2=0.999, eps=1e-8)を numpyで手書き実装。バッチサイズ32、1500ステップ(train setを約33周)
文章で「8×8のグレースケール画像」と書いてもどんなデータか想像しづらいと思うので、モデルが実際に見ているものを最初に目で確認しておく。
見ての通りかなり粗いが、人間の目でもぎりぎり読める程度には形が残っている。今回モデルにやらせるのは、この64個のピクセルの明るさの値を入力として受け取り、0〜9のどの数字が書かれているかを当てる10クラス分類だ。学習がうまくいったかどうかは、訓練に使っていないテスト用の360枚をどれだけ正しく分類できるか(テスト精度、以下test_acc)で測る。後で出てくる「崩壊」の判定も、このtest_accに対する閾値で決める。
BatchNormのような正規化層を入れないのは意図的だ。正規化があると学習率に対してかなり頑健になってしまい、そもそも「不安定になるLR」を作るのが難しくなる。今回はwarmupの効果を見たいので、不安定さが起きやすい「裸の」深いMLPを選んだ。
class MLP:
"""入力64 -> 128 -> 128 -> 128 -> 128 -> 10 の全結合ReLU MLP。numpyで手書き。"""
def __init__(self, seed, n_in=64, n_hid=128, n_layers=4, n_out=10):
r = np.random.default_rng(seed)
dims = [n_in] + [n_hid] * n_layers + [n_out]
self.W, self.b = [], []
for i in range(len(dims) - 1):
self.W.append(r.normal(0, np.sqrt(2.0 / dims[i]), (dims[i], dims[i + 1])))
self.b.append(np.zeros(dims[i + 1]))warmupの実装はシンプルで、warmupステップ数をwarmup_stepsとして、実際にAdamに渡す学習率cur_lrをステップごとに差し替えるだけだ。
if warmup_kind == "none" or t > warmup_steps:
cur_lr = lr
elif warmup_kind == "linear":
cur_lr = lr * t / warmup_steps
elif warmup_kind == "cosine":
frac = t / warmup_steps
cur_lr = lr * 0.5 * (1 - np.cos(np.pi * frac))warmup期間は全1500ステップの先頭10%(150ステップ)。linearは0からlrまで直線的に、cosineはコサインカーブでゆっくり立ち上がってから加速する。式だけだと立ち上がり方の違いがイメージしづらいので、実際にAdamに渡される学習率がステップに対してどう動くかを描いておく。なお、ここで言うcosineは学習率を徐々に下げていく有名な「cosine decay」とは逆向きで、cosineカーブで立ち上げるスケジュールのことだ。
違うのは黄色いwarmup区間の150ステップだけで、そこから先の1350ステップは3条件とも全く同じ学習率で走る。つまりこの実験は「最初の150ステップの入り方だけを変えると、その後の運命がどれだけ変わるか」を見るものだと言い換えられる。
まずLRを探す: 崩壊率は0.012→0.018の間で0%から50%へ跳ね上がった
「不安定になるLR」を最初から知っていたわけではない。恣意的に選ぶと説得力がないので、warmupなしのまま学習率を0.001から0.05まで動かして、8シードずつ回した。test_acc(テスト精度)が0.5を下回ったシードを「崩壊」と呼ぶことにする(10クラス分類のchanceレベルは0.1なので、0.5未満は訓練が実質的に破綻したとみなせる)。
lr=0.001 fails=0/8 mean_acc=0.970
lr=0.01 fails=0/8 mean_acc=0.961
lr=0.012 fails=0/8 mean_acc=0.945
lr=0.015 fails=2/8 mean_acc=0.761
lr=0.018 fails=4/8 mean_acc=0.577 <- 五分五分
lr=0.02 fails=4/8 mean_acc=0.502
lr=0.025 fails=5/8 mean_acc=0.374
lr=0.03 fails=8/8 mean_acc=0.255
lr=0.05 fails=8/8 mean_acc=0.1930.012までは8シード全部が安定して95%以上のテスト精度に収束する。0.015を超えたあたりから崩壊が混じり始め、0.018でちょうど8シード中4シードが崩壊する「五分五分」の境界に行き着いた。0.03以上になるともう全滅で、warmupがあってもなくても大差ないだろう極端な設定になってしまう。今回はこの0.018を本番実験のLRに採用した。「warmupが必要かどうか五分五分」の際どいラインを狙ったつもりだ。
スライダーで学習率を動かしながら、代表シード1本の損失曲線がどう変わっていくか確認できるようにしておいた。
本番: LR=0.018で20シード×3条件
warmupなし・linear warmup・cosine warmupのそれぞれについて、同じLR=0.018で20シード訓練した(モデル初期化とミニバッチの順序だけがシードで変わる)。
warmupなし : 崩壊 10/20 (50%) mean_acc=0.518 median_acc=0.514
linear warmup : 崩壊 12/20 (60%) mean_acc=0.435 median_acc=0.369
cosine warmup : 崩壊 11/20 (55%) mean_acc=0.498 median_acc=0.487正直、予想と違った。「warmupなしが一番崩壊しやすく、linear/cosineが目に見えて改善する」という絵を想像していたが、蓋を開けてみると崩壊率はほぼ同じ、むしろlinear warmupが数字の上ではわずかに悪い。20シードという規模でこの差(50% vs 60%)は誤差の範囲に収まるとみるのが妥当で、「warmupなし1強・warmupあり2強」のようなきれいな序列は出なかった。
3パネルを見ても、赤い線(崩壊したシード)と青い線(成功したシード)の混ざり具合はどのパネルでもほとんど同じに見える。warmup区間(黄色い帯)を過ぎたあとも、赤い線はコンスタントに湧き続けている。
学習の様子をアニメーションでも見てみる。
面白いのは、スパイクが上方向(損失が跳ね上がる)だけでなく下方向(損失がほぼ0まで落ちる)にも起きていることだ。Adamの更新幅は各パラメータでだいたい学習率のオーダーに正規化されるので、大きすぎる更新をすると、そのミニバッチに限ってたまたま出力がほぼ正解に一致してしまい損失が急落する、ということが起こりうる。次のステップではまた別のミニバッチで大きく外れて跳ね上がる。振動というより「毎回別の的に矢を放っている」ような暴れ方に見える。
warmupが効くのは前半だけ、後半には効いていない
崩壊率が変わらなかったので、warmupに全く意味がなかったのかというと、そうとも言い切れない。損失曲線の「ピーク値」を訓練の前半(1〜200ステップ、warmup区間150ステップを含む)と後半(201〜1500ステップ)に分けて平均すると、はっきり違う顔が見えてきた。
前半ピーク(平均) 後半ピーク(平均)
warmupなし 4.67 4.38
linear warmup 3.14 4.56
cosine warmup 3.40 4.32前半のピークはwarmupで明確に下がっている(4.67→3.14〜3.40)。ここだけ見れば「warmupは効く」という教えは正しい。ところが後半のピークはwarmupがあってもなくても変わらない(4.38 vs 4.56 vs 4.32)。むしろlinear warmupは後半の方が高い。
崩壊が起きるタイミングを調べてみると、それぞれの条件で損失が最大になったステップは訓練の最初から最後まで満遍なく散らばっていて、「序盤に集中している」という傾向は見られなかった。つまりこのモデルの不安定さは、Adamのモーメント推定が安定していない訓練開始直後だけの問題ではなく、LR=0.018という値そのものが、このアーキテクチャにとってはいつ踏んでも良い地雷のようなものになっている。warmupは最初の一歩を柔らかくするだけで、地雷そのものを取り除いてはくれない。だから最終的な崩壊率が変わらなかったのだと思う。
手を動かして意外だったこと
一番の発見は、「warmupは効かなかった」ではなく「warmupは効く場面と効かない場面がはっきり分かれた」ということだ。訓練の最初の一瞬だけを切り取れば、warmupは間違いなく仕事をしている。世間でよく紹介されるグラフや説明も、大抵は訓練序盤の挙動にフォーカスしている。でも今回のように、正規化なしの深いMLPを崩壊寸前のLRで長く回すと、危険は最初の一瞬だけでなく訓練全体に薄く広がっていて、warmupという「最初だけの処置」は根本原因には触れられない。
これは仕事の場面でも似たようなことがある気がする。プロジェクトの立ち上げ時に慎重に準備期間を置く(warmupに相当)のは確かに初動の失敗を減らすが、根っこにある無理な計画(高すぎるLR)自体を直さない限り、危うさは形を変えて後になって出てくる。導入部だけ丁寧にしても、土台そのものが危ういままなら、いつかどこかで踏み外す。「最初さえ乗り切れば大丈夫」という安心感は、意外と当てにならない。
もう1つ、20シードという規模で回してみて実感したのは、1〜2本の曲線だけを見て「warmupのおかげで綺麗に収束した」と結論づけるのがいかに危ういかということだ。今回もどの条件にも綺麗に収束するシードとそうでないシードが両方あって、都合の良い1本を選べばどちらの主張もできてしまう。分布を全部並べて初めて、「実はほとんど差がない」という地味だが正直な結論にたどり着けた。
まとめ
- 正規化なしの深さ4層ReLU MLP(numpy手書き、Adamも手書き)をsklearn digitsで訓練し、warmupなし/linear warmup/cosine warmupの3条件を比較した
- まずLRスイープ(8シード)で崩壊率が0.012の0%から0.018の50%へ急に立ち上がる境界を実測し、その境界のLR=0.018を本番実験に採用した
- LR=0.018で20シードずつ訓練した結果、崩壊率(test_acc<0.5)はwarmupなし10/20、linear warmup 12/20、cosine warmup 11/20とほぼ同じで、「warmupが崩壊率を下げる」という単純な図式は再現しなかった
- 訓練前半(1〜200ステップ)のピーク損失はwarmupで明確に下がった(4.67→3.14〜3.40)が、後半(201〜1500ステップ)のピークはほぼ変わらない、またはむしろ高かった(4.38→4.32〜4.56)
- 損失が最大になるステップは訓練全体に散らばっており、この設定での不安定さは訓練序盤だけの問題ではなかった。warmupは最初の一歩を柔らかくするが、LRそのものが高すぎるという根本原因には効かない
