Why ADAM Beats SGD for Attention Models	

1 INTRODUCTION . Stochastic gradient descent ( SGD ) is the canonical algorithm for training neural networks ( Robbins & Monro , 1951 ) . SGD iteratively updates model parameters in the negative gradient direction and thus seamlessly scales to large-scale settings . Though a well-tuned SGD outperforms adaptive methods ( Wilson et al. , 2017 ) in many traditional tasks including ImageNet classification ( see Figure 1b ) , certain tasks necessitate the use of adaptive variants of SGD ( e.g. , Adagrad ( Duchi et al. , 2011 ) , Adam ( Kingma & Ba , 2014 ) , AMSGrad ( Reddi et al. , 2019 ) ) , which employ adaptive perparameter learning rates . For instance , consider the task of training an attention model ( Vaswani et al. , 2017 ) using BERT ( Devlin et al. , 2018 ) . Figure 1a shows loss curves of BERT pretraining obtained from SGD with momentum as well as Adam . It can be seen that in spite of extensive hyperparameter tuning , SGD converges much slower than Adam during BERT training . The first significant hint to the performance of Adam on BERT comes from the distribution of the stochastic gradients . For Imagenet , the norms of the mini-batch gradients are typically quite small and well concentrated around their mean . On the other hand , the mini-batch gradient norms for BERT take a wide range of values and are sometimes much larger than their mean value . More formally , while the distribution of the stochastic gradients in Imagenet is well approximated by a Gaussian , the distribution for BERT seems to be heavy-tailed . In this work , we perform a rigorous theoretical and empirical study of the convergence of optimization methods under such heavy-tailed noise . In this setting , some of the stochastic gradients are much larger than the mean and can excessively influence the updates of SGD . This makes SGD unstable and leads to its poor performance . A natural strategy to stabilize the updates is to clip the magnitude of the stochastic gradients . We prove that indeed this is sufficient to ensure convergence even under heavy-tailed noise . Further , we show that adaptive methods in fact implicitly have such a clipping behavior , thereby providing a explanation for their superiority to SGD on BERT . More specifically , we make the following main contributions : • We empirically show that in tasks on which Adam outperforms SGD ( BERT pretraining ) , the noise in stochastic gradients is heavy-tailed . On the other hand , on tasks where traditionally SGD outperforms Adam ( ImageNet training ) , we show that the noise is well concentrated . • We study the convergence of gradient methods under heavy-tailed noise condition where SGD ’ s performance degrades and previous convergence proofs fail . We then establish the convergence of clipped gradient methods under the same condition and prove that they obtain theoretically optimal rates . • We show that Adam implicitly performs coordinate-wise gradient clipping and can hence , unlike SGD , tackle heavy-tailed noise . We prove that using such coordinate-wise clipping thresholds can be significantly faster than using a single global one . This can explain the superior performance of Adam on BERT pretraining . • Inspired by our theoretical analysis , we propose a novel adaptive-threshold coordinate-wise clipping algorithm and experimentally show that it outperforms Adam on BERT training tasks . 1.1 RELATED WORK . Adaptive step sizes . Adaptive step size during optimization has long been studied ( Armijo , 1966 ; Polyak , 1987 ) . More recently , Duchi et al . ( 2011 ) developed the Adagrad algorithm that benefits from the sparsity in stochastic gradients . Inspired by Adagrad , several adaptive methods have been proposed in the deep learning community ( Tieleman & Hinton , 2012 ; Kingma & Ba , 2014 ) . Recently , there has been a surge in interest to study the theoretical properties of these adaptive gradient methods due to ( Reddi et al. , 2019 ) , which pointed out the non-convergence of Adam and proposed an alternative algorithm , AMSGrad . Since then , many works studied different interesting aspects of adaptive methods , see ( Ward et al. , 2018 ; Li & Orabona , 2018 ; Zhou et al. , 2018a ; Staib et al. , 2019 ; Chen et al. , 2018 ; Zou & Shen , 2018 ; Zhou et al. , 2018b ; Agarwal et al. , 2018 ; Zou et al. , 2019 ) . Another direction of related work is normalized gradient descent , which has been studied for quasiconvex and non-convex settings ( Levy , 2016 ; Hazan et al. , 2015 ; Zhang et al. , 2019 ) . In contrast to our work , these prior works assume standard noise distributions that might not be applicable to key modern applications such as attention models , which exhibit heavy-tailed noise . Furthermore , mostly , the convergence rates of adaptive methods are worse than SGD . Noise in neural network . There has been little study of the actual stochastic gradient noise distributions in neural network training . To our knowledge , Simsekli et al . ( 2019 ) is the first work to focus on this topic and observe heavy tailed noise in network training . Our work differs in two important ways : first , we treat the noise as a high dimensional vector , while ( Simsekli et al. , 2019 ) treat deviations in each coordinate as scaler noises to estimate tail index . Second , we focus on convergence of optimization algorithm , the original paper focus on Langevin dynamics and consistence of the index estimator . More experiment comparison is in Appendix H . 2 HEAVY-TAILED NOISE IN STOCHASTIC GRADIENTS . To gain intuition about the difference between SGD and Adam , we start our discussion with the study of noise distributions of stochastic gradient that arise during neural network training . In particular , we focus on noise distributions while training two popular deep learning models — BERT and ResNet . Note that BERT and ResNet are typically trained with Adam and SGD ( with momentum ) respectively and can thus , provide insights about difference between these optimizers . We first investigate the distribution of the gradient noise norm ‖g −∇f ( x ) ‖ in the aforementioned neural network models , where g is the stochastic gradient computed from a minibatch sample . Figure 7 ( a ) and ( d ) show these distributions for ResNet50 on ImageNet and BERT on the Wikipedia and books dataset at model initialization respectively . For comparison , we plot distributions of a normalized sum of squared Gaussians , a well-concentrated distribution , and a Levy-α-stable distribution , a heavy-tailed distribution , in Figure 7 ( b ) and ( e ) respectively . We observe that the noise distribution for BERT appears heavy-tailed , while that of ResNet50 is well-concentrated . To support this observation , in Figure 7 ( c ) and ( f ) we further show the empirical variance of stochastic gradients with respect to the sample size used in the estimation . The results highlight that while the corresponding estimator converges for Imagenet , the empirical variance does not converge in BERT training even as the sample size approaches 107 . Figure 7 clearly shows that while the noise pattern is well-concentrated in some cases , it is also heavy tailed in other settings . We hypothesize that this is one major aspect that determines the performance of SGD and adaptive methods . More specifically , we argue and provide evidence that adaptive methods can be faster than SGD in scenarios where heavy-tailed noise distributions arise . More empirical details can be found in Section 5 . 2.1 OPTIMIZATION PROBLEM SETUP . We now formalize the optimization setup with heavy-tailed noise . Neural network training can be seen as minimizing a differentiable stochastic function f ( x ) = Eξ [ f ( x , ξ ) ] , where f : Rd → R can be potentially nonconvex . At each iteration , we assume access to an unbiased stochastic gradient g ( x ) = ∇f ( x , ξ ) corresponding to the parameters x , and also assume finite α moment . Assumption 1 ( Existence of α−moment ) . There exists positive real numbers α ∈ ( 1 , 2 ] andG > 0 such that for all x , E [ ‖g ( x ) ‖α ] ≤ Gα . Note that it is possible that the variance of g ( x ) is unbounded while simultaneously satisfying the above assumption for α < 2 , e.g . the Pareto distribution . Hence Assumption 1 is much weaker than the standard bounded second moment assumption . The possibility that the variance could be unbounded has a profound impact on the optimization process . In particular , previous convergence proofs for SGD fail and the algorithm may diverge ( see Appendix A ) . As we will illustrate over the next few sections , using a biased clipped stochastic gradient estimator allows us to circumvent the problem of unbounded variance , while ensuring convergence even under heavy-tailed noise . 3 CLIPPED GRADIENT DESCENT AND ACCLIP ALGORITHM . Our key idea is that by clipping the stochastic gradient , we can effectively control its variance . While this does introduce some additional bias , the clipping threshold can be ( possibly adaptively ) tuned to trade-off variance with the introduced bias . Empirically , this is a prevalent practice e.g . Merity et al . ( 2018 ) use clipped SGD to train AWDLSTM whereas Devlin et al . ( 2018 ) use clipped Adam to pretrain BERT . The general form of clipped gradient descent is outlined in Algorithm 1 . Algorithm 1 A general framework for clipped gradient descent 1 : x , mk ← x0 , 0 2 : for k = 1 , · , T do 3 : mk ← β1mk−1 + ( 1− β1 ) gk 4 : ĝk ← clip ( τk , mk ) 5 : xk ← xk−1 − ηkĝk return xK , where random variable K is supported on { 1 , · · · , T } . Perhaps the two most natural candidates for clipping mechanisms are Global CLIPping where a single global threshold is used and Coordinate-wise CLIPping using d coordinate-wise thresholds GClip ( τk , mk ) = min { τk ‖mk‖ , 1 } mk , for τk ∈ R≥0 or , ( GClip ) CClip ( τk , mk ) = min { τk |mk| , 1 } mk , for τk ∈ Rd≥0 . ( CClip ) where all operations for CClip are performed element-wise . Smaller threshold values of τk imply smaller variance , but also higher bias . If the noise distribution is very heavy-tailed ( or has a large variance ) , one could compensate by picking a smaller threshold . See Theorems 1 and 2 for a theoretical treatment of this for GClip . While GClip preserves the update direction and only scales its magnitude , CClip scales each coordinate individually and may not preserve the direction . However , if the noise distribution varies significantly across coordinates , CClip can take advantage of this and only clip those which are more heavy-tailed and consequently converge faster ( see Theorem 3 ) . However , to perform such a coordinate-wise variance-bias trade off optimally would require tuning all d thresholds , which could be extremely large in deep learning . Even if tuning were feasible , the noise distribution may be non-stationary and significantly vary as training progresses ( see Fig . 5 ) . To address this challenge , we propose Adaptive Coordinate-wise CLIPping , which adaptively sets the thresholds for each coordinate ACClip ( τk , mk ) = min { τk |mk|+ , 1 } mk , τ α k = β2τ α k−1 + ( 1− β2 ) |gk|α , ( ACClip ) where operations are element-wise , is a small number for numerical stability , and we set α = 1 by default . The use of a coordinate-wise α-moment ( as opposed to the traditional second moment ) is motivated by our theory and estimates B = E [ |gαk | ] 1/α in Assumption 2 . Since E [ |gαk | ] 1/α is increasing in α ≥ 1 , using α = 1 is a conservative bound on B when the true α is unknown .

Summary: This paper demonstrates empirically that the gradient noises of SGD with ResNet and Adam with Bert are different: one is well-concentrated, while the other one is heavy-tailed. The paper claims that this difference costs the failure of SGD on training Bert. Furthermore, the authors proposes gradient clipped SGD and its adaptive version ACClip. Experiments show that ACClip outperforms Adam on training Bert.