[论文解读] Large Batch Optimization for Deep Learning: Training BERT in 76 minutes

Yang You, Jing Li+8 more|arXiv (Cornell University)|Apr 1, 2019

Advanced Neural Network Applications参考文献 36被引用 164

一句话总结

引入 Lamb，一种逐层自适应的大批量优化器，能够在 Bert 训练中使用高达 32K 的批量大小，并将 Bert 预训练时间从 TPUv3 pod 的 3 天缩短到 76 分钟，同时对 Lamb 和 Lars 提供收敛性保证。

ABSTRACT

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

研究动机与目标

• 动机并解决使用大 minibatch 训练大规模深度网络的低效问题。
• 开发一种适用于大批量 SGD 的逐层自适应优化方法（Lamb）。
• 在非凸设置下给出 Lars 与 Lamb 的理论收敛性分析。
• 在大批量场景下，通过实证证明 Lamb 在 Bert 和 ResNet-50 上的优越性能。

提出的方法

• 提出一种通用策略：逐层将更新规范化为单位 L2 范数，并通过层范数的函数对学习率进行缩放，以稳定大批量训练。
• 在该策略下实现两种算法：Lars（基于动量）和 Lamb（基于 Adam 的逐层归一化）。
• 给出收敛性分析，表明在非凸目标下，Lars 与 Lamb 都收敛到驻点。
• 在 Bert（SQuAD，2.5B+0.8B 词汇语料）和 ImageNet/ResNet-50 上，用极大批量大小和最小化的超参数调整进行评估。
• 采用平方根学习率缩放和线性热身；在具有混合序列长度的多阶段 Bert 训练中使用重新热身。

实验结果

研究问题

• RQ1 逐层更新规范化加自适应学习率缩放能否在极大 minibatch 情况下实现稳定、高效的训练？
• RQ2 Lamb 和 Lars 在非凸设置下是否提供与 SGD 相似的理论收敛保证，并且在 Bert 与 ResNet-50 等任务上的经验比较如何？
• RQ3 Bert 与 ResNet-50 能容忍的批量大小是多少而不损失准确性，且在 TPUv3 Pods 上能实现多少的实际加速？
• RQ4 在多样化的架构和数据规模（语言模型对视觉模型）下，Lamb 是否比 Lars 更鲁棒？

主要发现

Solver	Batch size	Steps	F1 score (dev)	TPUs	Time
Baseline	512	1000k	90.395	16	81.4h
Lamb	512	1000k	91.752	16	82.8h
Lamb	1k	500k	91.761	32	43.2h
Lamb	2k	250k	91.946	64	21.4h
Lamb	4k	125k	91.137	128	693.6m
Lamb	8k	62500	91.263	256	390.5m
Lamb	16k	31250	91.345	512	200.0m
Lamb	32k	15625	91.475	1024	101.2m
Lamb	64k/32k	8599	90.584	1024	76.19m

• Lamb 在 Bert 的批量大小高达 32K 时实现稳定训练，并与基线相比在 SQuAD 上保持或提升 F1。
• 在混合批量训练和重新热身下，Bert 的训练时间从大约 3 天降至 76 分钟，使用 TPUv3 Pod（具备 64 倍扩展资源）。
• Lamb 在 ImageNet 上的 ResNet-50 达到具有竞争力或更高的 top-1 准确率，优于 Adam/AdamW，在大批量下接近或超越动量/Lars 基线。
• 在测试的批量大小下，Lars 通常不及 Lamb 在 Bert 上的表现，而 Lamb 随着批量大小增大仍保持较高的 F1（见 Table 1 和 Table 2）。
• 理论结果显示，在非凸设置下，Lars 和 Lamb 都收敛到驻点，Lamb 相对于 L_infty，在平均层 Lipschitz 常数 L_avg 上具有更有利的依赖性。