Year: 2019

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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND

heading: 2.1 NATURAL GRADIENT DESCENT
text: 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.

heading: 2.2 KRONECKER-FACTORED APPROXIMATE CURVATURE
text: We consider a MLP with L layers. At each layer i ∈ {1, . . . , L}, the MLP computation is given as: zi = Wiai−1 + bi
ai = φi(zi),
where ai−1 is an activation vector, zi is a pre-activation vector, Wi is a weight matrix, bi is a bias vector, and φi : R → R is an activation function. For convenience, we introduce homogeneous coordinates ā>i−1 = [a > i−1 1] > and W̄i = [Wi bi]. The above computation can be rewritten as
zi = W̄iāi−1 ai = φi(zi). (2.6)
We concatenate all of the network parameters W̄i into a single vector w,
w = [vec(W̄1) > vec(W̄2) > . . . vec(W̄L) >]>.
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(DW̄i) vec(DW̄j)>].
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 DW̄i = Dziā>i−1 and hence, we have vec(Dziā>i−1) = āi−1 ⊗Dzi. Then, F(w)i,i can be rewritten as:
F(w)i,i = E[vec(DW̄i) vec(DW̄i)>] = E[(āi−1⊗Dzi)(āi−1⊗Dzi)>] = E[āi−1ā>i−1⊗DziDz>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[āi−1ā>i−1]︸ ︷︷ ︸ :=Ai−1 ⊗E[DziDz>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(wk) −1 in Eqn. 2.3 by F̂(wk)−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 KFAC. Instead, we consider the class of transformations given by the following transformed network
z†i = W̄ † i ā † i−1 a†i = φ † i (z † i ) = Ωiφi(Φizi + τi) + γi,
(2.9)
where Ωi, Φi are invertible matrices and τi, γi are vectors. The transformed input is ā † 0 = Ω̄0ā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}Li=1. Suppose that we have activation functions {φ†i}Li=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}Li=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 gKFAC 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 gKFAC. More importantly, the invariance properties of the K-FAC algorithm are immmediately established in the same way as it was for exact natural gradient.

