提出无需固定训练时长的自适应学习率调度,提升大模型训练灵活性。
WSqD: A Horizon-Free Learning Rate Schedule for Large Model Training
- 用平方根衰减替代恒定阶段,实现学习率与训练时长解耦
- 在多个训练周期下表现稳定,仅需一个峰值学习率即可适配
- 理论证明收敛速率最优,适合需要灵活扩展训练的场景
标准学习率调度如余弦退火依赖固定训练周期,难以应对后期延长。温升稳定衰减(WSD)虽通过长平缓期缓解此问题,但其峰值学习率仍基于原周期设定,在训练延长后可能失效。受随机凸优化启发,本文提出WSqD(温升-平方根基-线性衰减)调度:以偏移的平方根函数替代WSD的恒定阶段,保留最终线性冷却。理论上,WSqD在随机凸设置下达到最优的 $O(1/\ oot\of{T})$ 最后迭代收敛速率。关键优势在于基础学习率调度与训练周期无关,仅需周期信息决定冷却启动时机。实验显示,在使用SlimPajama语料进行语言模型预训练时,WSqD在多个训练周期下表现优于或等同于精心调参的WSD及其他基线,且仅用单一峰值学习率即可复用。
原文摘要 · Abstract (English)
Standard learning rate schedules such as cosine annealing are tied to a fixed training horizon, limiting their ability to accommodate post hoc horizon extension. Warmup-stable-decay (WSD) partially addresses this issue by maintaining a long constant-rate phase before a short linear cooldown, allowing training to resume from a pre-decay checkpoint. However, its peak learning rate is still tuned based on the original training horizon and can become suboptimal when training is extended. Motivated by stochastic convex optimization, we propose WSqD (Warmup with Square-root base and linear Decay), a learning rate schedule that replaces WSD's constant stable phase with a shifted inverse-square-root base while retaining the final linear cooldown. In the stochastic convex setting, WSqD provably attains the minimax-optimal $O(1/\sqrt{T})$ last-iterate convergence rate. Importantly, its base learning rate schedule is horizon-independent, and the training horizon is needed only to determine when to begin the final cooldown. Empirically, on language-model pretraining using the SlimPajama corpus, WSqD matches or outperforms carefully tuned WSD and other baselines across multiple training horizons while reusing a single peak learning rate.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。