大规模批量优化用于深度学习：76 分钟训练 BERT

Large Batch Optimization for Deep Learning: Training BERT in 76 minutes

尤洋 Yang You · Google / UC Berkeley · 2019-04-01 · arXiv:1904.00962 ↗ · 被引 1213

摘要 · Abstract

在庞大数据集上训练大型深度神经网络在计算上极具挑战性。最近，使用大批量随机优化方法来解决这一问题引起了广泛关注。该研究领域最著名的算法是 LARS，它通过逐层自适应学习率在几分钟内训练了 ImageNet 上的 ResNet。然而，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 获取。

Training large deep neural networks on massive datasets is computationally very challenging. There has been recent surge in interest in using large batch stochastic optimization methods to tackle this issue. The most prominent algorithm in this line of research is LARS, which by employing layerwise adaptive learning rates trains ResNet on ImageNet in a few minutes. However, LARS performs poorly for attention models like BERT, indicating that its performance gains are not consistent across tasks. In this paper, we first study a principled layerwise adaptation strategy to accelerate training of deep neural networks using large mini-batches. Using this strategy, we develop a new layerwise adaptive large batch optimization technique called LAMB; we then provide convergence analysis of LAMB as well as LARS, showing convergence to a stationary point in general nonconvex settings. Our empirical results demonstrate the superior performance of LAMB across various tasks such as BERT and ResNet-50 training with very little hyperparameter tuning. In particular, for BERT training, our optimizer enables use of very large batch sizes of 32868 without any degradation of performance. By increasing the batch size to the memory limit of a TPUv3 Pod, BERT training time can be reduced from 3 days to just 76 minutes (Table 1). The LAMB implementation is available at https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py

核心贡献 · Key contributions

提出 LAMB 优化器，采用逐层自适应学习率进行大批量训练。
Proposes LAMB optimizer with layerwise adaptive learning rates for large-batch training.
在非凸设置下为 LAMB 和 LARS 提供收敛性分析。
Provides convergence analysis for LAMB and LARS in nonconvex settings.
使用高达 64K 的批量大小，将 BERT 训练时间从 3 天减少到 76 分钟。
Reduces BERT training time from 3 days to 76 minutes using batch size up to 64K.
证明 LAMB 在大批量下对 ResNet-50 达到最先进精度。
Demonstrates LAMB achieves state-of-the-art accuracy on ResNet-50 with large batches.
显示 LAMB 在 BERT 和 ResNet-50 任务上优于 LARS 和 AdamW。
Shows LAMB outperforms LARS and AdamW on BERT and ResNet-50 tasks.
引入混合批量训练与重新预热，实现稳定的大批量优化。
Introduces mixed-batch training with re-warm-up for stable large-batch optimization.

局限 · Limitations

收敛性分析假设每层梯度有界且光滑。
Convergence analysis assumes bounded gradients and smoothness per layer.
实证结果仅限于 BERT 和 ResNet-50，对其他模型的泛化性未验证。
Empirical results limited to BERT and ResNet-50; generalizability to other models unverified.
LAMB 需要极少调参，但最优超参数可能因任务而异。
LAMB requires minimal tuning but optimal hyperparameters may vary across tasks.
大批量训练需要大量硬件资源，如 TPUv3 Pod。
Large batch training requires significant hardware resources like TPUv3 Pod.
理论分析简化了β1=0 和λ=0 的情况，未完全覆盖一般情况。
Theoretical analysis simplified with β1=0 and λ=0; general case not fully covered.

论文章节 · Sections（共 13）

摘要 Abstract
1 引言 1 Introduction
1.1 相关工作 1.1 Related Work
2 预备知识 2 Preliminaries
3 算法 3 Algorithms
3.1 Lars 算法 3.1 Lars Algorithm
3.2 Lamb 算法 3.2 Lamb Algorithm
4 实验 4 Experiments
4.1 Bert 训练 4.1 Bert Training
4.2 使用 ResNet-50 进行 ImageNet 训练 4.2 ImageNet Training with ResNet-50.
4.3 批量大小缩放超参数 4.3 Hyperparameters for scaling the batch size
5 结论 5 Conclusion
6 致谢 6 Acknowledgement

阅读逐段中英对照全文 →