[论文解读] Optimal Linear Decay Learning Rate Schedules and Further Refinements

Aaron Defazio, Ashok Cutkosky+2 more|arXiv (Cornell University)|Oct 11, 2023

Stochastic Gradient Optimization Techniques被引用 5

一句话总结

本文提出了一种理论基础扎实、针对具体问题自适应的学习率调度框架，用于深度学习优化。通过从先前训练中提炼梯度范数序列，该方法推导出同时具备预热（warm-up）和线性衰减（linear decay）特性的调度策略——此前这些特性仅在实践中通过启发式方法实现。实验表明，在10项不同的深度学习任务中，线性衰减优于余弦衰减及其他调度策略，且通过进一步优化可获得更大性能提升。

ABSTRACT

Learning rate schedules used in practice bear little resemblance to those recommended by theory. We close much of this theory/practice gap, and as a consequence are able to derive new problem-adaptive learning rate schedules. Our main technical contribution is a refined analysis of learning rate schedules for a wide class of optimization algorithms (including SGD). When considering only worst-case analysis, our theory predicts that the optimal choice is the linear decay schedule where the step-size is set proportional to 1 - t/T, where t is the current iteration and T is the total number of steps. To go beyond this worst-case analysis, we use the observed gradient norms to derive schedules refined for any particular task. These refined schedules exhibit learning rate warm-up and rapid learning rate annealing near the end of training. Ours is the first systematic approach to automatically yield both of these properties. We perform the most comprehensive evaluation of learning rate schedules to date, evaluating across 10 diverse deep learning problems, a series of LLMs, and a suite of logistic regression problems. We validate that overall, the linear-decay schedule outperforms all commonly used default schedules including cosine annealing. Our adaptive schedule refinement method gives further improvements.

研究动机与目标

• 解决理论学习率调度与深度学习实践中选择之间的长期脱节问题。
• 弥合理论分析中对最终迭代点收敛性的研究与实际优化实践之间的差距。
• 提出一种系统化、问题自适应的方法，利用观测到的梯度统计特性来优化学习率调度。
• 证明在最坏情况假设下，线性衰减是最后迭代点的理论最优学习率调度。
• 表明经优化的调度策略自然产生预热和快速退火行为，从而解决长期存在的经验性启发式问题。

提出的方法

• 提出一个基于完整期望梯度范数序列的最终迭代点函数值的新理论界。
• 设计一种学习率调度策略，其中每一步的学习率大小由未来梯度范数的平方倒数乘积决定：$\eta_t = w_t \sum_{p=t+1}^T w_p $，其中 $ w_t = \hat{G}_t^{-2} $。
• 对原始梯度范数应用中值滤波，以平滑噪声并提升调度优化的鲁棒性。
• 利用前期训练运行收集梯度范数序列，进而为后续训练生成优化后的自适应调度。
• 将该方法推广至SGD以外的任意渐近遗憾趋于零的优化算法，使其可应用于Adam等方法。
• 对最终调度结果进行归一化，以确保最大学习率与标准基线一致。

Figure 1: Example gradient norm sequences (top row) and the resulting refined schedules given by Algorithm 1 (bottom row). Black dashed line at $y=1$ shown for reference. Percentage of runtime on the x-axis.

实验结果

研究问题

• RQ1 为何标准理论学习率调度无法匹配深度学习中的实际性能表现？
• RQ2 在最坏情况假设下，SGD最后迭代点的理论最优学习率调度是什么？
• RQ3 能否利用先前运行中获得的梯度范数序列，推导出自适应调度以提升泛化能力和收敛性？
• RQ4 理论是否自然产生预热和线性衰减行为，还是这些仅是经验性启发式方法？
• RQ5 所提出的优化方法相较于余弦退火等标准调度，在多样化深度学习任务中的表现如何？

主要发现

• 此前仅通过经验使用线性衰减学习率调度，在最坏情况假设下，被证明是SGD最后迭代点的理论最优解。
• 所提出的优化方法生成的调度策略同时具备预热和快速衰减特性，而这些在标准理论调度中并不存在。
• 在所有评估的10项深度学习任务中，包括视觉、自然语言处理和推荐系统，线性衰减作为非自适应调度，表现与余弦退火相当或更优。
• 经优化的调度策略在ResNet-50、ViT、GPT、RoBERTa和DLRM等模型上，持续优于标准调度，显著提升泛化性能。
• 该方法具有通用性：可应用于任意具有渐近遗憾趋于零特性的优化算法，包括Adam，从而实现理论基础坚实的自适应调度。
• 实证结果验证了：当基于观测到的梯度范数对理论推导的调度进行优化后，可在实际中实现可测量的性能提升。

Figure 2: Gradient Norm sequences and the resulting Refined schedules, generated using an initial linear decay schedule with warmup for the initial run. Log scale views are inset for scale reference. Fitted polynomial decay schedules are overlayed in blue.