学習率スケジュール — warmupとcosineの理由
学習率は固定の数字ではなく、訓練の最初から最後まで動かす曲線です。なぜ最初にわざと遅く走る(warmup)のか、なぜ余弦カーブで落とす(cosine)のか、バッチサイズを変えたら何を一緒に動かすのか。式・動く図・PyTorch実装・現場の落とし穴まで前提知識なしで解説します。
要点
- 学習率スケジュールとは、学習率を訓練ステップ数の関数として書き下したものです。
- 勾配は向きと急さしか教えず、進む距離を決めるのが学習率で、ステップごとに決める値です。
- warmupは学習率を0付近からピーク値まで直線的に上げ、Adam系で外すと初期に壊れやすくなります。
- cosine減衰は始点と終点で傾きが0で、ピーク値と総ステップ数を決めれば形が一意に定まります。
- バッチサイズをk倍にしたら学習率もk倍にする線形則は、臨界バッチサイズを超えると効かなくなります。
Attention Is All You Need
一次資料 — この記事の根拠
論文の発表 2017-06-12→この解説の公開 2026-08-229年2か月後
Attention Is All You NeedAshish Vaswani, Noam Shazeer, Niki Parmar ほか · 2017-06-12 · v7arXiv:1706.03762論文ページ·PDFSGDR: Stochastic Gradient Descent with Warm RestartsarXiv:1608.03983論文ページ·PDF
An Empirical Model of Large-Batch TrainingarXiv:1812.06162論文ページ·PDF
Training Compute-Optimal Large Language ModelsJordan Hoffmann, Sebastian Borgeaud, Arthur Mensch ほか · 2022-03-29 · v1arXiv:2203.15556論文ページ·PDF
原文の要旨(Abstract)を読む
The dominant sequence transduction models are based on complex recurrent or convolutional neural networks in an encoder-decoder configuration. The best performing models also connect the encoder and decoder through an attention mechanism. We propose a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely. Experiments on two machine translation tasks show these models to be superior in quality while being more parallelizable and requiring significantly less time to train. Our model achieves 28.4 BLEU on the WMT 2014 English-to-German translation task, improving over the existing best results, including ensembles by over 2 BLEU. On the WMT 2014 English-to-French translation task, our model establishes a new single-model state-of-the-art BLEU score of 41.8 after training for 3.5 days on eight GPUs, a small fraction of the training costs of the best models from the literature. We show that the Transformer generalizes well to other tasks by applying it successfully to English constituency parsing both with large and limited training data.
原文の要旨(Abstract)を読む
We investigate the optimal model size and number of tokens for training a transformer language model under a given compute budget. We find that current large language models are significantly undertrained, a consequence of the recent focus on scaling language models whilst keeping the amount of training data constant. By training over 400 language models ranging from 70 million to over 16 billion parameters on 5 to 500 billion tokens, we find that for compute-optimal training, the model size and the number of training tokens should be scaled equally: for every doubling of model size the number of training tokens should also be doubled. We test this hypothesis by training a predicted compute-optimal model, Chinchilla, that uses the same compute budget as Gopher but with 70B parameters and 4$\times$ more more data. Chinchilla uniformly and significantly outperforms Gopher (280B), GPT-3 (175B), Jurassic-1 (178B), and Megatron-Turing NLG (530B) on a large range of downstream evaluation tasks. This also means that Chinchilla uses substantially less compute for fine-tuning and inference, greatly facilitating downstream usage. As a highlight, Chinchilla reaches a state-of-the-art average accuracy of 67.5% on the MMLU benchmark, greater than a 7% improvement over Gopher.
比喩: 凍った駐車場から、白線の枠に停めるまで
雪の残る駐車場から車を出し、向こう側の枠にバックで停める。発進直後は路面が凍っているか分からないので、じわりと踏んでグリップを確かめます。枠まで距離がある区間は遠慮なく速度を出す。枠が近づいたら減速する。速度を保ったまま突っ込めば白線を通り過ぎ、切り返してまた通り過ぎ、いつまでも収まりません。
ニューラルネットの訓練も同じ3局面をたどります。「じわりと踏む」が warmup、「速度を出す」がピーク学習率、「減速して微調整」が decay(cosineなど)です。学習率スケジュールとは、このアクセルワークを訓練ステップ数の関数として書き下したものにすぎません。
学習率とは何を決めている数字か
ニューラルネットの学習は、重みを少しずつ書き換える作業の繰り返しです。1回の書き換えはこう書けます。
は ステップ目のパラメータ、 は損失 が下がる向きとその急さを表す矢印(勾配)、 が 学習率 です。式(1)は「いまいる場所から、坂を下る向きへ 倍の歩幅で1歩進む」と言っているだけです。
つまり、記号を日常語に戻すと「新しい重み = いまの重み −(歩幅)×(下り坂の向き)」。マイナス記号が「登らずに下る」ことを、掛け算が「向きは勾配が決め、距離は学習率が決める」という役割分担を担っています。勾配そのものの中身は損失関数と最適化で分解しています。
肝心なのは、勾配は向きと急さしか教えてくれない点です。「どれだけ進めば底に着くか」は入っていません。その距離を決めるのが学習率で、適正値は場所によって変わります。だから には添字 が付く。定数ではなく、ステップごとに決める値なのです。
同じ歩幅で走り続けると何が起きるか
学習率が小さすぎれば、正しい方向へ進んではいるのに一向に着きません。GPU時間は有限なので、これは事実上の失敗です。
大きすぎると、谷を1歩で飛び越えて反対斜面のより高い場所に着地します。そこでも勾配は谷を向いているのでまた飛び越える。往復のたびに振れ幅が広がり、損失は発散して NaN になります。
厄介なのは、その境目が事前には分からないことです。次の図で綱渡りを手で確かめてください。スライダーを上げると収束が速くなり、あるところを越えた瞬間に球が谷から飛び出します。
図から分かることは2つです。良い学習率には上限があること、そして谷底に近づくほど適正な歩幅は小さくなること。ならば「最初は大きく、後で小さく」すればいい — decayの発想はここから出てきます。
「では最初から小さめで通せば安全では」と思うかもしれません。実際それでも訓練は進みます。ただし同じ計算予算で到達できる損失は明らかに悪くなる。スケジュールは「壊れないため」ではなく「限られたGPU時間で最も低い損失に着くため」の道具です。
そしてもう一段ひねりがあります。現代的な訓練のほとんどは、素朴な右下がりではなく、わざと小さい学習率から始めて、いったん上げてから下げます。この不思議な山型が warmup + cosine です。
warmup — 最初の数百ステップをわざと遅く走る
warmup は単純です。学習率を 0 付近から始め、決めたステップ数 をかけてピーク値 まで直線的に引き上げます。
式(2)は「warmup区間では、進んだ割合ぶんだけ学習率を上げる」という意味です。1000ステップのwarmupなら、100ステップ目はピークの10%。それだけです。
つまり は「助走のうちどこまで来たか」を 0 から 1 で示す進捗メーターで、そのメーターの読みをそのままピーク値に掛けている、ということです。
問題は、なぜそんなことをするのか。理由は3つあり、どれか1つではなく重なっています。
1. 開始時の重みは乱数である。 損失は高く勾配は大きいのに、その向きが正しい保証はありません。この状態で大きな1歩を踏むと、乱数よりさらに悪い場所へ飛ぶことがあります。最初の数歩は、進むためというより情報を集めるための歩みだと考えてください。
2. Adamの内部統計が、まだ数個の勾配からの推定でしかない。 Adam は勾配の二乗平均を移動平均で推定し、その平方根で割ってパラメータごとの実効歩幅を決めます。訓練開始直後はこの推定値が小さくばらつき、小さい数で割るので実効学習率が跳ね上がります。RAdam の論文(Liu ら, 2020)はこの初期の分散を warmup の主因に挙げました。説明は研究者間で一本化されていませんが、Adam系で warmup を外すと初期に壊れやすい現象自体は広く再現されています。
3. バッチが大きいほど、最初の一歩が効きすぎる。 バッチを大きくすると勾配のノイズが減り、全パラメータが揃って同じ方向へ動きます。まだ正しいと決まっていない方向へ全員で大きく踏み出すのは危険です。
実例: Transformer原論文の山型スケジュール
上りと下りを1本の式にまとめた古典例が、Transformer 原論文(Vaswani ら, 2017)のスケジュールです。
はモデルの隠れ次元です。式(3)は「 までは に比例して上げ、過ぎたら に反比例して下げる」と読みます。2項の min を取ることで山型が1本の式になっています。原論文の設定は でした。
つまり先頭の は「モデルが太いほど全体の歩幅を一律に控えめにする」割引係数、min の中身は「上り坂の直線」と「下り坂の曲線」の2本で、そのステップで低いほうが採用される、ということです。
cosine — なぜ「余弦」で落とすのか
ピーク後の下げ方でいま最も標準的なのが cosine(余弦)減衰で、Loshchilov と Hutter の SGDR 論文(2017)の形が元になっています。
は最後に残す下限、 は訓練全体のステップ数です。式(4)は「warmup終了からゴールまでの進捗を角度 から に写し、その余弦で から へなめらかに落とす」と言っています。進捗が半分なら、学習率はちょうど中間値です。
つまり括弧の中の は、訓練が進むにつれて 2 から 0 へなめらかに下りていくダイヤルです。それを半分にして「上限と下限の差」に掛け、最後に下限を足し戻す — この3手でダイヤルの目盛りが学習率の値に変換される、ということです。
なぜ直線や階段でなく余弦なのか。始点と終点で傾きが0になり、中盤でいちばん急に落ちる形だからです。ピーク直後は速度を保って距離を稼ぎ、終盤は微調整のためゆっくり0へ近づく。階段状の step decay は同じことを段差で近似しますが、「いつ何倍に落とすか」を人間が決めなければなりません。cosine はピーク値と総ステップ数さえ決めれば形が一意に定まる — 調整すべき数が減ることが、実務で選ばれる理由です。
ただしその「総ステップ数で形が決まる」性質は落とし穴でもあります。Chinchilla 論文(Hoffmann ら, 2022)は、cosine の周期長を実際の訓練ステップ数に合わせるべきだと明示しました。周期を長く設定して途中で打ち切ると、学習率が下がりきらないまま終わるため最終損失が悪化します。同じモデル・同じデータでも設定ひとつで結論が変わる。この点がスケーリング則の実験でどう効いたかはスケーリング則を1から解説で扱っています。
近年は、長い定数区間を保って最後だけ短く落とす WSD(Warmup–Stable–Decay)型も使われます。「あと何ステップ回すか」を途中で変えたい継続事前学習では、総ステップ数に形が縛られない点が利点です。
バッチサイズとの関係 — 一緒に動かす2つのつまみ
式(1)の は、実際にはバッチ内サンプルの勾配の平均です。バッチを大きくすると、この平均のノイズが小さくなります。つまり大きなバッチでの1歩は「同じ歩幅でも、より確からしい方向へ踏む1歩」です。確からしいなら、もっと大きく踏んでよい。
これを最も単純に定式化したのが Goyal らの linear scaling rule(2017)で、「バッチサイズを 倍にしたら学習率も 倍にする」と述べます。彼らはこの規則と数エポックの gradual warmup を組み合わせ、ImageNet の ResNet-50 を非常に大きなバッチで訓練できることを示しました。ここでも warmup が相棒として登場します。
ただし比例は無限には続きません。McCandlish らの An Empirical Model of Large-Batch Training(2018)は、勾配ノイズの大きさから決まる臨界バッチサイズを導入しました。それより小さい領域ではバッチを増やすとステップ数を比例して減らせますが、超えるとステップ数はほとんど減らず、計算だけが無駄になります。ノイズがすでに十分小さく、平均をさらに精密にしても得るものがないからです。
実務上の含意は明快です。実効バッチが膨らんだら、ピーク学習率も上げる。分散データ並列では実効バッチが自動的に world_size 倍になるので、ノードを増やした瞬間にこのつまみが暗黙に動きます(分散学習を1から)。なお SGD では線形則、Adam では平方根に比例させる規則が経験的に語られますが、タスクや実装に依存し、確立した唯一の法則ではありません。小規模実験で確かめるのが安全です。
コードで書く: warmup + cosine
PyTorch なら LambdaLR に「ピーク学習率に対する倍率」を返す関数を渡すだけです。
import math
from torch.optim.lr_scheduler import LambdaLR
def warmup_cosine(step, warmup, total, min_ratio=0.1):
if step < warmup:
return step / max(1, warmup) # 0 → 1 へ直線
p = (step - warmup) / max(1, total - warmup) # 0 → 1 の進捗
return min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * p))
sched = LambdaLR(opt, lambda s: warmup_cosine(s, warmup=2000, total=100_000))
# opt.step() を呼んだ「直後」に sched.step()。勾配累積中の micro-batch では呼ばない
最後の行のコメントが、実際にいちばん事故を起こす箇所です。学習率とモーメンタムの相互作用も、次の図で一度見ておくと勘が働きます。
現場ではこう使う
誰がいつ触るか。 事前学習・継続事前学習・ファインチューニングのジョブを投げる ML エンジニアが、実行前の設定ファイルで決めます。学習率は途中で変えると以降の訓練が別物になるため、投げる前の10分がほぼすべてです。
実際に触る名前。 Hugging Face TrainingArguments なら learning_rate(=ピーク値)、lr_scheduler_type="cosine"、warmup_ratio か warmup_steps、max_steps、gradient_accumulation_steps。関数を直接使うなら get_cosine_schedule_with_warmup。素の PyTorch は LambdaLR、CosineAnnealingLR(T_max は「step() を呼ぶ回数」であってエポック数ではない)、OneCycleLR。DeepSpeed の設定JSONでは scheduler.params.warmup_num_steps と total_num_steps。
必ずログに出すもの。 損失だけでなく学習率そのものと勾配ノルムを W&B や TensorBoard に出します。スケジュールの事故は、損失カーブより学習率カーブに先に現れます。
事故になる落とし穴。
- 勾配累積で
sched.step()を micro-batch ごとに呼ぶ。gradient_accumulation_steps=8なら optimizer の更新は8回に1回です。毎回叩くとスケジュールは8倍速で走り切り、訓練の大半を下限学習率で過ごします。損失は下がるので気づきにくい。 - 再開時にスケジューラの
state_dictを復元し忘れる。 学習率がピークへ戻り、損失が跳ねます。プリエンプティブなGPUでは致命的です。 - 総ステップ数をあとから変える。
max_stepsを変えると cosine の形自体が変わり、周期長と実訓練長がずれた訓練は最終損失で損をします。 - cosine が下がりきる前のチェックポイントで良し悪しを判断する。 減衰の終盤で損失は目に見えて落ちるので、7割地点の値で切ると誤判定します。
- バッチサイズだけ変えて学習率を据え置く。 GPUを増やした、OOM回避で
per_device_train_batch_sizeを半分にした — どちらも実効バッチが変わっています。 - warmup を短くしすぎる。 序盤の損失スパイクや
NaNは、まず warmup の長さとピーク値を疑う。対症療法より先にそこを直します。
設計レビューで問われること。 「なぜ warmup が要るのか」は、初期パラメータの信頼性・Adam の2次モーメントの推定分散・大バッチでのノイズの少なさ、の3点で答えられれば十分です。
まとめ
- 学習率は定数ではなく、ステップ数の関数として設計する曲線
- warmup は、初期パラメータ・Adam の内部統計・大バッチのノイズの少なさという3つの不安定さへの保険
- cosine は、決めるべきハイパーパラメータをピーク値と総ステップ数の2つに減らす下げ方
- バッチサイズを変えたら学習率も動かす。ただし臨界バッチサイズを超えると増やす意味は薄れる
- 実装で最も事故るのは勾配累積と再開。学習率カーブを必ずログに出す
コメント
コメントにはログインが必要です