Large Batch Optimization for Deep Learning: Training BERT in 76 minutes — part 2/5
Cleaned-text code points [15559, 31254). Bibliography omitted here; source.md includes it.
for all layers i ∈ [ h ] i\in[h] and where x t ( i ) x^{(i)}_{t} and g t ( i ) g^{(i)}_{t} are the parameters and the gradients of the i th i^{\text{th}} layer at time step t t . The normalization modification is similar to one typically used in normalized gradient descent except that it is done layerwise. Note that the modification leads to a biased gradient update; however, in large-batch settings, it can be shown that this bias is small. It is intuitive that such a normalization provides robustness to exploding gradients (where the gradient can be arbitrarily large) and plateaus (where the gradient can be arbitrarily small). Normalization of this form essentially ignores the size of the gradient and is particularly useful in large batch settings where the direction of the gradient is largely preserved.
The scaling term involving ϕ \phi ensures that the norm of the update is of the same order as that of the parameter. We found that this typically ensures faster convergence in deep neural networks. In practice, we observed that a simple function of ϕ ( z ) = min { max { z , γ l } , γ u } \phi(z)=\min\{\max\{z,\gamma_{l}\},\gamma_{u}\} works well. It is instructive to consider the case where ϕ ( z ) = z \phi(z)=z . In this scenario, the overall change in the learning rate is ‖ x t ( i ) ‖ ‖ g t ( i ) ‖ \tfrac{\|x_{t}^{(i)}\|}{\|g_{t}^{(i)}\|} , which can also be interpreted as an estimate on the inverse of Lipschitz constant of the gradient (see equation 2 ). We now discuss different instantiations of the strategy discussed above. In particular, we focus on two algorithms: Lars ( 3.1 ) and the proposed method, Lamb ( 3.2 ).
### 3.1 Lars Algorithm
The first instantiation of the general strategy is Lars algorithm ( You et al., 2017 ) , which is obtained by using momentum optimizer as the base algorithm 𝒜 \mathcal{A} in the framework. Lars was earlier proposed for large batch learning for ResNet on ImageNet. In general, it is observed that the using (heavy-ball) momentum, one can reduce the variance in the stochastic gradients at the cost of little bias. The pseudocode for Lars is provide in Algorithm 1 .
We now provide convergence analysis for Lars in general nonconvex setting stated in this paper. For the sake of simplicity, we analyze the case where β 1 = 0 \beta_{1}=0 and λ = 0 \lambda=0 in Algorithm 1 . However, our analysis should extend to the general case as well. We will defer all discussions about the convergence rate to the end of the section.
###### Theorem 2 .
Let η t = η = 2 ( f ( x 1 ) − f ( x ∗ ) ) α u 2 ‖ L ‖ 1 T \eta_{t}=\eta=\sqrt{\tfrac{2(f(x_{1})-f(x^{*}))}{\alpha_{u}^{2}\|L\|_{1}T}} for all t ∈ [ T ] t\in[T] , b = T b=T , α l ≤ ϕ ( v ) ≤ α u \alpha_{l}\leq\phi(v)\leq\alpha_{u} for all v > 0 v>0 where α l , α u > 0 \alpha_{l},\alpha_{u}>0 . Then for x t x_{t} generated using Lars (Algorithm 1 ), we have the following bound ( 𝔼 [ 1 h ∑ i = 1 h ‖ ∇ i f ( x a ) ‖ ] ) 2 ≤ O ( ( f ( x 1 ) − f ( x ∗ ) ) L a v g T + ‖ σ ‖ 1 2 T h ) , \displaystyle\left(\mathbb{E}\left[\frac{1}{\sqrt{h}}\sum_{i=1}^{h}\|\nabla_{i}f(x_{a})\|\right]\right)^{2}\leq O\left(\frac{(f(x_{1})-f(x^{*}))L_{avg}}{T}+\frac{\|\sigma\|^{2}_{1}}{Th}\right), where x ∗ x^{*} is an optimal solution to the problem in equation 1 and x a x_{a} is an iterate uniformly randomly chosen from { x 1 , ⋯ , x T } \{x_{1},\cdots,x_{T}\} .
### 3.2 Lamb Algorithm
The second instantiation of the general strategy is obtained by using Adam as the base algorithm 𝒜 \mathcal{A} . Adam optimizer is popular in deep learning community and has shown to have good performance for training state-of-the-art language models like Bert . Unlike Lars , the adaptivity of Lamb is two-fold: (i) per dimension normalization with respect to the square root of the second moment used in Adam and (ii) layerwise normalization obtained due to layerwise adaptivity. The pseudocode for Lamb is provided in Algorithm 2 . When β 1 = 0 \beta_{1}=0 and β 2 = 0 \beta_{2}=0 , the algorithm reduces to be Sign Sgd where the learning rate is scaled by square root of the layer dimension ( Bernstein et al., 2018 ) .
The following result provides convergence rate for Lamb in general nonconvex settings. Similar to the previous case, we focus on the setting where β 1 = 0 \beta_{1}=0 and λ = 0 \lambda=0 . As before, our analysis extends to the general case; however, the calculations become messy.
###### Theorem 3 .
Let η t = η = 2 ( f ( x 1 ) − f ( x ∗ ) ) α u 2 ‖ L ‖ 1 T \eta_{t}=\eta=\sqrt{\tfrac{2(f(x_{1})-f(x^{*}))}{\alpha_{u}^{2}\|L\|_{1}T}} for all t ∈ [ T ] t\in[T] , b = T b=T , d i = d / h d_{i}=d/h for all i ∈ [ h ] i\in[h] , and α l ≤ ϕ ( v ) ≤ α u \alpha_{l}\leq\phi(v)\leq\alpha_{u} for all v > 0 v>0 where α l , α u > 0 \alpha_{l},\alpha_{u}>0 . Then for x t x_{t} generated using Lamb (Algorithm 2 ), we have the following bounds: 1. When β 2 = 0 \beta_{2}=0 , we have ( 𝔼 [ 1 d ‖ ∇ f ( x a ) ‖ 1 ] ) 2 \displaystyle\left(\mathbb{E}\left[\frac{1}{\sqrt{d}}\|\nabla f(x_{a})\|_{1}\right]\right)^{2} ≤ O ( ( f ( x 1 ) − f ( x ∗ ) ) L a v g T + ‖ σ ~ ‖ 1 2 T h ) , \displaystyle\leq O\left(\frac{(f(x_{1})-f(x^{*}))L_{avg}}{T}+\frac{\|\tilde{\sigma}\|^{2}_{1}}{Th}\right),
2. When β 2 > 0 \beta_{2}>0 , we have 𝔼 [ ‖ ∇ f ( x a ) ‖ 2 ] \displaystyle\mathbb{E}[\|\nabla f(x_{a})\|^{2}] ≤ O ( G 2 d h ( 1 − β 2 ) × [ 2 ( f ( x 1 ) − f ( x ∗ ) ) ‖ L ‖ 1 T + ‖ σ ~ ‖ 1 T ] ) , \displaystyle\leq O\left(\sqrt{\frac{G^{2}d}{h(1-\beta_{2})}}\times\left[\sqrt{\frac{2(f(x_{1})-f(x^{*}))\|L\|_{1}}{T}}+\frac{\|\tilde{\sigma}\|_{1}}{\sqrt{T}}\right]\right),
where x ∗ x^{*} is an optimal solution to the problem in equation 1 and x a x_{a} is an iterate uniformly randomly chosen from { x 1 , ⋯ , x T } \{x_{1},\cdots,x_{T}\} .
Discussion on convergence rates. We first start our discussion with the comparison of convergence rate of Lars with that of Sgd (Theorem 1 ). The convergence rates of Lars and Sgd differ in two ways: (1) the convergence criterion is ( 𝔼 [ ∑ i = 1 h ‖ ∇ i f ‖ ] ) 2 (\mathbb{E}[\sum_{i=1}^{h}\|\nabla_{i}f\|])^{2} as opposed to 𝔼 [ ‖ ∇ f ‖ 2 ] \mathbb{E}[\|\nabla f\|^{2}] in Sgd and (2) the dependence on L L and σ \sigma in the convergence rate. Briefly, the convergence rate of Lars is better than Sgd when the gradient is denser than curvature and stochasticity. This convergence rate comparison is similar in spirit to the one obtained in ( Bernstein et al., 2018 ) . Assuming that the convergence criterion in Theorem 1 and Theorem 2 is of similar order (which happens when gradients are fairly dense), convergence rate of Lars and Lamb depend on L a v g L_{avg} instead of L ∞ L_{\infty} and are thus, significantly better than that of Sgd . A more quantitative comparison is provided in Section C of the Appendix. The comparison of Lamb (with β 2 \beta_{2} = 0) with Sgd is along similar lines. We obtain slightly worse rates for the case where β 2 > 0 \beta_{2}>0 ; although, we believe that its behavior should be better than the case β 2 = 0 \beta_{2}=0 . We leave this investigation to future work.
## 4 Experiments
We now present empirical results comparing Lamb with existing optimizers on two important large batch training tasks: Bert and ResNet -50 training. We also compare Lamb with existing optimizers for small batch size ( < 1 K <1K ) and small dataset (e.g. CIFAR, MNIST) (see Appendix).
Experimental Setup. To demonstrate its robustness, we use very minimal hyperparameter tuning for the Lamb optimizer. Thus, it is possible to achieve better results by further tuning the hyperparameters. The parameters β 1 \beta_{1} and β 2 \beta_{2} in Algorithm 2 are set to 0.9 0.9 and 0.999 0.999 respectively in all our experiments; we only tune the learning rate. We use a polynomially decaying learning rate of η t = η 0 × ( 1 − t / T ) \eta_{t}=\eta_{0}\times(1-t/T) in Algorithm 2 ), which is the same as in Bert baseline. This setting also works for all other applications in this paper. Furthermore, for Bert and ResNet -50 training, we did not tune the hyperparameters of Lamb while increasing the batch size. We use the square root of LR scaling rule to automatically adjust learning rate and linear-epoch warmup scheduling. We use TPUv3 in all the experiments. A TPUv3 Pod has 1024 chips and can provide more than 100 petaflops performance for mixed precision computing. To make sure we are comparing with solid baselines, we use grid search to tune the hyper-parameters for Adam , Adagrad , AdamW ( Adam with weight decay), and Lars . We also tune weight decay for AdamW . All the hyperparameter tuning settings are reported in the Appendix. Due to space constraints, several experimental details are relegated to the Appendix.
### 4.1 Bert Training
We first discuss empirical results for speeding up Bert training. For this experiment, we use the same dataset as Devlin et al. (2018) , which is a concatenation of Wikipedia and BooksCorpus with 2.5B and 800M words respectively. We specifically focus on the SQuAD task 2 2 2 https://rajpurkar.github.io/SQuAD-explorer/ in this paper. The F1 score on SQuAD-v1 is used as the accuracy metric in our experiments. All our comparisons are with respect to the baseline Bert model by Devlin et al. (2018) . To train Bert , Devlin et al. (2018) first train the model for 900k iterations using a sequence length of 128 and then switch to a sequence length of 512 for the last 100k iterations. This results in a training time of around 3 days on 16 TPUv3 chips. The baseline Bert model 3 3 3 Pre-trained BERT model can be downloaded from https://github.com/google-research/bert achieves a F1 score of 90.395. To ensure a fair comparison, we follow the same SQuAD fine-tune procedure of Devlin et al. (2018) without modifying any configuration (including number of epochs and hyperparameters). As noted earlier, we could get even better results by changing the fine-tune configuration. For instance, by just slightly changing the learning rate in the fine-tune stage, we can obtain a higher F1 score of 91.688 for the batch size of 16K using Lamb . We report a F1 score of 91.345 in Table 1 , which is the score obtained for the untuned version. Below we describe two different training choices for training Bert and discuss the corresponding speedups.
For the first choice, we maintain the same training procedure as the baseline except for changing the training optimizer to Lamb . We run with the same number of epochs as the baseline but with batch size scaled from 512 to 32K. The choice of 32K batch size (with sequence length 512) is mainly due to memory limits of TPU Pod. Our results are shown in Table 1 . By using the Lamb optimizer, we are able to achieve a F1 score of 91.460 in 15625 iterations for a batch size of 32768 (14063 iterations for sequence length 128 and 1562 iterations for sequence length 512). With 32K batch size, we reduce Bert training time from 3 days to around 100 minutes. We achieved 49.1 times speedup by 64 times computational resources (76.7% efficiency). We consider the speedup is great because we use the synchronous data-parallelism. There is a communication overhead coming from transferring of the gradients over the interconnect. For ResNet -50, researchers are able to achieve 90% scaling efficiency because ResNet -50 has much fewer parameters (# parameters is equal to #gradients) than Bert (25 million versus 300 million).
To obtain further improvements, we use the Mixed-Batch Training procedure with Lamb . Recall that Bert training involves two stages: the first 9/10 of the total epochs use a sequence length of 128, while the last 1/10 of the total epochs use a sequence length of 512. For the second stage training, which involves a longer sequence length, due to memory limits, a maximum batch size of only 32768 can be used on a TPUv3 Pod. However, we can potentially use a larger batch size for the first stage because of a shorter sequence length. In particular, the batch size can be increased to 131072 for the first stage. However, we did not observe any speedup by increasing the batch size from 65536 to 131072 for the first stage, thus, we restrict the batch size to 65536 for this stage. By using this strategy, we are able to make full utilization of the hardware resources throughout the training procedure. Increasing the batch size is able to warm-up and stabilize the optimization process ( Smith et al., 2017 ) , but decreasing the batch size brings chaos to the optimization process and can cause divergence. In our experiments, we found a technique that is useful to stabilize the second stage optimization. Because we switched to a different optimization problem, it is necessary to re-warm-up the optimization. Instead of decaying the learning rate at the second stage, we ramp up the learning rate from zero again in the second stage (re-warm-up). As with the first stage, we decay the learning rate after the re-warm-up phase. With this method, we only need 8599 iterations and finish Bert training in 76 minutes (100.2% efficiency).
##### Comparison with AdamW and Lars .
To ensure that our approach is compared to a solid baseline for the Bert training, we tried three different strategies for tuning AdamW : (1) AdamW with default hyperparameters (see Devlin et al. (2018) ) (2) AdamW with the same hyperparameters as Lamb , and (3) AdamW with tuned hyperparameters. AdamW stops scaling at the batch size of 16K because it is not able to achieve the target F1 score (88.1 vs 90.4). The tuning information of AdamW is shown in the Appendix. For 64K/32K mixed-batch training, even after extensive tuning of the hyperparameters, we fail to get any reasonable result with AdamW optimizer. We conclude that AdamW does not work well in large-batch Bert training or is at least hard to tune. We also observe that Lamb performs better than Lars for all batch sizes (see Table 2 ).
### 4.2 ImageNet Training with ResNet-50.
ImageNet training with ResNet-50 is an industry standard metric that is being used in MLPerf 4 4 4 https://mlperf.org/ . The baseline can get 76.3% top-1 accuracy in 90 epochs ( Goyal et al., 2017 ) . All the successful implementations are based on momentum SGD ( He et al., 2016 ; Goyal et al., 2017 ) or Lars optimizer ( Ying et al., 2018 ; Jia et al., 2018 ; Mikami et al., 2018 ; You et al., 2018 ; Yamazaki et al., 2019 ) . Before our study, we did not find any paper reporting a state-of-the-art accuracy achieved by Adam , Adagrad , or AdamW optimizer. In our experiments, even with comprehensive hyper-parameter tuning, Adagrad / Adam / AdamW (with batch size 16K) only achieves 55.38%/66.04%/67.27% top-1 accuracy. After adding learning rate scheme of Goyal et al. (2017) , the top-1 accuracy of Adagrad / Adam / AdamW was improved to 72.0%/73.48%/73.07%. However, they are still much lower than 76.3%. The details of the tuning information are in the Appendix. Table 3 shows that Lamb can achieve the target accuracy. Beyond a batch size of 8K, Lamb ’s accuracy is higher than the momentum. Lamb ’s accuracy is also slightly better than Lars . At a batch size of 32K, Lamb achieves 76.4% top-1 accuracy while Lars achieves 76.3%. At a batch size of 2K, Lamb is able to achieve 77.11% top-1 accuracy while Lars achieves 76.6%.
### 4.3 Hyperparameters for scaling the batch size
For Bert and ImageNet training, we did not tune the hyperparameters of Lamb optimizer when increasing the batch size. We use the square root LR scaling rule and linear-epoch warmup scheduling to automatically adjust learning rate. The details can be found in Tables 4 and 5
## 5 Conclusion