2017 LARS:LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS (训练大Batch的卷积神经网络)

0. 写在前面

前一段时间调研AI芯片在尽量保证准确率的情况下,如何快速训练大规模数据集并分析其可行性。UC Berkeley的大佬尤洋[1]的这篇 提供了一个很好的保证在大Batch Size的分布式训练情况下,精度损失很小的策略。下面, 让我们开始~

1. 摘要 (Abstract)

加速训练大规模CNN的场景策略是分布式训练(即加多计算节点). 由于data-parallel的SGD的使用和计算节点的增大(成百上千个计算节点),训练的模型精度会剧烈降低。我们认为,目前对大batch的训练策略不够泛化而且可能导致训练无法收敛。

为了克服这个问题,作者提出了一种基于Layer-Wise Aaptive Rate Scaling (LARS) 的学习算法,通过使用LARS, 作者在Alexnet的batch size = 8K,Resnet-50 的batch size = 32K的时候可以让模型收敛而无精度损失。

2. 介绍 (Introduction)

本部分是删节版介绍,一些related work基本不会提及. 如感兴趣, 请查阅英文原版.

我们都知道,训练一个大型的CNN网络是非常耗时的,最直接的加速方式就是增加算力(比如更多的GPU节点),通过data-parallel 的SGD来加速训练。其中, 每个GPU节点(或者AI芯片)接收整体batch (比如8K,32K)中的一部分,为了充分利用计算资源,每个计算节点接收的batch应尽可能大。但是,天下没有免费的午餐,这种通过加计算节点分布式训练的策略,对模型的精度有显著的负面影响

容易理解,当Epoch保持固定(数据集也不变的时候), 大的batch size意味着权重调整的次数下降/减少. 一个直观的解决方式是(加大)学习率来compensate. 然而,使用更大的学习率会使得网络在初始阶段就发散(这也是为啥近几年warm up策略火的原因)。很多研究人员针对这个问题提出自己的思路,这里不详细介绍,总之:线性调整学习率 + warm up是大batch 训练的state-of-the-art (17年前最好的方法)

但是,作者使用这种策略来在ImageNet上训练Alexnet,发现batch size在2K后,随着batch size的增大,模型开始出现精度下降甚至发散的情况。为了解决这种情况,作者将Alexnet里面的归一化策略Local Response Normalization (LRN)替换为Batch Normalization (BN),通过将LRN替换为BN,作者发现精度的损失变小了。

为了分析在大学习率下训练的稳定性,作者分析了模型中每层权值范数和更新的梯度范数。一个有意思的现象被观察到:当这个比例很高时,训练就变得不稳定;当比例很低时,权值的更新变得很慢。因为这个比例在不同层之间差异很大,所以有必要对神经网络中的每一层都使用其独特的学习率。这就是LARS被提出的基础。

与ADAM和RMSProp等学习算法不同的是,LARS:

  • 网络中的每层都使用不同的学习率,使得训练更稳定。
  • 权值更新的量级与权值的范数有关,这样更好的控制训练的速度。

通过使用LARS,作者声称他们训练的Alexnet-BN和Resnet-50在batch size过万的情况下,精度不出现损失的情况。

3. 背景 (Background)

深度学习中,权值更新(或者说模型学习)的策略如下:

在每一步 t t t时, mini-batch个 x i x_i xi被从训练集中挑选出来,通过计算该mini-batch的损失函数的梯度 ∇ L ( x i , w ) \nabla L(x_i, w) L(xi,w), 网络的权值更新通过随机梯度法如下:
在这里插入图片描述
假设我们有 N N N个计算节点,总共的batch size为 B B B, 那么每个计算节点分到的mini-batch为 B N \frac{B}{N} NB个样本。但随着batch加大,我们发现训练就变得愈发困难(模型精度下降乃至发散)。之前的思路为了避免这种问题,会不断的调整一系列超参数(学习率,动量等)。

Krizhevsky[2]指出一种Linear Scaling策略:当我们对 B B B增大 k k k倍时,应该同样对学习率增大 k k k倍(并保持其它超参数不变)。

Linear Scaling背后的逻辑很直接:在保持epoch数量不变的情况下,如果你将batch size加大k倍, 那么更新权重的频次因此降低k倍。所以,加大学习率是一种很自然的方式。

但这种方式对batch size大于2K的情况, 训练变得很难收敛(training diverges), 将Alexnet中的LRN替换为BN后, 效果会好不少~,但是精度还是有损失的。

我们发现,主要的障碍在于高学习率带来的不稳定性,基于此,Goyal[3]等人提出LR+warm up的策略:

training starts with small LR, and then LR is gradually increased to
the target. 

这种策略,在 B B B=8K的时候,Resnet-50在无精度损失的情况下收敛,是目前的state-of-the-art。

