A Coordinate-Free Construction of Scalable Natural Gradient
Authors: Kevin Luk
Roger Grosse
Venue:  (2020)
Keywords: Natural gradient
second-order optimization
K-FAC
parameterization invariance
deep learning

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.

Under review as a conference paper at ICLR 2020
A COORDINATE-FREE CONSTRUCTION OF SCALABLE NATURAL GRADIENT
Anonymous authors Paper under double-blind review
ABSTRACT
Most neural networks trained using first-order optimization methods, which are sensitive to the parameterization of the model. Natural gradient descent is invariant smooth reparameterizations because it is defined coordinate-free way, but tractable approximations are typically defined in terms of coordinate sys- tems, and hence may lose the invariance properties. We analyze the invariance Approximate Curvature (K-FAC) algorithm properties of the Kronecker-Factee coordinate-free way. We explicitly construct a by constructing the algorithm Riemannian metric under which the natural gradient matches the K-FAC update; invariance affine transformations of the activations follows immediately. We extend our framework analyze the invariance properties of K-FAC appied convolutional networks and recurrent neural networks, as well as metrics other than the usual Fisher metric.
INTRODUCTION
Most neural networks are trained using stochastic gradient descent (SGD) (Bottou & Bousquet 2007). variants thereof which adapt step sizes for individual dimensions (Duchi al.) [001]
Kingma & Ba 2013). One well-known deficiency of SGD is that the updates are sensitive to the pa- rameterization 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 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]
loffe & Szegedy 2015). Such tricks can lead to large improvements in the speed of optimization.
Ideally, one would like to use an optimization algorithm which invariant to such transformations of 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 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 & Salakhut- dinov 2013 Martens & Grosse 2013 Desjardins et al. Schhllann et al.) 2015). The challenge that the exact natural gradient is impractical 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 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 diag- onal matrix with one block per layer of the network, and each block factorizes the Kronecker product of small matrices. Because of the Kronecker structure, the approximate natural gradient can be computed with low overhead relative ordinary SGD; K-FAC demonstrated significant

---

