A Coordinate-Free Construction of Scalable Natural Gradient

1 INTRODUCTION . Most neural networks are trained using stochastic gradient descent ( SGD ) ( Bottou & Bousquet , 2007 ) , or variants thereof which adapt step sizes for individual dimensions ( Duchi et al. , 2011 ; Kingma & Ba , 2015 ) . One well-known deficiency of SGD is that the updates are sensitive to the parameterization of the network . There are numerous tricks for reparameterizing network architectures so that they represent the same sets of functions , but in a friendlier coordinate system . Examples include replacing logistic activation functions with tanh ( Glorot & Bengio , 2010 ) , whitening the inputs or activations ( Desjardins et al. , 2015 ; Krizhevsky , 2009 ) , or centering the activations to have zero mean and/or unit variance ( Montavon & Müller , 2012 ; Cho et al. , 2011 ; Raiko et al. , 2012 ; Ioffe & Szegedy , 2015 ) . Such tricks can lead to large improvements in the speed of optimization . Ideally , one would like to use an optimization algorithm which is invariant to such transformations of a neural network , in order to avoid the pathologies which the transformations are meant to remedy . Natural gradient descent ( Amari , 1998 ) is a second-order optimization algorithm motivated by a key invariance property : to the first-order , its updates are invariant to smooth reparameterizations of a model . The natural gradient of a cost function can be seen as the gradient of the function on a Riemannian manifold ( typically using the Fisher information metric ( Amari & Nagaoka , 2000 ) ) , and the invariance properties of the algorithm follow directly from its definition in terms of differential geometric primitives . There have been many attempts to apply natural gradient descent , or approximations thereof , to training neural networks ( Amari et al. , 2000 ; Roux et al. , 2008 ; Martens , 2010 ; Grosse & Salakhutdinov , 2015 ; Martens & Grosse , 2015 ; Desjardins et al. , 2015 ; Schulman et al. , 2015 ) . The challenge is that the exact natural gradient is impractical to compute for large neural nets , because it requires solving a linear system whose dimension is the number of parameters ( which may be in the tens of millions for modern networks ) . Unfortunately , tractable approximations to the natural gradient are typically defined in terms of particular coordinate representations , and therefore may lose the invariance properties which motivated natural gradient in the first place . Kronecker-Factored Approximate Curvature ( K-FAC ) ( Martens & Grosse , 2015 ) is an approximate natural gradient optimizer where the Fisher information matrix F is approximated as a block diagonal matrix with one block per layer of the network , and each block factorizes as the Kronecker product of small matrices . Because of the Kronecker structure , the approximate natural gradient can be computed with low overhead relative to ordinary SGD ; K-FAC demonstrated significant speedups in training deep autoencoders ( Martens & Grosse , 2015 ) , classification convolutional networks ( Grosse & Martens , 2016 ; Ba et al. , 2017 ) , recurrent networks ( Martens et al. , 2018 ) and deep reinforcement learning ( Wu et al. , 2017 ) . The same Fisher matrix approximation has also led to significant improvements in modeling posterior uncertainty in Bayesian neural networks ( Ritter et al. , 2018b ; Zhang et al. , 2017 ) and avoiding catastrophic forgetting ( Ritter et al. , 2018a ) . Most recently , K-FAC was used to train a neural network to solve the many-electron Schrödinger equation ( Pfau et al. , 2019 ) . Although K-FAC does not satisfy the general invariance properties of natural gradient , it is still invariant to a broad and interesting class of reparameterizations : affine transformations of the activations in each layer ( Martens & Grosse , 2015 ) . This was verified through linear algebraic manipulation of the update rules , but unfortunately the proofs yielded little insight into the algorithm or advice about how it can be extended . In this paper , we take a different approach : we formulate K-FAC directly in terms of coordinate-free mathematical objects , so that the invariance properties follow immediately from the construction . Specifically , we view a neural network as a series of affine maps alternating with fixed nonlinear activation functions . The activations and pre-activations for each layer are viewed as abstract affine spaces , and the weights and biases of the network correspond to affine maps . The ordinary Fisher metric is a metric on this spaceW of affine maps . Our contribution is a recipe to convert a metric on W ( whose coordinate representation is extremely large ) into an approximate metric onW ( the “ KFAC metric ” ) , whose coordinate representation matches the K-FAC approximation . Hence , rather than view K-FAC as an approximation to the natural gradient under the Fisher metric , we view it as the exact natural gradient under the K-FAC metric . This entire construction is coordinate-free , so the invariance properties of K-FAC follow immediately . We can contrast K-FAC ’ s invariance properties with those of exact natural gradient descent . Since the exact natural gradient is derived in terms of a metric on a smooth manifold , the update is invariant to arbitrary smooth reparameterizations , but only up to the first-order . By contrast , we show global invariance to a more restricted class of reparameterizations . Our analysis imposes additional structure on the weight manifold W : the parameters are assumed to define affine maps between affine spaces . Choosing affine bases for the activations and pre-activations yields a natural affine basis for W . Therefore , the set of allowable reparameterizations for neural networks consists of affine change-of-basis transformations for the activations and pre-activations . This leaves out some unusual reparameterizations which exact natural gradient descent is invariant to , such as permuting the entries of the weight matrix . But it captures important classes of reparameterizations , such as whitening , normalization , and replacing the logistic activation function with tanh . And in exchange for imposing the affine structure , we obtain global invariance , not just first-order invariance . Our framework easily enables some generalizations of the basic result . Our construction is not limited to Fisher metrics , but rather it applies to general pullback metrics , where one places a metric on the network ’ s output space and pulls it back toW . Furthermore , we extend the invariance results to convolutional networks and recurrent networks through a straighforward application of our KFAC metric construction , highlighting the versatility of adopting the coordinate-free approach . Notations . Since our coordinate-free construction of K-FAC requires mathematical machinery from both abstract affine/linear algebra and differential geometry , we provide an overview of necessary concepts in Appendix B . As we go back-and-forth between coordinate-independent and coordinatedependent objects frequently in this paper , we provide a summary of all notations in Appendix A . The J·K notation throughout stands for choosing coordinates for an abstract mathematical object . 2 BACKGROUND . 2.1 NATURAL GRADIENT DESCENT . For simplicity , we define the natural gradient descent algorithm here in the context of multilayer perceptrons ( MLPs ) , i.e. , fully connected feed-forward networks . For an input-target pair ( x , y ) , let f ( x , w ) denote the output and w symbolizes the parameter vector of the MLP . Given a training set S of pairs ( xi , yi ) , we aim to minimize the empirical risk which is given by : h ( w ) = 1 |S| ∑ ( xi , yi ) ∈S L ( yi , f ( xi , w ) ) , ( 2.1 ) where L ( yi , f ( xi , w ) ) is the loss function measuring the disagreement between yi and f ( xi , w ) . Suppose that f ( x , w ) determines parameters z of the model ’ s predictive distribution Ry|z over y and furthermore , we reparameterize this as Py|x ( w ) = Ry|f ( x , w ) . Likewise , the density function of this distribution can be reparameterized as p ( y|x , w ) = r ( y|f ( x , w ) ) . In addition , we take the loss function here to be the negative log-likelihood L ( y , z ) = − log r ( y|z ) and denote log-likelihood gradients ∇wL ( y , f ( x , w ) ) by Dw ( D notation throughout remainder of the paper refers to loglikelihood gradients ) . The Fisher information matrix F ( w ) is defined as F ( w ) = Ex , y [ ( Dw ) ( Dw ) > ] , ( 2.2 ) where the expectation is taken over Py|x ( w ) for y and over the data distribution for x . Since F ( w ) is defined as the expectation of an outer product , F ( w ) is always guaranteed to be a positivesemidefinite ( PSD ) matrix . The natural gradient of the objective function h ( w ) in Eqn . 2.1 is F ( w ) −1∇h ( w ) . For a chosen learning rate > 0 , the natural gradient descent algorithm ( Amari , 1998 ) minimizes h ( w ) by using the natural gradient to update parameters of the network : wk+1 ← wk − F ( wk ) −1∇h ( wk ) . ( 2.3 ) Invariance Properties of Natural Gradient . Natural gradient descent possesses a key invariance property which does not hold for ordinary stochastic gradient descent ( SGD ) : given two equivalent networks which are parameterized differently , after applying the natural gradient descent update to each , the resulting networks will be equivalent up to the first-order . The reason for this is that the natural gradient admits an intrinsic coordinate-free construction in terms of differential geometric primitives . We give a detailed explanation of this by considering an abstract mathematical setting ; we defer to Appendix B for background on the mathematical terminology with which we use . LetM be a Riemannian manifold with nondegenerate metric given by g ( in this paper , our definition of “ metrics ” allows the possibility of it being degenerate ) . For a smooth function h : M→ R and a point p ∈M , the differential dh ( p ) is an abstract covector onM . We use the nondegenerate metric g to convert the covector dh ( p ) into a tangent vector . By definition , g ( p ) is a nondegenerate bilinear form which yields a linear isomorphism between the tangent space and the cotangent space : g ( p ) : TpM ∼=→ T ∗pM . This is commonly referred to as the musical isomorphism in mathematical literature . The inverse g ( p ) −1 gives a linear map the other way around , g ( p ) −1 : T ∗pM ∼=→ TpM . Applying this isomorphism to dh ( p ) yields the tangent vector g ( p ) −1dh ( p ) ∈ TpM . We call this tangent vector the natural gradient of h. We apply this mathematical framework to the objects of our interest . First , let W be a smooth manifold which characterizes the weight space of network parameters intrinsically . Let ω and ( ξ , υ ) be intrinsic versions of the parameter vector w and the input-target pair ( x , y ) respectively . ξ and υ belong to abstract input and output spaces which are also smooth manifolds . Now , W can be endowed with the Fisher metric gF which is defined as gF ( ω ) = Eξ , υ [ dLω ⊗ dLω ] , ( 2.4 ) where Lω = − log p ( υ|ξ , ω ) is the abstract log-likelihood loss function . The expectation above is taken over the abstract predictive distribution Pυ|ξ ( ω ) for υ and over the data distribution for ξ . Expressing this in a coordinate system , we obtain exactly the Fisher matrix as given in Eqn . 2.2 . AsW is a Riemannian manifold with the Fisher metric gF , the idealized gradient descent updates are given by Bonnabel et al . ( 2013 ) ωk+1 ← Expωk ( − gF ( ωk ) −1dh ( ωk ) ) , ( 2.5 ) where Expωk : TωkW → W is the exponential map . This update rule is exactly invariant to all smooth reparameterizations ofW since it is entirely coordinate-free . However , such an algorithm is infeasible in practice as computing the exponential map is typically an intractable problem . Instead , it is much easier to work with the following abstract natural gradient update rule which uses a firstorder approximation of the exponential map ωk+1 ← − gF ( ωk ) −1dh ( ωk ) . Since this is a first-order approximation of the update rule in Eqn . 2.5 , invariance to smooth reparameterizations holds only up to first-order . Additional approximations to the exponential map are necessary to obtain higher-order invariances ; we defer to Song & Ermon ( 2018 ) for a more detailed account of how this can be done .

Summary: Natural gradient (NG) has been proven efficient in statistical learning, and one of its attractive properties is being invariant under smooth transformations of the parameter space. Computing NG is often difficult as one has to derive the Fisher matrix and its inverse. K-FAC offers an approximate method for approximating the NG with the risk of losing the invariant property.