使用大batch优化深度学习：训练BERT仅需76分钟

在海量数据集上训练大型深度，是非常具有挑战性的。最近，有许多研究均使用大batch随机优化方法来解决此问题。在该研究领域中，目前最杰出的是LARS，它通过采用分层自适应学习率，可以在几分钟内在ImageNet上训练ResNet。但是，对于像BERT这样的注意力模型，LARS的表现很差，这说明它的性能提升在不同任务之间并不一致。在本文中，作者首先研究了一种有原则性的分层适应策略，这样就可以使用大的mini-batch来加速深度神经网络的训练。

 

使用这种策略，作者开发了一种称为LAMB的新分层自适应大批量优化技术。然后，作者提供了LAMB和LARS的收敛分析，展示了在一般非凸的要求中，可以收敛到固定点。

 

作者的实验结果证明了LAMB在各种任务（例如BERT和RseNet-50训练）中可以表现的非常出色，且仅需调整少量的超参数。重要的是，对于训练BERT，本文的优化程序可以使用非常大的32868 batch大小，而且不会降低性能。通过将批处理大小增加到TPUv3 Pod的内存限制，BERT的训练时间可以从3天减少到只有76分钟（可以见后面的表1）。LAMB的实现已经开源。

 

代码地址:

https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py

 

引言

随着大规模数据集的出现，即使使用诸如随机梯度下降（SGD）等有效的优化方法来训练深层的神经网络也变得特别难。例如，在16张TPUv3芯片上训练BERT和ResNet-50等模型需要3天，而在8张Tesla P100 gpu上则需要29小时。因此，研究者们对开发用于解决此问题的优化方法有浓厚的兴趣。

 

本文的目的是研究和开发优化技术，以加快训练大型深度神经网络的速度，主要是基于SGD方法的变种。基于SGD的方法通过在小batch上计算梯度的缩放方向来迭代更新模型参数。但是，SGD的可伸缩性受到其固有顺序性的限制。由于这一局限性，在深度学习的背景下改善SGD训练时间的方法很大程度上依赖于分布式异步的设置，但是，由于异步的引入隐性的限制了方法的并行化，这通常会导致性能下降。由于最近硬件的进步，看到了在大batch的数据上并行计算梯度的可行性。但是，单纯的增加批大小通常会导致泛化性能下降。

 

最近的研究表明，在达到某些mini-batch大小的情况下，学习率与mini-batch的线性比例可用于进一步加快训练速度。这些工作还阐明了两个可以在大batch同步SGD中使用线性缩放的内容：（i）学习速率的线性缩放在初始阶段是有害的；因此，最初需要使用手动调整的缓慢增加学习速度的warmup策略，并且（ii）学习速度的线性缩放可能会超出一定的批次大小。使用这些tricks，Goyal等人能够使用8192的batch大小将ResNet-50模型的训练时间从29小时大幅减少到1小时。

 

尽管这些工作证明了这种策略可以减少训练大型深层神经网络所需的时间，但他们还是强调了针对大batch学习的自适应学习率机制的需求。 

 

最近已经提出使用基于分层自适应学习率的SGD方法来解决这个问题。该研究领域中最成功的是LARS算法，该算法最初是为了训练RESNET提出的。使用LARS，只需几分钟即可在ImageNet上对ResNet-50完成训练！但是，已经观察到在各个任务中，其对于提升性能的效果不一致。例如，对于BERT等注意力模型，LARS的表现不佳。此外，在很大程度上缺乏了对LARS中采用的适应性的理论理解。因此，本文研究和开发了一种针对大batch设置的新方法。

 

更具体地说，本文主要有以下的贡献。

 

受LARS的启发，本文研究了专门针对大batch学习的通用适应策略，并为该策略提供了直觉性的方法。

基于适应策略，本文开发了一种新的优化算法（LAMB），以实现SGD中学习率的自适应。此外，本文为LARS和LAMB都提供了收敛分析，本文重点介绍了将这些方法用于大batch的好处。

 

