A C OORDINATE -F REE C ONSTRUCTION OF S CALABLE N ATURAL G RADIENT
Anonymous authors
Paper under double-blind review
A BSTRACT
Most neural networks are trained using ﬁrst-order optimization methods, which are sensitive to the parameterization of the model.
Natural gradient descent is invariant to smooth reparameterizations because it is deﬁned in a coordinate-free way, but tractable approximations are typically deﬁned 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 afﬁne 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.
I NTRODUCTION
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 deﬁciency 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 ﬁrst-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 deﬁnition 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 deﬁned in terms of particular coordinate representations, and therefore may lose the invariance properties which motivated natural gradient in the ﬁrst 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 signiﬁcant speedups in training deep autoencoders (Martens & Grosse, 2015), classiﬁcation 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 signiﬁcant 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: afﬁne transformations of the activations in each layer (Martens & Grosse, 2015).
This was veriﬁed 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.
Speciﬁcally, we view a neural network as a series of afﬁne maps alternating with ﬁxed nonlinear activation functions.
The activations and pre-activations for each layer are viewed as abstract afﬁne spaces, and the weights and biases of the network correspond to afﬁne maps.
The ordinary Fisher metric is a metric on this space W of afﬁne maps.
Our contribution is a recipe to convert a metric on
W (whose coordinate representation is extremely large) into an approximate metric on W (the “ 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 ﬁrst-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 deﬁne afﬁne maps between afﬁne spaces.
Choosing afﬁne bases for the activations and pre-activations yields a natural afﬁne basis for W .
Therefore, the set of allowable reparameterizations for neural networks consists of afﬁne 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 afﬁne structure, we obtain global invariance, not just ﬁrst-order invariance.
Our framework easily enables some generalizations of the basic result.
Our construction is not limited to Fisher metrics, but rather it applies to general pullback metrics , where one places a metric on the network’s output space and pulls it back to W .
Furthermore, we extend the invariance results to convolutional networks and recurrent networks through a straighforward application of our 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 afﬁne/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 (cid:74) · (cid:75) notation throughout stands for choosing coordinates for an abstract mathematical object.
B ACKGROUND 2. N ATURAL G RADIENT D ESCENT
For simplicity, we deﬁne the natural gradient descent algorithm here in the context of multilayer perceptrons (MLPs), i.e., fully connected feed-forward networks.
For an input-target pair ( x , y ) , let f ( x , w ) denote the output and w symbolizes the parameter vector of the MLP.
Given a training set
S of pairs ( x i , y i ) , we aim to minimize the empirical risk which is given by:
[Equation]
where L ( y i , f ( x i , w )) is the loss function measuring the disagreement between y i and f ( x i , w ) .
Suppose that f ( x , w ) determines parameters z of the model’s predictive distribution R y | z over y and furthermore, we reparameterize this as P y | x ( w ) =
R y | f
( x , w ) .
Likewise, the density function of this distribution can be reparameterized as p ( y | x , w ) =
r ( y | f ( x , w )) .
In addition, we take the loss function here to be the negative log-likelihood L ( y , z )
= − log r ( y | z ) and denote log-likelihood gradients ∇ w L ( y , f ( x , w )) by D w ( D notation throughout remainder of the paper refers to loglikelihood gradients).
The Fisher information matrix F ( w ) is deﬁned as
[Equation]
where the expectation is taken over P y | x ( w ) for y and over the data distribution for x . Since
F ( w ) is deﬁned as the expectation of an outer product, F ( w ) is always guaranteed to be a positivesemideﬁnite (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 (cid:15) > 0 , the natural gradient descent algorithm (Amari, 1998) minimizes h ( w ) by using the natural gradient to update parameters of the network:
[Equation]
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 ﬁrst-order.
The reason for this is that the natural gradient admits an intrinsic coordinate-free construction in terms of differential geometric primitives.
We give a detailed explanation of this by considering an abstract mathematical setting; we defer to Appendix B for background on the mathematical terminology with which we use.
Let M be a Riemannian manifold with nondegenerate metric given by g (in this paper, our deﬁnition of “metrics” allows the possibility of it being degenerate).
For a smooth function h : M → R and a point p ∈ M , the differential dh ( p ) is an abstract covector on M .
We use the nondegenerate metric g to convert the covector dh ( p ) into a tangent vector.
By deﬁnition, g ( p ) is a nondegenerate bilinear form which yields a linear isomorphism between the tangent space and the cotangent space:
[Equation]
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,
[Equation]
Applying this isomorphism to dh ( p ) yields the tangent vector g ( p ) − 1 dh (
p ) ∈ T p M .
We call this tangent vector the natural gradient of h .
We apply this mathematical framework to the objects of our interest.
First, let W be a smooth manifold which characterizes the weight space of network parameters intrinsically.
Let ω and ( ξ, υ ) be intrinsic versions of the parameter vector w and the input-target pair ( x , y ) respectively.
ξ and υ belong to abstract input and output spaces which are also smooth manifolds.
Now, W can be endowed with the Fisher metric g F which is deﬁned as
[Equation]
where L ω = − log p ( υ | ξ, ω ) is the abstract log-likelihood loss function.
The expectation above is taken over the abstract predictive distribution P υ | ξ ( ω ) for υ and over the data distribution for ξ .
Expressing this in a coordinate system, we obtain exactly the Fisher matrix as given in Eqn. 2.2.
As W is a Riemannian manifold with the Fisher metric g F , the idealized gradient descent updates are given by Bonnabel et al.
[Equation]
where Exp ω k : T ω k W
→ W is the exponential map.
This update rule is exactly invariant to all smooth reparameterizations of W since it is entirely coordinate-free.
However, such an algorithm is infeasible in practice as computing the exponential map is typically an intractable problem.
Instead, it is much easier to work with the following abstract natural gradient update rule which uses a ﬁrstorder approximation of the exponential map
[Equation]
Since this is a ﬁrst-order approximation of the update rule in Eqn.
2.5, invariance to smooth reparameterizations holds only up to ﬁrst-order.
Additional approximations to the exponential map are necessary to obtain higher-order invariances; we defer to Song & Ermon (2018) for a more detailed account of how this can be done.
2. K RONECKER -F ACTORED A PPROXIMATE C URVATURE
We consider a MLP with L layers.
At each layer i ∈ { 1 , . . .
, L } , the MLP computation is given as:
[Equation]
where a i − 1 is an activation vector,
z i is a pre-activation vector, W
i is a weight matrix, b i is a bias vector, and φ i : R → R is an activation function.
For convenience, we introduce homogeneous coordinates ¯ a (cid:62) i −
=
[ a (cid:62) i − 1 1] (cid:62) and ¯ W i =
[ W i b i ] .
The above computation can be rewritten as
[Equation]
[Equation]
[Equation]
We concatenate all of the network parameters
¯ W i into a single vector w ,
[Equation]
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
[Equation]
We now give an overview of the Kronecker-Factored Approximate Curvature (K-FAC) method (Martens & Grosse, 2015) of approximating the Fisher matrix.
Consider the diagonal ( i, i ) blocks of F ( w ) .
Using backpropagation, the log-likelihood gradient D ¯ W i
=
D
z i ¯ a (cid:62) i − 1 and hence, we have vec( D z i ¯ a (cid:62)
i − 1 ) =
¯ a i − 1 ⊗ D z i .
Then, F ( w ) i,i can be rewritten as:
[Equation]
where ⊗ denotes the Kronecker product of matrices.
If the activations and pre-activation derivatives are approximated as independent
[Equation]
this yields the following approximation ˆ F ( w ) i,i to F ( w ) i,i ,
[Equation]
The K-FAC approximation matrix ˆ F ( w ) to F ( w ) is deﬁned as
[Equation]
The K-FAC update rules are analogous to natural gradient descent updates; we simply replace
F ( w k ) − 1 in Eqn.
2.3 by ˆ F ( w k ) − 1 above.
Invariance Properties of K-FAC.
Since K-FAC uses the approximation ˆ F ( w ) rather than the Fisher matrix F ( w ) itself, the invariance properties of natural gradient do not necessarily carry over to KFAC.
Instead, we consider the class of transformations given by the following transformed network
[Equation]
where Ω i , Φ i are invertible matrices and τ i , γ i are vectors.
The transformed input is ¯ a † 0 = ¯ Ω 0 ¯ awhere
[Equation]
and the transformed output is a † L = f † ( x † , w † ) with w † deﬁned as
[Equation]
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 } L i =1 computes the same function as N .
Furthermore, if the Fisher matrix of the MLP is assumed to be positive-deﬁnite, then the K-FAC updates are equivalent, in the sense that the resulting networks compute the same function.
A proof of this result was given earlier in Martens & Grosse (2015); however, the proof was cumbersome and involved tedious manipulation of update rules.
In contrast, we like to provide a far more mathematically mature and elegant way to obtain this result.
That is, we construct an intrinsic
K-FAC metric g KFAC whose coordinate representation exactly matches the K-FAC approximation given in Eqn.
A K-FAC update can then be viewed as a natural gradient update with respect to the metric g KFAC .
More importantly, the invariance properties of the K-FAC algorithm are immmediately established in the same way as it was for exact natural gradient.
C OORDINATE -F REE K-FAC
We observe that MLPs consist of a sequence of afﬁne transformations and activation functions in alternation.
In order to capture this structure, we treat the spaces of activations and pre-activations as afﬁne 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 deﬁne afﬁne 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 afﬁne space
A
i − 1 .
• Pre-activations are taken to be elements ζ i in an afﬁne space Z i .
• Layerwise parameters are afﬁne transformations ω
i between A
i − 1 and Z i .
The collection of these transformations is an afﬁne space in its own right, which we denote by W i and refer to as the layerwise weight space.
• The weight space is given by the direct product W = W 1 × · · · × W L .
Elements in this space are written as ω =
( ω 1 , . .
. , ω L ) ∈ W .
• Input and outputs are denoted by ξ and f ( ξ, ω ) respectively.
The space of all inputs and outputs are afﬁne spaces denoted by X (= A 0 ) and
Y (=
A L ) respectively.
Moreover, the layerwise computation is given by
[Equation]
where ρ i : R → R is a ﬁxed nonlinear activation function which is assumed to be smooth throughout.
In Appendix D, we show that the MLP with computation deﬁned 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 deﬁne standard operations such as pushforwards and pullbacks.
While the commonly-used ReLU activation function is not smooth by deﬁnition, 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 functionsn 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
[Equation]
The following theorem (see Appendix E.1 for the proof) shows that by imposing afﬁne structure on the weight space, we can obtain global invariance instead of just ﬁrst-order invariance.
Theorem 3.2.
Let g be a nondegenerate metric on the weight space W of an abstract MLP.
For a chosen learning rate (cid:15) > 0 , the following update rule
[Equation]
is exactly invariant to all afﬁne 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 deﬁnes a metric on W .
The expected pullback metric over inputs, under a choice of coordinates around ω and α L , is given by
[Equation]
where G is the representation of g in these coordinates.
Example 3.3 (Fisher metric) .
Suppose that the outputs α L parameterize the model’s predictive distribution R υ | α L .
Let r ( υ | α L ) denote the density function of this distribution and furthermore, we take the loss function here to be the negative log-likelihood L α L = − log r
( υ | α L ) .
The output
Fisher metric g F, out on Y is deﬁned as
[Equation]
where the expectation is taken with respect to the predictive distribution R υ | α L .
Computing the expectation of Ψ ∗ ξ g F, out over the inputs ξ gives
[Equation]
This is exactly the Fisher metric deﬁned earlier in Eqn. 2.4.
We note that this construction quite general and not limited to Fisher metrics; for example, if we use the Euclidean metric on Y , then Eqn. 3.2 is the Gauss-Newton matrix (Martens, 2014).
A practical use case for the Gauss-Newton metric is when the outputs of the network do not have a natural probabilistic interpretation, e.g., the value network in an actor-critic architecture for reinforcement learning (Wu et al., 2017).
Another example is if we take the Riemannian metric induced by the
Bregman divergence associated with a convex functional on Y , then Eqn. 3.2 is the generalized
Gauss-Newton matrix (Martens, 2014).
3. I NDEPENDENCE M ETRIC
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 afﬁne maps between afﬁne spaces, but use notation which is suggestive of MLPs.
We assume the following:
• Afﬁne spaces A and Z • Afﬁne space W of afﬁne transformations between A and Z • Metric g on Z
Our ﬁrst task is to formulate a coordinate-free analogue of the outer product of homogenized activations, ¯ a i ¯ a (cid:62) i .
Consider the evaluation map ψ a : W → Z which is deﬁned 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 afﬁne spaces (see Corollary C. in Appendix C).
Let ∂ w be a tangent vector on W and f be a smooth function on Z .
Then,
[Equation]
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 deﬁne the tensor product of two elements in A as a mapping
[Equation]
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 , deﬁne g ind on W to be
[Equation]
where the ﬁrst expectation is over A and the second one is over Z .
Note that E [ g ( z )] is well deﬁned because the afﬁne structure of Z allows us to identify the cotangent spaces at all points z .
The following theorem, whose proof is given in Appendix E.2, shows that g ind is indeed a metric on W .
Theorem 3.4.
Let g be a metric on Z and ψ a : W → Z be the evaluation map.
Then, g ind as deﬁned in Eqn.
3.3 is a metric on W .
Moreover, if the expected pullback metric E a [ ψ ∗ a g ] is nondegenerate on W , then g ind is also nondegenerate.
From now on, we refer to g ind as the independence metric .
By direct computation, we obtain that the coordinate representation of g ind matches the K-FAC approximation of the layerwise Fisher blocks:
Proposition 3.5.
Suppose that we choose coordinate systems for the afﬁne spaces A , Z and in these coordinates,
[Equation]
Then the independence metric g ind can be expressed as
[Equation]
Remark 3.6 .
In the context of MLPs (which we explain in greater detail subsequently) where A =
A i − 1 , Z = Z i and W =
W
i , the matrix G ( z ) =
E y
[ D z i D z (cid:62) i ] where the expectation is taken over output space.
E [ G ( z )] means we furthermore take the expectation over Z i .
When we write E [ D z i D z (cid:62) i ] in Eqn. 2.7, we implicitly take this to mean E [ G ( z )] .
3. K-FAC M ETRIC
In this subsection, we formulate the layerwise independence approximation in a coordinate-free way, allowing us to deﬁne 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 deﬁne the K-FAC metric for MLPs.
Lastly, by viewing K-FAC as a metric on W , we show how invariances of the K-FAC algorithm can be obtained in a natural and straightforward manner.
Consider a MLP with L layers as described earlier in Section 3.
For every i ∈ { 1 , . . .
, L } , we deﬁne the following maps
• ψ iξ : W i → Z i which sends layerwise parameters ω
i to pre-activations ζ i by evaluation at activations α i − 1 .
• ϕ
iξ : Z i
→ Y which sends ζ i to network outputs α L = f ( ξ, ω ) .
The subscript ξ is used to highlight the fact that all of these maps have an implicit dependence on network inputs ξ .
Note that ψ iξ is a smooth map by deﬁnition.
Now, observe that ϕ
iξ is exactly the composition of network maps
[Equation]
Since all activation functions ρ i are assumed to be smooth maps, it follows immediately that ϕ
iξ is also a smooth map.
Moreover, we deﬁne the map Ψ iξ : W i
→ Y to be composition Ψ iξ
=
ϕ iξ ◦ ψ iξ .
Let g be a metric on Y .
Then, the pullback ( ϕ iξ ) ∗ g deﬁnes a metric on Z i .
Now, if we take A , Z , W in Section 3.1 to be
[Equation]
and the metric on Z
= Z i to be ( ϕ iξ ) ∗ g , the independence metric on W i here is
[Equation]
Now, given metrics g i ind on each W i , there is a natural “additive metric” deﬁned on the Cartesian product W = W 1 × · · · × W L .
We defer the formalities and the sum “+” operation on metrics to
Appendix B.5.
Deﬁnition 3.7.
The K-FAC metric on the weight space W of a MLP is deﬁned as
[Equation]
where the sum is as deﬁned in Eqn.
B.2 in Appendix B.5 and each g i ind is as given in Eqn. 3.4.
Theorem 3.8 (See Appendix E.3 for proof) .
Let g be a metric on Y .
Then, g KFAC given in Deﬁnition 3.7 is indeed a metric on the weight space W of an abstract MLP.
Moreover, if we assume that the expected pullback of g ,
[Equation]
under the map Ψ iξ : W i → Y is a nondegenerate metric on the layerwise weight space W i for every i , then g KFAC is also nondegenerate.
Coordinate-Free Proof of Theorem 2.1.
We are now ready to provide a natural and straightforward proof of Theorem 2.1.
In Appendix D, we show that networks N and N † correspond to two different choices of parameterizations for the same underlying abstract MLP.
Hence, they must compute the same function.
For the latter assertion, assume that the metric g on the output space Y in Theorem 3.8 is the output
Fisher metric g F, out in Example 3.3.
The pullback of this under ϕ iξ is given by
[Equation]
Let us choose coordinate systems on A i − 1 and Z i with
[Equation]
Then, by Proposition
[Equation] [Equation]
which is exactly ˆ F ( w ) i,i given earlier in Eqn. 2.7.
Furthermore,
[Equation]
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.
Now, observe that
[Equation]
and hence the K-FAC update rule is simply a natural gradient update rule with respect to the K-FAC metric g KFAC for abstract MLPs.
Suppose that g KFAC is a nondegenerate metric; which is true for example if the assumptions in the second assertion of Theorem 3.8 hold.
Applying Theorem 3.
shows that this update rule is invariant to any afﬁne reparameterizations of the model.
A PPLICATIONS TO C ONVOLUTIONAL AND R ECURRENT N ETWORKS
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 ﬂexbility of our analysis.
Due to space constraints, we give only sketches of the necessary constructions and defer the reader to Appendices F and G for full details.
4. C ONVOLUTIONAL N ETWORKS
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
[Equation]
where A exp l − 1 is the J | ∆ | × |T | matrix of expanded activations, W l is the
I × J | ∆ | weight matrix,
Z l is the I × |T | matrix of pre-activations, φ l is a nonlinear activation function and b l is the bias vector.
Adopting homogeneous coordinates for the various matrices, we can rewrite the above as
[Equation]
We brieﬂy introduce the concept of a transformed convolution layer.
For a convolution layer as deﬁned in the above equation, the parameters [ W l ] H and the transformed parameters [ W † l ] H are related in the following way
[Equation]
where Γ l and Υ l − 1 are invertible matrices.
The activation functions φ l and φ † l are related through a standard afﬁne change-of-basis as given in Eqn. 2.9.
For a convolutional network with L convolution layers, the Fisher matrix is given by
[Equation]
where w is the vector concatenating all convolution layer weights.
The K-FAC approximation is deﬁned by setting all off-diagonal blocks to be zero and approximating the diagonal blocks as
[Equation]
Here, a (: ,t ) l − 1 stands for the expanded activations at each spatial location t ∈ T (the J | ∆ |-dimensional column vectors of the matrix
A exp l − 1 in Eqn. 4.1) and z ( t ) l stands for the local pre
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-deﬁnite, the K-FAC updates are equivalent, in the sense that the resulting networks compute the same function.
Proof Sketch (Full proof in Appendix F.3).
A coordinate-dependent proof of this result was given previously in Grosse & Martens (2016).
Instead, we prove this intrinsically using the framework provided in Section 3.
Let A exp l − 1 be the afﬁne space of expanded pre-activations at t ∈ T , Z l be the afﬁne space of local pre-activations at t ∈ T , and W l be the layerwise weight space comprising of afﬁne transformations between A exp l − 1 and Z l .
Now, let us take A , Z , W in Section 3.1 to be A exp l − 1 ,
Z l , W l respectively and the metric on Z =
Z l to be the pullback of the output Fisher metric.
After summing over every spatial location t ∈ T , the independence metric g l ind on W l in coordinates, by Proposition 3.5, is exactly the expression given in Eqn. 4.4.
Taking the sum of g l ind ’s then gives exactly the K-FAC approximation for convolutional networks.
Finally, we know from Theorem 3.
that each g l ind indeed deﬁnes a metric and so by Theorem 3.2, we can conclude that the K-FAC updates are exactly invariant to all afﬁne reparameterizations of the model.
4. R ECURRENT N ETWORKS
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 ﬁxed length T .
We consider the recurrent computation step
[Equation]
where a t − 1 is an activation vector, z t is a pre-activation vector, and W is a recurrent weight matrix.
The transformed recurrent computation step is deﬁned as
[Equation]
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 deﬁned as
[Equation]
and the K-FAC approximation is deﬁned as
[Equation]
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-deﬁnite, 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 afﬁne space of local activations at each time step t , Z be the afﬁne space of local pre-activations at each time step t , and W be the parameter space consisting of all afﬁne transformations between A and Z .
Now, let us take A , Z ,
W in Section 3.1 to be A , Z , W respectively and the metric on Z
=
Z to be the pullback of the output Fisher mertric.
After summing over all time steps t ∈ T , the independence metric g ind on W in coordinates, by Proposition 3.5, is exactly the K-FAC approximation given in Eqn. 4.7.
Finally, we know from Theorem 3.4 that g ind indeed deﬁnes a metric and so by Theorem 3.2, we conclude that the K-FAC updates are exactly invariant to all afﬁne reparameterizations of the model.
A T ABLE OF N OTATIONS
[Table]
Coordinate
[Table]
ξ x
[Table]
υ y
[Table]
α i a i
[Table]
¯a i
[Table]
ζ i z i
[Table]
ω i
( W i b i )
[Table]
¯ W i
[Table]
ρ i φ i
[Table]
ω w
[Table]
f ( ξ, ω ) f
( x , w )
[Table]
L ( υ, f ( ξ, ω ))
L ( y , f ( x , w ))
[Table]
h ( ω ) h ( w )
[Table]
P υ | ξ ( ω )
P y | x ( w )
[Table]
p ( υ | ξ, ω ) p ( y | x , w )
Fisher metric/matrix g F ( ω )
F ( w ) layerwise log-likelihood differential/gradient d L ω
i
D ¯ W i log-likelihood differential/gradient d L ω
D w
K-FAC metric/
[Table]
g KFAC ( ω ) ˆ F ( w )
[Table]
⊗
[Table]
W
B M ATHEMATICAL B ACKGROUND B. V ECTOR S PACES AND T ENSOR A LGEBRA
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
[Equation]
for u 1 , u 2 ∈ U , v 1 , v 2 ∈ V and c ∈ R .
We now introduce tensors on vector spaces.
A k -tensor T on
V is a multilinear function
[Equation]
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 deﬁnition of a tensor product of vector spaces to
Appendix C.2.
A k -tensor T is symmetric if T is a symmetric multilinear function.
We work primarily with symmetric tensors in this paper.
B. C ANONICAL I SOMORPHISMS
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 deﬁne 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 ﬁnite dimensional vector space V .
There is an isomorphism between V and its dual space V ∗ : given a choice of basis e i for V , there is a dual basis e ∗ i for V ∗ and the map e i → e ∗ i is an isomorphism.
However, this is not canonical as it depends on the choice of basis e i for V .
On the other hand, consider the evaluation map ev v : V ∗ →
R deﬁned by ϕ (cid:55)→ ϕ ( v ) where ϕ ∈ V ∗ , v ∈ V .
The mapping v (cid:55)→ ev v then deﬁnes a canonical isomorphism from V to its double dual space V ∗∗ .
B. A FFINE A LGEBRA
A set A is an afﬁne space associated to the vector space V if there is a mapping A ×
A → V denoted by ( p, q ) ∈
A × A (cid:55)→ (cid:126)pq ∈ V satisfying the axioms
1. for any p, q, r ∈ A , (cid:126)pr = (cid:126)pq + (cid:126)qr 2.
for any p ∈ A and for any x ∈ V there is an unique q ∈
A such that x = (cid:126)pq .
Intuitively, an afﬁne space may be thought of as a vector space with no privileged origin.
Suppose that we choose an origin point o ∈ A and let { e 1 , . . .
, e n } be a basis for the associated vector space
V .
For any point p ∈ A , we can write (cid:126)op = (cid:80) ni
=1 x i ( p ) e i .
Here, { x 1 , . . .
, x n } is a set of coordinate functions, or more simply, a basis for A .
If we have two bases { x 1 , . .
, x n } and { y 1 , . . . , y n } , then they are related by y =
Bx + c where B =
[ b ij ] is an invertible n × n matrix and c
=
[ c i ] is a vector.
We now describe how to extend a change-of-basis on the afﬁne space A to the product space
A K =
A × · · · × A .
Let ι and κ be two choices of afﬁne bases on A , then
[Equation]
where homogeneous coordinates are used for a i and a † i .
Now, suppose that the change-of-basis from ι to κ is given by ( B c ) and denote [ B ] H
[Equation]
Then, we have
[Equation]
Thus, the induced change-of-basis on the product space
A K is given by the matrix I ⊗ [ B ] H .
B. D IFFERENTIALS , P USHFORWARDS AND P ULLBACKS
Let M be a smooth real manifold.
For p ∈ M , we denote the tangent space by T p M and the corresponding dual space, the cotangent space by T ∗ p M .
Given a smooth function f : M → R , the differential df ( p ) ∈ T ∗ p M is deﬁned by
[Equation]
Let ( x 1 , . . .
, x n ) be a coordinate system around p ∈ M , the differential df ( p ) can be expressed as
[Equation]
We observe that this coordinate representation corresponds to the gradient ∇ f (even though differentials and gradients are distinct objects for abstract manifolds).
For a smooth map ϕ : M 1 → M 2 of manifolds, the pushforward ϕ ∗ : T p M 1 → T ϕ ( p ) M 2 is deﬁned by
[Equation]
, where v ∈ T p M 1 and h : M 2 → R is a smooth function on M 2 .
If we suppose that M 1 = R n and
M 2 =
R m with ( x 1 , . . .
, x n ) a coordinate system around p and ( y 1 , . . .
, y m ) a coordinate system around ϕ ( p ) , the pushforward ϕ ∗ v can be represented as
[Equation]
where v = (cid:74) v (cid:75) .
This is exactly the Jacobian matrix J ϕ of ϕ and hence (cid:74)
ϕ ∗ v (cid:75) =
J ϕ v .
This Jacobian- vector product corresponds to the directional derivative, and can be computed using forward mode automatic differentiation (Schraudolph, 2002).
The dual notion of the pushforward, the pullback ϕ ∗ : T ∗ ϕ ( p ) M 2 → T ∗ p M 1 is deﬁned in the following way
[Equation]
With respect to the same coordinate systems chosen above, we can write (cid:74) ϕ
∗ u (cid:75) =
J (cid:62) ϕ u , where u = (cid:74) u (cid:75) .
Numerically, we can compute J (cid:62) ϕ
u efﬁciently using reverse mode auto-differentiation (i.e., backpropagation).
B. M ETRICS AND T HEIR P ROPERTIES
We introduce tensors on manifolds.
A symmetric k -tensor σ at the point p ∈ M is deﬁned as a symmetric k -tensor on the tangent space T p M .
Recall that this is a symmetric multilinear map on the k -fold product of T p M :
[Equation]
A metric on M is deﬁned as a smoothly varying symmetric 2 -tensor g which is positive-semideﬁnite 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 deﬁne how this works for general tensors.
Let ϕ : M 1 → M 2 be a smooth map of manifolds and σ be a symmetric k -tensor on M 2 at ϕ ( p ) .
The pullback ϕ ∗ σ of σ under ϕ is a symmetric k -tensor on M 1 deﬁned as
[Equation]
where v 1 , .
, v k ∈ T p M 1 .
In the case of metrics,
[Equation]
If we suppose that the metric g M 2 is given by G M 2 for a chosen coordinate system around ϕ ( p ) , then the pullback metric ϕ ∗ g M 2 on M 1 around p is given by
[Equation]
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 deﬁne metrics on product manifolds.
Given metrics g M 1 and g M 2 on
M 1 and M 2 respectively, we describe how to naturally deﬁne a metric on the product manifold
M 1 × M 2 .
For any point ( p, q ) ∈ M 1 × M 2 , there is a canonical isomorphism of tangent spaces:
[Equation]
The proof of this fact can be found in standard differential geometry literature (e.g. Lee (2003)) and so we do not elaborate further here.
Hence, any vector v ∈ T ( p,q ) ( M 1 × M 2 ) can be written as a pair ( v 1 , v 2 ) where v 1 ∈ T p M 1 and v 2 ∈ T q M 2 .
Now, we deﬁne the additive metric g M 1 + g Mon M 1 × M 2 as follows:
[Equation]
If we choose a coordinate system around ( p, q ) with the metrics g M 1 , g M 2 represented by matrices
G M 1 , G M 2 respectively, then we have
[Equation]
which is a matrix with block diagonals G M 1 , G M 2 and zero everywhere else.
This construction generalizes easily to sums of more than two terms.
C A DDITIONAL M ATHEMATICAL T ECHNICALITIES C. T ANGENT S PACE OF V ECTOR S PACES AND A FFINE S PACES
Theorem C.1.
Let V be a ﬁnite-dimensional vector space.
For each point p ∈ V , there is a canonical isomorphism V → T p V .
From now on, we suppress p in T p V and write TV when denoting tangent spaces of V .
Proof.
For any element v ∈ V , we can associate a tangent vector D v | p at p deﬁned by
[Equation]
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 afﬁne space and V be its associated vector space.
For each point a ∈ A , there is a canonical isomorphism between T a A and V .
From now on, we suppress a in T a A and write TA when denoting tangent spaces of A .
Proof.
Note that specifying a point a ∈ A naturally identiﬁes A with V .
Then, applying the above theorem gives the desired result.
C. T ENSOR P RODUCT OF V ECTOR S PACES
Let U and V be ﬁnite-dimensional vector spaces over the real numbers R .
Let R be the subspace of the free vector space R (cid:104) U
× V (cid:105) (set of all ﬁnite formal linear combinations of elements of U × V with real coefﬁcients) spanned by all elements of the following forms:
[Equation]
for u, u (cid:48) ∈ U , v, v (cid:48) ∈ V , and c ∈ R .
The tensor product U ⊗ V , is the quotient space R (cid:104)
U × V (cid:105) 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 (cid:55)→ ϕ ( 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 deﬁned by ( u,
[Equation]
(cid:55)→ ϕ (
[Equation] [Equation]
This induces an element on the tensor product ( U ⊗ V ) ∗ .
As such, we obtain an unique linear injection
[Equation]
Since all the vector spaces are ﬁnite-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. again, elements of this space are maps λ : T ∗ W → T ∗ Z .
D T WO P ARAMETERIZATIONS OF C OORDINATE -F REE MLP
We now show how the MLP with computation deﬁned 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 afﬁne bases for all of the activation spaces A 1 , . . .
A L − 1 , the pre-activation spaces Z 1 , . . .
, Z L , the input space X and the output space Y in the network.
Observe that a choice of bases for A i − 1 and Z i naturally induces a basis for each W i , and therefore also for the full weight space W .
Let ι, κ be two different choices of parameterizations for the network.
With respect to ι , we write (cid:74)
[Equation]
i
[Equation] [Equation]
=
[Equation]
i
[Equation]
and with respect to κ , we write
[Equation]
Hence, we can rewrite Eqn.
3.1 in the parameterizations ι, κ as
[Equation]
The parameters ( W i b i ) and ( W ‡ i b ‡ i ) are related as follows
[Equation]
where ( Ω i − 1 γ i − 1 ) is the change-of-basis from ι to κ on A i − 1 with Ω i − 1 an invertible matrix and γ i − 1 a vector.
Moreover, ( Φ i τ i ) is the change-of-basis from κ to ι on Z i with Φ i an invertible matrix and τ i a vector.
The activation functions φ i and φ ‡ i in Eqn.
D.1 are related in the following way:
[Equation]
where ( Ω i γ i ) is the change-of-basis from ι to κ on A i .
To verify Eqn.
D.2 and Eqn.
D.3, we consider the the following commutative diagram (the top horizontal arrow is equal to the composition of maps given by the other three arrows) which relates the two parameterizations ι and κ on the activation
afﬁne space
A i − 1 and the pre-activation afﬁne space
Z i :
[Equation]
Let a i − 1 ∈ (cid:74)
A i −
[Equation]
ι , this maps to z i
= W i a i − 1 + b i ∈ (cid:74) Z i (cid:75) ι under the top horizontal arrow in Eqn.
D.4.
Now, mapping a i − 1 under the composition of the other three arrows in Eqn.
D.4, we obtain
[Equation]
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):
[Equation]
Let z ‡ i ∈ (cid:74)
Z i (cid:75) κ , this maps to a ‡ i ∈ (cid:74)
A i (cid:75) κ under φ ‡ i .
Now, mapping z
i under the other three arrows in Eqn.
D.5, we obtain z ‡ i (cid:55)→ Φ i z
‡ i + τ i (apply left vertical arrow in Eqn.
D.5)
(cid:55)→
φ i ( Φ i z ‡ i + τ
i ) .
(apply upper horizontal arrow in Eqn.
D.5)
(cid:55)→
Ω i φ i ( Φ i z
‡ i + τ i )
+ γ
i
(apply right vertical arrow in Eqn.
D.5) .
This establishes Eqn.
D.3.
Lastly, we note that the equations in Eqn.
D.1 can be rewritten in homogeneous coordinates:
[Equation]
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.
This completes the proof.
E P ROOFS E. P ROOF OF T HEOREM 3.
First, note that g ( ω k ) − 1 dh ( ω k ) is an intrinsically deﬁned tangent vector on W .
The weight space
W of a MLP is an afﬁne 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 afﬁne basis for W , the algorithm is invariant to afﬁne reparameterizations.
E. P ROOF OF T HEOREM 3.
Before providing the proof of the theorem, we need the following lemma.
Lemma E.1.
Let g be a metric on Z and ψ a : W → Z be the evaluation map.
Then, the pullback metric ψ ∗ a g on W can be expressed as:
[Equation]
⊗ φ ] , where φ is a random covector.
Proof.
Given z ∈ Z , the metric g admits the rank-1 decomposition
[Equation]
where φ is a covector and the expectation is over Z .
This is akin to the more familiar case where any symmetric positive-semideﬁnite matrix admits a rank-1 spectral decomposition.
Computing the pullback of g under the map ψ a now gives
[Equation]
where z = ψ a ( w ) .
We analyze the pullback ψ ∗ a φ .
Let ∂ w be a tangent vector on W .
Then,
[Equation]
which shows that ψ ∗ a φ = a ⊗ φ .
Plugging this back into Eqn.
E.1, we obtain
[Equation]
which concludes the proof.
For the ﬁrst assertion of the theorem, we need to check that both components E
[ a ⊗ a ] and E [ g ( z )] deﬁne symmetric positive-semideﬁnite 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-semideﬁnite property,
[Equation]
where the latter inequality is due to the fact that a ⊗ a is positive-semideﬁnite.
Moreover, a ⊗ a is also symmetric and this property is preserved under expectations which implies that E [ a ⊗ a ] is both symmetric and positive-semideﬁnite.
For the second term E [ g ( z )] in g ind ( w ) , the fact that g is a metric on Z means that g ( z ) , by deﬁnition, is a symmetric positive-semideﬁnite 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-deﬁnite.
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. for further explanations) such that
[Equation]
and hence ( a ⊗ a )( λ ⊗ λ )
= 0 .
Now, consider the element λ ⊗ µ ∈ ( TW ⊗ T ∗ Z ) ⊗ TZ where µ ∈ TZ .
From the above lemma, we have
[Equation]
Then, evaluating E a [ ψ ∗ a g ( w )] at λ ⊗ µ :
[Equation]
This shows that E a [ ψ ∗ a g ] is not positive-deﬁnite which yields a contradiction as E a [ ψ ∗ a g ] was assumed to be a nondegenerate metric.
The exact same argument can be applied to show that E [ g ( z )] is positive-deﬁnite.
This gives us the desired result.
E. P ROOF OF T HEOREM 3.
From Theorem 3.4, we know that g i ind is a metric on W i .
Since g KFAC is deﬁned as the additive metric where each of the summands are g i ind , we can conclude that g KFAC is a metric.
For the second assertion of the theorem, recall that Ψ iξ
=
ϕ iξ ◦ ψ iξ .
By the functoriality property of pullback operations, we have
[Equation]
Since E ξ [(Ψ iξ ) ∗ g ] was assumed to be nondegenerate, this implies that
[Equation]
)] , is also nondegenerate.
Then, from the second assertion of Theorem 3.4, we obtain that g i ind is a nondegenerate metric on W i .
Consequently, g KFAC is nondegenerate which concludes the proof.
F C OORDINATE -F REE K-FAC FOR C ONVOLUTIONAL N ETWORKS F. C ONVOLUTIONAL N ETWORKS
We provide a complete account of the convolutional network case sketched out earlier in Section 4.1.
We begin by describing the convolution layer of a convolutional network in mathematical terms following Grosse & Martens (2016).
We only consider convolution layers as the pooling and response normalization layers of a convolutional network typically do not contain (many) trainable weights.
We then introduce the notion of a transformed convolution layer analogous to what was done in the case of MLPs.
Lastly, we use the abstract linear algebra machinery developed in Section 2 to give a coordinate-free description of convolution layers.
F.1. C ONVOLUTION L AYERS
We focus on a single convolution layer.
A convolution layer l takes as input a layer of activations a j,t , where j ∈ { 1 , . .
, J } indexes the input map and t ∈ T indexes the spatial location.
T here denotes the set of spatial locations, which we typically take to be a 2D-grid.
We assume that the convolution is performed with a stride of 1 and padding equal to the kernel radius R , so that the set of spatial locations is shared between the input and output feature maps.
This layer is parameterized by a set of weights w i,j,δ and biases b i , where i ∈ { 1 , . .
, I } indexes the output map and δ ∈ ∆ indexes the spatial offset.
The numbers of spatial locations and spatial offsets are denoted by |T | and | ∆ | respectively.
The computation at the convolution layer is given by
[Equation]
The pre-activations z i,t are then passed through a nonlinear activation function φ l .
Analogous to feed-forward networks, the weight derivatives are computed using backpropagation:
[Equation]
Following Grosse & Martens (2016), we represent the convolution layer computation in Eqn.
F. using matrix notation.
To do this, we write the activations a j,t as a J × |T | matrix A l − 1 , the preactivations z i,t as a I × |T | matrix Z l , the weights w i,j,δ as a
I × J | ∆ | matrix W l and the bias vector as b l .
For the activation matrix A l − 1 , if we extract the patches surrounding each spatial location t ∈ T and ﬂatten these patches into vectors where the vectors become columns of a matrix, we obtain a J | ∆ | × |T | matrix which we denote by A exp l − 1 .
From now on, we refer to this matrix as the expanded activations.
Finally, we can use these matrix notations to rewrite the computation in
Eqn. F.1 as
[Equation]
For convenience purposes later, we adopt homogeneous coordinates for various matrices:
[Equation]
Hence, Eqn.
F.2 can be rewritten as
[Equation]
where the activation function φ l here ignores the homogeneous coordinate.
We introduce the concept of a transformed convolution layer.
For a convolution layer as deﬁned in the above equation, the parameters [ W l ] H and the transformed parameters [ W † l ] H are related in the following way
[Equation]
where Γ l and Υ l − 1 are invertible matrices.
The activation functions φ l and φ † l are related through a standard afﬁne change-of-basis as given in Eqn. 2.9.
F
[Equation]
A BSTRACT C ONVOLUTION L AYERS
Just as in the coordinate-dependent case earlier, we focus on a single layer.
An abstract convolution layer l is deﬁned as follows:
• Local activations at each spatial location t ∈ T are taken to be elements α l − 1 in an afﬁne space
A l − 1 .
• Activations are taken to be elements α (:) l − 1 in A |T | l − 1 , (i.e. the direct product of A , |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 −. • Local pre-activations at t ∈ T are taken to be elements ζ ( t ) l in an afﬁne space Z l .
• Pre-activations are taken to be elements ζ (:) l in Z |T
| l . • Layerwise parameters are afﬁne transformations ω l between A | ∆ | l − 1 and Z l .
The collection of these
transformations
is an
afﬁne
space in its
own right
which
we denote
by
W l and
refer to as the
layerwise weight space
If
we apply ω
l
pointwise
, this
can
be extended to a map
[Equation]
The computation for this abstract layer is
[Equation]
where ρ l is a ﬁxed nonlinear activation function and α l are the l -th layer local activations deﬁned in exactly the same manner as α l − 1 .
We choose afﬁne bases on A l − 1 , Z l , and A l .
A basis on A l − 1 naturally induces a basis for A | ∆ | l − 1 .
Consequently, this gives a basis also for the layerwise parameter space W l .
Let ι , κ be two such choices.
With respect to ι , we write
[Equation]
and with respect to κ , we write
[Equation]
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 Z l and a l are J -dimensional column vectors of activations matrix A l .
Now, suppose that ( Ω l − 1 γ l − 1 ) is the change-of-basis from ι to κ on A l − 1 and ( Φ l τ l ) is the change-of-basis from κ to ι on Z l .
If we denote
[Equation]
then by the afﬁne change-of-basis formula for direct products (Eqn. B.1), I ⊗ [ Ω l − 1 ]
H deﬁnes the change-of-basis from ι to κ on A | ∆ | l − 1 .
The parameters [ W l ] H and [ W ‡ l ] H are related as follows:
[Equation]
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.
F. K RONECKER F ACTORS FOR C ONVOLUTION
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 ,
[Equation]
For an input-target pair ( x , y ) , the Fisher matrix for this network is
[Equation]
where D w is the log-likelihood gradient and the expectation is taken over the model’s predictive distribution P y | x ( w ) for y and over the data distribution for x .
The diagonal blocks F ( w ) l,l of
F ( w ) are
[Equation]
We are ready now to present the K-FAC approximation for convolutional networks.
For a particular layer l , we deﬁne the K-FAC approximation ˆ F ( w ) l,l to F ( w ) l,l as:
[Equation]
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.
F. C OORDINATE -F REE K-FAC FOR C ONVOLUTIONAL N ETWORKS
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 W l is the space of afﬁne transformations between A | ∆ | l − 1 and Z l .
The weight space of this network is the direct product of all layerwise weight spaces
[Equation]
Given an input ξ ∈ X and parameter ω
=
( ω 1 , . . .
, ω L ) ∈ W , denote the network output by f ( ξ, ω ) .
Now, for every l ∈ { 1 , . . .
, L } , deﬁne the following maps
• ψ lξ : W l → 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 deﬁnition.
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 deﬁne the map Ψ lξ :
W l → 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.
to be
[Equation]
and the metric on Z = Z l to be ( ϕ lξ,t ) ∗ g .
Summing over every spatial location t ∈ T , the independence metric on W l here is
[Equation]
Deﬁnition F.2.
The K-FAC metric on the weight space W of an abstract convolutional network is deﬁned as
[Equation]
where the sum above is as deﬁned in Eqn.
B.2 and each g l ind is as given in Eqn.
F.6.
Theorem F.3.
Let g be a metric on Y .
Then, g KFAC given in Deﬁnition F.2 is indeed a metric on the weight space W of an abstract convolutional network.
Moreover, if we assume that the expected pullback of g restricted to a single spatial location t ∈ T ,
[Equation]
under the map Ψ lξ,t : W l → Y is a nondegenerate metric on the layerwise weight space W l for every l , then g KFAC 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
[Equation]
is a metric on W l .
Since taking expectation over the set of spatial locations T and multiplying by the scale factor |T | preserves the metric properties, we obtain that g l ind in Eqn.
F.6 deﬁnes a metric on W l .
Consequently, g KFAC determines a metric on W .
To prove the latter assertion, note that by the functorial property of pullback operations,
[Equation]
is nondegenerate.
Using the second assertion of Theorem 3.4 yields that g l ind is nondegenerate which implies that this is true also for g KFAC .
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 g F, out .
For each spatial location t ∈ T , the pullback under ϕ lξ,t is
[Equation]
Now, choose coordinate systems on A l −
Using Proposition 3.5, the independence metric in Eqn.
F.6 can be expressed in coordinates as follows
[Equation]
which is exactly ˆ F ( w ) l,l given earlier in Eqn.
F.5.
Furthermore, (cid:74) g KFAC ( ω ) (cid:75) 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 g KFAC for abstract convolutional networks.
Lastly, if g KFAC is a nondegenerate metric; which is true for example if the assumptions in the second assertion of Theorem F.3 hold, then we can conclude that these updates are invariant to any afﬁne reparameterizations of the model.
G C OORDINATE -F REE K-FAC FOR R ECURRENT N ETWORKS
In this section, we provide a complete account of the recurrent network case sketched out earlier in Section 4.2.
We ﬁrst give a mathematical formulation of the recurrent computation step of these networks in both coordinate-dependent and coordinate-independent scenarios.
We proceed to give the Kronecker factorization of the Fisher matrix for recurrent networks and then state the invariance theorem for this optimization method.
Lastly, we prove the invariance theorem in the same way we did for MLPs and convolutional networks previously.
G. R ECURRENT N ETWORKS
As in the case of convolutional networks in Section F, it is not necessary to write out the full structure of a recurrent network.
Rather, we focus on the recurrent computation since the central object of our interest, the Fisher matrix for recurrent networks, only involves recurrent weights.
G.1. C OMPUTATIONAL S TEP
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 ﬁxed length T .
For an input x t at every t , the recurrent network maps this to an output o t .
Essentially, the network maps input sequences x
= ( x 1 , . .
, x T ) to output sequences o =
( o 1 , . .
, o T ) .
The computation, at every t , is
[Equation]
where a t − 1 is an activation vector, z t is a pre-activation vector, W is a recurrent weight matrix, V is a weight matrix, b is a recurrent bias vector, and φ is a ﬁxed nonlinear activation function.
For the remainder of this section, we focus on the ﬁrst equation
[Equation]
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 deﬁned as
[Equation]
The relationship between transformed parameters ( W † b † ) and original parameters ( W b ) is given by the standard change-of-basis formula as in Eqn.
D.2.
G.1.
A BSTRACT R ECURRENT N ETWORK
We now describe an abstract recurrent network formally.
• Local activations at each time step t are elements α t in an afﬁne space
A • Activations are elements α
= { α t } t ∈ T in the afﬁne space
A T • Local pre-activations at each t are elements ζ t in an afﬁne space
Z • Pre-activations are elements ζ = { ζ t } t ∈ T in the afﬁne space Z T
• Parameters are afﬁne transformations ω between A and Z .
The collection of these transformations is an afﬁne 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 afﬁne spaces X , Y respectively.
The input and output spaces are X T and Y T respectively; furthermore, elements here are written as ξ = { ξ t } t ∈ T and υ = { υ t } t ∈ T .
For every t , the abstract recurrent computation step is ζ t = ω ( α t − 1 ) .
A choice of parameterization for the abstract recurrent network consists of choosing afﬁne bases for A , Z , X and Y .
Since we have bases for A and Z , this induces a natural basis for W .
If we use exactly the same change-of-basis analysis given in Appendix D, then the recurrent network with computation given by Eqn.
G.1 and the transformed version in Eqn.
G.2 correspond to two different parameterizations of the same abstract recurrent network.
G. K-FAC FOR R ECURRENT N ETWORKS
We review the recent Kronecker factorization for recurrent networks method in Martens et al.
Recall that for every time step t , the recurrent computation can be written as
[Equation]
where ¯ W =
[ W b ] and ¯ a (cid:62) t − 1 =
[ a (cid:62) t − 1 1] (cid:62) .
Using backpropagation, the log-likelihood gradient is given by D z t ¯ a (cid:62) t − 1 .
The total contribution to the gradient across all t is the sum
D ¯ W =
T (cid:88) t =
D z t ¯ a (cid:62) t − 1 .
For an input-target pair ( x , y ) , the Fisher matrix F ( ¯ W ) for recurrent networks is deﬁned as
F ( ¯ W ) =
E x , y [vec( D ¯ W ) vec( D ¯ W ) (cid:62) ] .
Finally, the K-FAC approximation ˆ F ( ¯ W ) to F ( ¯ W ) for recurrent networks is deﬁned as
[Equation]
Remark G.1 .
As in the case of convolution layers, there is weight sharing in recurrent networks (across time here instead of spatial locations) and so it is not enough to just assume statistical independence between activations and pre-activation derivatives to make the K-FAC approximation here.
We defer the reader to Martens et al.
(2018) for detailed explanations on how the K-FAC approximation is derived for recurrent networks.
We now proceed to the last section of this paper to give a coordinate-free proof of Theorem 4.2.
The method of proof mirrors exactly the proofs given previously for Theorem 2.1 and Theorem 4.1.
G. C OORDINATE -F REE K-FAC FOR R ECURRENT N ETWORKS
Given an input ξ = { ξ t } t ∈ T ∈ X T and parameter ω ∈ W , denote the network output by f ( ξ, ω ) .
For a speciﬁc time step t , consider the following maps:
• ψ ξ,t : W → Z which sends parameters ω to pre-activations ζ t by evaluation at activations α t −• ϕ ξ,t : Z → Y which sends ζ t to outputs
In addition, we deﬁne the map Ψ ξ,t : W → Y as the composition Ψ ξ,t = ϕ ξ,t ◦ ψ ξ,t .
Let g be a metric on Y .
The pullback ϕ ∗ ξ,t g then deﬁnes a metric on Z .
Now, we take A , Z , W in
Section 3.1 to be
A =
[Equation] [Equation]
Z , W = W , and the metric on Z
= Z to be ϕ ∗ ξ,t g .
Summing over all time steps t ∈ T , we make the following deﬁnition which arises from the independence metric in Section 3.2:
Deﬁnition G.2.
The K-FAC metric on the weight space W of an abstract recurrent network is deﬁned as
[Equation]
Theorem G.3.
Let g be a metric on Y .
Then, g KFAC given in Deﬁnition G.2 is a metric on the weight space W of an abstract recurrent network.
Moreover, if we assume that the expected pullback of g ,
[Equation]
under the smooth map Ψ ξ,t : W → Y is a nondegenerate metric, then g KFAC 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
[Equation]
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 g KFAC is a metric on W .
For the nondegeneracy statement, using the functorial property of pullbacks, we know that
[Equation]
is nondegenerate.
Then, g KFAC 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 g F, out .
Then, the pullback under ϕ ξ,t is
[Equation]
Now, choose coordinate systems for A and Z .
We can write
[Equation]
By Proposition 3.5, the K-FAC metric in Eqn.
G.4 can be represented in these chosen coordinates as
[Equation]
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 g KFAC for abstract recurrent networks.
If we suppose that g KFAC is nondegenerate, then these updates are invariant to any afﬁne reparameterizations of the model.