另一个由大batch带来的问题是:“generalization gap” (泛化差距). Keskar[4]等人指出:large-batch的方法倾向于收敛于训练函数的精确最小化. 但是对泛化性能并没有什么提升…

4. 分析Alexnet在大batch下的训练效果

下表是作者在不同batch size, 不同的基准学习率下,对Alexnet原版和Alexnet-BN版本进行试验的结果统计。
2017 LARS:LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS (训练大Batch的卷积神经网络)_第1张图片
由右表(Alexnet-BN)可以发现, B B B=512和 B B B=8192时,精度掉了2.2个百分点。为了检查这个是不是"generalization gap" . 根据下图的loss图,我们并没有发现 B B B=512和 B B B=8192有显著的差别。

2017 LARS:LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS (训练大Batch的卷积神经网络)_第2张图片
我们的结论是,在这种情况下,精度损失与泛化差距无关,这是由低训练造成的。

5. Layer-wise Adaptive Rate Scaling LARS

2017 LARS:LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS (训练大Batch的卷积神经网络)_第3张图片
首先, 在第一次迭代后,分析每层权值的L2范数和梯度更新的L2范数以及其相应的比值:
∣ ∣ w ∣ ∣ / ∣ ∣ ∇ L ( w t ) ∣ ∣ ||w|| / ||\nabla L(w_t)|| w/L(wt).

接着,对每层 l l l使用其独特的学习率 λ l \lambda^{l} λl (Local LR),则其权值更新的值从原来的 Δ w t l = λ ∗ ∇ L ( w t ) \Delta w_{t}^{l} = \lambda * \nabla L(w_t) Δwtl=λL(wt)变为 Δ w t l = γ ∗ λ l ∗ ∇ L ( w t l ) \Delta w_{t}^{l} =\gamma *\lambda^{l} *\nabla L(w_t^{l}) Δwtl=γλlL(wtl). 其中, γ \gamma γ是全局学习率 (global LR),Local LR的定义为:
λ l = η ∗ ∣ ∣ w l ∣ ∣ ∣ ∣ ∇ L ( w l ) ∣ ∣ \lambda^{l} = \eta * \frac {||w^l||} {||\nabla L(w^{l})||} λl=ηL(wl)wl

η \eta η这个参数决定了一种置信度:即我们有多相信当前层(第 l l l层)会在一次更新中改变它的权重。

需要注意的是,现在,模型中每层权值更新的幅度与梯度的幅度无关,所以这有助于部分消除梯度消失和梯度爆炸等现象。我们的LARS可以扩展到SGD (平衡local LR与weight decay, β \beta β是weight decay系数):

λ l = η × ∣ ∣ w l ∣ ∣ ∣ ∣ ∇ L ( w l ) ∣ ∣ + β ∗ ∣ ∣ w l ∣ ∣ \lambda^{l} = \eta \times \frac {||w^l||} {||\nabla L(w^{l})|| + \beta * ||w^l||} λl=η×L(wl)+βwlwl

下表就是LARS + SGD的算法流程图:
2017 LARS:LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS (训练大Batch的卷积神经网络)_第4张图片

6. 试验结果

训练在Nvidia DGX1上进行,可以看出,试验结果非常不错,比起之前的结果好上不少。
2017 LARS:LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS (训练大Batch的卷积神经网络)_第5张图片
由下图可以看出,带LARS的训练策略, B B B=8K的表现远远比不带LARS的测试集要好。
2017 LARS:LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS (训练大Batch的卷积神经网络)_第6张图片
虽然带LARS的SGD比普通SGD的效果更好,但是,也没有达到baseline(即达到 B B B=512的Alexnet的水平)。 但是,这可以通过训练更长的时间来达到。

作者认为: 当batch非常大的时候, 随机梯度将与真实梯度非常接近. 所以, bs大的时候相比bs小的时候并没有更多的额外梯度信息. 下表的试验结果验证了这一观点。

2017 LARS:LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS (训练大Batch的卷积神经网络)_第7张图片

7. 结论

大Batch Size是扩大训练卷积网络的核心.

现有的方法通常是采用更大的学习率, 但这不可避免的带来了网络的不收敛. (即便使用warm up策略)

为了解决这个优化问题, 作者提出了LARS, 通过使用LARS, 作者成功的将Resnet50和Alexnet-BN扩展到BS=32k的情况(训练ImageNet).

参考资料

[1]: You Yang, Igor Gitman, Boris Ginsburg:
[2]: Alex Krizhevsky:
[3]: Priya Goyal, Piotr Dollár and et al. in 1 hour >
[4]: Nitish Shirish Keskar, Dheevatsa Mudigere, and et al:

你可能感兴趣的:(算法,机器学习,深度学习)