Title: A COORDINATE-FREE CONSTRUCTION OF SCALABLE NATURAL GRADIENT

Abstract: Most neural networks are trained using first-order optimization methods, which are sensitive to the parameterization of the model. Natural gradient descent is invariant to smooth reparameterizations because it is defined in a coordinate-free way, but tractable approximations are typically defined in terms of coordinate systems, and hence may lose the invariance properties. We analyze the invariance properties of the Kronecker-Factored Approximate Curvature (K-FAC) algorithm by constructing the algorithm in a coordinate-free way. We explicitly construct a Riemannian metric under which the natural gradient matches the K-FAC update; invariance to affine transformations of the activations follows immediately. We extend our framework to analyze the invariance properties of K-FAC appied to convolutional networks and recurrent neural networks, as well as metrics other than the usual Fisher metric.

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 space W of affine maps. Our contribution is a recipe to convert a metric on W (whose coordinate representation is extremely large) into an approximate metric on W (the "K-FAC 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 to W. Furthermore, we extend the invariance results to convolutional networks and recurrent networks through a straighforward application of our K-FAC 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 • 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 (x i , y i ), we aim to minimize the empirical risk which is given by:

h(w) = 1 |S| (xi,yi)∈S L(y i , f (x i , w)), (2.1)

where L(y i , f (x i , w)) is the loss function measuring the disagreement between y i and f (x i , w).

Suppose that f (x, w) determines parameters z of the model's predictive distribution R y|z over y and furthermore, we reparameterize this as P y|x (w) = R y|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 ∇ w L(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) = E x,y [(Dw)(Dw) ], (2.2)

where the expectation is taken over P y|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:

w k+1 ← w k -F(w k ) -1 ∇h(w k ). (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.

Let M 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 on M. 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) : T p M ∼ = → T * p M

. 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 * p M ∼ = → T p M.

Applying this isomorphism to dh(p) yields the tangent vector g(p) -1 dh(p) ∈ T p M. 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 g F which is defined as

g F (ω) = 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.

As W is a Riemannian manifold with the Fisher metric g F , the idealized gradient descent updates are given by Bonnabel et al. ( 2013)

EQUATION

where Exp ω k : T ω k W → W is the exponential map. This update rule is exactly invariant to all smooth reparameterizations of W 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 ← -g F (ω k ) -1 dh(ω 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.

2.2. KRONECKER-FACTORED APPROXIMATE CURVATURE:

We consider a MLP with L layers. At each layer i ∈ {1, . . . , L}, the MLP computation is given as:

z i = W i a i-1 + b i a i = φ i (z i ),

where a i-1 is an activation vector, z i is a pre-activation vector, W i is a weight matrix, b i is a bias vector, and φ i : R → R is an activation function. For convenience, we introduce homogeneous coordinates

ā i-1 = [a i-1 1] and Wi = [W i b i ].

The above computation can be rewritten as

z i = Wi āi-1 a i = φ i (z i ). (2.6)

We concatenate all of the network parameters Wi into a single vector w,

w = [vec( W1 ) vec( W2 ) . . . vec( WL ) ] .

Here, vec denotes the vectorization operator which stacks the columns of a matrix together to form a vector. The Fisher matrix for the MLP is a L × L block matrix F(w) where each (i, j)-th block is

F(w) i,j = E[vec(D Wi ) vec(D Wj ) ].

We now give an overview of the Kronecker-Factored Approximate Curvature (K-FAC) method (Martens & Grosse, 2015) of approximating the Fisher matrix. Consider the diagonal (i, i) blocks of F(w). Using backpropagation, the log-likelihood gradient D Wi = Dz i ā i-1 and hence, we have vec(Dz i ā i-1 ) = āi-1 ⊗ Dz i . Then, F(w) i,i can be rewritten as:

F(w) i,i = E[vec(D Wi ) vec(D Wi ) ] = E[(ā i-1 ⊗Dz i )(ā i-1 ⊗Dz i ) ] = E[ā i-1 ā i-1 ⊗Dz i Dz i ],

where ⊗ denotes the Kronecker product of matrices. If the activations and pre-activation derivatives are approximated as independent, this yields the following approximation

F(w) i,i to F(w) i,i , F(w) i,i = E[ā i-1 ā i-1 ] :=Ai-1 ⊗ E[Dz i Dz i ] :=Gi

(2.7)

The K-FAC approximation matrix F(w) to F(w) is defined as

F(w) =     A0 ⊗ G1 0 A1 ⊗ G2 . . . 0 AL-1 ⊗ GL     .

(2.8)

The K-FAC update rules are analogous to natural gradient descent updates; we simply replace F(w k ) -1 in Eqn. 2.3 by F(w k ) -1 above.

Invariance Properties of K-FAC. Since K-FAC uses the approximation F(w) rather than the Fisher matrix F(w) itself, the invariance properties of natural gradient do not necessarily carry over to K-FAC. Instead, we consider the class of transformations given by the following transformed network

z † i = W † i ā † i-1 a † i = φ † i (z † i ) = Ω i φ i (Φ i z i + τ i ) + γ i ,

(2.9)

where Ω i , Φ i are invertible matrices and τ i , γ i are vectors. The transformed input is ā † 0 = Ω0 ā0 where Ω0 = Ω 0 γ 0 0 1 , and the transformed output is a † L = f † (x † , w † ) with w † defined as

w † = [vec( W † 1 ) vec( W † 2 ) . . . vec( W † L ) ]

. The original and transformed network are equivalent in terms of the functions they compute. We observe that the transformations given in Eqn. 2.9 encompasses a wide range of transformations. These include common deep learning tricks such as centering the activations to have zero mean and/or unit variance and replacing logistic sigmoid activation functions with tanh. While K-FAC may not be invariant under smooth parameterizations of the model as in the case of natural gradient, the following theorem shows that it is invariant to the class of transformations given in Eqn. 2.9. Theorem 2.1. Let N be the network with parameter vector w and activation functions {φ i } L i=1 . Suppose that we have activation functions {φ † i } L i=1 as given in Eqn. 2.9. Then, there exists a parameter vector w † such that the transformed network N † with parameter vector w † and activation functions {φ † i } L i=1 computes the same function as N . Furthermore, if the Fisher matrix of the MLP is assumed to be positive-definite, then the K-FAC updates are equivalent, in the sense that the resulting networks compute the same function.

A proof of this result was given earlier in Martens & Grosse (2015) ; however, the proof was cumbersome and involved tedious manipulation of update rules. In contrast, we like to provide a far more mathematically mature and elegant way to obtain this result. That is, we construct an intrinsic K-FAC metric g KFAC whose coordinate representation exactly matches the K-FAC approximation given in Eqn. 2.8. A K-FAC update can then be viewed as a natural gradient update with respect to the metric g KFAC . More importantly, the invariance properties of the K-FAC algorithm are immmediately established in the same way as it was for exact natural gradient.

3. COORDINATE-FREE K-FAC:

We observe that MLPs consist of a sequence of affine transformations and activation functions in alternation. In order to capture this structure, we treat the spaces of activations and pre-activations as affine spaces. Note that this introduces more structure than was assumed when we discussed the exact natural gradient in Section 2.1; in that subsection, we treated the space of network parameters as a general smooth manifold. Here, the network weights and biases are assumed to define affine transformations. The set of allowable reparameterizations (and hence, the desired set of invariances) is correspondingly more limited (though still very broad). We now present the coordinate-free MLP formally. For i ∈ {1, . . . , L}, we have

• Activations are taken to be elements α i-1 in an affine space A i-1 .

• Pre-activations are taken to be elements ζ i in an affine space Z i .

• Layerwise parameters are affine transformations ω i between A i-1 and Z i . The collection of these transformations is an affine space in its own right, which we denote by W i and refer to as the layerwise weight space. • The weight space is given by the direct product

W = W 1 × • • • × W L . Elements in this

space are written as ω = (ω 1 , . . . , ω L ) ∈ W. • Input and outputs are denoted by ξ and f (ξ, ω) respectively. The space of all inputs and outputs are affine spaces denoted by X (= A 0 ) and Y(= A L ) respectively.

Moreover, the layerwise computation is given by

EQUATION

where ρ i : R → R is a fixed nonlinear activation function which is assumed to be smooth throughout.

In Appendix D, we show that the MLP with computation defined in Eqn. 2.6 and the transformed version in Eqn. 2.9 correspond to two different choices of parameterizations of the same underlying abstract MLP.

Remark 3.1. The smoothness condition on the activation functions ρ i is necessary as our later analysis is built upon the framework of differential geometry. Smoothness is required to define standard operations such as pushforwards and pullbacks. While the commonly-used ReLU activation function is not smooth by definition, it is locally smooth with probability 1 if the weights are sampled from a continuous distribution. Furthermore, the ReLU function is a limit of the softplus functions

1 n log(1 + e nx )

, which are all individually smooth.

Optimization. The optimization problem in the abstract setting is analogous to the coordinatedependent one. Let (ξ, υ) be an abstract input-target pair and L(υ, f (ξ, ω)) be the loss function measuring the disagreement between outputs f (ξ, ω) of the abstract MLP and targets υ. Given a training set S of abstract input-target pairs (ξ i , υ i ), the objective function we wish to minimize here is

h(ω) = 1 |S| (ξi,υi)∈S L(υ i , f (ξ i , ω)).

The following theorem (see Appendix E.1 for the proof) shows that by imposing affine structure on the weight space, we can obtain global invariance instead of just first-order invariance. Theorem 3.2. Let g be a nondegenerate metric on the weight space W of an abstract MLP. For a chosen learning rate > 0, the following update rule

ω k+1 ← ω k -g(ω k ) -1 dh(ω k ),

is exactly invariant to all affine reparameterizations of the model.

Pullback of Output Metrics to Weight Space. Consider a metric g on the output space Y of the MLP. Let Ψ ξ : W → Y be the smooth map which sends parameters ω to outputs α L = f (ξ, ω)

given an input ξ. The pullback Ψ * ξ g defines a metric on W. The expected pullback metric over inputs, under a choice of coordinates around ω and α L , is given by

EQUATION

where G is the representation of g in these coordinates.

Example 3.3 (Fisher metric). Suppose that the outputs α L parameterize the model's predictive distribution R υ|α L . Let r(υ|α L ) denote the density function of this distribution and furthermore, we take the loss function here to be the negative log-likelihood L α L = -log r(υ|α L ). The output Fisher metric g F,out on Y is defined as

g F,out (α L ) = E υ [dL α L ⊗ dL α L ],

where the expectation is taken with respect to the predictive distribution R υ|α L . Computing the expectation of Ψ * ξ g F,out over the inputs ξ gives

E ξ [Ψ * ξ g F,out (ω)] = E ξ [Ψ * ξ E υ [dL α L ⊗ dL α L ]] = E ξ,υ [Ψ * ξ (dL α L ⊗ dL α L )] = E ξ,υ [Ψ * ξ (dL α L ) ⊗ Ψ * ξ (dL α L )] = E ξ,υ [dL ω ⊗ dL ω ].

This is exactly the Fisher metric defined earlier in Eqn. 2.4. We note that this construction quite general and not limited to Fisher metrics; for example, if we use the Euclidean metric on Y, then Eqn. 3.2 is the Gauss-Newton matrix (Martens, 2014) . A practical use case for the Gauss-Newton metric is when the outputs of the network do not have a natural probabilistic interpretation, e.g., the value network in an actor-critic architecture for reinforcement learning (Wu et al., 2017) . Another example is if we take the Riemannian metric induced by the Bregman divergence associated with a convex functional on Y, then Eqn. 3.2 is the generalized Gauss-Newton matrix (Martens, 2014) .

3.1. INDEPENDENCE METRIC:

We now come to the heart of our paper: the construction of a metric inspired by the K-FAC approximation. Recall that K-FAC makes two approximations to obtain a tractable Fisher matrix: (1) it assumes independence of activations and pre-activation derivatives in order to push the expectation inside the Kronecker product (Eqn. 2.7), and (2) it keeps only the diagonal blocks corresponding to individual layers. In this section, we develop a coordinate-free way to push the expectation inside the Kronecker product, thereby obtaining an approximate metric we term the independence metric. We later use this construction to develop approximate metrics for MLPs. In Section 3.2, we develop a coordinate-free version of the block-diagonal approximation. Combining both approximations yields the K-FAC metric, an intrinsic metric whose coordinate representation matches the K-FAC approximate Fisher matrix.

We begin by setting up the mathematical framework. To avoid tying ourselves to MLPs, we consider the more general setting of metrics on affine maps between affine spaces, but use notation which is suggestive of MLPs. We assume the following:

• Affine spaces A and Z

• Affine space W of affine transformations between A and Z

• Metric g on Z Our first task is to formulate a coordinate-free analogue of the outer product of homogenized activations, āi ā i . Consider the evaluation map ψ a : W → Z which is defined by evaluating w ∈ W at a ∈ A. We compute the pushforward ψ a * : T W → T Z. Note that there is no need to specify particular points for the tangent spaces here since we are working with affine spaces (see Corollary C.2 in Appendix C). Let ∂ w be a tangent vector on W and f be a smooth function on Z. Then,

ψ a * (∂ w )(f ) = ∂ w (f • ψ a )(w) = (∂ w f )(ψ a (w)) • ψ a (w) = (∂ w f )(z) • a

This shows that the pushforward ψ a * is exactly multiplication by the element a. Hence, we can identify any a ∈ A with its linear map T W → T Z. Thus, this enables us to define the tensor product of two elements in A as a mapping a 1 ⊗ a 2 : T W × T W → T Z ⊗ T Z:

(a 1 ⊗ a 2 )(∂ w1 , ∂ w2 ) = a 1 (∂ w1 ) ⊗ a 2 (∂ w2 ).

We now introduce the central object of our study, inspired by the independence assumption for activations and pre-activation derivatives which led to Eqn. 2.7. For w ∈ W , define g ind on W to be

EQUATION

where the first expectation is over A and the second one is over Z. Note that E[g(z)] is well defined because the affine structure of Z allows us to identify the cotangent spaces at all points z. The following theorem, whose proof is given in Appendix E.2, shows that g ind is indeed a metric on W .

Theorem 3.4. Let g be a metric on Z and ψ a : W → Z be the evaluation map. Then, g ind as defined in Eqn. 3.3 is a metric on W . Moreover, if the expected pullback metric E a [ψ * a g] is nondegenerate on W , then g ind is also nondegenerate. From now on, we refer to g ind as the independence metric.

By direct computation, we obtain that the coordinate representation of g ind matches the K-FAC approximation of the layerwise Fisher blocks: Proposition 3.5. Suppose that we choose coordinate systems for the affine spaces A, Z and in these coordinates, a = a, z = z, g(z) = G(z).

Then the independence metric g ind can be expressed as

g ind (w) = E[āā ] ⊗ E[G(z)].

Remark 3.6. In the context of MLPs (which we explain in greater detail subsequently) where

A = A i-1 , Z = Z i and W = W i , the matrix G(z) = E y [Dz i Dz i ]

where the expectation is taken over output space. E[G(z)] means we furthermore take the expectation over Z i . When we write E[Dz i Dz i ] in Eqn. 2.7, we implicitly take this to mean E[G(z)].

3.2. K-FAC METRIC:

In this subsection, we formulate the layerwise independence approximation in a coordinate-free way, allowing us to define the K-FAC metric, whose coordinate representation matches the K-FAC approximation to the Fisher matrix. To do this, we use the independence metric developed in Section 3.1 to define the K-FAC metric for MLPs. Lastly, by viewing K-FAC as a metric on W, we show how invariances of the K-FAC algorithm can be obtained in a natural and straightforward manner.

Consider a MLP with L layers as described earlier in Section 3. For every i ∈ {1, . . . , L}, we define the following maps

• ψ i ξ : W i → Z i which sends layerwise parameters ω i to pre-activations ζ i by evaluation at activations α i-1 . • ϕ i ξ : Z i → Y which sends ζ i to network outputs α L = f (ξ, ω).

The subscript ξ is used to highlight the fact that all of these maps have an implicit dependence on network inputs ξ. Note that ψ i ξ is a smooth map by definition. Now, observe that ϕ i ξ is exactly the composition of network maps

ρ L • ω L • • • • • ω i+1 • ρ i : Z i → Y.

Since all activation functions ρ i are assumed to be smooth maps, it follows immediately that ϕ i ξ is also a smooth map. Moreover, we define the map

Ψ i ξ : W i → Y to be composition Ψ i ξ = ϕ i ξ • ψ i ξ .

Let g be a metric on Y. Then, the pullback (ϕ i ξ ) * g defines a metric on Z i . Now, if we take A, Z, W in Section 3.1 to be

EQUATION

Now, given metrics g i ind on each W i , there is a natural "additive metric" defined on the Cartesian product W = W 1 × • • • × W L . We defer the formalities and the sum "+" operation on metrics to Appendix B.5. Definition 3.7. The K-FAC metric on the weight space W of a MLP is defined as

g KFAC (ω) = g 1 ind (ω 1 ) + • • • + g L ind (ω L )

, where the sum is as defined in Eqn. B.2 in Appendix B.5 and each g i ind is as given in Eqn. 3.4. Theorem 3.8 (See Appendix E.3 for proof). Let g be a metric on Y. Then, g KFAC given in Definition 3.7 is indeed a metric on the weight space W of an abstract MLP. Moreover, if we assume that the expected pullback of g, E ξ [(Ψ i ξ ) * g], under the map Ψ i ξ : W i → Y is a nondegenerate metric on the layerwise weight space W i for every i, then g KFAC is also nondegenerate.

Coordinate-Free Proof of Theorem 2.1. We are now ready to provide a natural and straightforward proof of Theorem 2.1. In Appendix D, we show that networks N and N † correspond to two different choices of parameterizations for the same underlying abstract MLP. Hence, they must compute the same function.

For the latter assertion, assume that the metric g on the output space Y in Theorem 3.8 is the output Fisher metric g F,out in Example 3.3. The pullback of this under ϕ i ξ is given by

(ϕ i ξ ) * g F,out (ζ i ) = E[dL ζi ⊗ dL ζi ].

Let us choose coordinate systems on A i-1 and Z i with

α i-1 = a i-1 , ζ i = z i , (ϕ i ξ ) * g F,out (ζ i ) = E[dL ζi ⊗ dL ζi ] = E[Dz i Dz i ].

Then, by Proposition 3.5,

g i ind (ω i ) = E[ā i-1 ā i-1 ] ⊗ E[Dz i Dz i ],

which is exactly F(w) i,i given earlier in Eqn. 2.7. Furthermore,

g KFAC (ω) = g 1 ind (ω 1 ) + • • • + g L ind (ω L )

is the matrix with diagonal blocks F(w) i,i and zeros everywhere else (see Appendix B.5). This is precisely F(w) in Eqn. 2.8. Now, observe that

g KFAC (ω) -1 dh(ω) = g KFAC (ω) -1 dh(ω) = F(w) -1 ∇h(w),

and hence the K-FAC update rule is simply a natural gradient update rule with respect to the K-FAC metric g KFAC for abstract MLPs. Suppose that g KFAC is a nondegenerate metric; which is true for example if the assumptions in the second assertion of Theorem 3.8 hold. Applying Theorem 3.2 shows that this update rule is invariant to any affine reparameterizations of the model.

4. APPLICATIONS TO CONVOLUTIONAL AND RECURRENT NETWORKS:

In this section, we extend the preceding analysis to both convolutional and recurrent networks. Both cases are straightforward applications of our results in Section 3, highlighting the flexbility of our analysis. Due to space constraints, we give only sketches of the necessary constructions and defer the reader to Appendices F and G for full details.

4.1. CONVOLUTIONAL NETWORKS:

We focus on a single convolution layer following the exposition given in Grosse & Martens (2016) .

Let J be the number of input maps, I be the number of output maps, ∆ be the set of spatial offsets and T be the number of spatial locations. We write the convolution layer as a matrix multiplication

Z l = W l A exp l-1 + b l A l = φ l (Z l ), (4.1)

where A exp l-1 is the J|∆| × |T | matrix of expanded activations, W l is the I × J|∆| weight matrix, Z l is the I × |T | matrix of pre-activations, φ l is a nonlinear activation function and b l is the bias vector. Adopting homogeneous coordinates for the various matrices, we can rewrite the above as

[Z l ] H = [W l ] H [A exp l-1 ] H [A l ] H = φ l ([Z l ] H ). (4.2)

We briefly introduce the concept of a transformed convolution layer. For a convolution layer as defined in the above equation, the parameters [W l ] H and the transformed parameters [W † l ] H are related in the following way

EQUATION

where Γ l and Υ l-1 are invertible matrices. The activation functions φ l and φ † l are related through a standard affine change-of-basis as given in Eqn. 2.9.

For a convolutional network with L convolution layers, the Fisher matrix is given by

F(w) = E x,y [(Dw)(Dw) ],

where w is the vector concatenating all convolution layer weights. The K-FAC approximation is defined by setting all off-diagonal blocks to be zero and approximating the diagonal blocks as

F(w) l,l = |T |(E T [ā (:,t) l-1 (ā (:,t) l-1 ) ] ⊗ E T [Dz (t) l (Dz (t) l ) ]).

(4.4)

Here, a (:,t)

l-1 stands for the expanded activations at each spatial location t ∈ T (the J|∆|-dimensional column vectors of the matrix A exp l-1 in Eqn. 4.1) and z

(t)

l stands for the local pre-activations at t ∈ T (the I-dimensional column vectors of the matrix Z l in Eqn. 4.1).

Theorem 4.1. Let N be a convolutional network with parameter vector w and activation functions {φ l } L l=1 . Suppose that we have activation functions {φ † l } L l=1 which are related to {φ l } L l=1 by standard change-of-basis transformations. Then, there exists a parameter vector w † such that the transformed network N † with parameter vector w † and activation functions {φ l } L l=1 computes the same function as N . Furthermore, suppose that the Fisher matrix for the convolutional network is positive-definite, the K-FAC updates are equivalent, in the sense that the resulting networks compute the same function.

Proof Sketch (Full proof in Appendix F.3). A coordinate-dependent proof of this result was given previously in Grosse & Martens (2016) . Instead, we prove this intrinsically using the framework provided in Section 3. Let A exp l-1 be the affine space of expanded pre-activations at t ∈ T , Z l be the affine space of local pre-activations at t ∈ T , and W l be the layerwise weight space comprising of affine transformations between A exp l-1 and Z l . Now, let us take A, Z, W in Section 3.1 to be A exp l-1 , Z l , W l respectively and the metric on Z = Z l to be the pullback of the output Fisher metric. After summing over every spatial location t ∈ T , the independence metric g l ind on W l in coordinates, by Proposition 3.5, is exactly the expression given in Eqn. 4.4. Taking the sum of g l ind 's then gives exactly the K-FAC approximation for convolutional networks. Finally, we know from Theorem 3.4 that each g l ind indeed defines a metric and so by Theorem 3.2, we can conclude that the K-FAC updates are exactly invariant to all affine reparameterizations of the model.

4.2. RECURRENT NETWORKS:

As in the case of convolutional networks, we need not write out the full structure of a recurrent network. Rather, we focus solely on the recurrent computation since the object of our interest, the Fisher matrix, only involves recurrent weights. Let T be the number of different time steps and T = {1, . . . , T }. We use t to index the time step. Additionally, we assume that all sequences are of fixed length T . We consider the recurrent computation step

EQUATION

where a t-1 is an activation vector, z t is a pre-activation vector, and W is a recurrent weight matrix. The transformed recurrent computation step is defined as

z † t = W † a † t-1 + b † . (4.6)

The relationship between transformed parameters (W † b † ) and original parameters (W b) is given by the standard change-of-basis formula as in Eqn. D.2.

The Fisher matrix for recurrent networks is defined as

F( W) = E x,y [vec(D W) vec(D W) ],

and the K-FAC approximation is defined as

F( W) = T (E T [ā t-1 ā t-1 ] ⊗ E T [Dz t Dz t ]). (4.7)

Theorem 4.2. Let N be a recurrent network with recurrent parameters (W b). Suppose that we have a recurrent network N † with recurrent parameters (W † b † ) and the relationship between (W b) and (W † b † ) is a standard change-of-basis transformation. Then, the networks N and N † compute the same function. Furthermore, if the Fisher matrix for recurrent networks is assumed to be positive-definite, then the K-FAC updates are equivalent, in the sense that the resulting networks compute the same function.

Proof Sketch (Full proof in Appendix G.3). Let A be the affine space of local activations at each time step t, Z be the affine space of local pre-activations at each time step t, and W be the parameter space consisting of all affine transformations between A and Z. Now, let us take A, Z, W in Section 3.1 to be A, Z, W respectively and the metric on Z = Z to be the pullback of the output Fisher mertric. After summing over all time steps t ∈ T, the independence metric g ind on W in coordinates, by Proposition 3.5, is exactly the K-FAC approximation given in Eqn. 4.7. Finally, we know from Theorem 3.4 that g ind indeed defines a metric and so by Theorem 3.2, we conclude that the K-FAC updates are exactly invariant to all affine reparameterizations of the model.

A 1 : Since this paper involves the interplay of many coordinate-independent and coordinatedependent mathematical objects, we summarize these notations here. Note that these notations are primarily for the setting of MLPs, the focus of the main body of the paper. The notations for convolutional networks and recurrent networks in Appendix F and G will be self-contained. As a general rule of thumb, we use boldface to symbolize coordinate-dependent objects and standard math font for coordinate-independent ones.

B MATHEMATICAL BACKGROUND:

B.1 VECTOR SPACES AND TENSOR ALGEBRA Let V be a vector space. The dual space V * of V is the set of all linear functionals on V and this space itself admits the structure of a vector space. The direct sum U ⊕ V of two vector spaces U and V is a vector space where the set structure is the Cartesian product U × V and the addition and multiplication is given by

(u 1 , v 1 ) + (u 2 , v 2 ) = (u 1 + v 1 , u 2 + v 2 ) c(u 1 , v 1 ) = (cu 1 , cv 1 ), for u 1 , u 2 ∈ U , v 1 , v 2 ∈ V and c ∈ R.

We now introduce tensors on vector spaces. A k-tensor T on V is a multilinear function

T : V × • • • × V k copies → R,

We may think of T as an element of the vector space (V * ) ⊗k , the tensor product of the vector space V * with itself k-times. We delegate the definition of a tensor product of vector spaces to Appendix C.2. A k-tensor T is symmetric if T is a symmetric multilinear function. We work primarily with symmetric tensors in this paper.

B.2 CANONICAL ISOMORPHISMS:

We describe the distinction between an isomorphism and a canonical isomorphism of vector spaces. An isomorphism between two vector spaces U and V is a bijection between U and V which preserves addition and scalar multiplication. A canonical isomorphism is a stronger concept, it is an isomorphism of vector spaces which is natural, in the sense that it does not depend on any choice of bases to define the isomorphism. For example, any two vector spaces of the same dimension are isomorphic to one another but the isomorphism may not be canonical. Consider a finite dimensional vector space V . There is an isomorphism between V and its dual space V * : given a choice of basis e i for V , there is a dual basis e * i for V * and the map e i → e * i is an isomorphism. However, this is not canonical as it depends on the choice of basis e i for V . On the other hand, consider the evaluation map ev v : V * → R defined by ϕ → ϕ(v) where ϕ ∈ V * , v ∈ V . The mapping v → ev v then defines a canonical isomorphism from V to its double dual space V * * .

B.3 AFFINE ALGEBRA:

A set A is an affine space associated to the vector space V if there is a mapping A × A → V denoted by (p, q) ∈ A × A → pq ∈ V satisfying the axioms 1. for any p, q, r ∈ A, pr = pq + qr 2. for any p ∈ A and for any x ∈ V there is an unique q ∈ A such that x = pq.

Intuitively, an affine space may be thought of as a vector space with no privileged origin. Suppose that we choose an origin point o ∈ A and let {e 1 , . . . , e n } be a basis for the associated vector space V . For any point p ∈ A, we can write op = where homogeneous coordinates are used for a i and a † i . Now, suppose that the change-of-basis from ι to κ is given by (B c) and denote

[B] H = B c 0 1 . Then, we have    ā † 1 . . . ā † k    =    [B] H 0 . . . 0 [B] H       ā1 . . . āk    = (I ⊗ [B] H )    ā1 . . . āk    . (B.1)

Thus, the induced change-of-basis on the product space A K is given by the matrix I ⊗ [B] H .

B.4 DIFFERENTIALS, PUSHFORWARDS AND PULLBACKS:

Let M be a smooth real manifold. For p ∈ M, we denote the tangent space by T p M and the corresponding dual space, the cotangent space by T * p M. Given a smooth function f :

M → R, the differential df (p) ∈ T * p M is defined by df (p)(X p ) = X p (f ), X p ∈ T p M.

Let (x 1 , . . . , x n ) be a coordinate system around p ∈ M, the differential df (p) can be expressed as

df (p) =    ∂f ∂x1 . . . ∂f ∂xn    .

We observe that this coordinate representation corresponds to the gradient ∇f (even though differentials and gradients are distinct objects for abstract manifolds).

For a smooth map ϕ : M 1 → M 2 of manifolds, the pushforward ϕ * :

T p M 1 → T ϕ(p) M 2 is defined by ϕ * (v)h = v(h • ϕ), where v ∈ T p M 1 and h : M 2 → R is a smooth function on M 2 . If we suppose that M 1 = R n and M 2 = R m with (x 1 , . . . ,

x n ) a coordinate system around p and (y 1 , . . . , y m ) a coordinate system around ϕ(p), the pushforward ϕ * v can be represented as

ϕ * v =    ∂y1 ∂x1 . . . ∂y1 ∂xn . . . . . . . . . ∂ym ∂x1 . . . ∂ym ∂xn    v,

where v = v . This is exactly the Jacobian matrix J ϕ of ϕ and hence ϕ * v = J ϕ v. This Jacobianvector product corresponds to the directional derivative, and can be computed using forward mode automatic differentiation (Schraudolph, 2002) .

The dual notion of the pushforward, the pullback ϕ * :

T * ϕ(p) M 2 → T * p M 1 is defined in the follow- ing way (ϕ * u)(v) = u(ϕ * v), u ∈ T * ϕ(p) M 2 .

With respect to the same coordinate systems chosen above, we can write ϕ * u = J ϕ u, where u = u . Numerically, we can compute J ϕ u efficiently using reverse mode auto-differentiation (i.e., backpropagation).

B.5 METRICS AND THEIR PROPERTIES:

We introduce tensors on manifolds. A symmetric k-tensor σ at the point p ∈ M is defined as a symmetric k-tensor on the tangent space T p M. Recall that this is a symmetric multilinear map on the k-fold product of T p M:

σ(p) : T p M × • • • × T p M k copies → R.

A metric on M is defined as a smoothly varying symmetric 2-tensor g which is positive-semidefinite at every point p ∈ M. If g is nondegenerate, then this is just a usual Riemannian metric. However, in this paper, we use the term "nondegenerate" rather than "Riemannian" to describe such metrics.

Earlier in the paper, we pulled back metrics from the output space to the weight space of the network.

Here, we define how this works for general tensors. Let ϕ : M 1 → M 2 be a smooth map of manifolds and σ be a symmetric k-tensor on M 2 at ϕ(p). The pullback ϕ * σ of σ under ϕ is a symmetric k-tensor on M 1 defined as

ϕ * σ(p)(v 1 , . . . , v k ) = σ(ϕ(p))(ϕ * v 1 , . . . , ϕ * v k ),

where v 1 , . . . , v k ∈ T p M 1 . In the case of metrics,

ϕ * g M2 (p)(v 1 , v 2 ) = g M2 (ϕ(p))(ϕ * v 1 , ϕ * v 2 ).

If we suppose that the metric g M2 is given by G M2 for a chosen coordinate system around ϕ(p), then the pullback metric ϕ * g M2 on M 1 around p is given by

ϕ * g M2 (p) = J ϕ G M2 J ϕ ,

where J ϕ is the Jacobian of ϕ. While a metric always pulls back to a metric under a smooth map, the pullback of a nondegenerate metric can be degenerate as the pushforward map may have a non-trivial nullspace.

Lastly, we describe how to define metrics on product manifolds. Given metrics g M1 and g M2 on M 1 and M 2 respectively, we describe how to naturally define a metric on the product manifold M 1 × M 2 . For any point (p, q) ∈ M 1 × M 2 , there is a canonical isomorphism of tangent spaces:

T (p,q) (M 1 × M 2 ) ∼ = T p M 1 ⊕ T q M 2 .

The proof of this fact can be found in standard differential geometry literature (e.g. Lee ( 2003)) and so we do not elaborate further here. Hence, any vector v ∈ T (p,q) (M 1 × M 2 ) can be written as a pair (v 1 , v 2 ) where v 1 ∈ T p M 1 and v 2 ∈ T q M 2 . Now, we define the additive metric g M1 + g M2 on M 1 × M 2 as follows:

(g M1 + g M2 )(p, q)(u, v) = g M1 (p)(u 1 , v 1 ) + g M2 (q)(u 2 , v 2 ). (B.2)

If we choose a coordinate system around (p, q) with the metrics g M1 , g M2 represented by matrices G M1 , G M2 respectively, then we have

(g M1 + g M2 )(p, q) = G M1 0 0 G M2 ,

which is a matrix with block diagonals G M1 , G M2 and zero everywhere else. This construction generalizes easily to sums of more than two terms.

C ADDITIONAL MATHEMATICAL TECHNICALITIES C.1 TANGENT SPACE OF VECTOR SPACES AND AFFINE SPACES:

Theorem C.1. Let V be a finite-dimensional vector space. For each point p ∈ V , there is a canonical isomorphism V → T p V . From now on, we suppress p in T p V and write T V when denoting tangent spaces of V .

Proof. For any element v ∈ V , we can associate a tangent vector D v | p at p defined by

D v | p f = D v f (p) = d dt t=0 f (p + tv),

where f is a smooth function on V . This gives the desired canonical isomorphism since the above construction involves no choice of basis.

Corollary C.2. Let A be an affine space and V be its associated vector space. For each point a ∈ A, there is a canonical isomorphism between T a A and V . From now on, we suppress a in T a A and write T A when denoting tangent spaces of A.

Proof. Note that specifying a point a ∈ A naturally identifies A with V . Then, applying the above theorem gives the desired result.

C.2 TENSOR PRODUCT OF VECTOR SPACES:

Let U and V be finite-dimensional vector spaces over the real numbers R. Let R be the subspace of the free vector space R U × V (set of all finite formal linear combinations of elements of U × V with real coefficients) spanned by all elements of the following forms:

c(u, v) -(cu, v), c(u, v) -(u, cv), (u, v) + (u , v) -(u + u , v), (u, v) + (u, v ) -(u, v + v ), for u, u ∈ U , v, v ∈ V , and c ∈ R. The tensor product U ⊗ V , is the quotient space R U ×V R and the equivalence class of an element (u, v) in U ⊗ V is denoted by u ⊗ v.

We describe how the vector space of linear transformations between U and V , denoted by Hom(U, V ), may be thought of as tensor products. There is a canonical isomorphism

U * ⊗ V → Hom(U, V ) (C.1)

given by ϕ ⊗ v → ϕ(u)v where ϕ ∈ U * . Another isomorphism of interest to us is

U * ⊗ V * → (U ⊗ V ) * , (C.2)

which is again canonical. To derive this isomorphism, given

ϕ ∈ U * , φ ∈ V * , consider the bilinear map U × V → R defined by (u, v) → ϕ(u) • φ(v)

. This induces an element on the tensor product (U ⊗ V ) * . As such, we obtain an unique linear injection U * ⊗ V * → (U ⊗ V ) * . Since all the vector spaces are finite-dimensional, we can conclude that this is an isomorphism.

We now use these facts to explain several ingredients in the proof of Theorem 3.2 in greater detail. The element a is a linear map T W → T Z and by the above isomorphism, this means a ∈ T * W ⊗ T Z from Eqn. C.1. By Eqn. C.2, the dual space of (T * W ⊗ T Z) * ∼ = T W ⊗ T * Z. Using Eqn. C.1 again, elements of this are maps λ : T * W → T * Z.

D TWO PARAMETERIZATIONS OF COORDINATE-FREE MLP:

We now show how the MLP with computation defined in Eqn. 2.6 and the transformed version in Eqn. 2.9 can be viewed as two different choices of parameterizations for the same underlying abstract MLP.

First, a choice of parameterization, or a coordinate system, for the abstract MLP is a choice of affine bases for all of the activation spaces A 1 , . . . , A L-1 , the pre-activation spaces Z 1 , . . . , Z L , the input space X and the output space Y in the network. Observe that a choice of bases for A i-1 and Z i naturally induces a basis for each W i , and therefore also for the full weight space W. Let ι, κ be two different choices of parameterizations for the network. With respect to ι, we write

α i-1 ι = a i-1 , ζ i ι = z i , ω i ι = (W i b i ), ρ i ι = φ i , α i ι = a i ,

and with respect to κ, we write

α i-1 κ = a ‡ i-1 , ζ i κ = z ‡ i , ω i κ = (W ‡ i b ‡ i ), ρ i κ = φ ‡ i , α i κ = a ‡ i .

Hence, we can rewrite Eqn. 3.1 in the parameterizations ι, κ as

z i = W i a i-1 + b i z ‡ i = W ‡ i a ‡ i-1 + b ‡ i a i = φ i (z i ), a ‡ i = φ ‡ i (z ‡ i ).

(D.1)

The parameters (W i b i ) and (W ‡ i b ‡ i ) are related as follows

W i = Φ i W ‡ i Ω i-1 b i = Φ i W ‡ i γ i-1 + Φ i b ‡ i + τ i , (D.2)

where (Ω i-1 γ i-1 ) is the change-of-basis from ι to κ on A i-1 with Ω i-1 an invertible matrix and γ i-1 a vector. Moreover, (Φ i τ i ) is the change-of-basis from κ to ι on Z i with Φ i an invertible matrix and τ i a vector. The activation functions φ i and φ ‡ i in Eqn. D.1 are related in the following way:

a ‡ i = φ ‡ i (z ‡ i ) = Ω i φ i (Φ i z ‡ i + τ i ) + γ i , (D.

3) where (Ω i γ i ) is the change-of-basis from ι to κ on A i .

To verify Eqn. D.2 and Eqn. D.3, we consider the the following commutative diagram (the top horizontal arrow is equal to the composition of maps given by the other three arrows) which relates the two parameterizations ι and κ on the activation affine space A i-1 and the pre-activation affine space Z i :

A i-1 ι (Wi bi) / / (Ωi-1,γi-1) Z i ι A i-1 κ (W ‡ i b ‡ i ) / / Z i κ (Φi,τi) O O (D.4)

Let a i-1 ∈ A i-1 ι , this maps to z i = W i a i-1 + b i ∈ Z i ι under the top horizontal arrow in Eqn. D.4. Now, mapping a i-1 under the composition of the other three arrows in Eqn. D.4, we obtain

a i-1 → Ω i a i-1 + γ i-1

(apply left vertical arrow in Eqn. D.4)

→ W ‡ i (Ω i a i-1 + γ i-1 ) + b ‡ i

(apply bottom horizontal arrow in Eqn. D.4)

→ Φ i (W ‡ i Ω i a i-1 + W ‡ i γ i-1 ) + Φ i b ‡ i + τ i (apply right vertical arrow in Eqn. D.4) = Φ i W ‡ i Ω i-1 a i-1 + Φ i W ‡ i γ i-1 + Φ i b ‡ i + τ i .

This establishes the equality given in Eqn. D.2. For Eqn. D.3, we use the commutative diagram (this time, the bottom horizontal arrow is equal to the composition of the other three arrows):

Z i ι φi / / A i ι (Ωi,γi) Z i κ φ ‡ i / / (Φi,τi) O O A i κ (D.5) Let z ‡ i ∈ Z i κ

, this maps to a ‡ i ∈ A i κ under φ ‡ i . Now, mapping z ‡ i under the other three arrows in Eqn. D.5, we obtain

z ‡ i → Φ i z ‡ i + τ i (apply left vertical arrow in Eqn. D.5) → φ i (Φ i z ‡ i + τ i ).

(apply upper horizontal arrow in Eqn. D.5)

→ Ω i φ i (Φ i z ‡ i + τ i ) + γ i (apply right vertical arrow in Eqn. D.5).

This establishes Eqn. D.3.

Lastly, we note that the equations in Eqn. D.1 can be rewritten in homogeneous coordinates:

z i = Wi āi-1 z ‡ i = Wi ā ‡ i-1 a i = φ i (z i ), a ‡ i = φ ‡ i (z ‡ i ) = Ω i φ i (Φ i z ‡ i + τ i ) + γ i .

The left hand set of equations above is identical to the original MLP computation given in Eqn. 2.6. The right hand set of equations is identical to the transformed computation given in Eqn. 2.9. This completes the proof.

E PROOFS:

E.1 PROOF OF THEOREM 3.2

First, note that g(ω k ) -1 dh(ω k ) is an intrinsically defined tangent vector on W. The weight space W of a MLP is an affine space, and hence by Corollary C.2 in Appendix C, the tangent space of W at every point is canonically isomorphic to the vector space naturally associated to W. Thus, the exponential map Exp ω k corresponds to the above update rule. Since this construction did not require choosing an affine basis for W, the algorithm is invariant to affine reparameterizations.

E.2 PROOF OF THEOREM 3.4

Before providing the proof of the theorem, we need the following lemma.

Lemma E.1. Let g be a metric on Z and ψ a : W → Z be the evaluation map. Then, the pullback metric ψ * a g on W can be expressed as:

ψ * a g(w) = E[a ⊗ φ ⊗ a ⊗ φ],

where φ is a random covector.

Proof. Given z ∈ Z, the metric g admits the rank-1 decomposition

g(z) = E[φ ⊗ φ],

where φ is a covector and the expectation is over Z. This is akin to the more familiar case where any symmetric positive-semidefinite matrix admits a rank-1 spectral decomposition. Computing the pullback of g under the map ψ a now gives

ψ * a g(w) = ψ * a E[φ ⊗ φ] = E[ψ * a (φ ⊗ φ)] = E[ψ * a φ ⊗ ψ * a φ], (E.1)

where z = ψ a (w). We analyze the pullback ψ * a φ. Let ∂ w be a tangent vector on W . Then,

(ψ * a φ)(∂ w ) = φ((ψ a * )∂ w ) = φ(∂ w • a) = (a ⊗ φ)(∂ w ),

which shows that ψ * a φ = a ⊗ φ. Plugging this back into Eqn. E.1, we obtain

ψ * a g(w) = E[a ⊗ φ ⊗ a ⊗ φ],

which concludes the proof.

For the first assertion of the theorem, we need to check that both components E[a ⊗ a] and E[g(z)] define symmetric positive-semidefinite 2-tensors. Recall that a can be realized as a linear map from T W to T Z. Then, the dual element λ is a map from T * W to T * Z. We refer to Appendix C.2 for a formal explanation of this. To check the positive-semidefinite property,

E[a ⊗ a](λ, λ) = E[(a ⊗ a)(λ ⊗ λ)] ≥ 0,

where the latter inequality is due to the fact that a ⊗ a is positive-semidefinite. Moreover, a ⊗ a is also symmetric and this property is preserved under expectations which implies that E[a ⊗ a] is both symmetric and positive-semidefinite. For the second term E[g(z)] in g ind (w), the fact that g is a metric on Z means that g(z), by definition, is a symmetric positive-semidefinite 2-tensor on T * Z.

To establish the second assertion of the theorem, we need to show that both E[a ⊗ a] and E[g(z)] are positive-definite. Suppose to the contrary that this is not true for E[a ⊗ a]. Then there exists λ ∈ T W ⊗ T * Z (this is same as saying λ is a linear map from T * W to T * Z; refer to Appendix C.2 for further explanations) such that

0 = E[a ⊗ a](λ, λ) = E[(a ⊗ a)(λ ⊗ λ)],

and hence (a ⊗ a)(λ ⊗ λ) = 0. Now, consider the element λ ⊗ µ ∈ (T W ⊗ T * Z) ⊗ T Z where µ ∈ T Z. From the above lemma, we have

ψ * a g(w) = E[a ⊗ φ ⊗ a ⊗ φ] Then, evaluating E a [ψ * a g(w)] at λ ⊗ µ: E a [ψ * a g(w)](λ ⊗ µ, λ ⊗ µ) = E[a ⊗ φ ⊗ a ⊗ φ](λ ⊗ µ ⊗ λ ⊗ µ) = E[(a ⊗ a)(λ ⊗ λ) =0 •(φ ⊗ φ)(µ ⊗ µ)] = 0,

This shows that E a [ψ * a g] is not positive-definite which yields a contradiction as E a [ψ * a g] was assumed to be a nondegenerate metric. The exact same argument can be applied to show that E[g(z)] is positive-definite. This gives us the desired result.

E.3 PROOF OF THEOREM 3.8:

From Theorem 3.4, we know that g i ind is a metric on W i . Since g KFAC is defined as the additive metric where each of the summands are g i ind , we can conclude that g KFAC is a metric. For the second assertion of the theorem, recall that Ψ i ξ = ϕ i ξ • ψ i ξ . By the functoriality property of pullback operations, we have

(Ψ i ξ ) * = (ψ i ξ ) * • (ϕ i ξ ) * . Since E ξ [(Ψ i ξ )

* g] was assumed to be nondegenerate, this implies that

E ξ [(ψ i ξ ) * ((ϕ i ξ ) * g)],

is also nondegenerate. Then, from the second assertion of Theorem 3.4, we obtain that g i ind is a nondegenerate metric on W i . Consequently, g KFAC is nondegenerate which concludes the proof.

F COORDINATE-FREE K-FAC FOR CONVOLUTIONAL NETWORKS F.1 CONVOLUTIONAL NETWORKS:

We provide a complete account of the convolutional network case sketched out earlier in Section 4.1. We begin by describing the convolution layer of a convolutional network in mathematical terms following Grosse & Martens (2016) . We only consider convolution layers as the pooling and response normalization layers of a convolutional network typically do not contain (many) trainable weights. We then introduce the notion of a transformed convolution layer analogous to what was done in the case of MLPs. Lastly, we use the abstract linear algebra machinery developed in Section 2 to give a coordinate-free description of convolution layers.

F.1.1 CONVOLUTION LAYERS:

We focus on a single convolution layer. A convolution layer l takes as input a layer of activations a j,t , where j ∈ {1, . . . , J} indexes the input map and t ∈ T indexes the spatial location. T here denotes the set of spatial locations, which we typically take to be a 2D-grid. We assume that the convolution is performed with a stride of 1 and padding equal to the kernel radius R, so that the set of spatial locations is shared between the input and output feature maps. This layer is parameterized by a set of weights w i,j,δ and biases b i , where i ∈ {1, . . . , I} indexes the output map and δ ∈ ∆ indexes the spatial offset. The numbers of spatial locations and spatial offsets are denoted by |T | and |∆| respectively. The computation at the convolution layer is given by z i,t = δ∈∆ w i,j,δ a j,t+δ + b i .

(F.1)

The pre-activations z i,t are then passed through a nonlinear activation function φ l . Analogous to feed-forward networks, the weight derivatives are computed using backpropagation:

Dw i,j,δ = t∈T a j,t+δ Dz i,t .

Following Grosse & Martens (2016) , we represent the convolution layer computation in Eqn. F.1 using matrix notation. To do this, we write the activations a j,t as a J × |T | matrix A l-1 , the preactivations z i,t as a I × |T | matrix Z l , the weights w i,j,δ as a I × J|∆| matrix W l and the bias vector as b l . For the activation matrix A l-1 , if we extract the patches surrounding each spatial location t ∈ T and flatten these patches into vectors where the vectors become columns of a matrix, we obtain a J|∆| × |T | matrix which we denote by A exp l-1 . From now on, we refer to this matrix as the expanded activations. Finally, we can use these matrix notations to rewrite the computation in Eqn. F.1 as

Z l = W l A exp l-1 + b l A l = φ l (Z l ). (F.2)

For convenience purposes later, we adopt homogeneous coordinates for various matrices:

[A exp l-1 ] H = A exp l-1 1 , [Z l-1 ] H = Z l 1 , [W l ] H = W l b l 0 1 , [A l ] H = A l 1 .

Using Proposition 3.5, the independence metric in Eqn. F.6 can be expressed in coordinates as follows

g l ind (ω l ) = |T |(E T [ā (:,t) l-1 (ā (:,t) l-1 ) ] ⊗ E T [Dz (t) l (Dz (t) l ) ]

), which is exactly F(w) l,l given earlier in Eqn. F.5. Furthermore, g KFAC (ω) is exactly the K-FAC approximation F(w). Thus, the K-FAC update rule is a natural gradient update rule with respect to the K-FAC metric g KFAC for abstract convolutional networks. Lastly, if g KFAC is a nondegenerate metric; which is true for example if the assumptions in the second assertion of Theorem F.3 hold, then we can conclude that these updates are invariant to any affine reparameterizations of the model.

G COORDINATE-FREE K-FAC FOR RECURRENT NETWORKS:

In this section, we provide a complete account of the recurrent network case sketched out earlier in Section 4.2. We first give a mathematical formulation of the recurrent computation step of these networks in both coordinate-dependent and coordinate-independent scenarios. We proceed to give the Kronecker factorization of the Fisher matrix for recurrent networks and then state the invariance theorem for this optimization method. Lastly, we prove the invariance theorem in the same way we did for MLPs and convolutional networks previously.

G.1 RECURRENT NETWORKS:

As in the case of convolutional networks in Section F, it is not necessary to write out the full structure of a recurrent network. Rather, we focus on the recurrent computation since the central object of our interest, the Fisher matrix for recurrent networks, only involves recurrent weights. G.1.1 COMPUTATIONAL STEP Let T be the number of different time steps and T = {1, . . . , T }. We use t to index the time step. Throughout, we assume that all sequences are of fixed length T . For an input x t at every t, the recurrent network maps this to an output o t . Essentially, the network maps input sequences x = (x 1 , . . . , x T ) to output sequences o = (o 1 , . . . , o T ). The computation, at every t, is

z t = Wa t-1 + b z t = z t + Vx t a t = φ(z t )

where a t-1 is an activation vector, z t is a pre-activation vector, W is a recurrent weight matrix, V is a weight matrix, b is a recurrent bias vector, and φ is a fixed nonlinear activation function. For the remainder of this section, we focus on the first equation

z t = Wa t-1 + b, (G.1)

which represents the recurrent computation step. The latter two equations can be handled by the previous K-FAC analysis for MLPs given in Section 3. The transformed recurrent computation step is defined as

z † t = W † a † t-1 + b † . (G.

2)

The relationship between transformed parameters (W † b † ) and original parameters (W b) is given by the standard change-of-basis formula as in Eqn. D.2.

G.1.2 ABSTRACT RECURRENT NETWORK:

We now describe an abstract recurrent network formally.

• Local activations at each time step t are elements α t in an affine space A • Activations are elements α = {α t } t∈T in the affine space A T • Local pre-activations at each t are elements ζ t in an affine space Z • Pre-activations are elements ζ = {ζ t } t∈T in the affine space Z T

• Parameters are affine transformations ω between A and Z. The collection of these transformations is an affine space in its own right which we denote by W and refer to as the weight space • Network inputs and outputs at each t are elements ξ t , υ t in affine spaces X , Y respectively.

The input and output spaces are X T and Y T respectively; furthermore, elements here are written as ξ = {ξ t } t∈T and υ = {υ t } t∈T .

For every t, the abstract recurrent computation step is

ζ t = ω(α t-1 ).

A choice of parameterization for the abstract recurrent network consists of choosing affine bases for A, Z, X and Y. Since we have bases for A and Z, this induces a natural basis for W. If we use exactly the same change-of-basis analysis given in Appendix D, then the recurrent network with computation given by Eqn. G.1 and the transformed version in Eqn. G.2 correspond to two different parameterizations of the same abstract recurrent network.

G.2 K-FAC FOR RECURRENT NETWORKS:

We review the recent Kronecker factorization for recurrent networks method in Martens et al. (2018) .

Recall that for every time step t, the recurrent computation can be written as

z t = Wā t-1 ,

where W = [W b] and ā t-1 = [a t-1 1] . Using backpropagation, the log-likelihood gradient is given by Dz t ā t-1 . The total contribution to the gradient across all t is the sum

D W = T t=1

Dz t ā t-1 .

For an input-target pair (x, y), the Fisher matrix F( W) for recurrent networks is defined as Finally, the K-FAC approximation F( W) to F( W) for recurrent networks is defined as

F( W) = E x,

F( W) = T (E T [ā t-1 ā t-1 ] ⊗ E T [Dz t Dz t ]). (G.3)

Remark G.1. As in the case of convolution layers, there is weight sharing in recurrent networks (across time here instead of spatial locations) and so it is not enough to just assume statistical independence between activations and pre-activation derivatives to make the K-FAC approximation here. We defer the reader to Martens et al. (2018) for detailed explanations on how the K-FAC approximation is derived for recurrent networks.

We now proceed to the last section of this paper to give a coordinate-free proof of Theorem 4.2. The method of proof mirrors exactly the proofs given previously for Theorem 2.1 and Theorem 4.1.

G.3 COORDINATE-FREE K-FAC FOR RECURRENT NETWORKS:

Given an input ξ = {ξ t } t∈T ∈ X T and parameter ω ∈ W, denote the network output by f (ξ, ω).

For a specific time step t, consider the following maps:

• ψ ξ,t : W → Z which sends parameters ω to pre-activations ζ t by evaluation at activations α t-1 • ϕ ξ,t : Z → Y which sends ζ t to outputs

In addition, we define the map Ψ ξ,t : W → Y as the composition Ψ ξ,t = ϕ ξ,t • ψ ξ,t .

Let g be a metric on Y. The pullback ϕ * ξ,t g then defines a metric on Z. Now, we take A, Z, W in Section 3.1 to be A = A, Z = Z, W = W, and the metric on Z = Z to be ϕ * ξ,t g. Summing over all time steps t ∈ T, we make the following definition which arises from the independence metric in Section 3.2:

:



Figures and Tables:

i=1 x i (p)e i . Here, {x 1 , . . . , x n } is a set of coordinate functions, or more simply, a basis for A. If we have two bases {x 1 , . . . , x n } and {y 1 , . . . , y n }, then they are related by y = Bx + c where B = [b ij ] is an invertible n × n matrix and c = [c i ] is a vector.We now describe how to extend a change-of-basis on the affine space A to the product space A K = A × • • • × A. Let ι and κ be two choices of affine bases on A, then (a 1 , . . . , a k ) ι = (ā 1 , . . . , āk ), (a 1 , . . . , a k ) κ = (ā † 1 , . . . , ā † k ),

y [vec(D W) vec(D W) ].

TABLE OF NOTATIONS

Table content: <table><tr><td/><td colspan="2">Coordinate-independent Coordinate-dependent</td></tr><tr><td>inputs</td><td>ξ</td><td>x</td></tr><tr><td>targets</td><td>υ</td><td>y</td></tr><tr><td>activations</td><td>α i</td><td>a i</td></tr><tr><td>homogenized activations</td><td/><td>āi</td></tr><tr><td>pre-activations</td><td>ζ i</td><td>z i</td></tr><tr><td>layerwise parameters (and biases) homogenized layerwise parameters</td><td>ω i</td><td>(W i b i ) Wi</td></tr><tr><td>activation functions</td><td>ρ i</td><td>φ i</td></tr><tr><td>network parameter</td><td>ω</td><td>w</td></tr><tr><td>network output</td><td>f (ξ, ω)</td><td>f (x, w)</td></tr><tr><td>loss function</td><td>L(υ, f (ξ, ω))</td><td>L(y, f (x, w))</td></tr><tr><td>objective function</td><td>h(ω)</td><td>h(w)</td></tr><tr><td>predictive distribution</td><td>P υ|ξ (ω)</td><td>P y|x (w)</td></tr><tr><td>predictive distribution density function</td><td>p(υ|ξ, ω)</td><td>p(y|x, w)</td></tr><tr><td>Fisher metric/matrix layerwise log-likelihood differential/gradient</td><td>g F (ω) dL ωi</td><td>F(w) D Wi</td></tr><tr><td>log-likelihood differential/gradient</td><td>dL ω</td><td>Dw</td></tr><tr><td>K-FAC metric/matrix</td><td>g KFAC (ω)</td><td>F(w)</td></tr><tr><td>tensor product/Kronecker product</td><td>⊗</td><td>⊗</td></tr><tr><td>input space</td><td>X</td><td/></tr><tr><td>output space</td><td>Y</td><td/></tr><tr><td>space of activations</td><td>A i</td><td/></tr><tr><td>space of pre-activations</td><td>Z i</td><td/></tr><tr><td>layerwise weight space</td><td>W i</td><td/></tr><tr><td>weight space of network</td><td>W</td><td/></tr></table>

Peer review (iclr2020):
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. Why is the invariance important? How does the proposed framework help us develop better algorithms? Without empirical studies, it is not easy to see the significance of this work.