提出可解析求解的学习率调度理论,揭示深度学习训练的两种优化阶段。
Theory of Optimal Learning Rate Schedules and Scaling Laws for a Random Feature Model
- 基于最优控制理论推导出分阶段的最优学习率调度策略。
- 发现易阶段为多项式衰减,难阶段为预热-稳定-衰减模式,性能更优。
- 适用于模型结构与任务相关的学习率迁移,对ResNet和GPT-2均有指导意义。
设定深度学习模型的学习率是成功训练的关键。当前通常依赖经验试错。本文研究在随机特征模型上使用随机梯度下降(SGD)时,最优学习率调度 $η_T^/star(t)$ 的可解模型,其中 $t$ 为当前迭代步数,$T$ 为总训练步数。通过数值优化与最优控制理论,分析揭示两种阶段:易阶段中最优调度为多项式衰减 $η_T^/star(t) \simeq T^{-ξ} (1-t/T)^δ$,参数由特征与任务决定;难阶段则呈现预热-稳定-衰减形式,初始学习率恒定,衰减发生在极短比例训练期内。联合优化学习率与批次大小发现,在易阶段使用批次增长可加速收敛。进一步推导动量参数 $β(t)$ 的最优调度,证明其在难阶段能改善损失缩放指数。与恒定学习率 $η_T(t) ilde T^{-ξ}$ 及幂律调度 $η_T(t) ilde T^{-ξ} t^{-χ}$ 比较,本文方案取得更优收敛速率。理论表明学习率迁移性取决于模型与任务结构。在CIFAR-5M上的ResNet图像分类呈现难阶段行为,最优基础学习率在充分退火下保持恒定;而类似GPT-2的变换器在语言建模中表现为易阶段,即使退火也需调整学习率。
原文摘要 · Abstract (English)
Setting the learning rate (LR) for a deep learning model is a critical part of successful training. Choosing LRs is often done empirically with trial and error. In this work, we explore a solvable model of optimal LR schedules for a powerlaw random feature model trained with stochastic gradient descent (SGD). We consider the optimal schedule $η_T^\star(t)$ where $t$ is the current iterate and $T$ is the training horizon. This schedule is computed both as a numerical optimization problem and also analytically using optimal control theory. Our analysis reveals two regimes which we term the easy phase and hard phase. In the easy phase the optimal schedule is a polynomial decay $η_T^\star(t) \simeq T^{-ξ} (1-t/T)^δ$ where $ξ$ and $δ$ depend on the properties of the features and task. In the hard phase, the optimal schedule resembles warmup-stable-decay with constant initial LR and annealing performed over a vanishing fraction of training steps. We investigate joint optimization of LR and batch size and find batch ramps can improve the wall-clock time in the easy phase. Beyond SGD, we derive optimal schedules for momentum parameter $β(t)$ and show that it improves the loss-scaling exponent in the hard phase. We compare our optimal schedule to various benchmarks including (1) optimal constant learning rates $η_T(t) \sim T^{-ξ}$ (2) optimal power laws $η_T(t) \sim T^{-ξ} t^{-χ}$, finding that our schedule achieves better rates than either of these. Our theory suggests that LR transfer across training horizon depends on the structure of the model and task. For ResNet image classification on CIFAR-5M, the learning curves exhibit hard-phase behavior where optimal base LRs are constant under sufficient annealing. GPT-2 style transformers trained in language modeling exhibit easy-phase behavior where optimal LRs shift even under annealing.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。