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

0. 写在前面

前一段时间调研AI芯片在尽量保证准确率的情况下，如何快速训练大规模数据集并分析其可行性。UC Berkeley的大佬尤洋[1]的这篇<LARS：LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS> 提供了一个很好的保证在大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$ 时, mini-batch个 $x_i$ 被从训练集中挑选出来，通过计算该mini-batch的损失函数的梯度 $\nabla L(x_i, w)$ ，网络的权值更新通过随机梯度法如下：
在这里插入图片描述
假设我们有 $N$ 个计算节点，总共的batch size为 $B$ , 那么每个计算节点分到的mini-batch为 $\frac{B}{N}$ 个样本。但随着batch加大，我们发现训练就变得愈发困难（模型精度下降乃至发散）。之前的思路为了避免这种问题，会不断的调整一系列超参数（学习率，动量等）。

Krizhevsky[2]指出一种Linear Scaling策略：当我们对 $B$ 增大 $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$ =8K的时候，Resnet-50在无精度损失的情况下收敛，是目前的state-of-the-art。

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

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

下表是作者在不同batch size, 不同的基准学习率下，对Alexnet原版和Alexnet-BN版本进行试验的结果统计。
在这里插入图片描述
由右表(Alexnet-BN)可以发现， $B$ =512和 $B$ =8192时，精度掉了2.2个百分点。为了检查这个是不是"generalization gap" . 根据下图的loss图，我们并没有发现 $B$ =512和 $B$ =8192有显著的差别。

在这里插入图片描述
我们的结论是，在这种情况下，精度损失与泛化差距无关，这是由低训练造成的。

5. Layer-wise Adaptive Rate Scaling `LARS`

在这里插入图片描述
首先, 在第一次迭代后，分析每层权值的L2范数和梯度更新的L2范数以及其相应的比值:
$||\nabla L(w_t)||$ .

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

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

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

$\lambda^{l} = \eta \times \frac {||w^l||} {||\nabla L(w^{l})|| + \beta * ||w^l||}$

下表就是LARS + SGD的算法流程图：
在这里插入图片描述

6. 试验结果

训练在Nvidia DGX1上进行，可以看出，试验结果非常不错，比起之前的结果好上不少。
在这里插入图片描述
由下图可以看出，带LARS的训练策略， $B$ =8K的表现远远比不带LARS的测试集要好。

虽然带LARS的SGD比普通SGD的效果更好，但是，也没有达到baseline(即达到 $B$ =512的Alexnet的水平)。但是，这可以通过训练更长的时间来达到。