摘要

在大规模数据集上训练大型深度神经网络计算难度非常高。近期，使用大批次随机优化方法解决这一问题的兴趣激增。该领域的最著名算法是LARS，通过采用逐层自适应学习率，LARS能够在几分钟内完成ResNet在ImageNet上的训练。然而，LARS在处理类似BERT的注意力模型时表现不佳，这表明其性能提升在不同任务之间并不一致。本文中，我们首先研究了一种基于原理的逐层自适应策略，以加速使用大迷你批次训练深度神经网络的过程。利用这一策略，我们开发了一种新的逐层自适应大批次优化技术，称为LAMB；随后，我们对LAMB以及LARS进行了收敛性分析，在一般非凸设置下证明了它们能够收敛到一个稳定点。我们的实证结果展示了LAMB在各种任务中的优越性能，例如BERT和ResNet-50的训练，并且几乎不需要超参数调整。特别是对于BERT的训练，我们的优化器能够在不降低性能的情况下使用高达32868的大批次大小。通过将批次大小增加到TPUv3 Pod的内存极限，BERT的训练时间可以从3天缩短至仅76分钟（表1）。LAMB的实现代码可在以下链接获取：https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py