arXiv:2410.05192cs.LGcs.CL2024-10被引 72

提出WSD学习率策略,实现无需预设算力的持续训练与高效模型生成。

Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective

  • 采用暖身-稳定-衰减三段式学习率,稳定期高学习率加速路径探索。
  • 在简单二元语法数据上即可形成河谷型损失曲面,解释为何稳定期损失不降反升。
  • 新变体WSD-S可复用历史检查点,单次运行适配多种算力预算,性能更优。

当前语言模型训练需预先设定固定算力预算,因主流余弦学习率依赖总步数。而暖身-稳定-衰减(WSD)调度使用恒定学习率,使迭代点形成可无限延续的主分支,任意算力下均可从合适时间点分叉并快速衰减学习率以产出强模型。实验发现WSD产生非传统损失曲线:稳定期损失维持高位,衰减期骤降。为此,我们提出预训练损失具河谷景观假设——深谷中有一条河流。在此假设下,稳定期高学习率导致大振荡但沿河快速推进;衰减期低学习率抑制振荡,逼近河床,显现真实优化进展。因此,高学习率期负责河谷方向,快速衰减期负责山地方向,二者均关键。分析预测现象与实证一致,且该景观可由简单二元语法数据生成。受此启发,我们提出WSD-S,仅保留一个主分支,复用此前检查点的衰减阶段,从已衰减检查点恢复。WSD-S在参数量0.1B至1.2B下,单次运行即获得多个算力预算下的模型检查点,优于WSD与周期余弦策略。

原文摘要 · Abstract (English)

Training language models currently requires pre-determining a fixed compute budget because the typical cosine learning rate schedule depends on the total number of steps. In contrast, the Warmup-Stable-Decay (WSD) schedule uses a constant learning rate to produce a main branch of iterates that can in principle continue indefinitely without a pre-specified compute budget. Then, given any compute budget, one can branch out from the main branch at a proper time with a rapidly decaying learning rate to produce a strong model. Empirically, WSD generates a non-traditional loss curve: the loss remains elevated during the stable phase but sharply declines during the decay phase. Towards explaining this phenomenon, we conjecture that pretraining loss exhibits a river valley landscape, which resembles a deep valley with a river at its bottom. Under this assumption, we show that during the stable phase, the iterate undergoes large oscillations due to the high learning rate, yet it progresses swiftly along the river. During the decay phase, the rapidly dropping learning rate minimizes the iterate's oscillations, moving it closer to the river and revealing true optimization progress. Therefore, the sustained high learning rate phase and fast decaying phase are responsible for progress in the river and the mountain directions respectively, and are both critical. Our analysis predicts phenomenons consistent with empirical observations and shows that this landscape can emerge from pretraining on a simple bi-gram dataset. Inspired by the theory, we introduce WSD-S, a variant of WSD that reuses previous checkpoints' decay phases and keeps only one main branch, where we resume from a decayed checkpoint. WSD-S empirically outperforms WSD and Cyclic-Cosine in obtaining multiple language model checkpoints across various compute budgets in a single run for parameters scaling from 0.1B to 1.2B.

学习率调度河谷景观语言模型高效训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。