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

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

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

###### 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 online 1 1 1 https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py .

## 1 Introduction

With the advent of large scale datasets, training large deep neural networks, even using computationally efficient optimization methods like Stochastic gradient descent ( Sgd ) (\textsc{Sgd}) , has become particularly challenging. For instance, training state-of-the-art deep learning models like Bert and ResNet-50 takes 3 days on 16 TPUv3 chips and 29 hours on 8 Tesla P100 gpus respectively ( Devlin et al., 2018 ; He et al., 2016 ) . Thus, there is a growing interest to develop optimization solutions to tackle this critical issue. The goal of this paper is to investigate and develop optimization techniques to accelerate training large deep neural networks, mostly focusing on approaches based on variants of Sgd .

Methods based on Sgd iteratively update the parameters of the model by moving them in a scaled (negative) direction of the gradient calculated on a minibatch. However, Sgd ’s scalability is limited by its inherent sequential nature. Owing to this limitation, traditional approaches to improve Sgd training time in the context of deep learning largely resort to distributed asynchronous setup ( Dean et al., 2012 ; Recht et al., 2011 ) . However, the implicit staleness introduced due to the asynchrony limits the parallelization of the approach, often leading to degraded performance. The feasibility of computing gradient on large minibatches in parallel due to recent hardware advances has seen the resurgence of simply using synchronous Sgd with large minibatches as an alternative to asynchronous Sgd . However, naïvely increasing the batch size typically results in degradation of generalization performance and reduces computational benefits ( Goyal et al., 2017 ) .

Synchronous Sgd on large minibatches benefits from reduced variance of the stochastic gradients used in Sgd . This allows one to use much larger learning rates in Sgd , typically of the order square root of the minibatch size. Surprisingly, recent works have demonstrated that up to certain minibatch sizes, linear scaling of the learning rate with minibatch size can be used to further speed up the training Goyal et al. (2017) . These works also elucidate two interesting aspects to enable the use of linear scaling in large batch synchronous Sgd : (i) linear scaling of learning rate is harmful during the initial phase; thus, a hand-tuned warmup strategy of slowly increasing the learning rate needs to be used initially, and (ii) linear scaling of learning rate can be detrimental beyond a certain batch size. Using these tricks, Goyal et al. (2017) was able to drastically reduce the training time of ResNet-50 model from 29 hours to 1 hour using a batch size of 8192. While these works demonstrate the feasibility of this strategy for reducing the wall time for training large deep neural networks, they also highlight the need for an adaptive learning rate mechanism for large batch learning.

Variants of Sgd using layerwise adaptive learning rates have been recently proposed to address this problem. The most successful in this line of research is the Lars algorithm ( You et al., 2017 ) , which was initially proposed for training ResNet . Using Lars , ResNet-50 can be trained on ImageNet in just a few minutes! However, it has been observed that its performance gains are not consistent across tasks. For instance, Lars performs poorly for attention models like Bert . Furthermore, theoretical understanding of the adaptation employed in Lars is largely missing. To this end, we study and develop new approaches specially catered to the large batch setting of our interest.

Contributions. More specifically, we make the following main contributions in this paper.

• Inspired by Lars , we investigate a general adaptation strategy specially catered to large batch learning and provide intuition for the strategy.

• Based on the adaptation strategy, we develop a new optimization algorithm ( Lamb ) for achieving adaptivity of learning rate in Sgd . Furthermore, we provide convergence analysis for both Lars and Lamb to achieve a stationary point in nonconvex settings. We highlight the benefits of using these methods for large batch settings.

• We demonstrate the strong empirical performance of Lamb across several challenging tasks. Using Lamb we scale the batch size in training Bert to more than 32k without degrading the performance; thereby, cutting the time down from 3 days to 76 minutes. Ours is the first work to reduce Bert training wall time to less than couple of hours.

