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: The paper proposes adaptive gradient approaches where the step-size is not determined on the per-coordinate basis but rather for blocks of coordinates. Theoretical results are presented in terms of regret in online convex optimization, regarding convergence in non-convex optimization,  and with respect to uniform stability and generalization. These indicate that under certain conditions adaptivity at the block level could outperform coordinate-wise adaptivity. The approach is evaluated against alternatives on simulated and real-world problems.