[论文解读] Scaling SGD Batch Size to 32K for ImageNet Training.

Yang You, Igor Gitman+1 more|arXiv (Cornell University)|Aug 13, 2017

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

一句话总结

本文提出层自适应率缩放（LARS），一种基于权重和梯度范数按层调整学习率的学习率调度方法，实现了稳定的大批次训练。利用LARS，作者成功将ResNet-50和AlexNet在ImageNet-1k上的SGD批量大小扩展至32,768和8,192，实现了显著加速——例如，在批量大小4,096下训练速度比批量大小512快3倍，且准确率未下降。

ABSTRACT

The most natural way to speed-up the training of large networks is to use data-parallelism on multiple GPUs. To scale Stochastic Gradient (SG) based methods to more processors, one need to increase the batch size to make full use of the computational power of each GPU. However, keeping the accuracy of network with increase of batch size is not trivial. Currently, the state-of-the art method is to increase Learning Rate (LR) proportional to the batch size, and use special learning rate with "warm-up" policy to overcome initial optimization difficulty. <p>By controlling the LR during the training process, one can efficiently use large-batch in ImageNet training. For example, Batch-1024 for AlexNet and Batch-8192 for ResNet-50 are successful applications. However, for ImageNet-1k training, state-of-the-art AlexNet only scales the batch size to 1024 and ResNet50 only scales it to 8192. The reason is that we can not scale the learning rate to a large value. To enable large-batch training to general networks or datasets, we propose Layer-wise Adaptive Rate Scaling (LARS). LARS LR uses different LRs for different layers based on the norm of the weights and the norm of the gradients. By using LARS algoirithm, we can scale the batch size to 32768 for ResNet50 and 8192 for AlexNet. Large batch can make full use of the system's computational power. For example, batch-4096 can achieve 3x speedup over batch-512 for ImageNet training by AlexNet model on a DGX-1 station (8 P100 GPUs).

研究动机与目标

• 解决在ImageNet训练中将SGD批量大小扩展至8,192以上时保持模型准确率的挑战。
• 克服在批量大小超过某一阈值后通常出现的训练不稳定和准确率下降问题。
• 通过将批量大小扩展至32,768，实现多GPU系统的高效利用，同时保持收敛性和性能。
• 开发一种根据层特定权重和梯度范数自适应调整的学习率调度策略，以实现稳定的大批量优化。

提出的方法

• 提出层自适应率缩放（LARS），其计算方式为：将全局基础学习率与基于权重L2范数与梯度L2范数之比的缩放因子相乘，得到每层特定的学习率。
• 采用与批量大小成比例缩放的全局基础学习率，遵循常见的线性缩放规则。
• 为学习率引入预热阶段，以在早期训练阶段（尤其是大批次时）稳定训练过程。
• 引入权重衰减组件以正则化模型，并在大批次训练中提升泛化能力。
• 在模型（AlexNet、ResNet-50）间采用一致的训练协议，仅学习率调度策略除外，其余超参数保持固定。
• 在DGX-1系统上使用8块P100 GPU进行数据并行训练，以评估可扩展性和加速效果。

实验结果

研究问题

• RQ1 能否通过一种新颖的学习率调度策略，使ImageNet-1k上的大批次SGD训练（最高达32,768）实现稳定且准确？
• RQ2 在大批次设置下，层自适应学习率缩放（LARS）与标准线性缩放和预热策略相比，在收敛性和准确率方面表现如何？
• RQ3 在保持模型准确率的前提下，将批量大小从512提升至4,096，能实现多大程度的训练加速？
• RQ4 LARS是否能在极端批量大小下，为深层（ResNet-50）和浅层（AlexNet）网络均实现稳定训练？

主要发现

• LARS使ResNet-50在ImageNet-1k上实现32,768的大批量稳定训练，达到大规模训练的最先进性能。
• 对于AlexNet，LARS成功将批量大小扩展至8,192，显著超过此前最先进水平（1,024）。
• 在批量大小4,096时，使用LARS训练的AlexNet在8-GPU DGX-1系统上相比批量大小512实现了3倍加速。
• 基础学习率的线性缩放与LARS的层自适应机制相结合，即使在极端批量大小下也能保持模型准确率。
• LARS中的预热阶段有效缓解了大批次训练初期的优化不稳定性。
• LARS在大批次设置下对不同网络架构（包括ResNet-50和AlexNet）均表现出一致的性能。