• We also demonstrate the efficiency of Lamb for training state-of-the-art image classification models like ResNet . To the best of our knowledge, ours is first adaptive solver that can achieve state-of-the-art accuracy for ResNet -50 as adaptive solvers like Adam fail to obtain the accuracy of Sgd with momentum for these tasks.

### 1.1 Related Work

The literature on optimization for machine learning is vast and hence, we restrict our attention to the most relevant works here. Earlier works on large batch optimization for machine learning mostly focused on convex models, benefiting by a factor of square root of batch size using appropriately large learning rate. Similar results can be shown for nonconvex settings wherein using larger minibatches improves the convergence to stationary points; albeit at the cost of extra computation. However, several important concerns were raised with respect to generalization and computational performance in large batch nonconvex settings. It was observed that training with extremely large batch was difficult ( Keskar et al., 2016 ; Hoffer et al., 2017 ) . Thus, several prior works carefully hand-tune training hyper-parameters, like learning rate and momentum, to avoid degradation of generalization performance ( Goyal et al., 2017 ; Li, 2017 ; You et al., 2018 ; Shallue et al., 2018 ) .

( Krizhevsky, 2014 ) empirically found that simply scaling the learning rate linearly with respect to batch size works better up to certain batch sizes. To avoid optimization instability due to linear scaling of learning rate, Goyal et al. (2017) proposed a highly hand-tuned learning rate which involves a warm-up strategy that gradually increases the LR to a larger value and then switching to the regular LR policy (e.g. exponential or polynomial decay). Using LR warm-up and linear scaling, Goyal et al. (2017) managed to train ResNet -50 with batch size 8192 without loss in generalization performance. However, empirical study ( Shallue et al., 2018 ) shows that learning rate scaling heuristics with the batch size do not hold across all problems or across all batch sizes.

More recently, to reduce hand-tuning of hyperparameters, adaptive learning rates for large batch training garnered significant interests. Several recent works successfully scaled the batch size to large values using adaptive learning rates without degrading the performance, thereby, finishing ResNet -50 training on ImageNet in a few minutes ( You et al., 2018 ; Iandola et al., 2016 ; Codreanu et al., 2017 ; Akiba et al., 2017 ; Jia et al., 2018 ; Smith et al., 2017 ; Martens & Grosse, 2015 ; Devarakonda et al., 2017 ; Mikami et al., 2018 ; Osawa et al., 2018 ; You et al., 2019 ; Yamazaki et al., 2019 ) . To the best of our knowledge, the fastest training result for ResNet -50 on ImageNet is due to Ying et al. (2018) , who achieve 76+% top-1 accuracy. By using the Lars optimizer and scaling the batch size to 32K on a TPUv3 Pod, Ying et al. (2018) was able to train ResNet -50 on ImageNet in 2.2 minutes. However, it was empirically observed that none of these performance gains hold in other tasks such as BERT training (see Section 4 ).

## 2 Preliminaries

##### Notation.

For any vector x t ∈ ℝ d x_{t}\in\mathbb{R}^{d} , either x t , j x_{t,j} or [ x t ] j [x_{t}]_{j} are used to denote its j th j^{\text{th}} coordinate where j ∈ [ d ] j\in[d] . Let 𝕀 \mathbb{I} be the d × d d\times d identity matrix, and let 𝕀 = [ 𝕀 1 , 𝕀 2 , … , 𝕀 h ] \mathbb{I}=[\mathbb{I}_{1},\mathbb{I}_{2},...,\mathbb{I}_{h}] be its decomposition into column submatrices 𝕀 i = d × d h \mathbb{I}_{i}=d\times d_{h} . For x ∈ ℝ d x\in\mathbb{R}^{d} , let x ( i ) x^{(i)} be the block of variables corresponding to the columns of I i I_{i} i.e., x ( i ) = 𝕀 i ⊤ ​ x ∈ ℝ d i x^{(i)}=\mathbb{I}_{i}^{\top}x\in\mathbb{R}^{d_{i}} for i = { 1 , 2 , ⋯ , h } i=\{1,2,\cdots,h\} . For any function f : ℝ d → ℝ f:\mathbb{R}^{d}\rightarrow\mathbb{R} , we use ∇ i f ​ ( x ) \nabla_{i}f(x) to denote the gradient with respect to x ( i ) x^{(i)} . For any vectors u , v ∈ ℝ d u,v\in\mathbb{R}^{d} , we use u 2 u^{2} and u / v u/v to denote elementwise square and division operators respectively. We use ∥ . ∥ \|.\| and ∥ . ∥ 1 \|.\|_{1} to denote l 2 l_{2} -norm and l 1 l_{1} -norm of a vector respectively.

