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

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 ) , 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 1https : //github.com/tensorflow/addons/blob/master/tensorflow_addons/ optimizers/lamb.py 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 ( Reddi et al. , 2018 ; Zaheer et al. , 2018 ; Zhang et al. , 2019 ) . 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 xt ∈ Rd , either xt , j or [ xt ] j are used to denote its jth coordinate where j ∈ [ d ] . Let I be the d×d identity matrix , and let I = [ I1 , I2 , ... , Ih ] be its decomposition into column submatrices Ii = d× dh . For x ∈ Rd , let x ( i ) be the block of variables corresponding to the columns of Ii i.e. , x ( i ) = I > i x ∈ Rdi for i = { 1 , 2 , · · · , h } . For any function f : Rd → R , we use ∇if ( x ) to denote the gradient with respect to x ( i ) . For any vectors u , v ∈ Rd , we use u2 and u/v to denote elementwise square and division operators respectively . We use ‖.‖ and ‖.‖1 to denote l2-norm and l1-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∈Rd f ( x ) : = Es∼P [ ` ( x , s ) ] + λ 2 ‖x‖2 , ( 1 ) where ` is a smooth ( possibly nonconvex ) function and P is a probability distribution on the domain S ⊂ Rk . Here , x corresponds to model parameters , ` is the loss function and P is an unknown data distribution . We assume function ` ( x ) is Li-smooth with respect to ith block , i.e. , there exists a constant Li such that ‖∇i ` ( x , s ) −∇i ` ( x+ Iiδ , s ) ‖ ≤ Li‖δ‖ , ∀ x ∈ Rd , δ ∈ Rdi and s ∈ S , ( 2 ) for all i ∈ [ h ] . We use L = ( L1 , · · · , Lh ) > to denote the h-dimensional vector of Lipschitz constants . We use L∞ and Lavg to denote maxi Li and ∑ i Li h respectively . We assume the following bound on the variance in stochastic gradients : E‖∇i ` ( x , s ) −∇if ( x ) ‖2 ≤ σ2i for all x ∈ Rd and i ∈ [ h ] . Furthermore , we also assume E‖ [ ∇ ` ( x , s ) ] i − [ ∇f ( x ) ] i‖2 ≤ σ̃2i for all x ∈ Rd and i ∈ [ d ] . We use σ = ( σ1 , · · · , σh ) > and σ̃ = ( σ̃1 , · · · , σ̃d ) > 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 for all j ∈ [ d ] , x ∈ Rd and s ∈ S. Note that such assumptions are typical in the analysis of stochastic first-order methods ( cf . ( Ghadimi & Lan , 2013a ; Ghadimi et al. , 2014 ; Reddi et al. , 2016 ; 2018 ) ) . Stochastic gradient descent ( SGD ) is one of the simplest first-order algorithms for solving problem in Equation 1 . The update at the tth iteration of SGD is of the following form : xt+1 = xt − ηt 1 |St| ∑ st∈St ∇ ` ( xt , st ) + λxt , ( SGD ) where St is set of b random samples drawn from the distribution 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 and using appropriate learning rate , we have the following for the iterates of SGD : E [ ‖∇f ( xa ) ‖2 ] ≤ O ( ( f ( x1 ) − f ( x∗ ) ) L∞ T + ‖σ‖2 T ) . where x∗ is an optimal solution to the problem in equation 1 and xa is an iterate uniformly randomly chosen from { x1 , · · · , xT } . However , tuning the learning rate ηt in SGD , especially in large batch settings , is difficult in practice . Furthermore , the dependence on L∞ ( the maximum of smoothness across dimension ) can lead to significantly slow convergence . In the next section , we discuss algorithms to circumvent this issue .

Summary: This paper proposes a learning rate adaptation mechanism, called LAMB, for large-batch distributed training. The goal is to stabilize the training as the batch size increases. The idea is simple and straightforward -- there should be a layerwise learning rate adjusted by normalizing the layer weights and gradients at each layer so that layers with larger weights take larger learning steps, and vice versa. The authors perform empirical studies on BERT-large and ResNet to conclude that LAMB can scale up training batch size while still being able to converge in time with comparable accuracy.