[论文解读] Automated Learning Rate Scheduler for Large-batch Training

Chiheon Kim, Saehoon Kim+3 more|arXiv (Cornell University)|Jul 13, 2021

Advanced Neural Network Applications参考文献 29被引用 7

一句话总结

本文提出 AutoWU，一种用于大批次训练的自动学习率调度器，通过基于高斯过程的在线损失平滑技术，动态确定预热长度和峰值学习率。该方法通过自适应地从指数预热切换到余弦衰减，仅需极少的超参数调优，即在不同架构和批量大小下，实现了优于调优基线的最先进性能。

ABSTRACT

Large-batch training has been essential in leveraging large-scale datasets and models in deep learning. While it is computationally beneficial to use large batch sizes, it often requires a specially designed learning rate (LR) schedule to achieve a comparable level of performance as in smaller batch training. Especially, when the number of training epochs is constrained, the use of a large LR and a warmup strategy is critical in the final performance of large-batch training due to the reduced number of updating steps. In this work, we propose an automated LR scheduling algorithm which is effective for neural network training with a large batch size under the given epoch budget. In specific, the whole schedule consists of two phases: adaptive warmup and predefined decay, where the LR is increased until the training loss no longer decreases and decreased to zero until the end of training. Here, whether the training loss has reached the minimum value is robustly checked with Gaussian process smoothing in an online manner with a low computational burden. Coupled with adaptive stochastic optimizers such as AdamP and LAMB, the proposed scheduler successfully adjusts the LRs without cumbersome hyperparameter tuning and achieves comparable or better performances than tuned baselines on various image classification benchmarks and architectures with a wide range of batch sizes.

研究动机与目标

• 解决因优化步数不足和学习率调度不稳定导致的大批次训练性能不佳的问题。
• 消除在大批次设置中手动调优学习率调度的需求。
• 开发一种在线、自适应的学习率调度器，以数据驱动的方式检测最优预热时长和峰值学习率。
• 在训练过程中保持低计算开销的同时，对损失波动具有鲁棒性。
• 在多种模型和大批次设置下，实现与人工调优基线相当或更优的性能。

提出的方法

• 调度器分为两个阶段：第一阶段为自适应预热阶段，学习率从 $10^{-5}$ 指数增长至 1.0；第二阶段为预定义的衰减阶段。
• 使用高斯过程（GP）对训练损失进行实时平滑，以稳健检测损失最小点，从而触发从预热到衰减的过渡。
• 过渡条件定义为：损失达到最小值的概率超过置信度阈值 $c$，并设置耐心期 $p$ 以避免噪声引起的误报。
• 峰值学习率设定为检测到的最小损失点对应的学习率，衰减阶段从此值开始，采用余弦调度。
• 预热阶段采用指数调度，增长因子为 $\gamma = (\eta_{\max}/\eta_{\min})^{1/\lfloor\rho_w T\rfloor}$，其中 $\rho_w$ 控制分配给预热阶段的训练步数比例。
• 该方法与 AdamP 和 LAMB 等自适应优化器兼容，增强了大批次场景下的稳定性和性能。

实验结果

研究问题

• RQ1 在大批次训练中，自动化学习率调度器是否能在无需手动调优的情况下检测到最优预热长度和峰值学习率？
• RQ2 基于 GP 的在线损失平滑如何提升预热阶段检测对训练噪声的鲁棒性？
• RQ3 不同预热比例超参数 ($\rho_w$) 对不同批量大小下的最终模型准确率有何影响？
• RQ4 在大批次设置下，所提出的调度器与人工调优基线及 prior 自动化方法（如 SALSA）相比表现如何？
• RQ5 在 $\eta_{\max}$ 等超参数选择上，指数预热调度是否比线性预热更具鲁棒性？

主要发现

• 使用指数预热的 AutoWU 在批量大小为 16,384 的 ResNet-50 上实现了 75.22% 的 ImageNet 顶级-1 准确率，优于调优基线（75.02%）。
• AutoWU 中的线性预热对 $\eta_{\max}$ 极其敏感，当 $\eta_{\max}=1.0$ 时，准确率仅为 29.65%，凸显了指数调度的优越性。
• 当 $\rho_w = 0.25$ 时，AutoWU 在所有批量大小下均取得最佳性能，在 4K 批量大小下达到 76.60% 准确率，高于 $\rho_w = 0.125$ 时的 75.89%。
• 随着批量大小增加和预热速度加快（$\rho_w$ 较小），衰减阶段的起始学习率上升，但在极大批量下该趋势中断，表明存在临界批量大小效应。
• 该方法在多个架构和数据集上均实现了稳定的性能提升，证明其在 ResNet-50 之外的 ImageNet 上也具有良好的泛化能力。
• 设置耐心期 $p=3$ 在鲁棒性与训练稳定性之间提供了最佳平衡，降低了对局部损失波动的敏感性。