Large Batch Optimization for Deep Learning: Training BERT in 76 minutes — part 4/5

Cleaned-text code points [41220, 57092). Bibliography omitted here; source.md includes it.

Since the function f f is L L -smooth, we have the following: f ⁡ ( x t + 1 ) \displaystyle f(x_{t+1}) ≤ f ⁡ ( x t ) + ⟨ ∇ i f ​ ( x t ) , x t + 1 ( i ) − x t ( i ) ⟩ + ∑ i = 1 h L i 2 ​ ‖ x t + 1 ( i ) − x t ( i ) ‖ 2 \displaystyle\leq f(x_{t})+\langle\nabla_{i}f(x_{t}),x_{t+1}^{(i)}-x_{t}^{(i)}\rangle+\sum_{i=1}^{h}\frac{L_{i}}{2}\|x_{t+1}^{(i)}-x_{t}^{(i)}\|^{2} = f ⁡ ( x t ) ​ − η t ∑ i = 1 h ∑ j = 1 d i ϕ ( ∥ x t ( i ) ∥ ) × ( [ ∇ i f ( x t ) ] j × r t , j ( i ) ‖ r t ( i ) ‖ ) ⏟ T 1 + ∑ i = 1 h L i ​ α u 2 ​ η t 2 2 \displaystyle=f(x_{t})\underbrace{-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\phi(\|x_{t}^{(i)}\|)\times\left([\nabla_{i}f(x_{t})]_{j}\times\frac{r_{t,j}^{(i)}}{\|r_{t}^{(i)}\|}\right)}_{T_{1}}+\sum_{i=1}^{h}\frac{L_{i}\alpha_{u}^{2}\eta_{t}^{2}}{2} (7) The above inequality simply follows from the lipschitz continuous nature of the gradient. We bound term T 1 T_{1} in the following manner: T 1 \displaystyle T_{1} ≤ − η t ∑ i = 1 h ∑ j = 1 d i ϕ ( ∥ x t ( i ) ∥ ) × ( [ ∇ i f ( x t ) ] j × r t , j ( i ) ‖ r t ( i ) ‖ ) \displaystyle\leq-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\phi(\|x_{t}^{(i)}\|)\times\left([\nabla_{i}f(x_{t})]_{j}\times\frac{r_{t,j}^{(i)}}{\|r_{t}^{(i)}\|}\right) ≤ − η t ∑ i = 1 h ∑ j = 1 d i 1 − β 2 G 2 ​ d i ( ϕ ( ∥ x t ( i ) ∥ ) × [ ∇ i f ( x t ) ] j × g t , j ( i ) ) \displaystyle\leq-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\sqrt{\frac{1-\beta_{2}}{G^{2}d_{i}}}\left(\phi(\|x_{t}^{(i)}\|)\times[\nabla_{i}f(x_{t})]_{j}\times g_{t,j}^{(i)}\right) − η t ∑ i = 1 h ∑ j = 1 d i ( ϕ ( ∥ x t ( i ) ∥ ) × [ ∇ i f ( x t ) ] j × r t , j ( i ) ‖ r t ( i ) ‖ ) 𝟙 ( s i g n ( ∇ i f ( x t ) ] j ) ≠ s i g n ( r t , j ( i ) ) ) \displaystyle\qquad\qquad-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\left(\phi(\|x_{t}^{(i)}\|)\times[\nabla_{i}f(x_{t})]_{j}\times\frac{r_{t,j}^{(i)}}{\|r_{t}^{(i)}\|}\right)\mathds{1}(sign(\nabla_{i}f(x_{t})]_{j})\neq sign(r_{t,j}^{(i)})) (8) This follows from the fact that ‖ r t ( i ) ‖ ≤ d i 1 − β 2 \|r_{t}^{(i)}\|\leq\sqrt{\frac{d_{i}}{1-\beta_{2}}} and v t ≤ G \sqrt{v_{t}}\leq G . If β 2 = 0 \beta_{2}=0 , then T 1 T_{1} can be bounded as follows: T 1 \displaystyle T_{1} ≤ − η t ∑ i = 1 h ∑ j = 1 d i 1 d i ( ϕ ( ∥ x t ( i ) ∥ ) × | [ ∇ i f ( x t ) ] j | ) \displaystyle\leq-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\sqrt{\frac{1}{d_{i}}}\left(\phi(\|x_{t}^{(i)}\|)\times|[\nabla_{i}f(x_{t})]_{j}|\right) − η t ∑ i = 1 h ∑ j = 1 d i ( ϕ ( ∥ x t ( i ) ∥ ) × [ ∇ i f ( x t ) ] j × r t , j ( i ) ‖ r t ( i ) ‖ ) 𝟙 ( s i g n ( ∇ i f ( x t ) ] j ) ≠ s i g n ( r t , j ( i ) ) ) \displaystyle\qquad\qquad-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\left(\phi(\|x_{t}^{(i)}\|)\times[\nabla_{i}f(x_{t})]_{j}\times\frac{r_{t,j}^{(i)}}{\|r_{t}^{(i)}\|}\right)\mathds{1}(sign(\nabla_{i}f(x_{t})]_{j})\neq sign(r_{t,j}^{(i)})) The rest of the proof for β 2 = 0 \beta_{2}=0 is similar to argument for the case β 2 > 0 \beta_{2}>0 , which is shown below. Taking expectation, we have the following: 𝔼 ⁡ [ T 1 ] \displaystyle\mathbb{E}[T_{1}] ≤ − η t ∑ i = 1 h ∑ j = 1 d i 1 − β 2 G 2 ​ d i 𝔼 [ ϕ ( ∥ x t ( i ) ∥ ) × ( [ ∇ i f ( x t ) ] j × g t , j ( i ) ) ] \displaystyle\leq-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\sqrt{\frac{1-\beta_{2}}{G^{2}d_{i}}}\mathbb{E}\left[\phi(\|x_{t}^{(i)}\|)\times\left([\nabla_{i}f(x_{t})]_{j}\times g_{t,j}^{(i)}\right)\right] − η t ∑ i = 1 h ∑ j = 1 d i 𝔼 [ ϕ ( ∥ x t ( i ) ∥ ) × ( [ ∇ i f ( x t ) ] j × r t , j ( i ) ‖ r t ( i ) ‖ ) 𝟙 ( s i g n ( ∇ i f ( x t ) ] j ) ≠ s i g n ( g t , j ( i ) ) ) ] \displaystyle\qquad\qquad-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\mathbb{E}\left[\phi(\|x_{t}^{(i)}\|)\times\left([\nabla_{i}f(x_{t})]_{j}\times\frac{r_{t,j}^{(i)}}{\|r_{t}^{(i)}\|}\right)\mathds{1}(sign(\nabla_{i}f(x_{t})]_{j})\neq sign(g_{t,j}^{(i)}))\right] ≤ − η t ∑ i = 1 h ∑ j = 1 d i 1 − β 2 G 2 ​ d i 𝔼 [ ( ϕ ( ∥ x t ( i ) ∥ ) × [ ∇ i f ( x t ) ] j × g t , j ( i ) ) ] \displaystyle\leq-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\sqrt{\frac{1-\beta_{2}}{G^{2}d_{i}}}\mathbb{E}\left[\left(\phi(\|x_{t}^{(i)}\|)\times[\nabla_{i}f(x_{t})]_{j}\times g_{t,j}^{(i)}\right)\right] + η t ∑ i = 1 h ∑ j = 1 d i 𝔼 [ α u | [ ∇ i f ( x t ) ] j | 𝟙 ( s i g n ( ∇ i f ( x t ) ] j ) ≠ s i g n ( g t , j ( i ) ) ) ] \displaystyle\qquad\qquad+\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\mathbb{E}\left[\alpha_{u}|[\nabla_{i}f(x_{t})]_{j}|\mathds{1}(sign(\nabla_{i}f(x_{t})]_{j})\neq sign(g_{t,j}^{(i)}))\right] ≤ − η t ∑ i = 1 h ∑ j = 1 d i 1 − β 2 G 2 ​ d i 𝔼 [ ϕ ( ∥ x t ( i ) ∥ ) × ( [ ∇ i f ( x t ) ] j × g t , j ( i ) ) ] \displaystyle\leq-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\sqrt{\frac{1-\beta_{2}}{G^{2}d_{i}}}\mathbb{E}\left[\phi(\|x_{t}^{(i)}\|)\times\left([\nabla_{i}f(x_{t})]_{j}\times g_{t,j}^{(i)}\right)\right] − η t ∑ i = 1 h ∑ j = 1 d i α u | [ ∇ i f ( x t ) ] j | ℙ ( s i g n ( ∇ i f ( x t ) ] j ) ≠ s i g n ( g t , j ( i ) ) ) \displaystyle\qquad\qquad-\eta_{t}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\alpha_{u}|[\nabla_{i}f(x_{t})]_{j}|\mathbb{P}(sign(\nabla_{i}f(x_{t})]_{j})\neq sign(g_{t,j}^{(i)})) Using the bound on the probability that the signs differ, we get: 𝔼 ⁡ [ T 1 ] ≤ − η t ​ α l ​ h ⁡ ( 1 − β 2 ) G 2 ​ d ​ ‖ ∇ f ​ ( x t ) ‖ 2 + η t ​ α u ​ ∑ i = 1 h ∑ j = 1 d i σ i , j b . \displaystyle\mathbb{E}[T_{1}]\leq-\eta_{t}\alpha_{l}\sqrt{\frac{h(1-\beta_{2})}{G^{2}d}}\|\nabla f(x_{t})\|^{2}+\eta_{t}\alpha_{u}\sum_{i=1}^{h}\sum_{j=1}^{d_{i}}\frac{\sigma_{i,j}}{\sqrt{b}}. Substituting the above bound on T 1 T_{1} in equation 7 , we have the following bound: 𝔼 ⁡ [ f ⁡ ( x t + 1 ) ] \displaystyle\mathbb{E}[f(x_{t+1})] ≤ f ⁡ ( x t ) − η t ​ α l ​ h ⁡ ( 1 − β 2 ) G 2 ​ d ​ ‖ ∇ f ​ ( x t ) ‖ 2 + η t ​ α u ​ ‖ σ ~ ‖ 1 b + η t 2 ​ α u 2 ​ ‖ L ‖ 1 2 \displaystyle\leq f(x_{t})-\eta_{t}\alpha_{l}\sqrt{\frac{h(1-\beta_{2})}{G^{2}d}}\|\nabla f(x_{t})\|^{2}+\eta_{t}\alpha_{u}\frac{\|\tilde{\sigma}\|_{1}}{\sqrt{b}}+\frac{\eta_{t}^{2}\alpha_{u}^{2}\|L\|_{1}}{2} (10) Summing the above inequality for t = 1 t=1 to T T and using telescoping sum, we have the following inequality: 𝔼 ⁡ [ f ⁡ ( x T + 1 ) ] \displaystyle\mathbb{E}[f(x_{T+1})] ≤ f ⁡ ( x 1 ) − η t ​ α l ​ h ⁡ ( 1 − β 2 ) G 2 ​ d ​ ∑ t = 1 T 𝔼 ⁡ [ ‖ ∇ f ​ ( x t ) ‖ 2 ] + η ​ T ​ α u ​ ‖ σ ~ ‖ 1 b + η 2 ​ α u 2 ​ T 2 ​ ‖ L ‖ 1 . \displaystyle\leq f(x_{1})-\eta_{t}\alpha_{l}\sqrt{\frac{h(1-\beta_{2})}{G^{2}d}}\sum_{t=1}^{T}\mathbb{E}[\|\nabla f(x_{t})\|^{2}]+\eta T\alpha_{u}\frac{\|\tilde{\sigma}\|_{1}}{\sqrt{b}}+\frac{\eta^{2}\alpha_{u}^{2}T}{2}\|L\|_{1}.

