76分钟完成训练BERT！谷歌大批量优化研究取得重大突破！

全球最热门的人工智能前沿资讯平台

全文共1742字，阅读时长需要6分钟

在海量数据集上训练大型深层神经网络在计算上具有很大的挑战性。最近，人们对使用大批量随机优化方法来解决这个问题的兴趣激增。在这方面最突出的算法是LARS，它通过采用分层自适应学习率在几分钟内在ImageNet上训练RESNET。然而，LARS对于像BERT这样的注意模型表现不佳，这表明它的性能增益在任务之间并不一致。最近，Google与UCBerkeley的一项研究成功解决了大批量训练的问题。

BEGIN

本文首先研究了一种原则性的分层自适应策略，以加速大批量、小批量的深层神经网络训练。利用该策略，本文发展了一种新的分层自适应大批量优化技术LAMB，并给出了LAMB和LARS的收敛性分析，比较了一般非凸情形下LAMB和LARS的收敛性。

LAMB是什么？

LAMB的全称是Layer-wiseAdaptive Moments optimizer for Batching training。BERT训练的baseline使用权重衰减的Adam 作为优化器，这是Adam优化器的一个变体。另一个自适应优化器LARS，早前被提出用于ImageNet上RESNET的大批量学习。为了得到进一步的改进，研究团队对LAMB使用混合批处理训练。BERT训练包括两个阶段：前9/10的epochs使用128的序列长度，而后1/10的epochs使用512的序列长度。由于内存限制，第二阶段的训练需要更长的序列长度，因此在TPUv3Pod上最多只能使用32768个批处理大小。

LAMB 算法的概述如下所示：

图1 LAMB算法流程图

ADAM优化器在深度学习社区中很受欢迎，并且在训练BERT这样的最新语言模型方面表现出了良好的性能。与LARS不同，LAMB的适应性有两个方面：（i）ADAM中使用的二阶矩平方根的每一维归一化；（ii）由于分层自适应性而获得的分层归一化。算法2给出了LAMB的伪码。当β1=0和β2=0时，算法简化为符号SGD，其中学习率按层维度的平方根缩放。

实验结果

本研究在两个重要的大批量训练任务：BERT和RESNET-50训练上，比较LAMB与现有优化器的实证结果。还比较了在小批量（

BERT训练:

首先是加速BERT训练的实验结果。本实验使用由Wikipedia和BooksCorpus拼接而成的数据集。在实验中，以SQuAD-v1上的F1分数作为精度指标，所有的比较都是基于Devlin等人的baselineBERT模型。为了训练BERT,Devlin等人首先使用序列长度为128的900k迭代训练模型，然后在最后的100k迭代中转换为512的序列长度，这导致了在16个TPUv3芯片上大约需要3天的训练时间。本文在表中报告F1分数为91.345，这是未微调版本的分数。

表1 使用SQuAD-v1的F1score作为精度指标

F1的baseline成绩是由BERT的公共知识库提供的预训练模型(BERT-large)实现的(截止到2019年2月1日)。本文在实验中使用TPUv3s。本文使用了与baseline相同的设置：前9/10的epochs使用序列长度128，最后1/10的epochs使用序列长度512。所有的实验运行相同数量的epochs。值得注意的是，本文可以通过手动调整超参数来获得更好的结果。

本文的实验，除了将训练优化器更改为LAMB之外，保持与baseline相同的训练过程，使用与baseline相同数量的epochs运行，但batchsize从512扩展到32K（选择32K大小（序列长度512）主要是由于TPUPod的内存限制）。通过使用LAMB优化器，能够在batchsize为32768的15625次迭代（序列长度为128的14063次迭代和序列长度为512的1562次迭代）中获得91.460的F1分数。对于批量大小为32K，本文将BERT训练时间从3天缩短到约100分钟。

主要贡献

（1）受LARS的启发，研究了一种特别适合大批量学习的通用适应策略，并为该策略提供了实验验证。

（2）基于自适应策略，提出了一种新的优化算法（LAMB）来实现SGD中学习速率的自适应。此外，还提供了LARS和LAMB的收敛分析以获得非凸情况下的一个驻点。

（3）展示了LAMB在多个挑战性任务中的强大实证性能。使用LAMB，将训练BERT的批量大小扩展到32k以上，而不会降低性能，并且所需时间从3天减少到76分钟。这是目前将BERT训练经过时间减少到几个小时以内的第一项工作。

总结

大批量处理技术是加快深度神经网络训练的关键。本文提出了LAMB优化器，它支持自适应的单元更新(Adaptive Elementwise Updating)和逐层学习率(Layerwise Learning Rates)。此外，LAMB是一个通用优化器，适用于小批量和大批量。本文还为LAMB优化器提供了理论分析，重点介绍了其性能优于标准SGD的情况。对于广泛的应用程序，LAMB的性能要优于现有的优化器。通过使用LAMB，我们可以将BERT预训练的批处理规模扩展到64K而不失准确性，从而将BERT训练时间从3天缩短到76分钟左右。LAMB也是第一个能够在RESNET-50的ImageNet训练中获得最先进精度的大批量自适应解算器。

参考文献：

1.https://arxiv.org/pdf/1904.00962.pdf

2.https://tech.sina.com.cn/csj/2020-01-13/doc-iihnzhha2131900.shtml

喜欢本篇内容请给我们点个在看哟