学習率スケジュール — warmupとcosineの理由
学習率は固定の数字ではなく、訓練の最初から最後まで動かす曲線です。なぜ最初にわざと遅く走る(warmup)のか、なぜ余弦カーブで落とす(cosine)のか、バッチサイズを変えたら何を一緒に動かすのか。式・動く図・PyTorch実装・現場の落とし穴まで前提知識なしで解説します。
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. バッチが大きいほど、最初の一歩が効きすぎる。 バッチを大きくすると勾配のノイズが減り、全パラメータが揃って同じ方向へ動きます。まだ正しいと決まっていない方向へ全員で大きく踏み出すのは危険です。
コメント
コメントにはログインが必要です