Rearranging the terms of the above inequality, and dividing by η ​ T ​ α l \eta T\alpha_{l} , we have: h ⁡ ( 1 − β 2 ) G 2 ​ d ​ 1 T ​ ∑ t = 1 T 𝔼 ⁡ [ ‖ ∇ f ​ ( x t ) ‖ 2 ] \displaystyle\sqrt{\frac{h(1-\beta_{2})}{G^{2}d}}\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla f(x_{t})\|^{2}] ≤ f ⁡ ( x 1 ) − 𝔼 ⁡ [ f ⁡ ( x T + 1 ) ] T ​ η ​ α l + α u ​ ‖ σ ~ ‖ 1 α l ​ b + η 2 ​ ‖ L ‖ 1 \displaystyle\leq\frac{f(x_{1})-\mathbb{E}[f(x_{T+1})]}{T\eta\alpha_{l}}+\frac{\alpha_{u}\|\tilde{\sigma}\|_{1}}{\alpha_{l}\sqrt{b}}+\frac{\eta}{2}\|L\|_{1} ≤ f ⁡ ( x 1 ) − f ⁡ ( x ∗ ) T ​ η ​ α l + α u ​ ‖ σ ~ ‖ 1 α l ​ b + η ​ α u 2 2 ​ α l ​ ‖ L ‖ 1 . \displaystyle\leq\frac{f(x_{1})-f(x^{*})}{T\eta\alpha_{l}}+\frac{\alpha_{u}\|\tilde{\sigma}\|_{1}}{\alpha_{l}\sqrt{b}}+\frac{\eta\alpha_{u}^{2}}{2\alpha_{l}}\|L\|_{1}. ∎

