Preconditioner on Matrix Lie Group for SGD

We study two types of preconditioners and preconditioned stochastic gradient descent ( SGD ) methods in a unified framework . We call the first one the Newton type due to its close relationship to the Newton method , and the second one the Fisher type as its preconditioner is closely related to the inverse of Fisher information matrix . Both preconditioners can be derived from one framework , and efficiently estimated on any matrix Lie groups designated by the user using natural or relative gradient descent minimizing certain preconditioner estimation criteria . Many existing preconditioners and methods , e.g. , RMSProp , Adam , KFAC , equilibrated SGD , batch normalization , etc. , are special cases of or closely related to either the Newton type or the Fisher type ones . Experimental results on relatively large scale machine learning problems are reported for performance study . 1 INTRODUCTION . This paper investigates the use of preconditioner for accelerating gradient descent , especially in large scale machine learning problems . Stochastic gradient descent ( SGD ) and its variations , e.g. , momentum ( Rumelhart et al. , 1986 ; Nesterov , 1983 ) , RMSProp and Adagrad ( John et al. , 2011 ) , Adam ( Kingma & Ba , 2015 ) , etc. , are popular choices due to their simplicity and wide applicability . These simple methods do not use well normalized step size , could converge slow , and might involve more controlling parameters requiring fine tweaking . Convex optimization is a well studied field ( Boyd & Vandenberghe , 2004 ) . Many off-the-shelf methods there , e.g. , ( nonlinear ) conjugate gradient descent , quasi-Newton methods , Hessian-free optimizations , etc. , can be applied to small and middle scale machine learning problems without much modifications . However , these convex optimization methods may have difficulty in handling gradient noise and scaling up to problems with hundreds of millions of free parameters . For a large family of machine learning problems , natural gradient with the Fisher information metric is equivalent to a preconditioned gradient using inverse of the Fisher information matrix as the preconditioner ( Amari , 1998 ) . Natural gradient and its variations , e.g. , Kronecker-factored approximate curvature ( KFAC ) ( Martens & Grosse , 2015 ) and the one in ( Povey et al. , 2015 ) , all use such preconditioners . Other less popular choices are the equilibrated preconditioner ( Dauphin et al. , 2015 ) and the one proposed in ( Li , 2018 ) . Momentum or the heavy-ball method provides another independent way to accelerate converge ( Nesterov , 1983 ; Rumelhart et al. , 1986 ) . Furthermore , momentum and preconditioner can be combined to further accelerate convergence as shown in Adam ( Kingma & Ba , 2015 ) . This paper groups the above mentioned preconditioners and preconditioned SGD methods into two classes , the Newton type and the Fisher type . The Newton type is closely related to the Newton method , and is suitable for general purpose optimizations . The Fisher type preconditioner relates to the inverse of Fisher information matrix , and is limited to a large subclass of stochastic optimization problems where the Fish information metric can be well defined . Both preconditioners can be derived from one framework , and estimated on any matrix Lie groups designated by the user with almost the same natural or relative gradient descent methods minimizing specific preconditioner estimation criteria . 2 BACKGROUND . 2.1 NOTATIONS . We consider the minimization of cost function f ( θ ) = Ez [ ` ( θ , z ) ] ( 1 ) whereEz takes expectation over random variable z , ` is a loss function , and θ is the model parameter vector to be optimized . For example , in a classification problem , ` could be the cross entropy loss , z is a pair of input feature vector and class label , vector θ consists of all the trainable parameters in the classification model , and Ez takes average over all samples from the training data set . By assuming second order differentiable model and loss , we could approximate ` ( θ , z ) as a quadratic function of θ within a trust region around θ , i.e. , ` ( θ , z ) = bTz θ + 0.5θ TH zθ + az , where az is the sum of higher order approximation errors and constant term independent of θ , H z is a symmetric matrix , and subscript z in bz , H z and az reminds us that these three terms depend on z . Clearly , these three terms depend on θ as well , although we do not explicitly show this dependence to simplify our notations since we just consider parameter updates in the same trust region . Now , we may rewrite ( 1 ) as f ( θ ) = bTθ + 0.5θTHθ + a , where b = Ez [ bz ] , H = Ez [ H z ] , and a = Ez [ az ] . We do not impose any assumption , e.g. , positive definiteness , on H except for being symmetric . Thus , the quadratic surface in the trust region could be non-convex . To simplify our notations , we no longer consider the higher order approximation errors included in a , and simply assume that f ( θ ) is a quadratic function of θ in the trust region . 2.2 PRECONDITIONER . Let us consider a certain iteration . Preconditioned SGD updates θ as θ ← θ − µP ∂f̂ ( θ ) /∂θ ( 2 ) where µ > 0 is the step size , f̂ ( θ ) is an estimate of f ( θ ) obtained by replacing expectation with sample average , and positive definite matrix P could be a fixed or adaptive preconditioner . By letting θ ′ = P−0.5θ , we can rewrite ( 2 ) as θ ′ ← θ ′ − µ∂f̂ ( θ ) /∂θ ′ ( 3 ) where P−0.5 denotes the principal square root of P . Hence , ( 3 ) suggests that preconditioned SGD is equivalent to SGD in a transformed parameter domain . Within the considered trust region , let us write the stochastic gradient , ∂f̂ ( θ ) /∂θ , explicitly as ∂f̂ ( θ ) /∂θ = Ĥθ + b̂ ( 4 ) where Ĥ and b̂ are estimates of H and b , respectively . Combining ( 4 ) and ( 2 ) gives the following linear system θ ← ( I − µPĤ ) θ − µP b̂ ( 5 ) for updating θ within the assumed trust region , where I is the identity matrix . A properly determined P could significantly accelerate convergence of the locally linear system in ( 5 ) . We review a few facts shown in ( Li , 2018 ) before introducing our main contributions . Let δθ be a random perturbation of θ , and be small enough such that θ+ δθ still resides in the same trust region . Then , ( 4 ) suggests the following resultant perturbation of stochastic gradient , δĝ def = ∂f̂ ( θ + δθ ) /∂θ − ∂f̂ ( θ ) /∂θ = Ĥδθ = Hδθ + ε ( 6 ) where ε accounts for the error due to replacing Ĥ with H . Note that by definition , δĝ is a random vector dependent on both z and δθ . The preconditioner in ( Li , 2018 ) is pursued by minimizing criterion c ( P ) = Ez , δθ [ δĝ TP δĝ + δθTP−1δθ ] ( 7 ) where subscript δθ in Ez , δθ denotes taking expectation over δθ . Under mild conditions , criterion ( 7 ) determines a unique positive definite P , which is optimal in the sense that it preconditions the stochastic gradient such that PEz , δθ [ δĝδĝ T ] P = Eδθ [ δθδθ T ] ( 8 ) which is comparable to relationship H−1δgδgTH−1 = δθδθT , where δg = Hδθ is the perturbation of noiseless gradient , and we assume that H is invertible , but not necessarily positive definite . Clearly , this preconditioner is comparable toH−1 . It perfectly preconditions the gradient such that the amplitudes of parameter perturbations match that of the associated preconditioned gradient , regardless of the amount of gradient noise . Naturally , preconditioned SGD with this preconditioner inherits the scale-invariance property of the Newton method . Note that in the presence of gradient noise , the optimal P and P−1 given by ( 8 ) are not unbiased estimates of H−1 and H , respectively . Actually , even if H is positive definite and available , H−1 may not always be a good preconditioner when H is ill-conditioned since it could significantly amplify the gradient noise along the directions of the eigenvectors of H associated with small eigenvalues , and lead to divergence . More specifically , it is shown in ( Li , 2018 ) that H−1Ez , δθ [ δĝδĝ T ] H−1 ≥ Eδθ [ δθδθT ] , whereA ≥ B means thatA −B is nonnegative definite . 3 TWO PRECONDITIONER ESTIMATION CRITERIA . 3.1 THE NEWTON TYPE CRITERION . Preconditioner estimation criterion ( 7 ) requires δθ to be small enough such that θ and θ + δθ reside in the same trust region . In practice , numerical error might be an issue when handling small numbers with floating point arithmetic . This concern becomes more grave with the popularity of single and even half precision math in large scale neural network training . Luckily , ( 6 ) relates δĝ to the Hessianvector product , which can be efficiently evaluated with automatic differentiation software tools . Let v be a random vector with the same dimension as θ . Then , ( 4 ) suggests the following method for Hessian-vector product evaluation , ∂ ∂θ { [ ∂f̂ ( θ ) /∂θ ] T v } = ∂2f̂ ( θ ) ∂θ∂θT v = Ĥv ( 9 ) Now , replacing ( δθ , δĝ ) in ( 7 ) with ( v , Ĥv ) leads to our following new preconditioner estimation criterion , cn ( P ) = Ez , v [ v TĤPĤv + vTP−1v ] ( 10 ) where the subscript v in Ez , v suggests taking expectation over v . We no longer have the need to assume v to be an arbitrarily small vector . It is important to note that this criterion only requires the Hessian-vector product . The Hessian itself is not of interest . We call ( 10 ) the Newton type preconditioner estimation criterion as the resultant preconditioned SGD method is closely related to the Newton method . 3.2 THE FISHER TYPE CRITERION . We consider the machine learning problems where the empirical Fisher information matrix can be well defined by F = Ez [ ∂ ` ( θ , z ) ∂θ ( ∂ ` ( θ , z ) ∂θ ) T ] . Replacing ( δθ , δĝ ) in ( 7 ) with ( v , ĝ + λv ) leads to criterion cf ( P ) = Ez , v [ ( ĝ + λv ) TP ( ĝ + λv ) + vTP−1v ] ( 11 ) where ĝ = ∂f̂ ( θ ) /∂θ is a shorthand for stochastic gradient , and λ ≥ 0 is a damping factor . Clearly , v is independent of ĝ . Let us further assume that v is drawn from standard multivariate normal distribution N ( 0 , I ) , i.e. , Ev [ v ] = 0 and Ev [ vvT ] = I . Then , we could simplify cf ( P ) as cf ( P ) = tr { PEz , v [ ( ĝ + λv ) ( ĝ + λv ) T ] +P−1Ev [ vvT ] } = tr { P [ Ez [ ĝĝT ] + λ2I ] +P−1 } ( 12 ) By letting the derivative of cf ( P ) with respect to P be zero , the optimal positive definite solution for cf ( P ) is readily shown to be P = { Ez [ ĝĝT ] + λ2I } −0.5 ( 13 ) When ĝ is a gradient estimation obtained by taking average over B independent samples , Ez [ ĝĝ T ] is related to the Fisher information matrix by Ez [ ĝĝ T ] = F /B + ( B − 1 ) ggT /B ( 14 ) We call this preconditioner the Fisher type one due to its close relationship to the Fisher information matrix . One can easily modify this preconditioner to obtain an unbiased estimation of F−1 . Let s be an exponential moving average of ĝ . Then , after replacing the ĝ in ( 13 ) with ĝ − s + s/ √ B and setting λ = 0 , P 2/B will be an unbiased estimation of F−1 . Generally , it might be acceptable to keep the bias term , ( B−1 ) ggT /B , in ( 14 ) for two reasons : it is nonnegative definite and regularizes the inversion in ( 13 ) ; it vanishes when the parameters approach a stationary point . Actually , the Fisher information matrix could be singular for many commonly used models , e.g. , finite mixture models , neural networks , hidden Markov models . We might not be able to inverse F for these singular statistical models without using regularization or damping . A Fisher type preconditioner with λ > 0 loses the scale-invariance property of a Newton type preconditioner . BothP andP 2 can be useful preconditioners when the step size µ and damping factor λ are set properly .

Summary: This paper proposes a preconditioned SGD method where the preconditioner is adapted by performing some type of gradient descent on some secondary objective "c".  The preconditioner lives in one of a restricted class of invertible matrices (e.g. symmetric, diagonal, Kronecker-factored) constituting a Lie group (which is where the title comes from).