Under review as a conference paper at ICLR 2020
classification convolutional net- works (Grosse & Martens] 2016 Ba et al. 2017). recurrent networks (Martens et al.) 2018 and deep reinforcement learning (Wu et a.. (017). The same Fisher matrix approximation has also led to sig- nificant improvements in modeling posterior uncertainty in Bayesian neural networks (Ritter et al.) 2018h Zhang al. 2017) and avoiding catastrophic forgetting (Ritter al. (0188). Most recently, K-FAC was used to train a neural network to solve the many-electron Schrödinger equation (Plau
speedups in training deep autoencoders Martens & Grosse 2015),
al. 2019).
Although K-FAC does not satisfy the general invariance properties of natural gradient, it is still invariant to broad and interesting class of reparameterization:: affine transformations of the acti- vations in each layer (Martens & Grosse (0015). This was verified through linear algebraic manip- ulation 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 recipe to convert metric on W (whose coordinate representation 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 in- variant arbitrary smooth reparameterizations, but only up the first-order. By contrast, we show global invariance to 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 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 reparameterizatioos,, such as whitening, normalization, and replacing the logistic activation function with tanh. And in exchange for imposing the affine structure, obtain global invariance, not just first-order invariance.
Our framework easily enables some generalizations of the basic result. Our construction is not limited Fisher metrics, but rather applies to general pullback metrics, where one places metric on the network's 's output space and pulls it back to W. Furthermore, we extend the invariance results convolutional networks and recurrent networks through 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-ing indinattt-ddeeedeent and coordinate- dependent 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.
BACKGROUND
2.1 NATURAL GRADIENT DESCENT
For simplicity, 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 training set
2

---

Under review as a conference paper at ICLR 2020
S of pairs Ys), aim to minimize the empirical risk which is given by:
h h(w)
1
Σ
Cly (x,, w)).
(2.1)
S]
(x..)))S
where C(yi: fx,, w)) is the loss function measuring the disagreement between y; and f(;,, w).
Suppose that f(x, w) determines parameters z of the model's predictive distribution Ryjx over and furthermore, we reparameterize this as Pyix( (w) = fx) Likewise, the density function of this distribution can be reparameterized as p(y [x, w) = r(ylf(x, w)). In addition, we take the loss function here to be the negative log-likelihood [(y, z) = -log r(y/z) and denote log-likelihood gradients V C(y, f w)) by Dw (D notation throughout remainder of the paper refers to log- likelihood gradients). The Fisher information matrix F(w) is defined
F(w)= = Ex (PDw)))(
(2.2)
where the expectation is taken over (w) for y and over the data distribution for Since F(w) is defined as the expectation of an outer product, F F(w) is always guaranteed to be positive- semidefinite (PSD) matrix.
The natural gradient of the objective function l(w) in Eqn. 2.1 is F(w) -Wh))-VVWWww For chosen learning rate 6 > O, the natural gradient descent algorithm (Amari 1998) minimizes l(w) by using
the natural gradient update parameters of the network:
Wk+1 < W: - =F((w))
(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-fit free construction in terms of differential geometric primitives. We give detailed explanation of this by considering abstract mathematical setting; we defer to Appendix for background on the mathematical terminology with which we use.
Let M be Riemannian manifold with nondegenerate metric given by (in this paper, our definition of "metrics" allows the possibility of it being degenerate). For a smooth function h : . + R and a point p M, the differential dp(p) is abstract covector on M. We use the nondegenerate metric to convert the covector dh(p) into a tangent vector. By definition, g(p) is nondegenerate bilinear form which yields linear isomorphism between the tangent space and the cotangent space:
(p)) Ty.M $ M.
This is commonly referred to as the musical isomorphism in mathematical literature. The inverse g(p)-1 gives linear map the other way around,
g(p)~ -1 T.M T.M.
Applying this isomorphism to dh(p) yields the tangent vector -'d(pp) E TpMM 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 w and (g, v) be intrinsic versions of the parameter vector w and the input-target pair (x respectively. & and belong to abstract input and output spaces which are also smooth manifolds. Now, W can be endowed with the Fisher metric ge which is defined as
gr (w) = E.g. @ dCw),
(2.4)
where L. = log (() is the abstract log- likeliiel loss function. The expectation above taken over the abstract predictive distribution (w 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 a Riemannian manifold with the Fisher metric gr, the idealized gradient descent updates are given by Bonnabel et al. (0033)
Wk+1 Expry (-egr(wk) -
(2.5)
3

---

Under review as a conference paper at ICLR 2020
where Exp Tw, W W is the exponential map. This update rule is exactly invariant to all smooth reparameterizations of W since it entirely coordimate-fres However, such an algorithm is infeasible in practice computing the exponential map is typically an intractable problem. Instead, is much easier to work with the following abstract natural gradient update rule which uses a first- order approximation of the exponential map
Wk++1 (rggrwwx>>ditie
Since this is first-order approximation of the update rule in Eqn. 2.5 invariance smooth repa- rameterizations holds only up to first-order. Additional approximations to the exponential map are necessary to obtain higher-order invariances; defer to |Song & Ermon (2088) for more detailed account of how this can be done.
2.2 SAAECCKKER-FCTORRE APPROXIMATE CURVATURE
We consider a MLP with layers. each layer E {l,..., L], the MLP computation is given as:
z: W a, b
a: (z:).
where a;-1 is an activation vector, z, is pre-activation vector, W. is a weight matrix, b, is a bias vector, and d4 IR - is activation function. For convenience, introduce homogeneous coordinates __
and W, -W; b,]. The above computation can be rewritten as
2: W
(2.6)
(z,).
We concatenate all of the network parameters W; into single vector w,
w = [vec(W]) vec(Wz)
Here, vec denotes the vectorization operator which stacks the columns of matrix together to form vector. The Fisher matrix for the MLP is L x L block matrix F F(w) where each (i,h) block is
F( ((w) J.. (Dv,)(DW,) veo(DW,))
We now give an overview of the Kroneeker-Facttreed Approximate Curvature (K-FAC) method (Martens & Grosse 2015) of approximating the Fisher matrix. Consider the diagonal (i, blocks of F(w). Using backpropagation, the log-likelihooo gradient DW, = Dz,a,_ and hence, we have vec(Dz,,a,1) = a @ Dz,. Then, F(w),,s can be rewritten as:
F(w),,: = (evvvcDDW.)) veo(DW.) .) "] 1 ( (ā.-1 "] E -aa-ja__ &D, Dz,
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),, F F(w)....
F (w)),
Ela,-
X
E [Dz
(2.7)
:
The K-FAC approximation matrix F(w) to F(w) is defined as
Ao @ Gi A &G
0
F(w)=
(2.8)
A--1 GL
The K-FAC update rules are analogous to natural gradient descent updates; we simply replace
F in Eqn. 2.3 by ((wx) 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, consider the class of transformations given by the following transformed network
z Wia!1
(2.9)
a (z si (,0 + T.) + Yt
4

---

Under review as a conference paper at ICLR 2020
where S2, Φ, invertible matrices and Ts. 7 vectors. The transformed input is aj Sgas where
Go = So
and the transformed output is at, = (xt, wT) with w defined as
w [vec W (TT vec(WW))
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 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 {o: Suppose that we have activation functions {o] Ji as given in Equ. 2.9 Then, there exists a pa- rameter vector w| such that the transformed network A with parameter vector wi and activation functions 1- computes the same function as N. Furthermore, if the Fisher matrix of the MLP is assumed to be positive-defiite, 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 cum- bersome and involved tedious manipulation of update rules. 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 9KFAС whose coordinate representation exactly matches the K-FAC approximation given in Eqn. 2.8 A K-FAC update can then be viewed as natural gradient update with respect to the metric OKFAC- More importantly, the invariance properties the K-FAC algorithm are immme- diately established in the same way as it was for exact natural gradient.
3 COORDINATE-PRRT K-FAC
We observe that MLPs consist of a sequence of affine transformations and activation functions in alternation. In order to capture this structure, 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.] in that subsection, 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- MLP formally. For i (1,..., L), we have
Activations are taken to be elements o;-1 in an affine space AA-1-
Pre-activations are taken be elements & in an affine space Z,.
Layerwise parameters are affine transformations wi between A;- and 2,. The collection these transformations is an affine space own right, which we denote by W, and refer to the layerwise weight space.
The weight space is given by the direct product W = W1 x ... x Wr. Elements this space are written as (Wi,..., Wy.) W.
Input and outputs are denoted by & and f(E, w) respectively. The space of all inputs and outputs affine spaces denoted by x²[= An) and D(= AL) respectively.
Moreover, the layerwise computation is given by = wi(a;=1)
(3.1)
a; = P:(Gi,
where Ps IR + R is fixed nonlinear activation function which is assumed to be smooth throughout. In Appendix D 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 MI.P
5

---

Under review as a conference paper at ICLR 2020
Remark 3.1. The smoothness condition the activation functions P: is necessary ysis 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 func- tion not smooth by definition, is locally smooth with probability 1 if the weights are sampled from continuous distribution. Furthermore, the ReL.U function is limit of the softplus functions 1 log(1 + ex which are all individually smooth.
later anal-
Optimization. The optimization problem in the abstract setting is analogous to the coordinate- dependent one. Let (ç.v) be an abstract input-target pair and L(v. f(E,W) be the loss functioon measuring the disagreement between outputs f(ç. w) of the abstract MLP and targets v. Given training set of abstract input-target pairs (Ea, 14), the objective function we wish to minimize here is
h(w)
Σ (ELRA)ES
f(i,
S
The following theorem (see Appendix E. for the proof) shows that by imposing affine structure on the weight space, we can obtain global invariance instead just first-order invariance.
Theorem 3.2. Let g be nondegenerate metric on the weight space W of abstract MLP. For chosen learning rate > O, the following update rule
Wk+1 Wk -g(wu)- "hh(w))-d(((w))
exactly invariant to all affine reparameterizzations of the model.
Pullback of Output Metrics to Weight Space. Consider metric on the output space ) of the MLP. Let Ve W > y the smooth map which sends parameters w to outputs ar [(E, w) given an input g. The pullback V2g defines a metric on W. The expected pullback metric over inputs, under a choice of coordinates around w and az, is given by
[E [Be Ex [J GJ
(3.2)
where G is the representation of 9 in these coordinates.
Example 3.3 (Fisher metric). Suppose that the outputs as parameterize the model's predictive distribution Rejoz. Let r(vlar.) denote the density function of this distribution and furthermore, take the loss function here to be the negative logelikeld far. log (vlaa)). The output Fisher metric out on y is defined as
JF (F = [dCa & dhorl.
where the expectation taken with respect the predictive distribution Ryjor Computing the expectation of OF over the inputs & gives
Bel (w)]
= E. @ dCaull
IE (ddCaa @ dCar)l
LE ( 8 (d(Caa))
LE IdC 2 dCu).
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 V, then Eqn. 3.2 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 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 convex functional on v, 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 approx- imation. Recall that K-FAC makes two approximations obtain tractable Fisher matrix: (1) it

---

Under review as a conference paper at ICLR 2020
assumes independence of activations and pre-activation derivatives in order to push the expectation inside the Kronecker product (Eqn. (.7)). and (2) it keeps only the diagonal blocks corresponding to individual layers. In this section, develop a coordinate-fin 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 develop approximate metrics for MLPs. In Section [3.2] 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 MLPs, we consider the more general setting of metrics on affine maps between affine spaces, but use notation which is suggestive of Ml.Ps. We assume the following:
Affine spaces and Z
Affine space W of affine transformations between and Z
Metric g Z
Our first task is to formulate a coordinate-free analogue of the outer product of homogenized acti- vations, a.a." Consider the evaluation map vla - which defined by evaluating w E W at E A. We compute the pushforward Vor TW + TZ. Note that there is need to specify par- ticular points for the tangent spaces here since we are working with affine spaces (see Corollary C.2
in Appendix C). Let 0n be tangent vector W and f be a smooth function on Z. Then,
Vas (Ow) = 0m o vhw)(w)
(0 ((baf((b((w) 2% (w)
()) a
This shows that the pushforward Va. is exactly multiplication by the element Hence, we identify any a E with its linear map TW - TZ. Thus, this enables us to define the tensor product of two elements in A mapping a @ 02 TW TW TZ & TZ:
(a1 X (0w; Ow,)) 01 (8w,)) 2 02 (Owz).
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 be
Iind (w) Ela @ a) B[g(=)].
(3.3)
where the first expectation is over A and the second one is over Z. Note that B[g(=)] 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 Jind is indeed a metric on W.
Theorem 3,4. Let be metric on Z and Via : W + Z be the evaluation map. Then, Jind as defined in Eqn. 3.3 is metric on W. Moreover, if the expected pullback metric [Vagg nondegenerate on W. then Jind is also nondegenerate. From now on, we refer to Jind 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 the affine spaces A, Z and in these coordinates,
[a]
[2] z, [[{())] -G(z.
Then the independence metric Gind can be expressed as
[y E aa 1 (())))
Remark 3.6. In the context of MLPs (which we explain in greater detail subsequently) where A = Ar-1, Z Z, and W W. the matrix G(2) Ey [Dz,Dz,1 where the expectation is taken over output space. [[)(2) means we furthermore take the expectation over Z, When we write E(D2, Dz, in Eqn. 2.7 implicitly take this mean E[G(z)).

---

Under review as a conference paper at ICLR 2020
3.2 K- K-FAC METRIC
In this subsection, we formulate the layerwise independence approximation in a coordinate-free way, allowing
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 Sec- tion 3.1 to define the K-FAC metric for MLPs. Lastly, by viewing K-FAC a metric on W, we show how invariances of the K-FAC algorithm can be obtained in natural and straightforward manner.
Consider MLP with L layers as described earlier in Section For every i the following maps
L., we define
V : W; + 2, which sends layerwise parameters w; pre-activations ç by evaluation at activations a:-1-
p/ 2; + y which sends G to network outputs az = (()
The subscript & is used to highlight the fact that all these maps have an implicit dependence on network inputs Ç. Note that v/ is smooth map by definition. Now, observe that go is exactly the composition network maps
PL
wi+1 Pr 2; +
Since all activation functions Pr assumed to be smooth maps, it follows immediately that 9o is also a smooth map. Moreover, we define the map V W; - y to composition V = A o v
Let be metric on V. Then, the pullback (+p)) defines metric on Z, Now, if we take A, Z, W in Section 3.1 to be
A = A,-1, Z = 2,, W = W; and the metric on Z 2, to (vp)) g. the independence metric on W; here
(ws) = Ela;-1 @ a4-1] @ El(vE)" g(()])
(3.4)
Now, given metrics Hind on each Wi, there is natural "additive metric" defined on the Cartesian product W = Wi x ... x Wr. 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
IKFAC (w) = Dina (w1) + +... (wn(
where the sum is defined in Eqn. B.2 in Appendix B.S and each glnd is as given in Eqn. 3.4
Theorem 3.8 (See Appendix E.3 for proof). Let g be a metric on V. Then, JKFAC given in Defini- tion 3.7 is indeed metric the weight space W of an abstract MLP. Moreover, if assume that the expected pullback of g.
[V((( g],
under the map V i, then gKFAC is also nondegenerate.
y nondegenerate metric on the layerwise weight space Wafor every
Coordinate- 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 A 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 on the output space y in Theorem 3.8 the output Fisher metric guo in Example 3.3 The pullback of this under &ó given by
(vt) out (Gi) EldGç, @ dC
Let us choose coordinate systems on A and Z, with
[a:-1] a;-1, [Ga] = Zi, [((t)) out (G:) [Edde dC E(D2,D2]]
Then, by Proposition 3.5
Ioina (w:)] E Ela;
E D2,
8

---

This is Under review as a conference paper at ICLR 2020
which is exactly F F(w),, given earlier in Eqn. 2.7 Furthermore,
[ΦKFAс[W]] = -no + ... Find (wr))
is the matrix with diagonal blocks F(w),,) and zeros everywhere else (see Appendix B.5 precisely F (ww in Eqn. 2.8 Now, observe that
B.5
[UKFAA(())
- ([XFFAc()""I[(hlaa)
F(w) Vh((w,
and hence the K-FAC update rule is simply a natural gradient update rule with respect the K-FAC metric 9KFAC for abstract MLPs. Suppose that 9KFAC is nondegenerate metric; which is true for example the assumptions the second assertion of Theorem [.8 hold. Applying Theorem 3.2 shows that this update rule is invariant to any affine reparameterizations of the model.
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 Appendices F and for full details.
4.1 CONVOLUTIONAAL NETWORKS
We focus on a single convolution layer following the exposition given in Grosse & Martens 2016 Let J be the number of input maps, be the number of output maps, A the spatial offsets and T be the number of spatial locations. We write the convolution layer as matrix multiplication
Z W,A, - + b
(4.1)
A, (Z,),
where A is the J| x [T| matrix of expanded activations, Wi is the / x J/A/ weight matrix, Zi is the IT matrix of pre-activations, di is nonlinear activation function and by is the bias vector. Adopting homogeneous coordinates for the various matrices, we can rewrite the above
[Z]H = [W [AS
(4.2)
[Ai] = (Z.l))
We briefly introduce the concept of transformed convolution layer. For a convolution layer as defined in the above equation, the parameters [Wi] and the transformed parameters [W]] are related in the following way
Wil D (W/ T--1).
(4.3)
where F2 and Ti are invertible matrices. The activation functions o and of are related through standard affine change-of-basts given in Eqn. 2.9
For convolutional network with L convolution layers, the Fisher matrix is given by
F(w)= ((Dw)
where w is the vector concatenating all convolution layer weights. The K-FAC approximation defined by setting all off-diagonal blocks to be zero and approximating the diagonal blocks as
F F(w),. [ |a
E, (Dz)" (4)
(D22))
(4.4)
Here, a stands for the expanded activations at each spatial location / E T (the J Alddimmenso column vectors of the matrix A in Eqn. 4.1) and z/" stands for the local pre-activations at E T (the 7-dimensional column vectors of the matri Z, in Eqn. 4.1)

---

Under review as a conference paper at ICLR 2020
Theorem 4.1. Let N be convolutional network with parameter vector and activation func- tions (or)fer- Suppose that we have activation finctions {o) which are related to f6t) by standard change-of taagg--o-bbsis-- transformations. Then, there exists parameter vector w such that the transformed network N with parameter vector w and activation functions {dr){=1 computes the same function as N. Furthermore, suppose that the Fisher matrix for the convolutional network is the K-FAC updates are equivalent, in the sense that the resulting networks compute the same function.
Proof Sketch (Full proof in Appendix F33 A coordinate-dependee proof of this result was given previously in Grosse & Martens (2016). Instead, we prove this intrinsically using the framework provided in Section Let A be the affine space of expanded pre-activations at T, Z₁ be the affine space of local antevactittions at & E T, and W; be the layerwise weight space comprising of affine transformations between A and Z,. Now, let us take A, Z, W in Section 3.1 to be A 21 W/ respectively and the metric on Z = Z: to be the pullback of the output Fisher metric. After summing over every spatial location / E T, the independence metric Fnnd on W: coordinates, by Proposition 3.5 is exactly the expression given Eqn. 4.4 Taking the sum of gind then gives exactly the K-FAC approximation for convolutional networks. Finally, we know from Theorem 3.4 that each g'nd indeed defines a metric and so by Theorem 3.2 we can conclude that the K-FAC updates are exactly invariant to all affine reparameterizatioos of the model.
4.2 RECURRENT NETWORKS
As in the case of convolutional networks, need not write out the full structure of 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....,7) We use I to index the time step. Additionally, we assume that all sequences are of fixed length T. We consider the recurrent computation step
Z Wa
(4.5)
where a is an activation vector, z, is a pre-activation vector, and W is a recurrent weight matrix. The transformed recurrent computation step is defined as
z w a|_ +b.
(4.6)
The relationship between transformed parameters (WT bt) and original parameters (W b) is given
by the standard change-of- -aasis formula as in Eqn. D.2
The Fisher matrix for recurrent networks is defined
F F(W)= Ex.y/
and the K-FAC approximation is defined as
F (W) = ((E 2 Er Dz,,
(4.7)
Theorem 4.2. Let N be recurrent network with recurrent parameters (W b), Suppose that we have recurrent network Nt with recurrent parameters (Wi bi) and the relationship between (W b) and (Wt bt) is standard change-of basis transformation. Then, the networks N and N't 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.
Let A be the affine space of local activations at each time step f, Z be the affine space of local pre-activations at each time step f, and W be the parameter space consisting of all affine transformations between and Z. Now, let us take A, Z, W in Section 3.] to be A, Z, W respectively and the metric Z = Z to be the pullback of the output Fisher mertric. After summing over all time steps E T, the independence metric Iind 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 gind indeed defines a metric and so by Theorem 3.2 we conclude that the K-FAC updates exactly invariant all affine reparameterizations of the model.
Proof Sketch (Full proof in Appendix
10

---

Under review as a conference paper at ICLR 2020
REFERENCES
S. Amari and H. Nagaoka. Methods of Information Geometry. American Mathematical Society and Oxford University Press, 2000.
S. Amari, H. Park, and Fukumizu. Adaptive method of realizing natural gradient learning for multilayer perceptrons, Neural Computation, 2000.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251- 276, 1998.
Ba, R. Grosse, and J. Martens. Distributed second-order optimization using Kronecker-factored approximations, In ICLR, 2017.
Silvere Bonnabel et al. Stochastic gradient descent riemannian manifolds. 2013.
L. Bottou and O. Bousquet. The tradeoff's of large scale learning. In NIPS, 2007.
K. Cho, T. Raiko, and A. Ilin. Enhanced gradient and adaptive learning rate for training restricted Boltzmann machines. In ICML, 2011.
Guillaume Desjardins, Karen Simonyan, Razvan Pascanu, et al. Natural neural networks. In Ad- vances Neural Information Processing Systems, pp. 2071-2079 2015.
Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. JMLR, 2011.
X. Glorot and Y. Bengio. Understanding the difficulty of training deep feedforward neural networks. In AISTATS, 2010.
R. Grosse and Martens. A Kronecker-factored approximate Fisher matrix for convolution layers. In ICML, 2016.
R. B. Grosse and R. Salakhutdinov. Scaling up natural gradient by sparsely factorizing the inverse Fisher matrix. In JCMI., 2015.
S. loffe and C. Szegedy. Batch normalization: accelerating deep network training by reducing internal covariate shift. In ICML, 2015.
D. Kingma and J. Ba. Adam: a method for stochastic optimization. ICLR, 2015.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
John M Lee. Smooth manifolds. In Introduction to Smooth Manifolds, pp. 1-29. Springer, 2003.
J. Martens. Deep learning via hessian-free optimization. In /CML, 2010.
James Martens. New insights and perspectives the natural gradient method. arXiv preprint arXiv: 1412. 193, 2014.
James Martens and Roger B Grosse. Optimizing neural networks with kronecker-factored approxi- mate curvature. In ICML, 2015.
James Martens, Jimmy Ba, and Matt Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In ICLR, 2018.
G. Montavon and K.-R. Müller. Deep Boltzmann machines and the centering trick. In Neural Networks: Tricks of the Trade. Springer, 2012.
David Pfau, James S Spencer, Alexander G de Matthews, and WMC Foulkes. Ab-initio solu- tion of the many-clectron schri." odinger equation with deep neural networks. arXiv preprint 1XXi:19999002887. 2019.
T. Raiko, H. Valpola, and Y. LeCun. Deep learning made easier by linear transformations in percep- trons. In AISTATS, 2012.
11

---

Under review as a conference paper at ICLR 2020
Hippolyt Ritter, Aleksandar Botev, and David Barber. Online structured laplace approximations for overcoming catastrophic forgetting, arXiv preprint
2018a.
Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. 2018b.
N. Le Roux, P.-.. Manzagol, and Bengio. Topmoumoute online natural gradient algorithm. In NIPS, 2008.
Nicol N Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 14(7):1723-1738, 2002.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region
policy optimization. In International Conference on Machine Learning, pp. 1889-1897, 2015.
Yang Song and Stefano Ermon. Accelerating natural gradient with higher-order invariance. arXiv preprint
2018.
Yuhuai Wu, Elman Mansimov, Shun Liao, Roger Grosse, and Jimmy Ba. Scalable trust-region method for deep reinforcement learning using kronecker-factored approximation. arXiv preprint arXiv:1708.055 14,, 2017.
Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. arXiv preprint 17Xiv:1712.02399, 2017.
12

---

Under review as a conference paper at ICLR 2020
A
TABLE OF NOTATIONS
Coordinate-indepeneenl
Coordinate-dependent
inputs
X
targets
y
activations
homogenized activations
a,
are-activations
layerwise parameters (and biases)
Wi
(W; b,)
homogenized layerwise parameters
W
activation functions
dz
network parameter
w
network output
s(E,w) ((,,((E,W) h(w)
f(x, w)
loss function
C(y, S(y, w))
objective function
h(w)
predictive distribution
Pate(w)
Pyl((W)
predictive distribution density function
p(ulE,w)
plyix, w)
Fisher metric/matrix
gF(w)
F(w)
layerwise log-likelihood differential/graddent
dhw.
DW;
log-liked likelihood differentiallgradientt
dew
Dw
K-FAC metric/matrix
[KFAC(W)
F(w)
tensor product/Kroneckee product input space output space
60
X
y
space of activations space of arevactivssions layerwise weight space weight space of network
A
Z
W
W
Table I: Since this paper involves the interplay of many pooddinate-nndpendeen and coordinate- dependent 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 general rule of thumb, we use boldface to symbolize coordinate-dependeet objects and standard math font for coordinate-indeeendeet ones.
B MATHEMATICAL BACKGROUND
B. VECTOR SPACES AND TENSOR ALGEBRA
Let V be a vector space. The dual space V of is the set of all linear functionals on V and this space itself admits the structure of 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 x V and the addition and multiplication is given by
(u), + (u2, v2) (u) + Ui, 02 + v2) c(u,, (cu,, cu1),
for U1, U2 E U. v2 V and c R. We now introduce tensors vector spaces, A k-tensor T on is multilinear function
T: T:VX...VV'
copies
We may think of as an element of the vector space (V*)&8 the tensor product of the vector space V with itself k-times. We delegate the definition of tensor product of vector spaces to Appendix C.2 A k-tensor T is symmetric T is a symmetric multilincar function. work primarily with symmetric tensors in this paper.
13

---

Under review as a conference paper at ICLR 2020
B.2 CANONICAL ISOMORPHISMS
We describe the distinction between an isomorphism and a canonical isomorphism of vector spaces. isomorphism between two vector spaces and V is a bijection between U and V which pre- serves addition and scalar multiplication. A canonical isomorphism is stronger concept, it is an isomorphism 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 isomorphic to one another but the isomorphism may not be canonical. Consider a finite dimensional vector space V. There is isomorphism between V and its dual space V given choice of basis e, for V. there a dual basis o for V and the map e₁ - o is an isomorphism. However, this is canonical as it depends tion map ev V* + IR defined by & + (p(v) where E V*, v E V. The mapping v + ev, then defines a canonical isomorphism from V to double dual space V**
the choice of basis e; for V. On the other hand, consider the evalua-
B.3 AFFINE ALGEBRA
set A is an affine space associated to the vector space V if there is mapping A x + V denoted by (p, q) A x A pę satisfying the axioms
1. for any p, q, A, pt pi + qr
2. for any p E and for any x E there an unique a such that x pq.
Intuitively, an affine space may be thought of as vector space with no privileged origin. Suppose that we choose origin point A and let (ei.....e. ) be basis for the associated vector space V. For any point p E A. we can write op = Σi= x, (p)e,, Here, {xI,..., x is a set of coordinate functions, more simply, basis for A. If we have two bases (x)...., Xn and {yi..... ..n. then they are related by = Bx + where vector.
b.y] is an invertible n x n matrix and = [c,] is a
We now describe how to extend a change-of- on the affine space to the product space A A. Let and be two choices of affine bases A, then
A
[(ai,..... }]. = (a....., an), [(a.,...,a.))
(a)
where homogeneous coordinates are used for a, and a.. Now, suppose that the change basis from i. given by (B e) and denote
B 0 1
B =
Then, have
B H
a
o
B
ag
(B.1)
aj
(I &
B)n)
Thus, the induced change-of-basin on the product space K given by the matrix I & B
B.4 DIFFERENTIALS, PUSHFORWARDE AND PULLBACKS
Let M be smooth real manifold. For p M, we denote the tangent space by T,M and the corresponding dual space, the cotangent space by T M. Given smooth function f M - R. the differential (p) e T M is defined by
df df()) Xx()), X, T.M.
14

---

Under review as a conference paper at ICLR 2020
Let (X1...., X,) be coordinate system around p M, the differential (p) can be expressed as
0f
ax1
[d[(p]]
0f
8x.
We observe that this coordinate representation corresponds to the gradient V (even though differ- entials and gradients are distinct objects for abstract manifolds).
For a smooth map yp MM - M2 of manifolds, the pushforward p. : TyMi + To(p).M2 is defined by
(v(v) = v(h o p), a smooth function M2. If we suppose that M₁ RD and
where u M, and h M2 M2 = R' with (xI....> Xn) ) a coordinate system around and (y1...., .m) coordinate system around ppp), the pushforward yo, v can be represented as
Oy
By 8x.
[
v,
0x1
Ux..
where [v]] This is exactly the Jacobian matrix J. of yp and hence [s] = J.v. This Jacobian- vector product corresponds to the directional derivative, and can be computed using forward mode automatic differentiation (Schraudolph) 20022..
The dual notion of the pushforward, the pullback sp" Th (o + M. is defined in the follow- ing way
(p" = u(p.v), u E T M2.
With respect to the same coordinate systems chosen above, we can write [I*" J u, where [u]. Numerically,
compute J efficiently using reverse mode auto-diifereeniaaion (i.e., backpropagation).
B.5 METRICS AND THEIR PROPERTIES
We introduce tensors manifolds. A symmetric k-tensor at the point p E M defined symmetric k-tensor on the tangent space T..M Recall that this is symmetric multilinear map on the k-fold product of T TrMM:
(p) TTMx ... TrM IR.
J: copics
A metric on M is defined as smoothly varying symmetric 2-tensor which is positive-semidefinite at every point E M. g is nondegenerate, then this is just usual Riemannian metric. However, in this paper, 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 the network. Here, define how this works for general tensors. Let ip M1 - M2 be smooth map of manifolds and a be a symmetric k-tensor on M2 at ypp). The pullback soo a a under is a symmetric k-tensor on Mi defined as
x0 (p) un) = (p. (pukk)
where 01 o1...., Up Mp. In the case of metrics,
4" (p) = gM2 (Mx)))()))(( v2.2)
If we suppose that the metric gM, is given by GM, for chosen coordinate system around p(p), then the pullback metric 4p gM, Mi around is given by
I (p)) GM
where J is the Jacobian of p. While metric always pulls back to metric under smooth map, the pullback a nondegenerate metric can be degenerate as the pushforward map may have non-trivial nullspace.
15

---

Under review as a conference paper at ICLR 2020
Lastly, we describe how define metries product manifolds. Given metrics 9.M, and JM, Mi and M2 respectively, we describe how to naturally define a metric on the product manifold M1 x M2. For any point (p.9) E M₁ x M2. there is a canonical isomorphism of tangent spaces: T(p.)) x M2) M, M,
The proof of this fact can be found in standard differential geometry literature (e.g. Lee (2003)) and we do not elaborate further here. Hence, any vector E (xp.))(M .M2) can be written pair (V1, U2) where E Tp and U2 E M2. Now, we define the additive metric 9M, + 9M,
on Mi M2 as follows:
(gM, + 9m, (p. (u, u) = g.M, (p) 01) + 9.M, (g)(()) 02).
(B.2)
we choose coordinate system around (p, 4) with the metrics JM 9M2 represented by matrices G. GM,, respectively, then we have
GM, 0
M[a, + gp, 40
GM²
which is matrix with block diagonals GM, GM, 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 fmite-dimenssin vector space. For each point E V, there is canonical isomorphism + T.V. V. From now on, we suppress p T,V and write TV when denoting langent spaces V.
Proof. For any element
V. we can associate tangent vector Dulp at defined by
De = Dof(p)
f(p tv),
dt 1=0
where i 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 E A, there a canonical isomorphism between T.A and V. From now on, we suppress in T and write T A when denoting tangent spaces of A.
Proof. Note that specifying a point a E A naturally identifies with Then, applying the above theorem gives the desired result.
C.2 TENSOR PRODUCT OF VECTOR SPACES
Let U and be miite-dimensional vector spaces over the real numbers R. Let R be the subspace of the free vector space IR 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 + uu, v), (u, v) (u, v') (u, v u v'),
for u, u' E U, U, E V, and c R. The tensor product & V, is the quotient space the equivalence class of an element (u, v) in U @ V is denoted by 8 v.
R(UxV) R
and
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 canonical isomorphism
U 3 V
Hom(U, V)
(C.1)
16

---

Under review as a conference paper at ICLR 2020
given by 4 p(u)u where &p U*. Another isomorphism of interest us is
o V' - (U & V)'
(C.2)
which is again canonical.
To derive this isomorphism, given E U", V*, consider the bilinear R defined by
map U
(u, + p(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,
conclude that this is isomorphism.
We now use these facts to explain several ingredients the proof of Theorem 3.2 in greater detail. The element is linear map TW -+ TZ and by the above isomorphism, this means E T Ø TZ from Eqn. [C. By Eqn. C.2 the dual space of (T W @ T2)" TW & T Z. Using Eqn. C.1
again, elements of this space are maps À= T W T Z.
D TWO PARAMETERIZATIONS OF COORDINATE-FRES 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 parametenization for the same underlying abstract ML..
First, a choice of parameterization, or coordinate system, for the abstract MLP a choice of affine
bases for all of the activation spaces : the pre-activation spaces Z1,..., Zz, the input space X' and the output space y in the network. Observe that a choice of bases for A and Z; naturally induces basis for each W,, and therefore also for the full weight space W. Let e, be two different choices of parameterizations for the network. With respect to we write
[a-- D|. a;-1; [G.. 2i, [ws]. (W. b.), [o. Φi [o:]. a;
and with respect to K, we write
[o:- a [G] 2 |+:] = (w! b)), I0. = 0/. [ail = a.
Hence, we can rewrite Eqn.
3.1 in the parameterizations K as
= W, + b; a, (2(2)),
z wia b
(D.1)
a o
The parameters (W; b.) and (W b!) are related as follows
W, = क w
(D.2)
b, の W 7- - -0,1+b//-- +
D,b
where (D2-1 7t-1) is the change-of-basit from to A:-1 with 82-1 an invertible matrix and 7- vector. Moreover, (æ, T.) is the cf from x to z on Z, with &, an invertible matrix and T; a vector. The activation functions d, and of in Eqn. D.I are related in the following way:
af = = 2 [ ( + T;) +
(D.3)
where (I2; 7:) is the ofange-of-bassee from to A..
To verify Eqn. D.2 and Eqn. D.3 we consider the the following commutative diagram (the top horizontal arrow equal to the composition of maps given by the other three arrows) which relates the two parameterizations and к on the activation affine space A;-1 and the pre-activation affine space Z,;
(W. b,)
[A
1].
[
(D.4)
(0-1177-1)
(..T))
(W? b!)
[A
[2
17

---

Under review as a conference paper at ICLR 2020
Let a, [[l, this maps to z, W,a,-1 + b, [2 under the top horizontal arrow in Eqn. D.4 Now, mapping a,-1 under the composition the other three arrows in Eqn. D.4 we obtain
a--1
+ -, 7/-1
(apply left vertical arrow in Eqn. D.4
- ( + 7--1) + by
(apply bottom horizontal arrow in Eqn. D.4
P. + w 7 b
(apply right vertical arrow in Eqn. D.4
4, W. + W + B,b,
+ T--
we use the commutative diagram (this time, the bottom horizontal arrow equal composition the other three arrows):
This establishes the equality given in Eqn. D.2
For Eqn. D.3
[
[A..
(D.5)
(+)))))
(2..%)
[2
[Ail]
Let 2 12 this maps to af E [A.]x under of Now, mapping zद under the other three arrows Eqn. D.5] we obtain
2f 2 +
(apply left vertical arrow in Eqn. D.5
+ ( + Ti).
(apply upper horizontal arrow in Eqn. D.5
1.0 Ti) +
(apply right vertical arrow in Eqn. D.5
This establishes Eqn. D.3 Lastly, we note that the equations in Eqn. D.I can be rewritten in homogeneous coordinates:
2; W
z a = (z) 0.
(z,),
22 + T.) + %.
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 Eqn. 2.9 This
completes the proof.
E PROOFS
E. PROOF OF THEOREM 3.2
is an intrinsically defined tangent vector on W. The weight space W of 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 Expw,, corresponds the above update rule. Since this construction did not require choosing an affine basis for W, the algorithm is invariant to affine reparameterizations.
First, note that ((w
E.2 PROOF THEOREM 3.4
Before providing the proof of the theorem, we need the following lemma.
Lemma E.1. Let g be a metric Z and Vh W Z be the evaluation map. Then, the pullback metric V on W can be expressed as:
(ig(w)
15
Ø),
where d is random covector.
18

---

Under review as a conference paper at ICLR 2020
Proof. Given Z, the metric g admits the rank-1 decomposition
((2) = E & 6].
where é is covector and the expectation is over Z. This is akin to the familiar case where any symmetric positive-semidefiite matrix admits a rank-1 spectral decomposition. Computing the pullback of y under the map V now gives
(ag(w) B @ ø)
El @ 6]
(E.1)
Elpg @ [.0),
where 2 = (V)).. We analyze the pullback 02 Let an be a tangent vector on W. Then,
(+)))((w)) (Wa)) = @(ww a)
(a @ (3w),
which shows that 6/ = a @ Φb Plugging this back into Eqn. we obtain
Vig(w) Ela
which concludes the proof.
For the first assertion of the theorem, we need to check that both components E[a @ a] and E(g(=)) define symmetric sesitive-semidefiinii
2-tensors. Recall that can be realized a linear map from TW to TZ. Then, the dual element À is map from T" to T*Z. We refer to Appendix C.2 for formal explanation of this. To check the positive-semidefinte property,
Ela 2 () E[(a & a) A) 2 0,
where the latter inequality is due to the fact that a & a is positive-semidefnni Moreover, @ a also symmetric and this property is preserved under expectations which implies that Efaa @ al is both symmetric and positive-semidefintt. For the second term E[g(=)] in gina(w) the fact that is metric Z means that g(=), by definition, is symmetric nositive-semidefniite 2-tensor on 7 Z.
To establish the second assertion of the theorem, we need to show that both IE(a 6 a] and Rig(=)]
positive-definite Suppose to the contrary that this not true for Ela a). Then there exists E TW @ T Z (this is same as saying À is linear map from 1 W to T 2; refer to Appendix C.2
for further explanations) such that
0 = E a
(a x)]:
and hence (a & () & X) 0. Now, consider the element A (TW @ T* @ TZ where Iu TZ. From the above lemma, we have
( = Eladsaad
Then, evaluating Ea [x(Vig(w)] at & μr:
= E (A@)))))))
= 0,
This shows that E.[vi] is not dositive-definit which yields contradiction E, [ was sumed to nondegenerate metric. The exact same argument can be applied to show that Eig(=)] is positive-definit, This gives the desired result.
19

---

Under review as a conference paper at ICLR 2020
E.3 PROOF OF THEOREM 3.8
From Theorem 3.4 we know that Find is a metric on W, Since 9keaC is defined as the additive metric where each of the summands are ging. we can conclude that 9KFAC is a metric. For the second assertion of the theorem, recall that V = pp vig. By the functoriality property of pullback operations, we have
(w = (w (+t))
Since Ec(( was assumed to be nondegenerate, this implies that
EE((v)) ((pE)* g)l,
also nondegenerate. Then, from the second assertion of Theorem 3.4 we obtaín that Find is nondegenerate metric W,. Consequently, JKFAC is nondegenerate which concludes the proof.
F COORDINATE-PREE K-FAC FOR CONVOLUTIONAL NETWORKS
F.1 CONVOLUTIONAL NETWORKS
We provide complete account convolutional network case sketched out earlier in Section We begin by describing the convolution layer of convolutional network in mathematical terms fol- lowing 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 the case of MLPs. Lastly, coordinate- free description convolution layers.
the abstract linear algebra machinery developed in Section 2 to give
F.1.1 CONVOLUTION LAYERS
We focus on single convolution layer. convolution layer I takes as input a layer of activations aj., where E (1.....J] J} indexes the input map and / E T indexes the spatial location. 'T here denotes the set of spatial locations, which we typically take to be 2D-grid. We assume that the convolution is performed with a stride of I and padding equal to the kernel radius R, so that the set of spatial locations shared between the input and output feature maps. This layer is parameterized by set of weights Wa,.5 and biases b, where E f1,..., T} indexes the output map and A 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
Σ
Wi,, b,.
(F.1)
The pre-activations a are then passed through a nonlinear activation function Φr. Analogous to feed-forward networks, the weight derivatives are computed using backpropagation:
Dw
Σ aj.448Dzi.
Following Grosse & Martens (2016), we represent the convolution layer computation in Eqn. F. using matrix notation. To do this, we write the activations ay.: as a [T] matrix A;-1, the pre- activations Zi,e as a / x [T] matrix Z, the weights Wi,.. as a J| matrix W; and the bias vector as b,. For the activation matrix A,-1. location t E T and flatten these patches into vectors where the vectors become columns of matrix, we obtain / x [T] matrix which we denote by A From now on, we refer to this matrix as the expanded activations. Finally, we can use these matrix notations to rewrite the computation in
extract the patches surrounding each spatial
Eqn. [.T
Z WiA EXP b: A, oi (Z.).
(F.2)
For convenience purposes later, we adopt homogeneous coordinates various matrices:
Z:
Wi
b 1
A
A exp
[Z-1l =
W =
[A
1-1
=
1
0
=
20

---

Under review as a conference paper at ICLR 2020
Hence, Eqn. F.2 can be rewritten
Zi] / [Wi] [A [A--) [Ai] (Z))
(F3)
where the activation function Φ here ignores the homogeneous coordinate. We introduce the concept of transformed convolution layer. For convolution layer defined in the above equation, the parameters [ and the transformed parameters [W]] are related in the following way
[iil
ri [ H (I Tı-1),
(F.4)
where F and T are invertible matrices. The activation functions Φi and of are related through standard affine change-of-bass as given in Eqn. 2.9
F.1.2 ABSTRACT CONVOLUTION LAYERS
Just as in the nordinnteeddependae case earlier, we focus on single layer. An abstract convolution layer 7 is defined as follows:
Local activations at each spatial location t. E T are taken to be elements 0t-1 in an affine space At-
Activations are taken to be elements 2 in A (i.e. the direct product of A, [T] times).
[7]
(The superscripts are meant to be suggestive of Python slicing notation.)
Expanded activations at E T are taken to be elements (:.0) ai-1 in AAA The full expanded
activations are taken to be elements (64) ai-1 in [A
Local pre-activations at E T are taken to be elements (t) in affine space Z,
Pre-activations are taken to be elements c in 1T!
Layerwise parameters are affine transformations wy between A and 2,. The collection of these transformations is an affine space in its right which denote by Wi and refer to as the layerwise weight space. If we apply wi pointwise, this can be extended to map
A | + 2 IT
The computation for this abstract layer is
(1) wi (attti))
on - pr())))
where Pi fixed nonlinear activation function and at are the I-th layer local activations defined in exactly the same manner 01-1-
We choose affine bases on A/-1, 2,, and Ar. A basis on A naturally induces a basis for A|A1 Consequently, this gives a basis also for the layerwise parameter space W.. Let n be two such choices. With respect to i, we write
Ta."011 = a Is
(144
21) [w]].
[p:l. di, [[olle = a,.
and with respect to K, we write
[o [o[20] (a(:))) 1c =
(2,0==(2)))))))) (2) [w]] = (w) b/), [pd] = 08, [[o] = a7-
(:.0) Note that a a are J| -dimensionnl column vectors of the expanded activations matrix A z(() are I-dimensional column vectors of pre-activations matrix Zi and ai are J-dimensional column vectors of activations matrix A!-
Now, suppose that (Ω 7-1) is the change-of-baase basis from to R on A and (Φ: T)) is the change-of- basis from on Zj. If we denote
Ω 12114
7-11
[ =
$
0
1
21

---

Under review as a conference paper at ICLR 2020
then by the affine sfaagge-ff-basee formula for direct products (Eqn. B.1 & Mq-1l defines the change-of-basia from to A on A The parameters [Wi]H and [W]]æ are related as follows:
[ # = [ [|] @ (Q=1)=)
By taking T and Ti in Eqn. F.4 to be Y [Ωt=1] and T2 (
I 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.
F.2 KRONECKER FACTORS FOR CONVOLUTION
We review the Kronecker Factors for Convolution method due to Grosse & Martens (2016) which a version of K-FAC for convolutional networks. The network architecture to consider is a con- volutional network with L convolution layers. First, let be the concatenation of all trainable parameters W
[vec(Wi) (1) vec (W2) ... (2)
For an input-tare target pair (x, y). the Fisher matrix for this network is
F(w)= ((xy)(()))
where Dw is the log-likelihood gradient and the expectation is taken the model's predictive distribution (y)x(w) for and over the data distribution for x. The diagonal blocks F(w)L, of F(w) are
F F(w)L. E[vec(DW,) vee (De.)
We ready now to present the K-FAC approximation for convolutional networks. For particular layer we define the K-FAC approximation F(w)r, to F(w)), as:
F F(w),, = [T)
a D
E, E, [D (D2, (T)
(F.5)
where a|_? is the homogeneous notation for a The K-FAC approximation F(w) to F(w) is the matrix with diagonal blocks F( (w),,, 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. i 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 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.
F.3 COORDINATE- FREE K-FAC FOR CONVOLUTIONAL NETWORKS
We begin by considering an abstract convolutional network with convolution layers. Let X' and y denote the input and output spaces of this network respectively. Recall that the layerwise weight space W; is the space of affine transformations between A/A| and Z, The weight space of this network is the direct product of all layerwise weight spaces
W = Wi x WL.
Given input & E X' and parameter w = (w....,, wL) E W, denote the network output by f({,w), w)).
Now, for every 1 ( ......... L). define the following maps
Vif Wi + z which sends layerwise parameters wi to pre-activations c.0) by evaluating local activations a1-1 across every spatial location E T
p 2 which sends pre-activations 50 ((E.w)
22

---

Under review as a conference paper at ICLR 2020
Again, d/g is trivially smooth map from its definition. The map I' 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 ReLUU and max-pooling, we conjecture that the non-smooth case can be addressed by taking limits of smooth functions.) Finally, define the map V Wi y the composition V = 4g vig
Let be a metric on y and consider the pullback (fgg) restricted to single spatial location which we denote by (vEE,) " "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 be
A A Z = 2,, W WW,
and the metric on Z = Z₁ to be (vpE,)) Summing over every spatial location E T, the indepen- dence metric on Wi here is
Iina = [T| & (1:0) & BT((P))) " g(c)
(F.)
Definition F.2. The K-FAC metric on the weight space W of an abstract convolutional network defined as
(KFAC(W) = Gind (w1) + .... gind (WL),
where the sum above is as defined in Eqn. B.2 and each Flnd is as given in Eqn. F.6
Theorem F.3. Let be a metric on y. Then, 9KFAC given in Definition F2 is indeed a metric on the weight space W of an abstract convolutional network. Moreover we assume that the expected pullback of restricted to single spatial location t E T,
Be(( 2.
under the map V Wi + y is nondegenerate metric on the layerwise weight space Wi for every I, then 9KFAC 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
(:.6) Ela X a & Bl((Ve))) ([(F),,)"(((")) "(({(")))
a metric on Wi. Since taking expectation over the set of spatial locations T and multiplying by the scale factor [T] preserves the metric properties, obtain that gnd in Eqn. F.6 defines metric on Wi. Consequently, ЧкFАС determines a metric on W. To prove the latter assertion, note that by the functorial property of pullback operations,
Bel(ve (1 g)]
nondegenerate. Using the second assertion of Theorem 3.4 yields that gind nondegenerate which implies that this is true also for JKFAC-
We conclude this section with the proof of Theorem 4.1 Our proof is coordinate-frt and given in exactly the same manner as the proof Theorem 2.1 at the end Scction 3
Coordinate-Free Proof of Theorem 4.1 As shown earlier in Section F1.22 each convolution layer of N' and W correspond to two different choices of parameterizations for the same underlying abstract convolution layer. Hence, A and ' must compute the same function.
Assume that the metric on the output space y in Theorem F33 is the output Fisher metric gFo For each spatial location E T, the pullback under yp is
(PEE)) gF (c( EldC to dC (.).
Now, choose coordinate systems on Ar-1 and Z, This induces coordinates for A/A1 and we write
[a (s.0) = 2 11
1c 20, [((E,)) (
(C (c [B|dC @ dC ddc E|D2 (Dz( )
23

---

Under review as a conference paper at ICLR 2020
Using Proposition 3.5 the independence metric in Eqn. F6 can be expressed in coordinates as follows
[mna (T) (a ( @ E, (22,
(1) ( )
which is exactly F(w)r, given earlier in Eqn. F.5 Furthermore, [[к] is exactly the K- FAC approximation F(w). Thus, the K-FAC update rule is simply natural gradient update rule with respect to the K-FAC metric 9KFAC for abstract convolutional networks. Lastly, if 9KFAC is nondegenerate metric; which true for example if the assumptions in the second assertion of Theo- rem 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 complete account of the recurrent network case sketched out earlier in Section 4.2. We first give mathematical formulation of the recurrent computation step of these networks 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. is not necessary to write out the full structure of recurrent network. Rather, 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 different time steps and = (1....,T), We use to index the time step. Throughout, assume that all sequences are of fixed length T. For input x, at every I, the recurrent network maps this to an output o,, Essentially, the network maps input sequences (Xi,..., XT) ) to output sequences ( OT). The computation, at every I, is
+
z = + Vx; a, = o(Z))
where a;-1 an activation vector, 2, is pre-activation vector, W is a recurrent weight matrix, V is weight matrix, b is recurrent bias vector, and á is fixed nonlinear activation function. For the remainder this section, we focus on the first equation
Z Wa 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
w a! +b²
(G.2)
The relationship between transformed parameters (Wi bt) and original parameters (W b) given
by the standard change-of-basis formula in Eqn. D.2
G. 1.2
ABSTRACT RECURRENT NETWORK
We now describe abstract recurrent network formally.
Local activations at each time step f are elements in an affine space A
Activations are elements a = (aeher in the affine space AT
Local pre-activations at each are elements & in an affine space
Pre-activations elements < (G: in the affine space 2T
24

---

Under review as a conference paper at ICLR 2020
Parameters are affine transformations between A and 2. The collection of these trans- formations 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 are elements &,, in affine spaces X, y respectively. The input and output spaces are X² and )" respectively; furthermore, elements here are written as Ç = {& heT and = fu, hem.
For every t, the abstract recurrent computation step is
C w((
choice of parameterization for the abstract recurrent network consists of choosing affine bases for A. Z, X' and V. Since we have bases for A and Z, this induces a natural basis for W. If we use exactly the same change-of- -fabiisss 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 Recall that for every time step t, the recurrent computation can written as
in Martens et al. 2018
2 Wa;-1
where W W bj and a!_ = [a_1 TT. Using backpropagation, the log- log-likelihood gradient is given by D The total contribution to the gradient across all is the sum
DW Σ
For an input-target pair yx, the Fisher matrix F (W) for recurrent networks is defined as F(W) Ex,y (voo[W) vec(DW)
Finally, the K-FAC approximation F(W) F(W) for recurrent networks is defined
F(W)= = TEE ā a @ (22,
(G.3)
Remark G.I. As in the case of convolution layers, there is weight sharing in recurrent networks (across time here instead spatial locations) and so is not enough to just assume statistical in- dependence between activations and pre-activation derivatives to make the K-FAC approximation here. We defer the reader 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 coordinate-free proof of Theorem 4.2
The
method of proof mirrors exactly the proofs given previously for Theorem 2. and Theorem 4.1
G.3 COORDINATE- K-FAC FOR RECURRENT NETWORKS
Given input { her X and parameter W. denote the network output by f(E, For specific time step f, consider the following maps:
Veg W + Z which sends parameters to pre-activations & by evaluation at activations
Q--1
iPe.s Z + y which sends & to outputs
In addition, we define the map V W - 3 as the composition V = SPE,t o the,t-
Let be metric on y. The pullback ot, then defines metric 2. Now, take A, 7, W in Section to be
A Z, W W, and the metric on = to be f2.9 Summing over all time steps t E T, we make the following definition which arises from the independence metric in Section 3.2
25

---

Under review as a conference paper at ICLR 2020
Definition G.2. The K-FAC metric the weight space W of abstract recurrent network is defined as
9KFAC(W) T(Br(ar- & at-1) & Erlpg,
(G.4)
Theorem G.3. Let 9 be metric on y. Then, 9KFAC given in Desinition G. a metric on the weight space W of an abstract recurrent network, Moreover, if we assume that the expected pullbuck of g.
under the smooth map ge W + y is nondegenerate metric, then 9KFAC 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 have that
Eia,-1 @ ax-1ه @ ElpE)
is 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 (KFAC metric W. For the nondegeneracy statement, using the functorial property of pullbacks, we know that
Eg (g(1)E.
nondegenerate. Then, JKFAC is nondegenerate by the second assertion of Theorem 3.4
Coordinate-Free Proof of Theorem 4.2, We conclude this paper with coordinate-free proof of Theorem 4.2 As mentioned at the end of Section G.1.2 N and A'l correspond to two different choices parameterizations for the same underlying abstract recurrent network and so they must compute the same function.
Assume that the metric g on y Theorem G.3 above is the output Fisher metric Futtt Then, the pullback under &Pg.t. is
Pnn (G) = R[dfc, @ dcc.).
Now, choose coordinate systems for A and Z. We can write
[ |] at [[]]] = Zz, [ (4+) = D(Du,
By Proposition 3.5 the K-FAC metric in Eqn. G.4 be represented in these chosen coordinates as [як = T(E a (r,Dz,Dz,]]
which is exactly the K-FAC approximation F(W) in Eqn. G.3 Thus, the K-FAC update rule is natural gradient update with respect to the K-FAC metric 9KFAC for abstract recurrent networks. suppose that gKFAC is nondegenerate, then these updates invariant any affine reparameter- izations of the model.
26


Reviews:
[{"review_id": "Skx1SGmx5r", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nWhy 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.", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "N/A", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "BJlVUPfMKr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper is concerned with tractable (approximate) forms of natural gradient updates for neural networks, in particular with the recent K-FAC approximation, which applies a set of approximation (layer-wise independence, Kronecker structure for affine maps) in order to obtain a Hessian that can be computed and inverted efficiently. K-FAC has been introduced for MLPs, and has previously been generalized to convolutional and certain recurrent NNs.\n\nThe stated goal of this paper is to provide a mathematical re-formulation of K-FAC in terms of Riemannian metrics. While K-FAC has been developed as approximation to the exact natural gradient update, they come up with a different Riemannian metric, definition of space, etc., such that in the end, K-FAC is the exact natural gradient for that. The authors here also obtain a more precise answer to invariance properties and, given some heavy maths, what they claim to be more elegant proofs of previously known properties of K-FAC.\n\nThe paper uses very heavy math, well \"over my head\" and likely most ICLR attendees. Along with me, they'll ask the obvious question of what this is good for. As far as I can see, there is nothing really new being proposed here in terms of practical consequences. The authors also do not make much effort to explain why their viewpoint is useful, say to obtain practically relevant insights in future work. So, as far as I am concerned, I do not see why this work should be of much relevance to ICLR, which is not an abstract maths conference.\n\nA final comment is that people have for a very long time tried to use second-order optimization for MLPs. The aspect that always was tricky there, is that SGD is *stochastic*, and the second-order info is hard to estimate from a mini-batch. The sets of approximations of K-FAC are pretty extreme, but they may just be needed to make things work in the end, because they may stabilize that \"stochastic inverse Fisher info matrix\" enough to not make the optimization process fail altogether. Now, all theoretical arguments, like \"invariance to this and that\", always ignore the crucial fact that you are talking about a stochastic estimate over a mini-batch, and your theory is always for E_x[...] \"being the truth\". It is not, it is just over a small mini-batch. I am not saying the additional theoretical insight from this work here (over previous K-FAC work), whatever it may be in the end, is not useful. I am just saying I'd be a lot more confident if the authors would specifically address the stochastic property.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory.", "review_assessment:_checking_correctness_of_experiments": "N/A", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "B1l2T3Ue5B", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThis present paper offers a different approach: rather than approximating the exact NG directly as K-FAC does, it views the approximate K-FAC natural gradient as the exact \"natural gradient\" under the K-FAC metric. This is an interesting idea and the paper goes on to show that the new \"exact\" NG under the K-FAC metric (rather than the Fisher-Rao metric) is invariant under certain affine transformations.\n\nI find the paper extremely hard to follow. My main concern is that many concepts and math objects are at risk of being not mathematically rigorous or well defined. For example, equation (2.5) is not correct, this is not a genuine update based on exponential map. Similarly, the update equation in the first graph of page 4 doesn't make sense as \\mathcal{W} is an abstract manifold and the subtract here is not defined. The way paper is written has a high risk of causing confusion between coordinate-dependent and coordinate-free objects. On page 3, line -5, I have trouble understanding the abstract distribution P_{v|\\psi}(\\omega). On page 7, I don't understand what \"evaluating $w$ at $a$\" means, and even $a$ here isn't defined yet I think. In equation (2.3), $-w_k$ should be $w_k$ (i.e. no minus sign). There are many other typos and confusions that I opt not to  point out here.\n\nIn summary, the main issue of the paper in its current form is that its presentation isn't clear with many typos. I suggest the authors to re-write the paper carefully so that it is more accessible.\n\nI know that this paper focuses on theoretical properties of NG, but it would also be good if the authors offer some discussion on the use of their method in practice.\n\n  \n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "N/A", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors analyze the natural gradient algorithm for training a neural net from a theoretical perspective and prove connections to the K-FAC algorithm. The paper is poorly written and contains no experimental evaluation or well established implications wrt practical significance of the results.
Decision: Reject