## Appendix C Comparison of Convergence Rates of Lars and Sgd

Inspired by the comparison used by ( Bernstein et al., 2018 ) for comparing SIGN Sgd with Sgd , we define the following quantities: ( ∑ i = 1 h ‖ ∇ i f ​ ( x t ) ‖ ) 2 \displaystyle\left(\sum_{i=1}^{h}\|\nabla_{i}f(x_{t})\|\right)^{2} = ψ ⁡ ( ∇ f ​ ( x t ) ) ​ d ​ ‖ ∇ f ​ ( x t ) ‖ 2 h ≥ ψ g ​ d ​ ‖ ∇ f ​ ( x t ) ‖ 2 h \displaystyle=\frac{\psi(\nabla f(x_{t}))d\|\nabla f(x_{t})\|^{2}}{h}\geq\frac{\psi_{g}d\|\nabla f(x_{t})\|^{2}}{h} ‖ L ‖ 1 2 \displaystyle\|L\|_{1}^{2} ≤ ψ L ​ d 2 ​ ‖ L ‖ ∞ 2 h 2 \displaystyle\leq\frac{\psi_{L}d^{2}\|L\|_{\infty}^{2}}{h^{2}} ‖ σ ‖ 1 2 \displaystyle\|\sigma\|_{1}^{2} = ψ σ ​ d ​ ‖ σ ‖ 2 h . \displaystyle=\frac{\psi_{\sigma}d\|\sigma\|^{2}}{h}. Then Lars convergence rate can be written in the following manner: ( 𝔼 [ ∥ ∇ f ( x a ) ∥ ) 2 ≤ O ( ( f ⁡ ( x 1 ) − f ⁡ ( x ∗ ) ) ​ L ∞ T ψ L ψ g 2 + ‖ σ ‖ 2 T ψ σ 2 ψ g 2 ) . \displaystyle\left(\mathbb{E}[\|\nabla f(x_{a})\|\right)^{2}\leq O\left(\frac{(f(x_{1})-f(x^{*}))L_{\infty}}{T}\frac{\psi_{L}}{\psi_{g}^{2}}+\frac{\|\sigma\|^{2}}{T}\frac{\psi_{\sigma}^{2}}{\psi_{g}^{2}}\right). If ψ L ≪ ψ g 2 \psi_{L}\ll\psi_{g}^{2} and ψ σ ≪ ψ g 2 \psi_{\sigma}\ll\psi_{g}^{2} then Lars (i.e., gradient is more denser than curvature or stochasticity), we gain over Sgd . Otherwise, Sgd ’s upper bound on convergence rate is better.

## Appendix D N-LAMB: Nesterov Momentum for LAMB

Sutskever et al. (2013) report that Nesterov’s accelerated gradient (NAG) proposed by Nesterov (1983) is conceptually and empirically better than the regular momentum method for convex, non-stochastic objectives. Dozat (2016) incorporated Nesterov’s momentum into Adam optimizer and proposed the Nadam optimizer. Specifically, only the first moment of Adam was modified and the second moment of Adam was unchanged. The results on several applications (Word2Vec, Image Recognition, and LSTM Language Model) showed that Nadam optimizer improves the speed of convergence and the quality of the learned models. We also tried using Nesterov’s momentum to replace the regular momentum of LAMB optimizer’s first moment. In this way, we got a new algorithm named as N-LAMB (Nesterov LAMB). The complete algorithm is in Algorithm 3 . We can also Nesterov’s momentum to replace the regular momentum of LAMB optimizer’s second moment. We refer to this algorithm as NN-LAMB (Nesterov’s momentum for both the first moment and the second moment). The details of NN-LAMB were shown in Algorithm 4 .

Dozat (2016) suggested the best performance of Nadam was achieved by β 1 \beta_{1} = 0.975, β 2 \beta_{2} = 0.999, and ϵ \epsilon = 1e-8. We used the same settings for N-LAMB and NN-LAMB. We scaled the batch size to 32K for ImageNet training with ResNet-50. Our experimental results show that N-LAMB and NN-LAMB can achieve a comparable accuracy compared to LAMB optimizer. Their performances are much better than momentum solver (Figure 1 ).

## Appendix E LAMB with learning rate correction

There are two operations at each iteration in original Adam optimizer (let us call it adam-correction): m t = m t / ( 1 − β 1 t ) m_{t}=m_{t}/(1-{\beta}_{1}^{t}) v t = v t / ( 1 − β 2 t ) v_{t}=v_{t}/(1-{\beta}_{2}^{t}) It has an impact on the learning rate by η t := η t ∗ ( 1 − β 2 t ) / ( 1 − β 1 t ) {\eta}_{t}:={\eta}_{t}*\sqrt{(1-{\beta}_{2}^{t})/(1-{\beta}_{1}^{t})} . According to our experimental results, adam-correction essentially has the same effect as learning rate warmup (see Figure 2 ). The warmup function often was implemented in the modern deep learning system. Thus, we can remove adam-correction from the LAMB optimizer. We did not observe any drop in the test or validation accuracy for BERT and ImageNet training.

## Appendix F LAMB with different norms

We need to compute the matrix/tensor norm for each layer when we do the parameter updating in the LAMB optimizer. We tried different norms in LAMB optimizer. However, we did not observe a significant difference in the validation accuracy of ImageNet training with ResNet-50. In our experiments, the difference in validation accuracy is less than 0.1 percent (Figure 3 ). We use L2 norm as the default.

## Appendix G Regular Batch Sizes for Small Datasets: MNIST and CIFAR-10.

According to DAWNBench, DavidNet (a custom 9-layer Residual ConvNet) is the fastest model for CIFAR-10 dataset (as of April 1st, 2019) 5 5 5 https://dawn.cs.stanford.edu/benchmark/CIFAR10/train.html . The baseline uses the momentum SGD optimizer. Table 6 and Figure 4 show the test accuracy of CIFAR-10 training with DavidNet. The PyTorch implementation (momentum SGD optimizer) on GPUs was reported on Standford DAWNBench’s website, which achieves 94.06% in 24 epochs. The Tensorflow implementation (momentum SGD optimizer) on TPU achieves a 93.72% accuracy in 24 epochs 6 6 6 https://github.com/fenwickslab/dl_tutorials/blob/master/tutorial3_cifar10_davidnet_fix.ipynb . We use the implementation of TensorFlow on TPUs. Lamb optimizer is able to achieve 94.08% test accuracy in 24 epochs, which is better than other adaptive optimizers and momentum SGD. Even on the smaller tasks like MNIST training with LeNet, Lamb is able to achieve a better accuracy than existing solvers (Table 7 ).

## Appendix H Implementation Details and Additional Results

There are several hyper-parameters in Lamb optimizer. Although users do not need to tune them, we explain them to help users to have a better understanding. β 1 {\beta}_{1} is used for decaying the running average of the gradient. β 2 {\beta}_{2} is used for decaying the running average of the square of gradient. The default setting for other parameters: weight decay rate λ \lambda =0.01, β 1 {\beta}_{1} =0.9, β 2 {\beta}_{2} =0.999, ϵ \epsilon =1e-6. We did not tune β 1 {\beta}_{1} and β 2 {\beta}_{2} . However, our experiments show that tuning them may get a higher accuracy.

Based on our experience, learning rate is the most important hyper-parameter that affects the learning efficiency and final accuracy. Bengio (2012) suggests that it is often the single most important hyper-parameter and that it always should be tuned. Thus, to make sure we have a solid baseline, we carefully tune the learning rate of Adam , AdamW , Adagrad , and momentum Sgd

In our experiments, we found that the validation loss is not reliable for large-batch training. A lower validation loss does not necessarily lead to a higher validation accuracy (Figure 5 ). Thus, we use the test/val accuracy or F1 score on dev set to evaluate the optimizers.

#### H.0.1 BERT

Table 8 shows some of the tuning information from BERT training with AdamW optimizer. AdamW stops scaling at the batch size of 16K. The target F1 score is 90.5. Lamb achieves a F1 score of 91.345. The table shows the tuning information of AdamW . In Table 8 , we report the best F1 score we observed from our experiments.

The loss curves of BERT training by Lamb for different batch sizes are shown in Figure 6 . We observe that the loss curves are almost identical to each other, which means our optimizer scales well with the batch size.

The training loss curve of BERT mixed-batch pre-training with LAMB is shown in Figure 7 . This figure shows that LAMB can make the training converge smoothly at the batch size of 64K.

Figure 8 shows that we can achieve 76.8% scaling efficiency by scaling the batch size (49.1 times speedup by 64 times computational resources) and 101.8% scaling efficiency with mixed-batch (65.2 times speedup by 64 times computational resources)

#### H.0.2 ImageNet

Figures 9 - 14 show the LAMB trust ratio at different iterations for ImageNet training with ResNet-50. From these figures we can see that these ratios are very different from each other for different layers. LAMB uses the trust ratio to help the slow learners to train faster.

### H.1 Baseline tuning details for ImageNet training with ResNet-50

If you are not interested in the baseline tuning details, please skip this section.

Goyal et al. (2017) suggested a proper learning rate warmup and decay scheme may help improve the ImageNet classification accuracy. We included these techniques in Adam/AdamW/AdaGrad tuning. Specifically, we use the learning rate recipe of Goyal et al. (2017) : (1) 5-epoch warmup to stablize the initial stage; and (2) multiply the learning rate by 0.1 at 30th, 60th, and 80th epoch. The target accuracy is around 76.3% ( Goyal et al., 2017 ) . There techniques help to improve the accuracy of Adam/AdamW/AdaGrad to around 73%. However, even with these techniques, Adam/AdamW/AdaGrad stil can not achieve the target validation accuracy.