We start our discussion by formally stating the problem setup. In this paper, we study nonconvex stochastic optimization problems of the form min x ∈ ℝ d ⁡ f ⁡ ( x ) := 𝔼 s ∼ ℙ ​ [ ℓ ⁡ ( x , s ) ] + λ 2 ​ ‖ x ‖ 2 , \displaystyle\min_{x\in\mathbb{R}^{d}}f(x):=\mathbb{E}_{s\sim\mathbb{P}}[\ell(x,s)]+\frac{\lambda}{2}\|x\|^{2}, (1) where ℓ \ell is a smooth (possibly nonconvex) function and ℙ \mathbb{P} is a probability distribution on the domain 𝒮 ⊂ ℝ k \mathcal{S}\subset\mathbb{R}^{k} . Here, x x corresponds to model parameters, ℓ \ell is the loss function and ℙ \mathbb{P} is an unknown data distribution.

We assume function ℓ ⁡ ( x ) \ell(x) is L i L_{i} - smooth with respect to x ( i ) x^{(i)} , i.e., there exists a constant L i L_{i} such that ‖ ∇ i ℓ ​ ( x , s ) − ∇ i ℓ ​ ( y , s ) ‖ ≤ L i ​ ‖ x ( i ) − y ( i ) ‖ , ∀ x , y ∈ ℝ d , and ​ s ∈ 𝒮 , \|\nabla_{i}\ell(x,s)-\nabla_{i}\ell(y,s)\|\leq L_{i}\|x^{(i)}-y^{(i)}\|,\quad\forall\ x,y\in\mathbb{R}^{d},\text{ and }s\in\mathcal{S}, (2) for all i ∈ [ h ] i\in[h] . We use L = ( L 1 , ⋯ , L h ) ⊤ L=(L_{1},\cdots,L_{h})^{\top} to denote the h h -dimensional vector of Lipschitz constants. We use L ∞ L_{\infty} and L a ​ v ​ g L_{avg} to denote max i ⁡ L i \max_{i}L_{i} and ∑ i L i h \sum_{i}\tfrac{L_{i}}{h} respectively. We assume the following bound on the variance in stochastic gradients: 𝔼 ​ ‖ ∇ i ℓ ​ ( x , s ) − ∇ i f ​ ( x ) ‖ 2 ≤ σ i 2 \mathbb{E}\|\nabla_{i}\ell(x,s)-\nabla_{i}f(x)\|^{2}\leq\sigma_{i}^{2} for all x ∈ ℝ d x\in\mathbb{R}^{d} and i ∈ [ h ] i\in[h] . Furthermore, we also assume 𝔼 ​ ‖ [ ∇ ℓ ​ ( x , s ) ] i − [ ∇ f ​ ( x ) ] i ‖ 2 ≤ σ ~ i 2 \mathbb{E}\|[\nabla\ell(x,s)]_{i}-[\nabla f(x)]_{i}\|^{2}\leq\tilde{\sigma}_{i}^{2} for all x ∈ ℝ d x\in\mathbb{R}^{d} and i ∈ [ d ] i\in[d] . We use σ = ( σ 1 , ⋯ , σ h ) ⊤ \sigma=(\sigma_{1},\cdots,\sigma_{h})^{\top} and σ ~ = ( σ ~ 1 , ⋯ , σ ~ d ) ⊤ \tilde{\sigma}=(\tilde{\sigma}_{1},\cdots,\tilde{\sigma}_{d})^{\top} to denote the vectors of standard deviations of stochastic gradient per layer and per dimension respectively. Finally, we assume that the gradients are bounded i.e., [ ∇ l ​ ( x , s ) ] j ≤ G [\nabla l(x,s)]_{j}\leq G for all i ∈ [ d ] i\in[d] , x ∈ ℝ d x\in\mathbb{R}^{d} and s ∈ 𝒮 s\in\mathcal{S} . Note that such assumptions are typical in the analysis of stochastic first-order methods (cf. ( Ghadimi & Lan, 2013a ; Ghadimi et al., 2014 ) ).