heading: 3 COORDINATE-FREE K-FAC
text: 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 Ai−1. • Pre-activations are taken to be elements ζi in an affine space Zi. • Layerwise parameters are affine transformations ωi between Ai−1 and Zi. The collection
of these transformations is an affine space in its own right, which we denote by Wi and refer to as the layerwise weight space.
• The weight space is given by the direct productW = W1 × · · · × WL. 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 (= A0) and Y(= AL) respectively.
Moreover, the layerwise computation is given by ζi = ωi(αi−1)
αi = ρi(ζi), (3.1)
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 spaceW of an abstract MLP. For a chosen learning rate > 0, the following update rule
ωk+1 ← ωk − g(ωk)−1dh(ω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
JEξ[Ψ∗ξg(ω)]K = Ex[J>ΨξGJΨξ ], (3.2) 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 gF,out on Y is defined as
gF,out(αL) = Eυ[dLαL ⊗ dLαL ], where the expectation is taken with respect to the predictive distribution Rυ|αL . Computing the expectation of Ψ∗ξgF,out over the inputs ξ gives
Eξ[Ψ∗ξgF,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).

heading: 3.1 INDEPENDENCE METRIC
text: 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, āiā>i . Consider the evaluation map ψa : W → Z which is defined by evaluating w ∈W at a ∈ A. We compute the pushforward ψa∗ : TW → TZ. 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) = (∂wf)(ψa(w)) · ψ′a(w) = (∂wf)(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 TW → TZ. Thus, this enables us to define the tensor product of two elements in A as a mapping a1 ⊗ a2 : TW × TW → TZ ⊗ TZ:
(a1 ⊗ a2)(∂w1 , ∂w2) = a1(∂w1)⊗ a2(∂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 gind on W to be
gind(w) := E[a⊗ a]⊗ E[g(z)], (3.3)
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 gind is indeed a metric on W .
Theorem 3.4. Let g be a metric onZ and ψa : W → Z be the evaluation map. Then, gind as defined in Eqn. 3.3 is a metric on W . Moreover, if the expected pullback metric Ea[ψ∗ag] is nondegenerate on W , then gind is also nondegenerate. From now on, we refer to gind as the independence metric.
By direct computation, we obtain that the coordinate representation of gind 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,
JaK = a, JzK = z, Jg(z)K = G(z). Then the independence metric gind can be expressed as
Jgind(w)K = E[āā>]⊗ E[G(z)].
Remark 3.6. In the context of MLPs (which we explain in greater detail subsequently) where A = Ai−1, Z = Zi and W = Wi, the matrix G(z) = Ey[DziDz>i ] where the expectation is taken over output space. E[G(z)] means we furthermore take the expectation over Zi. When we write E[DziDz>i ] in Eqn. 2.7, we implicitly take this to mean E[G(z)].

heading: 3.2 K-FAC METRIC
text: 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 onW , 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ξ : Wi → Zi which sends layerwise parameters ωi to pre-activations ζi by evaluation at activations αi−1.
• ϕiξ : Zi → 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 : Zi → 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ξ :Wi → Y to be composition Ψiξ = ϕiξ ◦ ψiξ.
Let g be a metric on Y . Then, the pullback (ϕiξ)∗g defines a metric on Zi. Now, if we take A, Z, W in Section 3.1 to be A = Ai−1, Z = Zi, W =Wi, and the metric on Z = Zi to be (ϕiξ)∗g, the independence metric onWi here is
giind(ωi) = E[αi−1 ⊗ αi−1]⊗ E[(ϕiξ)∗g(ζi)]. (3.4)
Now, given metrics giind on each Wi, there is a natural “additive metric” defined on the Cartesian productW = W1 × · · · × WL. We defer the formalities and the sum “+” operation on metrics to Appendix B.5. Definition 3.7. The K-FAC metric on the weight spaceW of a MLP is defined as
gKFAC(ω) = g 1 ind(ω1) + · · ·+ gLind(ωL),
where the sum is as defined in Eqn. B.2 in Appendix B.5 and each giind is as given in Eqn. 3.4. Theorem 3.8 (See Appendix E.3 for proof). Let g be a metric on Y . Then, gKFAC given in Definition 3.7 is indeed a metric on the weight spaceW of an abstract MLP. Moreover, if we assume that the expected pullback of g, Eξ[(Ψiξ)∗g], under the map Ψiξ :Wi → Y is a nondegenerate metric on the layerwise weight spaceWi for every i, then gKFAC 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 networksN andN † 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 gF,out in Example 3.3. The pullback of this under ϕiξ is given by
(ϕiξ) ∗gF,out(ζi) = E[dLζi ⊗ dLζi ].
Let us choose coordinate systems on Ai−1 and Zi with
Jαi−1K = ai−1, JζiK = zi, J(ϕiξ)∗gF,out(ζi)K = JE[dLζi ⊗ dLζi ]K = E[DziDz>i ]. Then, by Proposition 3.5,
Jgiind(ωi)K = E[āi−1ā>i−1]⊗ E[DziDz>i ],
which is exactly F̂(w)i,i given earlier in Eqn. 2.7. Furthermore,
JgKFAC(ω)K = Jg1ind(ω1) + · · ·+ gLind(ωL)K
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
JgKFAC(ω)−1dh(ω)K = JgKFAC(ω)−1KJdh(ω)K = 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 gKFAC for abstract MLPs. Suppose that gKFAC 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.

heading: 4 APPLICATIONS TO CONVOLUTIONAL AND RECURRENT NETWORKS
text: 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.

heading: 4.1 CONVOLUTIONAL NETWORKS
text: 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
Zl = WlA exp l−1 + bl Al = φl(Zl), (4.1)
where Aexpl−1 is the J |∆| × |T | matrix of expanded activations, Wl is the I × J |∆| weight matrix, Zl is the I × |T | matrix of pre-activations, φl is a nonlinear activation function and bl is the bias vector. Adopting homogeneous coordinates for the various matrices, we can rewrite the above as
[Zl]H = [Wl]H [A exp l−1]H [Al]H = φl([Zl]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 [Wl]H and the transformed parameters [W † l ]H are related in the following way
[Wl]H = Γl[W † l ]H(I⊗Υl−1), (4.3)
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) = Ex,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 |(ET [ā(:,t)l−1 (ā (:,t) l−1 ) >]⊗ ET [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 Aexpl−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 Zl in Eqn. 4.1).
Theorem 4.1. Let N be a convolutional network with parameter vector w and activation functions {φl}Ll=1. Suppose that we have activation functions {φ † l }Ll=1 which are related to {φl}Ll=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}Ll=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 Aexpl−1 be the affine space of expanded pre-activations at t ∈ T , Zl be the affine space of local pre-activations at t ∈ T , andWl be the layerwise weight space comprising of affine transformations between Aexpl−1 and Zl. Now, let us take A, Z, W in Section 3.1 to be A exp l−1, Zl,Wl respectively and the metric on Z = Zl to be the pullback of the output Fisher metric. After summing over every spatial location t ∈ T , the independence metric glind on Wl in coordinates, by Proposition 3.5, is exactly the expression given in Eqn. 4.4. Taking the sum of glind’s then gives exactly the K-FAC approximation for convolutional networks. Finally, we know from Theorem 3.4 that each glind 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.

heading: 4.2 RECURRENT NETWORKS
text: 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
zt = Wat−1 + b, (4.5)
where at−1 is an activation vector, zt 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̄) = Ex,y[vec(DW̄) vec(DW̄)>],
and the K-FAC approximation is defined as
F̂(W̄) = T (ET[āt−1ā>t−1]⊗ ET[DztDz>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 gind onW 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 gind 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.

heading: B MATHEMATICAL BACKGROUND
text: 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
(u1, v1) + (u2, v2) = (u1 + v1, u2 + v2)
c(u1, v1) = (cu1, cv1),
for u1, u2 ∈ U , v1, v2 ∈ 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.

heading: B.2 CANONICAL ISOMORPHISMS
text: 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 ei for V , there is a dual basis e∗i for V
∗ and the map ei → e∗i is an isomorphism. However, this is not canonical as it depends on the choice of basis ei for V . On the other hand, consider the evaluation map evv : V ∗ → R defined by ϕ 7→ ϕ(v) where ϕ ∈ V ∗, v ∈ V . The mapping v 7→ evv then defines a canonical isomorphism from V to its double dual space V ∗∗.

heading: B.3 AFFINE ALGEBRA
text: A setA is an affine space associated to the vector space V if there is a mappingA×A→ V denoted by (p, q) ∈ A×A 7→ ~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 {e1, . . . , en} be a basis for the associated vector space V . For any point p ∈ A, we can write ~op = ∑n i=1 xi(p)ei. Here, {x1, . . . ,xn} is a set of coordinate functions, or more simply, a basis for A. If we have two bases {x1, . . . ,xn} and {y1, . . . ,yn}, then they are related by y = Bx + c where B = [bij ] is an invertible n × n matrix and c = [ci] is a vector.
We now describe how to extend a change-of-basis on the affine space A to the product space AK = A× · · · ×A. Let ι and κ be two choices of affine bases on A, then
J(a1, . . . , ak)Kι = (ā1, . . . , āk), J(a1, . . . , ak)Kκ = (ā†1, . . . , ā†k),
where homogeneous coordinates are used for ai 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  ā † 1 ...
ā†k
 =  [B]H 0. . .
0 [B]H
  ā1...
āk  = (I⊗ [B]H)
 ā1... āk  . (B.1)
Thus, the induced change-of-basis on the product space AK is given by the matrix I⊗ [B]H .

heading: B.4 DIFFERENTIALS, PUSHFORWARDS AND PULLBACKS
text: Let M be a smooth real manifold. For p ∈ M, we denote the tangent space by TpM and the corresponding dual space, the cotangent space by T ∗pM. Given a smooth function f :M→ R, the differential df(p) ∈ T ∗pM is defined by
df(p)(Xp) = Xp(f), Xp ∈ TpM.
Let (x1, . . . ,xn) be a coordinate system around p ∈M, the differential df(p) can be expressed as
Jdf(p)K =  ∂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 ϕ : M1 → M2 of manifolds, the pushforward ϕ∗ : TpM1 → Tϕ(p)M2 is defined by ϕ∗(v)h = v(h ◦ ϕ), where v ∈ TpM1 and h :M2 → R is a smooth function onM2. If we suppose thatM1 = Rn and M2 = Rm with (x1, . . . ,xn) a coordinate system around p and (y1, . . . ,ym) a coordinate system around ϕ(p), the pushforward ϕ∗v can be represented as
Jϕ∗vK =  ∂y1 ∂x1 . . . ∂y1∂xn ... . . . ...
∂ym ∂x1 . . . ∂ym∂xn v, where v = JvK. This is exactly the Jacobian matrix Jϕ of ϕ and hence Jϕ∗vK = 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)M2 → T ∗ pM1 is defined in the following way (ϕ∗u)(v) = u(ϕ∗v), u ∈ T ∗ϕ(p)M2.
With respect to the same coordinate systems chosen above, we can write Jϕ∗uK = J>ϕu, where u = JuK. Numerically, we can compute J>ϕu efficiently using reverse mode auto-differentiation (i.e., backpropagation).

heading: B.5 METRICS AND THEIR PROPERTIES
text: 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 TpM. Recall that this is a symmetric multilinear map on the k-fold product of TpM:
σ(p) : TpM× · · · × TpM︸ ︷︷ ︸ k copies → R.
A metric onM 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 ϕ : M1 → M2 be a smooth map of manifolds and σ be a symmetric k-tensor on M2 at ϕ(p). The pullback ϕ∗σ of σ under ϕ is a symmetric k-tensor onM1 defined as
ϕ∗σ(p)(v1, . . . , vk) = σ(ϕ(p))(ϕ∗v1, . . . , ϕ∗vk),
where v1, . . . , vk ∈ TpM1. In the case of metrics, ϕ∗gM2(p)(v1, v2) = gM2(ϕ(p))(ϕ∗v1, ϕ∗v2).
If we suppose that the metric gM2 is given by GM2 for a chosen coordinate system around ϕ(p), then the pullback metric ϕ∗gM2 onM1 around p is given by
Jϕ∗gM2(p)K = J>ϕGM2Jϕ, 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 gM1 and gM2 on M1 and M2 respectively, we describe how to naturally define a metric on the product manifold M1 ×M2. For any point (p, q) ∈M1 ×M2, there is a canonical isomorphism of tangent spaces: T(p,q)(M1 ×M2) ∼= TpM1 ⊕ TqM2. 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)(M1 ×M2) can be written as a pair (v1, v2) where v1 ∈ TpM1 and v2 ∈ TqM2. Now, we define the additive metric gM1 + gM2 onM1 ×M2 as follows:
(gM1 + gM2)(p, q)(u, v) = gM1(p)(u1, v1) + gM2(q)(u2, v2). (B.2)
If we choose a coordinate system around (p, q) with the metrics gM1 , gM2 represented by matrices GM1 , GM2 respectively, then we have
J(gM1 + gM2)(p, q)K = [
GM1 0 0 GM2
] ,
which is a matrix with block diagonals GM1 ,GM2 and zero everywhere else. This construction generalizes easily to sums of more than two terms.

heading: C ADDITIONAL MATHEMATICAL TECHNICALITIES

heading: C.1 TANGENT SPACE OF VECTOR SPACES AND AFFINE SPACES
text: Theorem C.1. Let V be a finite-dimensional vector space. For each point p ∈ V , there is a canonical isomorphism V → TpV . From now on, we suppress p in TpV and write TV when denoting tangent spaces of V .
Proof. For any element v ∈ V , we can associate a tangent vector Dv|p at p defined by
Dv|pf = Dvf(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 TaA and V . From now on, we suppress a in TaA and write TA 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.

heading: C.2 TENSOR PRODUCT OF VECTOR SPACES
text: Let U and V be finite-dimensional vector spaces over the real numbers R. LetR 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 7→ ϕ(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) 7→ ϕ(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 TW → TZ and by the above isomorphism, this means a ∈ T ∗W ⊗ TZ from Eqn. C.1. By Eqn. C.2, the dual space of (T ∗W ⊗ TZ)∗ ∼= TW ⊗ T ∗Z. Using Eqn. C.1 again, elements of this space are maps λ : T ∗W → T ∗Z.

heading: D TWO PARAMETERIZATIONS OF COORDINATE-FREE MLP
text: 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 A1, . . . ,AL−1, the pre-activation spaces Z1, . . . ,ZL, the input space X and the output space Y in the network. Observe that a choice of bases for Ai−1 and Zi naturally induces a basis for each Wi, and therefore also for the full weight space W . Let ι, κ be two different choices of parameterizations for the network. With respect to ι, we write
Jαi−1Kι = ai−1, JζiKι = zi, JωiKι = (Wi bi), JρiKι = φi, JαiKι = ai, and with respect to κ, we write
Jαi−1Kκ = a‡i−1, JζiKκ = z ‡ i , JωiKκ = (W ‡ i b ‡ i ), JρiKκ = φ ‡ i , JαiKκ = a ‡ i .
Hence, we can rewrite Eqn. 3.1 in the parameterizations ι, κ as
zi = Wiai−1 + bi z ‡ i = W ‡ ia ‡ i−1 + b ‡ i ai = φi(zi), a ‡ i = φ ‡ i (z ‡ i ).
(D.1)
The parameters (Wi bi) and (W ‡ i b ‡ i ) are related as follows
Wi = ΦiW ‡ iΩi−1
bi = ΦiW ‡ iγi−1 + Φib ‡ i + τi,
(D.2)
where (Ωi−1 γi−1) is the change-of-basis from ι to κ on Ai−1 with Ωi−1 an invertible matrix and γi−1 a vector. Moreover, (Φi τi) is the change-of-basis from κ to ι on Zi 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(Φiz ‡ i + τi) + γi, (D.3) where (Ωi γi) is the change-of-basis from ι to κ on Ai. 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 Ai−1 and the pre-activation affine space Zi:
JAi−1Kι (Wi bi) //
(Ωi−1,γi−1)
JZiKι
JAi−1Kκ (W‡i b ‡ i )// JZiKκ
(Φi,τi)
OO (D.4)
Let ai−1 ∈ JAi−1Kι, this maps to zi = Wiai−1 + bi ∈ JZiKι under the top horizontal arrow in Eqn. D.4. Now, mapping ai−1 under the composition of the other three arrows in Eqn. D.4, we obtain
ai−1 7→ Ωiai−1 + γi−1 (apply left vertical arrow in Eqn. D.4) 7→W‡i (Ωiai−1 + γi−1) + b ‡ i (apply bottom horizontal arrow in Eqn. D.4)
7→ Φi(W‡iΩiai−1 + W ‡ iγi−1) + Φib ‡ i + τi (apply right vertical arrow in Eqn. D.4) = ΦiW ‡ iΩi−1ai−1 + ΦiW ‡ iγi−1 + Φib ‡ 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):
JZiKι φi // JAiKι
(Ωi,γi)
JZiKκ φ‡i //
(Φi,τi)
OO
JAiKκ
(D.5)
Let z‡i ∈ JZiKκ, this maps to a ‡ i ∈ JAiKκ under φ ‡ i . Now, mapping z ‡ i under the other three arrows in Eqn. D.5, we obtain
z‡i 7→ Φiz ‡ i + τi (apply left vertical arrow in Eqn. D.5)
7→ φi(Φiz‡i + τi). (apply upper horizontal arrow in Eqn. D.5) 7→ Ωiφi(Φiz‡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:
zi = W̄iāi−1 z ‡ i = W̄iā ‡ i−1 ai = φi(zi), a ‡ i = φ ‡ i (z ‡ i ) = Ωiφi(Φiz ‡ 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.

heading: E PROOFS

heading: E.1 PROOF OF THEOREM 3.2
text: First, note that g(ωk)−1dh(ω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 forW , the algorithm is invariant to affine reparameterizations.

heading: E.2 PROOF OF THEOREM 3.4
text: 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 ψ∗ag on W can be expressed as:
ψ∗ag(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
ψ∗ag(w) = ψ ∗ aE[φ⊗ φ]
= 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
ψ∗ag(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 TW to TZ. 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 gind(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 λ ∈ TW ⊗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 λ ⊗ µ ∈ (TW ⊗ T ∗Z) ⊗ TZ where µ ∈ TZ. From the above lemma, we have
ψ∗ag(w) = E[a⊗ φ⊗ a⊗ φ]
Then, evaluating Ea[ψ∗ag(w)] at λ⊗ µ:
Ea[ψ∗ag(w)](λ⊗ µ, λ⊗ µ) = E[a⊗ φ⊗ a⊗ φ](λ⊗ µ⊗ λ⊗ µ) = E[(a⊗ a)(λ⊗ λ)︸ ︷︷ ︸
=0
·(φ⊗ φ)(µ⊗ µ)]
= 0,
This shows that Ea[ψ∗ag] is not positive-definite which yields a contradiction as Ea[ψ∗ag] 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.

heading: E.3 PROOF OF THEOREM 3.8
text: From Theorem 3.4, we know that giind is a metric on Wi. Since gKFAC is defined as the additive metric where each of the summands are giind, we can conclude that gKFAC 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 giind is a nondegenerate metric onWi. Consequently, gKFAC is nondegenerate which concludes the proof.

heading: F COORDINATE-FREE K-FAC FOR CONVOLUTIONAL NETWORKS

heading: F.1 CONVOLUTIONAL NETWORKS
text: 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.

heading: F.1.1 CONVOLUTION LAYERS
text: We focus on a single convolution layer. A convolution layer l takes as input a layer of activations aj,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 wi,j,δ and biases bi, 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
zi,t = ∑ δ∈∆ wi,j,δaj,t+δ + bi. (F.1)
The pre-activations zi,t are then passed through a nonlinear activation function φl. Analogous to feed-forward networks, the weight derivatives are computed using backpropagation:
Dwi,j,δ = ∑ t∈T aj,t+δDzi,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 aj,t as a J × |T | matrix Al−1, the preactivations zi,t as a I × |T | matrix Zl, the weights wi,j,δ as a I × J |∆| matrix Wl and the bias vector as bl. For the activation matrix Al−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 Aexpl−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
Zl = WlA exp l−1 + bl Al = φl(Zl). (F.2)
For convenience purposes later, we adopt homogeneous coordinates for various matrices:
[Aexpl−1]H =
[ Aexpl−1
1
] , [Zl−1]H = [ Zl 1 ] , [Wl]H = [ Wl bl 0 1 ] , [Al]H = [ Al 1 ] .
Hence, Eqn. F.2 can be rewritten as
[Zl]H = [Wl]H [A exp l−1]H [Al]H = φl([Zl]H), (F.3)
where the activation function φl here ignores the homogeneous coordinate. We introduce the concept of a transformed convolution layer. For a convolution layer as defined in the above equation, the parameters [Wl]H and the transformed parameters [W † l ]H are related in the following way
[Wl]H = Γl[W † l ]H(I⊗Υl−1), (F.4)
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.

heading: F.1.2 ABSTRACT CONVOLUTION LAYERS
text: Just as in the coordinate-dependent case earlier, we focus on a single layer. An abstract convolution layer l is defined as follows:
• Local activations at each spatial location t ∈ T are taken to be elements αl−1 in an affine space Al−1.
• Activations are taken to be elements α(:)l−1 inA |T | l−1, (i.e. the direct product ofA, |T | times).
(The superscripts are meant to be suggestive of Python slicing notation.)
• Expanded activations at t ∈ T are taken to be elements α(:,t)l−1 in A |∆| l−1. The full expanded
activations are taken to be elements α(:,:)l−1 in A |∆|⊗|T | l−1 .
• Local pre-activations at t ∈ T are taken to be elements ζ(t)l in an affine space Zl. • Pre-activations are taken to be elements ζ(:)l in Z |T | l . • Layerwise parameters are affine transformations ωl between A|∆|l−1 and Zl. The collection of these transformations is an affine space in its own right which we denote byWl and refer to as the layerwise weight space. If we apply ωl pointwise, this can be extended to a map
A|∆|⊗|T |l−1 → Z |T | l .
The computation for this abstract layer is
ζ (t) l = ωl(α (:,t) l−1 )
αl = ρl(ζ (t) l ),
where ρl is a fixed nonlinear activation function and αl are the l-th layer local activations defined in exactly the same manner as αl−1.
We choose affine bases on Al−1, Zl, and Al. A basis on Al−1 naturally induces a basis for A|∆|l−1. Consequently, this gives a basis also for the layerwise parameter space Wl. Let ι, κ be two such choices. With respect to ι, we write
Jα(:,t)l−1 Kι = a (:,t) l−1 , Jζ (t) l Kι = z (t) l , JωlKι = (Wl bl), JρlKι = φl, JαlKι = al,
and with respect to κ, we write
Jα(:,t)l−1 Kκ = (a (:,t) l−1 ) ‡, Jζ(t)l Kκ = (z (t) l ) ‡, JωlKκ = (W‡l b ‡ l ), JρlKκ = φ ‡ l , JαlKκ = a ‡ l .
Note that a(:,t)l−1 are J |∆|-dimensional column vectors of the expanded activations matrix A exp l−1, z (t) l are I-dimensional column vectors of pre-activations matrix Zl and al are J-dimensional column vectors of activations matrix Al.
Now, suppose that (Ωl−1 γl−1) is the change-of-basis from ι to κ on Al−1 and (Φl τl) is the change-of-basis from κ to ι on Zl. If we denote
[Ωl−1]H =
[ Ωl−1 γl−1
0 1
] , [Φl]H = [ Φl τl 0 1 ] ,
then by the affine change-of-basis formula for direct products (Eqn. B.1), I ⊗ [Ωl−1]H defines the change-of-basis from ι to κ on A|∆|l−1. The parameters [Wl]H and [W ‡ l ]H are related as follows:
[Wl]H = [Φl]H [W ‡ l ]H(I⊗ [Ωl−1]H)
By taking Υl−1 and Γl in Eqn. F.4 to be Υl−1 = [Ωl−1]H and Γl = [Φl]H , we can conclude that a convolution layer and its transformed version simply correspond to two different choices of parameterizations for the same underlying abstract convolution layer.

heading: F.2 KRONECKER FACTORS FOR CONVOLUTION
text: We review the Kronecker Factors for Convolution method due to Grosse & Martens (2016) which is a version of K-FAC for convolutional networks. The network architecture to consider is a convolutional network with L convolution layers. First, let w be the concatenation of all trainable parameters W̄l,
w = [vec(W̄1) > vec(W̄2) > . . . vec(W̄L) >]>.
For an input-target pair (x,y), the Fisher matrix for this network is
F(w) = Ex,y[(Dw)(Dw)>],
where Dw is the log-likelihood gradient and the expectation is taken over the model’s predictive distribution Py|x(w) for y and over the data distribution for x. The diagonal blocks F(w)l,l of F(w) are
F(w)l,l = E[vec(DW̄l) vec(DW̄l)>].
We are ready now to present the K-FAC approximation for convolutional networks. For a particular layer l, we define the K-FAC approximation F̂(w)l,l to F(w)l,l as:
F̂(w)l,l = |T |(ET [ā(:,t)l−1 (ā (:,t) l−1 ) >]⊗ ET [Dz(t)l (Dz (t) l ) >]), (F.5)
where ā(:,t)l−1 is the homogeneous notation for a (:,t) l−1 . The K-FAC approximation F̂(w) to F(w) is the matrix with diagonal blocks F̂(w)l,l as given above and zeros everywhere else. Remark F.1. Unlike MLPs where K-FAC is derived from assuming only the statistical independence of activations and pre-activation derivatives, convolution layers admit weight sharing and additional assumptions are necessary to derive the approximation F̂(w)l,l in Eqn. F.5. We refer to Grosse & Martens (2016) for extensive details on how these approximations are derived.
We devote the remainder of this section to deriving invariance properties of K-FAC for convolutional networks. We refer the reader to Grosse & Martens (2016) for other aspects of the K-FAC algorithm on convolutional networks, such as implementation details and experimental results.

heading: F.3 COORDINATE-FREE K-FAC FOR CONVOLUTIONAL NETWORKS
text: We begin by considering an abstract convolutional network with L convolution layers. Let X and Y denote the input and output spaces of this network respectively. Recall that the layerwise weight space Wl is the space of affine transformations between A|∆|l−1 and Zl. The weight space of this network is the direct product of all layerwise weight spaces
W =W1 × · · · ×WL.
Given an input ξ ∈ X and parameter ω = (ω1, . . . , ωL) ∈ W , denote the network output by f(ξ, ω). Now, for every l ∈ {1, . . . , L}, define the following maps
• ψlξ :Wl → Z |T | l which sends layerwise parameters ωl to pre-activations ζ (:) l by evaluating
local activations αl−1 across every spatial location t ∈ T
• ϕlξ : Z |T | l → Y which sends pre-activations ζ (:) l to f(ξ, ω)
Again, ψlξ is trivially a smooth map from its definition. The map ϕ l ξ includes all operations in convolutional networks such as max-pooling and response normalization. We make an assumption here that all these operations are smooth. (While this is not the case for common operations such as ReLU and max-pooling, we conjecture that the non-smooth case can be addressed by taking limits of smooth functions.) Finally, we define the map Ψlξ :Wl → Y as the composition Ψlξ = ϕlξ ◦ ψlξ.
Let g be a metric on Y and consider the pullback (ϕlξ)∗g restricted to a single spatial location t which we denote by (ϕlξ,t)
∗g. More concretely, this metric is computed by assuming components of the tangent vector at all other spatial locations are zero. Now, let us take A, Z, W in Section 3.1 to be
A = A|∆|l−1, Z = Zl, W =Wl,
and the metric on Z = Zl to be (ϕlξ,t)∗g. Summing over every spatial location t ∈ T , the independence metric onWl here is
glind(ωl) = |T |(ET [α (:,t) l−1 ⊗ α (:,t) l−1 ]⊗ ET [(ϕ l ξ,t) ∗g(ζ (t) l )]). (F.6)
Definition F.2. The K-FAC metric on the weight spaceW of an abstract convolutional network is defined as
gKFAC(ω) = g 1 ind(ω1) + · · ·+ gLind(ωL),
where the sum above is as defined in Eqn. B.2 and each glind is as given in Eqn. F.6.
Theorem F.3. Let g be a metric on Y . Then, gKFAC given in Definition F.2 is indeed a metric on the weight spaceW of an abstract convolutional network. Moreover, if we assume that the expected pullback of g restricted to a single spatial location t ∈ T ,
Eξ[(Ψlξ,t)∗g],
under the map Ψlξ,t : Wl → Y is a nondegenerate metric on the layerwise weight space Wl for every l, then gKFAC is also nondegenerate.
Proof. The proof of this theorem mirrors the proof given earlier for Theorem 3.8. By Theorem 3.4, we know that
E[α(:,t)l−1 ⊗ α (:,t) l−1 ]⊗ E[(ϕ l ξ,t) ∗g(ζ (t) l )],
is a metric onWl. Since taking expectation over the set of spatial locations T and multiplying by the scale factor |T | preserves the metric properties, we obtain that glind in Eqn. F.6 defines a metric onWl. Consequently, gKFAC determines a metric onW . To prove the latter assertion, note that by the functorial property of pullback operations,
Eξ[(ψlξ,t)∗((ϕlξ,t)∗g)]
is nondegenerate. Using the second assertion of Theorem 3.4 yields that glind is nondegenerate which implies that this is true also for gKFAC.
We conclude this section with the proof of Theorem 4.1. Our proof is coordinate-free and given in exactly the same manner as the proof of Theorem 2.1 at the end of Section 3.
Coordinate-Free Proof of Theorem 4.1. As shown earlier in Section F.1.2, each convolution layer of N and N † correspond to two different choices of parameterizations for the same underlying abstract convolution layer. Hence, N and N † must compute the same function. Assume that the metric g on the output space Y in Theorem F.3 is the output Fisher metric gF,out. For each spatial location t ∈ T , the pullback under ϕlξ,t is
(ϕlξ,t) ∗gF,out(ζ (t) l ) = E[dLζ(t)l ⊗ dLζ(t)l ].
Now, choose coordinate systems on Al−1 and Zl. This induces coordinates for A|∆|l−1 and we write
Jα(:,t)l−1 K = a (:,t) l−1 , Jζ (t) l K = z (t) l , J(ϕlξ,t)∗gF,out(ζ (t) l )K = JE[dLζ(t)l ⊗ dLζ(t)l ]K = E[Dz (t) l (Dz (t) l ) >].
Using Proposition 3.5, the independence metric in Eqn. F.6 can be expressed in coordinates as follows
Jglind(ωl)K = |T |(ET [ā (:,t) l−1 (ā (:,t) l−1 ) >]⊗ ET [Dz(t)l (Dz (t) l ) >]),
which is exactly F̂(w)l,l given earlier in Eqn. F.5. Furthermore, JgKFAC(ω)K is exactly the KFAC approximation F̂(w). Thus, the K-FAC update rule is simply a natural gradient update rule with respect to the K-FAC metric gKFAC for abstract convolutional networks. Lastly, if gKFAC 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.

heading: G COORDINATE-FREE K-FAC FOR RECURRENT NETWORKS
text: 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.

heading: G.1 RECURRENT NETWORKS
text: 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.

heading: G.1.1 COMPUTATIONAL STEP
text: 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 xt at every t, the recurrent network maps this to an output ot. Essentially, the network maps input sequences x = (x1, . . . ,xT ) to output sequences o = (o1, . . . ,oT ). The computation, at every t, is
zt = Wat−1 + b z′t = zt + Vxt at = φ(z ′ t)
where at−1 is an activation vector, zt 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
zt = Wat−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.

heading: G.1.2 ABSTRACT RECURRENT NETWORK
text: 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 AT
• 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 ZT
• 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 YT 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.

heading: G.2 K-FAC FOR RECURRENT NETWORKS
text: 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
zt = W̄āt−1,
where W̄ = [W b] and ā>t−1 = [a > t−1 1] >. Using backpropagation, the log-likelihood gradient is given by Dztā>t−1. The total contribution to the gradient across all t is the sum
DW̄ = T∑ t=1 Dztā>t−1.
For an input-target pair (x,y), the Fisher matrix F(W̄) for recurrent networks is defined as
F(W̄) = Ex,y[vec(DW̄) vec(DW̄)>].
Finally, the K-FAC approximation F̂(W̄) to F(W̄) for recurrent networks is defined as
F̂(W̄) = T (ET[āt−1ā>t−1]⊗ ET[DztDz>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.

heading: G.3 COORDINATE-FREE K-FAC FOR RECURRENT NETWORKS
text: 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 ϕ∗ξ,tg 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 ϕ∗ξ,tg. Summing over all time steps t ∈ T, we make the following definition which arises from the independence metric in Section 3.2:
Definition G.2. The K-FAC metric on the weight space W of an abstract recurrent network is defined as gKFAC(ω) = T (ET[αt−1 ⊗ αt−1]⊗ ET[ϕ∗ξ,tg(ζt)]). (G.4) Theorem G.3. Let g be a metric onY . Then, gKFAC given in Definition G.2 is a metric on the weight spaceW of an abstract recurrent network. Moreover, if we assume that the expected pullback of g,
Eξ[Ψ∗ξ,tg],
under the smooth map Ψξ,t :W → Y is a nondegenerate metric, then gKFAC is also nondegenerate.
Proof. The proof of this theorem is analogous to the proofs of Theorem 3.8 and Theorem F.3. From Theorem 3.4, we have that E[αt−1 ⊗ αt−1]⊗ E[ϕ∗ξ,tg(ζt)] is a metric on W . Since this remains true after taking expectation over the set of time steps T and multiplying by the scale factor T , we can conclude that gKFAC is a metric on W . For the nondegeneracy statement, using the functorial property of pullbacks, we know that
Eξ[ψ∗ξ,t(ϕ∗ξ,tg)]
is nondegenerate. Then, gKFAC is nondegenerate by the second assertion of Theorem 3.4.
Coordinate-Free Proof of Theorem 4.2. We conclude this paper with a coordinate-free proof of Theorem 4.2. As mentioned at the end of Section G.1.2, N and N † correspond to two different choices of parameterizations for the same underlying abstract recurrent network and so they must compute the same function.
Assume that the metric g on Y in Theorem G.3 above is the output Fisher metric gF,out. Then, the pullback under ϕξ,t is ϕ∗ξ,tgF,out(ζt) = E[dLζt ⊗ dLζt ]. Now, choose coordinate systems for A and Z . We can write
Jαt−1K = at−1, JζtK = zt, Jϕ∗ξ,tgF,out(ζt)K = E[DztDz>t ]. By Proposition 3.5, the K-FAC metric in Eqn. G.4 can be represented in these chosen coordinates as
JgKFAC(ω)K = T (ET[āt−1ā>t−1]⊗ ET[DztDz>t ]),
which is exactly the K-FAC approximation F̂(W̄) in Eqn. G.3. Thus, the K-FAC update rule is a natural gradient update with respect to the K-FAC metric gKFAC for abstract recurrent networks. If we suppose that gKFAC is nondegenerate, then these updates are invariant to any affine reparameterizations of the model.

Review summaries:
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.

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.