Blockwise Adaptivity:  Faster Training and Better Generalization in Deep Learning

1 INTRODUCTION . Deep networks have achieved excellent performance in a variety of domains such as computer vision ( He et al. , 2016 ) , language modeling ( Zaremba et al. , 2014 ) , and speech recognition ( Graves et al. , 2013 ) . The most popular optimizer is stochastic gradient decent ( SGD ) ( Robbins & Monro , 1951 ) , which is simple and has low per-iteration complexity . Its convergence rate is also well-established ( Ghadimi & Lan , 2013 ; Bottou et al. , 2018 ) . However , vanilla SGD is sensitive to the choice of stepsize , and requires careful tuning . To improve the efficiency and robustness of SGD , many variants based on coordinate-wise adaptive stepsize ( Almeida et al. , 1999 ; Duchi et al. , 2011 ; Schaul et al. , 2013 ; Kingma & Ba , 2015 ; Tieleman & Hinton , 2012 ; Zeiler , 2012 ; Zheng & Kwok , 2017 ; Reddi et al. , 2018 ; Zaheer et al. , 2018 ; Zou & Shen , 2018 ; Chen et al. , 2019 ; Zou et al. , 2019 ) have been proposed . Though this is effective in accelerating convergence , its generalization performance is often worse than SGD ( Wilson et al. , 2017 ) . Recently , it is shown that coordinate-wise adaptive gradient descent is closely related to signbased gradient descent ( Balles & Hennig , 2018 ; Bernstein et al. , 2018 ) . Theoretical arguments and empirical evidence suggest that the gradient sign impedes generalization ( Balles & Hennig , 2018 ) . To contract the generalization gap , a partial adaptive parameter for the second-order momentum is proposed ( Chen & Gu , 2018 ) . By using a smaller partial adaptive parameter , the adaptive gradient algorithm behaves less like sign descent and more like SGD . Based on the similar motivation , Luo et al . ( 2019 ) proposed to gradually transform Adam Kingma & Ba ( 2015 ) to SGD . Moreover , to avoid numerical problems in practical implementation , a small ( = 10−8 ) parameter is typically used in these methods . This parameter controls adaptivity of the algorithm ( Zaheer et al. , 2018 ) . A larger value ( say , = 10−3 ) can reduce adaptivity and empirically helps Adam to match its generalization performance with SGD . This implies that coordinate-wise adaptivity may be too aggressive for good generalization performance . To improve generalization performance , attempts have been made to use a layer-wise stepsize ( Singh et al. , 2015 ; You et al. , 2017 ; Yu et al. , 2017 ; Zhou et al. , 2018 ) , which assign different stepsizes to different layers or normalize the layer-wise gradient . However , there has been no theoretical analysis for its empirical success . In this paper , we consider the more general case in which model parameters are partitioned into blocks , instead of simply into layers . By revisiting the derivation of Adagrad , we propose the use of a blockwise stepsize which depends on the corresponding gradient block . Such blockwise adaptivity is less aggressive than coordinate-wise adaptivity , as it adapts to parameter blocks instead of to individual parameters . This can thus have a better balance between adaptivity and generalization . Moreover , unlike coordinate-wise adaptivity , blockwise adaptivity is not sign-based gradient descent , and so does not suffer from its performance deterioration . As in ( Bernstein et al. , 2018 ; Ghadimi & Lan , 2013 ; Ward et al. , 2018 ; Zaheer et al. , 2018 ; Zou & Shen , 2018 ; Zou et al. , 2019 ) , we will focus on the expected risk minimization problem : min θ F ( θ ) = Ez [ f ( θ ; z ) ] , ( 1 ) where f is some possibly nonconvex loss function , z is a random sample , θ is the model parameter , and the expectation is taken w.r.t . the underlying sample distribution . The expected risk measures the generalization performance on unseen data ( Bottou et al. , 2018 ) , and reduces to the empirical risk when a finite training set is considered . We show theoretically that the proposed blockwise adaptive gradient descent can be faster than its coordinate-wise counterpart . Using tools on uniform stability ( Bousquet & Elisseeff , 2002 ; Hardt et al. , 2016 ) , we also show that blockwise adaptivity has potentially lower generalization error than coordinate-wise adaptivity . Empirically , blockwise adaptive gradient descent converges faster and obtains better generalization performance than coordinatewise descent ( Adam ) and Nesterov ’ s accelerated gradient ( NAG ) ( Sutskever et al. , 2013 ) . Notations . For an integer n , [ n ] = { 1 , 2 , . . . , n } . For a vector x , xT denotes its transpose , Diag ( x ) is a diagonal matrix with x on its diagonal , √ x is the element-wise square root of x , x2 is the coordinate-wise square of x , ‖x‖2 = √ xTx , ‖x‖∞ = maxi |xi| , ‖x‖2Q = xTQx , where Q is a positive semidefinite ( psd ) matrix , and x ≥ 0 means xi ≥ 0 for all i . For two vectors x and y , x/y , and 〈x , y〉 denote the element-wise division and dot product , respectively . For a square matrix X , X−1 is its inverse , and X 0 means that X is psd . Moreover , 1d = [ 1 , 1 , . . . , 1 ] T ∈ Rd . 2 RELATED WORK . 2.1 ADAGRAD . Adagrad ( shown in Algorithm 1 ) is an adaptive gradient method in online convex learning with coordinate-wise stepsize ( Duchi et al. , 2011 ; McMahan & Streeter , 2010 ) . It is particularly useful for sparse learning , as parameters for the rare features can take large steps . Recently , Ward et al . ( 2018 ) established its convergence properties with a global adaptive stepsize in nonconvex optimization . It is shown that Adagrad converges to a stationary point at the optimal O ( 1/ √ T ) rate ( up to a factor log ( T ) ) , where T is the total number of iterations . Recall that the SGD iterate is the solution to the problem : θt+1 = arg minθ〈gt , θ〉 + 12η‖θ − θt‖ 2 2 , where gt is the gradient of the loss ft at iteration t , and θt ∈ Rd is the parameter vector . To incorporate information about the curvature of sequence { ft } , the ` 2-norm in the SGD update can be replaced by the Mahalanobis norm , leading to ( Duchi et al. , 2011 ) : θt+1 = arg min θ 〈gt , θ〉+ 1 2η ‖θ − θt‖2Diag ( st ) −1 , ( 2 ) where st ≥ 0 . This is an instance of mirror descent ( Nemirovski & Yudin , 1983 ) . Its regret bound has a gradient-related term ∑T t=1 ‖gt‖2Diag ( st ) −1 . Adagrad ’ s stepsize can be obtained by examining a similar objective ( Duchi et al. , 2011 ) : min s∈S T∑ t=1 ‖gt‖2Diag ( s ) −1 , ( 3 ) where S = { s : s ≥ 0 , 〈s , 1〉 ≤ c } , and c is some constant . At optimality , s∗ , i = c‖g1 : T , i‖2/ ∑d j=1 ‖g1 : T , j‖2 , where g1 : T , i = [ g1 , i , . . . , gT , i ] T . As st can not depend on gj ’ s with j > t , this suggests st , i ∝ ‖g1 : t , i‖2 . Theoretically , this choice of st leads to a regret bound that is competitive with the best post-hoc optimal bound ( McMahan & Streeter , 2010 ) . To solve the expected risk minimization problem in ( 1 ) , an Adagrad variant called weighted AdaEMA is recently proposed in ( Zou et al. , 2019 ) . It employs weighted averaging of g2t , i ’ s for stepsize . Moreover , weighted AdaEMA is a general coordinate-wise adaptive method and includes many Adagrad variants , including Adam and RMSprop , as special cases . 2.2 UNIFORM STABILITY . Given n samples S = { zi } ni=1 drawn i.i.d . from an underlying data distribution D , one often learns the model by minimizing the empirical risk : minθ ΦS ( θ ) ≡ 1n ∑n i=1 f ( θ ; zi ) . Let M ( S ) be the output of a possibly randomized algorithm M ( e.g. , SGD ) running on data S. Definition 1 . ( Hardt et al. , 2016 ) Let S and S′ be two data sets of size n that differ in only one sample . Algorithm M is u-uniformly stable if stab ≡ supS , S′ supz∈D EM [ f ( M ( S ) ; z ) − f ( M ( S′ ) ; z ) ] ≤ u . The generalization error is defined as gen ≡ ES , M [ ΦS ( M ( S ) ) −F ( M ( S ) ) ] , where the expectation is taken w.r.t . the set S and randomness of M ( Hardt et al. , 2016 ) . It is shown that the generalization error is bounded by the uniform stability of M , i.e. , | gen| ≤ stab ( Hardt et al. , 2016 ) . In other words , the more uniformly stable an algorithm is , the lower is its generalization error . 3 BLOCKWISE ADAPTIVE DESCENT . Consider the under-determined least squares problem : min θ ‖Xθ − y‖22 , ( 4 ) whereX ∈ Rn×d is the input matrix ( with sample size n , and dimensionality d > n ) , and output y ∈ Rn . We assume that XXT is invertible . As pointed out in ( Zhang et al. , 2017 ; Wilson et al. , 2017 ) , any stochastic gradient descent method on problem ( 4 ) with a global stepsize outputs a trajectory with iterates lying in the span of the rows of X . One solution of ( 4 ) is XT ( XXT ) −1y , which happens to be the solution with minimum ` 2-norm among all possible global minimizers . This minimum-norm solution has the largest margin , and maximizing margin typically leads to lower generalization error ( Boser et al. , 1992 ) . It is known that SGD converges to the minimum ` 2-norm solution of problem ( 4 ) ( Zhang et al. , 2017 ) . On the other hand , coordinate-wise adaptive methods ( such as Adagrad , RMSprop , and Adam ) fail to find the minimum ` 2-norm solution , but converge to solutions with low ` ∞-norm instead ( Wilson et al. , 2017 ) . Examples in ( Wilson et al. , 2017 ) show that solutions obtained by these adaptive methods can generalize arbitrarily poorly , while the SGD solution makes no error . In Section 3.1 , we first show that blockwise adaptivity , unlike coordinate-wise adaptivity , can find the minimum ` 2-norm solution of a nonlinear least squares problem in layer-wise training of a neural network . This motivates us to further exploit blockwise adaptivity in end-to-end neural network training ( Section 3.2 ) . To provide more analysis , Section 3.3 studies the proposed algorithm in the online convex learning setting as in ( Duchi et al. , 2011 ; Kingma & Ba , 2015 ; Reddi et al. , 2018 ) .

Summary: In this paper, the authors propose a generalization of AdaGrad, called BAG, that operates on blocks of parameters instead of each individual parameter. The authors also propose a momentum version of BAG called BAGM. Convergence rate results are proved for the algorithm, and some uniform stability results show situations where BAG would generalize better than previous adaptive gradient methods.