本文展示了LAMB在多个任务中的性能。使用LAMB，本文在训练BERT时将batch大小扩展到超过32k，而不会降低性能。因此，训练时间从3天减少到76分钟。本文是第一个将BERT训练时间减少到几个小时以内的工作。

本文还展示了LAMB在训练像RESNET这样的图像分类模型方面的效率。据作者所知，本文是第一个可以为RESNET-50达到SOTA精度的自适应求解器。

 

方法

作者首先讨论了在大Batch设置中适应学习率的一般策略。然后是使用此策略的两种特定算法。重点讨论深度学习中的表现。

 

一般策略：假设本文在具有以下分层更新规则的小batch设置中使用迭代基础算法A（例如SGD或ADAM）：

 

 

其中u_t是在t时刻由A算法更新的参数。在大batch中，主要做以下两个修改。

 

1、更新的参数会进行L2范数的归一化。且在每层网络能做归一化的操作。

 

2、学习率会由一个ф函数进行缩放。也是在每层网络都做这一操作。

 

假设基础算法A是SGD，因此上述修改的结果如下：

 

 

其中x_t和g_t分别是第i层的参数和梯度。在实践中，作者观察到一个非常简单的ф函数ф=min{max(,),}。具体的理论分析，详见论文。

 

LARS vs LAMB算法

 

LARS和本文提出的LAMB算法都是上述一般策略的实例。LARS最初是为在ImageNet使用大batch训练RESNET而提出的。通常，可以观察到，使用动量优化器可以以很小的偏差为代价减小随机梯度的方差。与LARS不同，LAMB使用的基础算法是adam。且在自适应的设置上有两个方面不同，一是对ADAM中第二个moment平方根的每维进行归一化。二是由于逐层自适应性可以获得逐层归一化。上图是两种算法的伪代码。

 

实验

本文将提供将使用现有优化器的LAMB方法在两个大batch训练任务上进行比较的实验结果：BERT和RESNET-50训练。本文还将LAMB与现有的针对mini-batch（<1K大小）和小型数据集（例如CIFAR，MNIST）的优化器进行了比较。为了阐述该方法的鲁棒性，本文使用了非常的超参数来调优LAMB优化器，具体的详情可以参见论文。

 

BERT训练实验

本文采用的数据集包含wiki和bookscorpus词条，主要关注SQuAD任务。评价指标是F1分数。均是与BERT baseline进行对比。具体的训练细节参见论文。结果如下：

 

 

作者还将LAMB与LARS算法进行了对比，可见LAMB算法在16k batch大小以上还可以稳定收敛。

 

 

此外，作者还对比了在IMAGENET上面的实验结果，可以发现在大batch下面，lamb算法相较于adam/adagrad等算法可以收敛的更好。

 

 

对于BERT和ImageNet训练，在增加batch大小时，作者并没有调整LAMB优化器的超参数。作者使用了LR平方根缩放规则和线性warmup来自动调整学习率。如下表格所示:

 

 

 

结论

大batch处理技术对于加快深度神经网络的训练至关重要。在本文中，作者提出了LAMB优化器，它支持自适应的像素更新和分层的学习率。此外，LAMB是适用于小batch和大batch的通用优化器。

 

此外还为LAMB优化器提供了理论分析，重点介绍了其性能优于标准SGD的情况。通过使用LAMB，本文能够将BERT预训练的batch大小扩展到64K，而不会损失准确性，从而将BERT训练时间从3天减少到大约76分钟。LAMB还是第一个可以使用RESNET-50在ImageNet训练上达到SOTA结果的大batch自适应优化器。

 

原文链接：

https://static.aminer.cn/upload/pdf/program/5e5e18b793d709897ce2a20b_0.pdf

 

声明：文章收集于网络，版权归原作者所有，为传播信息而发，如有侵权，请联系小编删除，谢谢！