Open1
Warmup-Stable-Decay (WSD)
bilzardpaper: Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations
Summary
- Cosine LR decayは学習ステップ数を最初から決めておく必要がある。またステップ数を変える場合は学習を最初からやり直す必要がある。
- 学習中の大部分は一定の学習率をキープし、終盤で一気にDecayするスケジューリング方式(以降cooldownと呼称)が注目されている。この手法は学習率を変えないのでcooldown前のcheckpointから学習を再開できる。
- 本論文では1) 8Bまでのモデルでのcosine lrとのPPLの比較、2) SWAやSFOとの比較、3) スケール則(チンチラ則)が保持されることの確認を行った。
結果
- 8Bまでのモデルでcooldownは一貫してcosine lrを上回った
- SWA(Stochastic Weight Averaging)でもほぼcooldownと同等の効果が得られたが、cooldownには及ばなかった→SWAはcooldownの代替の選択肢として有望(高いLRをキープすることによる不安定性を回避したい場合など)
- SFO(Scheduler-Free Optimization)はdecayパラメータを最適に設定すればcooldownと同等の性能が得られるが、誤ったパラメータを選択すると性能が著しく劣化する→パラメータにロバストなcooldownの方が優れている
- 360Mまでのモデル、10B tokensまでの実験でチンチラ即はcosine lrの場合と完全にマッチした→学習のやり直しが不要なcooldownの方がコスパでは圧倒的に優れている(O(N) vs O(N^2))
Limitation
- 高い学習率をキープすることによる学習の不安定性の問題。8Bモデルで20 stepまでは安定して学習できたが、ステップ数を増やしても安定するか？→QK-normalizationなどで回避できるかも？
- 1B~100Bまでの現実的なパラメータ数でもスケールするか？は未検証