Stochastic gradient descent ( Sgd ) is one of the simplest first-order algorithms for solving problem in Equation 1 . The update at the t th t^{\text{th}} iteration of Sgd is of the following form: x t + 1 = x t − η t ​ 1 | 𝒮 t | ​ ∑ s t ∈ 𝒮 t ∇ ℓ ​ ( x t , s t ) + λ ​ x t , x_{t+1}=x_{t}-\eta_{t}\frac{1}{|\mathcal{S}_{t}|}\sum_{s_{t}\in\mathcal{S}_{t}}\nabla\ell(x_{t},s_{t})+\lambda x_{t}, ( Sgd ) where S t S_{t} is set of b b random samples drawn from the distribution ℙ \mathbb{P} . For very large batch settings, the following is a well-known result for Sgd .

###### Theorem 1 ( ( Ghadimi & Lan, 2013b ) ) .

With large batch b = T b=T and using appropriate learning rate, we have the following for the iterates of Sgd : 𝔼 ⁡ [ ‖ ∇ f ​ ( x a ) ‖ 2 ] ≤ O ⁡ ( ( f ⁡ ( x 1 ) − f ⁡ ( x ∗ ) ) ​ L ∞ T + ‖ σ ‖ 2 T ) . \displaystyle\mathbb{E}\left[\|\nabla f(x_{a})\|^{2}\right]\leq O\left(\frac{(f(x_{1})-f(x^{*}))L_{\infty}}{T}+\frac{\|\sigma\|^{2}}{T}\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}\} .

However, tuning the learning rate η t \eta_{t} in Sgd , especially in large batch settings, is difficult in practice. Furthermore, the dependence on L ∞ L_{\infty} (the maximum of smoothness across dimension) can lead to significantly slow convergence. In the next section, we discuss algorithms to circumvent this issue.

## 3 Algorithms

In this section, we first discuss a general strategy to adapt the learning rate in large batch settings. Using this strategy, we discuss two specific algorithms in the later part of the section. Since our primary focus is on deep learning, our discussion is centered around training a h h -layer neural network.

General Strategy. Suppose we use an iterative base algorithm 𝒜 \mathcal{A} (e.g. Sgd or Adam ) in the small batch setting with the following layerwise update rule: x t + 1 = x t + η t ​ u t , \displaystyle x_{t+1}=x_{t}+\eta_{t}u_{t}, where u t u_{t} is the update made by 𝒜 \mathcal{A} at time step t t . We propose the following two changes to the update for large batch settings: 1. The update is normalized to unit l 2 l_{2} -norm. This is ensured by modifying the update to the form u t / ‖ u t ‖ u_{t}/\|u_{t}\| . Throughout this paper, such a normalization is done layerwise i.e., the update for each layer is ensured to be unit l 2 l_{2} -norm.

2. The learning rate is scaled by ϕ ⁡ ( ‖ x t ‖ ) \phi(\|x_{t}\|) for some function ϕ : ℝ + → ℝ + \phi:\mathbb{R}^{+}\rightarrow\mathbb{R}^{+} . Similar to the normalization, such a scaling is done layerwise.

Suppose the base algorithm 𝒜 \mathcal{A} is Sgd , then the modification results in the following update rule: