[论文解读] AdaScale SGD: A User-Friendly Algorithm for Distributed Training

Tyler B. Johnson, Pulkit Agrawal+2 more|arXiv (Cornell University)|Jul 9, 2020

Advanced Neural Network Applications被引用 20

一句话总结

AdaScale SGD 是一种用户友好的算法，通过基于梯度方差动态缩放学习率，自动适应大批次训练，实现近乎完美的线性加速，且无需超参数调优。它在不同任务和批量大小下均保持模型质量，优于线性或平方根缩放等固定缩放规则，并自然地产生类似预热的行为，而无需显式的预热阶段。

ABSTRACT

When using large-batch training to speed up stochastic gradient descent, learning rates must adapt to new batch sizes in order to maximize speed-ups and preserve model quality. Re-tuning learning rates is resource intensive, while fixed scaling rules often degrade model quality. We propose AdaScale SGD, an algorithm that reliably adapts learning rates to large-batch training. By continually adapting to the gradient's variance, AdaScale automatically achieves speed-ups for a wide range of batch sizes. We formally describe this quality with AdaScale's convergence bound, which maintains final objective values, even as batch sizes grow large and the number of iterations decreases. In empirical comparisons, AdaScale trains well beyond the batch size limits of popular "linear learning rate scaling" rules. This includes large-batch training with no model degradation for machine translation, image classification, object detection, and speech recognition tasks. AdaScale's qualitative behavior is similar to that of "warm-up" heuristics, but unlike warm-up, this behavior emerges naturally from a principled mechanism. The algorithm introduces negligible computational overhead and no new hyperparameters, making AdaScale an attractive choice for large-scale training in practice.

研究动机与目标

• 为解决在大批次分布式训练中保持模型质量的挑战，因为固定学习率缩放规则常导致性能下降。
• 消除大规模训练中对手动超参数调优或启发式预热调度的需求。
• 开发一种基于随机梯度内在方差的系统性、自适应方法，以动态缩放学习率。
• 在包括图像分类、目标检测、机器翻译和语音识别在内的多种机器学习任务中，实现可靠且用户友好的大批次训练。
• 提供一种理论基础坚实的替代方案，以替代线性或平方根学习率缩放，确保在大规模训练中保持收敛性和模型准确性。

提出的方法

• AdaScale 在每次迭代中通过将学习率反比于随机梯度估计方差进行动态调整，确保在不同批量大小下期望更新幅度保持稳定。
• 该算法保持了量 $\eta_t \mathbb{E}[\|\bar{\mathbf{g}}_t\|^2]$ 在不同尺度下的不变性，这是其收敛界的核心，确保了优化进度的一致性。
• 它不引入任何新超参数，且计算开销可忽略不计，使其适用于大规模训练系统。
• 该方法自然地产生类似预热的学习率调度行为，而无需显式设置预热阶段，其行为源于自适应缩放机制。
• AdaScale 与标准学习率调度（如指数衰减）兼容，可自动适应不同批量大小和尺度。
• 该算法源自一个理论收敛界，保证了即使批量大小增加、迭代次数减少，最终目标值也能被保持。

实验结果

研究问题

• RQ1 能否设计一种学习率自适应策略，在无需手动重新调优的情况下，保持在广泛批量大小范围内的模型质量？
• RQ2 基于梯度方差的自适应学习率缩放与线性或平方根缩放等固定规则相比，在模型准确性和训练速度方面表现如何？
• RQ3 是否存在一种自适应机制，能自然产生类似预热的行为，而无需显式设置预热超参数？
• RQ4 AdaScale 在大批次分布式训练中能在多大程度上实现线性加速，同时保持最终模型性能？
• RQ5 AdaScale 能否在包括图像分类、目标检测、机器翻译和语音识别在内的多种机器学习任务中有效应用？

主要发现

• AdaScale 在多个基准测试中实现了近乎完美的线性加速，包括使用 32,000 批量大小的 ImageNet 和每批次最大 262,000 个 token 的 Transformer 模型。
• 对于 CIFAR-10，AdaScale 在规模 16 下对 13×13 的指数衰减学习率调度网格均保持了模型质量，优于对缩放 SGD 进行的直接超参数搜索。
• 与需要预热的线性缩放不同，后者在批量大小增加时会降低模型准确率，而 AdaScale 即使在极端规模下也能保持验证准确率。
• AdaScale 的收敛界确保了随着批量大小增加、迭代次数减少，最终目标值仍能被保持，为其鲁棒性提供了理论依据。
• 该算法自然地产生类似预热的学习率行为，而无需显式设置预热超参数或修改调度策略。
• 实证结果表明，AdaScale 在多种任务和规模下系统性地保持了模型质量，证明了其在实际分布式训练中的可靠性与实用性。