Cosine Annealing(コサインアニーリング)
コサイン関数に従って学習率を緩やかに減衰させるスケジュール手法。LLM訓練における事実上の標準Decay方式で、滑らかな減衰により安定した収束を実現する。
Cosine Annealing(コサインアニーリング)とは
Cosine Annealingは、学習率をコサイン関数の半周期(πラジアン)に沿って減衰させるスケジューリング手法です。2017年にLoshchilovとHutterが提案した「SGDR: Stochastic Gradient Descent with Warm Restarts」で導入され、現在のLLM訓練では最も広く使用されるDecay方式となっています。
数学的定式化
Cosine Annealingの学習率は以下の式で計算されます:
η(t) = η_min + 0.5 × (η_max - η_min) × (1 + cos(π × t / T))
- η(t): 現在のステップtでの学習率
- η_max: 最大学習率(Warmup後のピーク値)
- η_min: 最小学習率(通常はη_maxの1/10〜1/100)
- T: 総Decay期間(Warmup後のステップ数)
この関数は、初期に緩やかに減衰し、中盤で最も急峻に減衰し、終盤で再び緩やかになるS字状の減衰曲線を描きます。
Linear Decayとの比較
| 特性 | Cosine Annealing | Linear Decay |
|---|---|---|
| 減衰曲線 | 非線形(S字型) | 一定勾配 |
| 序盤の減衰 | 緩やか | 一定 |
| 中盤の減衰 | 急峻 | 一定 |
| 終盤の減衰 | 緩やか | 一定 |
| 収束安定性 | 高い | やや低い |
| 採用例 | GPT-3/4、LLaMA | BERT、RoBERTa |
Cosine Annealingが優位な理由は、訓練序盤で高い学習率を長く維持し、終盤では十分に低い学習率で精密な最適化を行えることにあります。この特性がLLMの大規模パラメータ空間の探索に適しています。
Cosine Annealing with Warm Restarts(SGDR)
元論文で提案されたWarm Restart版は、コサイン曲線が最小値に達した時点で学習率を再びピーク値にリセットし、新しいサイクルを開始します。各サイクルの長さを倍増させる(T_mult=2)ことで、徐々に長い探索期間を確保します。ただし、LLM訓練ではRestart版はあまり使用されず、単一サイクルのCosine Decayが主流です。
最小学習率の設定
最小学習率η_minの設定は最終性能に影響します。多くの大規模LLMでは最大学習率の10%(η_min = 0.1 × η_max)が使用されます。
| モデル | η_max | η_min | 比率 |
|---|---|---|---|
| GPT-3 | 6e-5 | 6e-6 | 10% |
| LLaMA | 3e-4 | 3e-5 | 10% |
| Chinchilla | 1e-4 | 1e-5 | 10% |
| PaLM | 1e-2 | 1e-5 | 0.1% |
実装上の注意点
PyTorchではtorch.optim.lr_scheduler.CosineAnnealingLRが標準実装として利用可能です。Warmupとの組み合わせにはSequentialLRを使用するか、HuggingFace Transformersのget_cosine_schedule_with_warmupを利用するのが一般的です。
FAQ
Q1: なぜLLM訓練ではLinear DecayよりCosine Annealingが好まれるのですか?
Cosine Annealingは序盤に高い学習率を長く維持するため、loss landscapeの広い範囲を探索できます。LLMの巨大なパラメータ空間では、この探索フェーズが最終的な汎化性能に大きく影響します。
Q2: サイクル数を増やすWarm Restart版はLLMで使わないのですか?
LLMの事前学習は通常1 epoch以下で完了するため、複数サイクルのRestartは不要です。Restartが有効なのは、同じデータを複数epoch学習するCV(コンピュータビジョン)タスクが中心です。
Q3: Cosine Annealing中にloss spikeが発生した場合はどうすべきですか?
まずGradient Clippingの閾値を確認し、η_maxを下げるか、Warmup期間を延長します。中盤のspikeはデータ品質の問題であることも多く、データパイプラインの検証も推奨されます。