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: This paper analyzes the invariance properties of the K-FAC algorithm by reconstructing the algorithm in a coordinate-free way where the neural network is viewed as a series of affine mappings alternating with nonlinear activation functions. It converts the original metric into an approximate metric, whose coordinate representation matches the K-FAC approximation. So K-FAC can be viewed as the exact natural gradient under the new metric rather than an approximation under the Fisher metric.