Title: KRONECKER-FACTORED CURVATURE APPROXIMA-TIONS FOR RECURRENT NEURAL NETWORKS

Abstract: Kronecker-factor Approximate Curvature (Martens & Grosse, 2015) (K-FAC) is a 2nd-order optimization method which has been shown to give state-of-the-art performance on large-scale neural network optimization tasks (Ba et al., 2017) . It is based on an approximation to the Fisher information matrix (FIM) that makes assumptions about the particular structure of the network and the way it is parameterized. The original K-FAC method was applicable only to fully-connected networks, although it has been recently extended by Grosse & Martens (2016) to handle convolutional networks as well. In this work we extend the method to handle RNNs by introducing a novel approximation to the FIM for RNNs. This approximation works by modelling the statistical structure between the gradient contributions at different time-steps using a chain-structured linear Gaussian graphical model, summing the various cross-moments, and computing the inverse in closed form. We demonstrate in experiments that our method significantly outperforms general purpose state-of-the-art optimizers like SGD with momentum and Adam on several challenging RNN training tasks.

1. INTRODUCTION:

As neural networks have become ubiquitous in both research and applications the need to efficiently train has never been greater. The main workhorses for neural net optimization are stochastic gradient descent (SGD) with momentum and various 2nd-order optimizers that use diagonal curvature-matrix approximations, such as RMSprop (Tieleman & Hinton, 2012) and Adam (Ba & Kingma, 2015) . While the latter are typically easier to tune and work better out of the box, they unfortunately only offer marginal performance improvements over well-tuned SGD on most problems.

Because modern neural networks have many millions of parameters it is computationally too expensive to compute and invert an entire curvature matrix and so approximations are required. While early work on non-diagonal curvature matrix approximations such as TONGA (Le Roux et al., 2008) and the Hessian-free (HF) approach (Martens, 2010; Martens & Sutskever, 2011; 2012; Desjardins et al., 2013; Sainath et al., 2013) demonstrated the potential of such methods, they never achieved wide adoption due to issues of scalability (to large models in the case of the former, and large datasets in the case of the latter).

Motivated in part by these older results and by the more recent success of centering and normalization methods (e.g. Schraudolph, 1998; Vatanen et al., 2013; Ioffe & Szegedy, 2015) a new family of methods has emerged that are based on non-diagonal curvature matrix approximations the rely on the special structure of neural networks. Such methods, which include Kronecker-factored approximated curvature (K-FAC) (Martens & Grosse, 2015) , Natural Neural Nets (Desjardins et al., 2015) , Practical Riemannian Neural Networks (Marceau-Caron & Ollivier, 2016) , and others (Povey et al., 2015) , have achieved state-of-the-art optimization performance on various challenging neural network training tasks and benchmarks.

While the original K-FAC method is applicable only to standard feed-forward networks with fully connected layers, it has recently been extended to handle convolutional networks (Grosse & Martens, 2016) through the introduction of the "Kronecker Factors for Convolution" (KFC) approximation. Ba et al. (2017) later developed a distributed asynchronous version which proposed additional approximations to handle very large hidden layers.

In this work we develop a new family of curvature matrix approximations for recurrent neural networks (RNNs) within the same design space. As in the original K-FAC approximation and the KFC approximation, we focus on the Fisher information matrix (a popular choice of curvature matrix), and show how it can be approximated in different ways through the adoption of various approximating assumptions on the statistics of the network's gradients. Our main novel technical contribution is an approximation which uses a chain-structured linear Gaussian graphical model to describe the statistical relationship between gradient contributions coming from different time-steps. Somewhat remarkably, it is possible to sum the required cross-moments to obtain a Fisher approximations which has enough special algebraic structure that it can still be efficiently inverted. In experiments we demonstrate the usefulness of our approximations on several challenging RNN training tasks.

2. NOTATION AND BACKGROUND:



2.1. NETWORK, LOSS, AND OBJECTIVE FUNCTION:

We denote by f (x, θ) the neural network function associated evaluated on input x, where θ are the parameters. We will assume a loss function of the form L(y, z) = -log r(y|z), where r is the density function associated with a predictive distribution R. The loss associated with a single training case is then given by L(y, f (x, θ)) ≡ -log r(y|f (x, θ)). Throughout the rest of this document we will use the following special notation for derivatives of the single-case loss w.r.t. some arbitrary variable Z (possibly matrix-valued): DZ = dL(y, f (x, θ)) dZ

The objective function which we wish to minimize is the expected loss h(θ) = E Q [L(y, f (x, θ))]

over the training distribution Q on x and y.

2.2. THE FISHER, THE NATURAL GRADIENT, AND 2ND-ORDER OPTIMIZATION:

The Fisher information matrix (aka "the Fisher") associated with the model's predictive distribution P y|x (θ) is given by

F = E DθDθ = cov(Dθ, Dθ).

Note that here, and for the remainder of this paper, y is taken to be distributed according to the model's predictive conditional distribution P y|x (θ), so that E[DZ] = 0 for any variable Z that is conditionally independent of y given the value of f (x, θ) (this includes Dθ). All expectations and covariances are defined accordingly. This is done because the expectation that defines the Fisher information matrix uses P y|x (θ). If we were to instead use the training distribution Q y|x on y, we would essentially be computing to the "empirical Fisher" (or approximations thereof), which as argued by Martens (2014) is a less appropriate choice for a curvature matrix than the true Fisher.

The natural gradient is defined as F -1 ∇h, and is the update direction used in natural gradient descent. As argued by Amari (1998) , natural gradient descent has the two key advantages: it is invariant to the parameterization of the model, and has "Fisher efficient" convergence 1 . However, as shown by Martens (2014) these two facts have several important caveats. First, the parameterization invariance only holds approximately in practice when non-infinitesimal step-sizes are used. Second, Fisher efficiency is actually a weak property possessed by simpler methods like SGD with Polyak/parameter averaging (Polyak & Juditsky, 1992) , and even then will only be achieved when the method converges to a global minimizer and the model is capable of perfectly capturing the true distribution of y given x.

An alternative explanation for the empirical success of the natural gradient method is that it is a 2ndorder method, whose update minimizes the following local quadratic approximation to the objective h(θ + δ):

EQUATION

This is similar to the 2nd-order Taylor series approximation of h(θ + δ), but with the Fisher substituted in for the Hessian. This substitution can be justified by the observation that the Fisher is a kind of PSD approximation to the Hessian (Pascanu & Bengio, 2014; Martens, 2014) . And as argued by Martens (2014) , while stochastic 2nd-order methods like natural gradient descent cannot beat the asymptotically optimal Fisher efficient convergence achieved by SGD with Polyak averaging, they can enjoy better pre-asymptotic convergence rates. Moreover, insofar as gradient noise can be mitigated through the use of large mini-batches -so that stochastic optimization starts to resemble deterministic optimization -the theoretical advantages of 2nd-order methods become further pronounced, which agrees with the empirical observation that the use of large-minibatches speeds up 2nd-methods much more than 1st-order methods (Martens & Grosse, 2015; Ba et al., 2017) .

In addition to providing an arguably better theoretical argument for the success of natural gradient methods, their interpretation as 2nd-order methods also justifies the common practice of computing the natural gradient as (F + λI) -1 ∇h instead of F -1 ∇h. In particular, this practice can be viewed as a type of "update damping/regularization", where one encourages δ to lie within some region around δ = 0 where eqn. 1 remains a trustworthy approximation (e.g. Nocedal & Wright, 2006; Martens & Sutskever, 2012) .

2.3. KRONECKER-FACTORED APPROXIMATE CURVATURE (K-FAC):

Because modern neural network have millions (or even billions) of parameters it is computationally too expensive to compute and invert the Fisher. To address this problem, the K-FAC method of Martens & Grosse (2015) uses a block-diagonal approximation of the Fisher (where the blocks correspond to entire layers/weight matrices), and where the blocks are further approximated as Kronecker products between much smaller matrices. The details of this approximation are given in the brief derivation below.

Let W be a weight matrix in the network which computes the mapping s = W a, where a and s are vector-valued inputs and outputs respectively and denote g = Ds.

As in the original K-FAC paper we will assume that a includes a homogeneous coordinate with value 1 so that the bias vector may be folded into the matrix W .

Here and throughout the rest of this document, F will refer to the block of the Fisher corresponding to this particular weight-matrix W .

The Kronecker product of matrices B and C, denoted by B ⊗ C for matrices B ∈ R m×n and C of arbitrary dimensions, is a block matrix defined by

B ⊗ C ≡    [B] 1,1 C • • • [B] 1,n C . . . . . . . . . [B] m,1 C • • • [B] m,n C   

Note that the Kronecker product has many convenient properties that we will make use of in this paper. (See Van Loan (2000) for a good discussion of the Kronecker product and its properties.)

A simple application of the chain rule gives DW = ga . If we approximate g and a as statistically independent, we can write F as

F = E[vec(DW) vec(DW ) ] = E[vec(ga ) vec(ga ) ] = E[(a ⊗ g)(a ⊗ g) ] = E[(aa ) ⊗ (gg )] = E[aa ] ⊗ E[gg ] = A ⊗ G,

where we have defined

A = E[aa ] and G = E[gg ].

The matrices A and G can be estimated using simple Monte Carlo methods, and averaged over lots of data by taking an exponentially decaying average across mini-batches. This is the basic Kronecker factored approximation (Heskes, 2000; Martens & Grosse, 2015; Povey et al., 2015) , which is related to the approximation made in the Natural Neural Nets approach (Desjardins et al., 2015) . It is shown by Martens & Grosse (2015) that the approximation is equivalent to neglecting the higher-order cumulants of the as and gs, or equivalently, assuming that they are Gaussian distributed. To see why this approximation is useful, we observe that inversion and multiplication of a vector by F amounts to inverting the factor matrices A and G and performing matrix-matrix multiplications with them, due to the following two basic identities:

EQUATION

The required inversion and matrix multiplication operations are usually computational feasible because the factor matrices have dimensions equal to the size of the layers, which is typically just a few thousand. And when they are not, additional approximations can be applied, such as approximate/iterative inversion (Povey et al., 2015) , or additional Kronecker-factorization applied to either A or G (Ba et al., 2017) . Moreover, the computation of the inverses can be amortized across iterations of the optimizer at the cost of introducing some staleness into the estimates.

3. APPROXIMATING F FOR RNNS:

The basic Kronecker-factored approximation to the Fisher block F described in the previous section assumed that the weight matrix W was used to compute a single mapping of the form s = W a.

When W is used to compute multiple such mappings, as is often the case for RNNs, or a mapping of a different flavor, as is the case for convolutional networks (CNNs), the approximation is not applicable, strictly speaking. Grosse & Martens (2016) recently showed that by making additional approximating assumptions, the basic Kronecker-factored approximation can be extended to convolutional layers. This new approximation, called "KFC", is derived by assuming that gradient contributions coming from different spatial locations are uncorrelated, and that their intra and inter-location statistics are spatially homogeneous, in the sense that they look the same from all reference locations. These assumptions are referred to "spatially uncorrelated derivatives" and "spatial homogeneity," respectively.

In this section we give the main technical contribution of this paper, which is a family of Kroneckerbased approximations of F that can be applied to RNNs. To build this we will apply various combinations of the approximating assumptions used to derive the original K-FAC and KFC approaches, along with several new ones, including an approximation which works by modelling the statistical structure between the gradient contributions from time-steps using a chain-structured linear Gaussian graphical model.

3.1. PRELIMINARIES:

Let W be some weight matrix which is used at T different time-steps (or positions) to compute the mapping

s t = W a t ,

where t indexes the time-step. T is allowed to vary between different training cases.

Defining g t = Ds t , the gradient of the single-case loss with respect to W can be written as

DW = T t=1 g t a t = T t=1 D t W,

where D t W = g t a t denotes the contribution to the gradient from time-step t. When it is more convenient to work with the vector-representations of the matrix-valued variables D t W we will use the notation

w t = vec(D t W ),

so that vec(DW ) = T t=1 w t . Let F T denote the conditional Fisher of DW for a particular value of T . We have

EQUATION

Observe that F can be computed from

F T via F = E T [F T ].

To proceed with our goal of obtaining a tractable approximation to F we will make several approximating assumptions, as discussed in the next section.

3.2. BASIC INITIAL APPROXIMATIONS:



3.2.1. INDEPENDENCE OF T:

One simplifying approximation we will make immediately is that T is independent of the w t 's, so

that E[w t w s |T ] = E[w t w s ].

In this case eqn. 3 can be written as

EQUATION

where we have defined

V t,s = E[w t w s ].

Independence of T and the w t 's is a reasonable approximation assumption to make because 1) for many datasets T is constant (which formally implies independence), and 2) even when T varies substantially, shorter sequences will typically have similar statistical properties to longer ones (e.g. short paragraphs of text versus longer paragraphs).

3.2.2. TEMPORAL HOMOGENEITY:

Another convenient and natural approximating assumption we will make is that the w t 's are temporally homogeneous, which is to say that the statistical relationship between any w t and w s depends only on their distance in time (d = ts). This is analogous to the "spatial homogeneity" assumption of KFC. Under this assumption the following single-subscript notation is well-defined:

V t-s = E[w t w s ]. We note that V -d = V d .

Applying this notation to eqn. 4 we have

EQUATION

where we have used the fact that there are T -|d| ways to write d as ts for t, s ∈ {1, 2, . . . , T }.

Temporal homogeneity is a pretty mild approximation, and is analogous to the frequently used "steady-state assumption" from dynamical systems. Essentially, it is the assumption that the Markov chain defined by the system "mixes" and reaches its equilibrium distribution. If the system has any randomness, and its external inputs reach steady-state, the steady-state assumption is quite accurate for states sufficiently far from the beginning of the sequence (which will be most of them).

3.2.3. INDEPENDENCE BETWEEN THE a t 'S AND THE g t 'S:

If we have that a t and g s are pair-wise independent for each t and s, which is the obvious generalization of the basic approximation used to derive the K-FAC approach, then following a similar derivation to the one from Section 2.3 we have

V t,s = E[(a t a s ) ⊗ (g t g s )] = E[(a t a s ) ⊗ (Ds t Ds s )] = A t,s ⊗ G t,s ,

where we have defined

A t,s = E[a t a s ] and G t,s = E[g t g s ].

Extending our temporal homogeneity assumption from the w t 's to the a t 's and g t 's (which is natural to do since w t = vec(g t a t )), the following notation becomes well-defined:

A t-s = A t,s and G t-s = G t,s ,

which allows us to write

V d = A d ⊗ G d .

3.3. AN INITIAL ATTEMPT TO OBTAIN A TRACTABLE FISHER APPROXIMATION:

Given the approximating assumptions made in the previous subsections we have

F T = T d=-T (T -|d|)V d = T d=-T (T -|d|)(A d ⊗ G d ).

Assuming for the moment that all of the training sequences have the same length, so that F = F T0 for some T 0 , we have that F will be the sum of 2T 0 + 1 Kronecker products.

Without assuming any additional structure, such as a relationship between the various A d 's or G d 's, there doesn't appear to be any efficient way to invert such a sum. One can use the elementary identity (B⊗C) -1 = B -1 ⊗C -1 to invert a single Kronecker product, and there exists decomposition-based methods to efficiently invert sums of two Kronecker products (see Martens & Grosse (2015) ), however there is no known efficient algorithm for inverting sums of three or more Kronecker products.

Thus is appears that we must make additional approximating assumptions in order to proceed.

3.4. ASSUMING INDEPENDENCE OF THE w t 'S ACROSS TIME:

If we assume that the contributions to the gradient (the w t 's) are independent across time, or at least uncorrelated, this means that V d = 0 for d = 0. This is analogous to the "spatially uncorrelated derivatives" assumption of KFC.

In this case eqn. 5 simplifies to

F T = T d=-T (T -|d|)V d = (T -0)V 0 = T V 0 , so that F = E T [F T ] = E T [T V 0 ] = E T [T ]V 0 .

Using the identities in eqn. 2, and the symmetry of A 0 and G 0 , we can thus efficiently multiply F -1 by a vector z = vec(Z) using the formula

EQUATION

This is, up to normalization by E T [T ], identical to the inverse multiplication formula used in the original K-FAC approximation for fully-connected layers.

We note that E T [T ] = i ω i T i , where T i are the different values of T , and ω i 0 are normalized weights (with i ω i = 1) that measure their proportions in the training set.

3.5. MODELING THE RELATIONSHIPS BETWEEN THE w t 'S USING AN LGGM:

As we saw in Section 3.3, the approximation assumptions made in Section 3.2 (independence of T , temporal homogeneity, and independence between the a t 's and the g t 's), aren't sufficient to yield a tractable formula for F -1 . And while additionally assuming independence across time of the w t 's is sufficient (as shown in Section 3.4), it seems like an overly severe approximation to make.

In this section we consider a less severe approximation which we will show still produces a tractable F -1 . In particular, we will assume that the statistical relationship of the w t 's is described by a simple linear Gaussian graphical model (LGGM) with a compact parameterization (whose size is independent of T ). Such an approach to computing a tractable Fisher approximations was first explored by Grosse & Salakhutdinov (2015) for RBMs, although our use of it here is substantially different, and requires additional mathematical machinery.

The model we will use is a fairly natural one. It is a linear Gaussian graphical model with a onedimensional chain structure corresponding to time. The graphical structure of our model is given by the following picture:

… … w t 1 w t w t+1 ✏ t 1 ✏ t ✏ t+1

Variables in the model evolve forward in time according to the following equation:

w t = Ψw t-1 + t

where Ψ is a square matrix and t are i.i.d. from N (0, Σ) for some positive definite matrix Σ (which is the conditional covariance of w t given w t-1 ).

Due to the well-known equivalence between directed and undirected Gaussian graphical models for tree-structured graphs like this one, the decision of whether to make the edges directed or undirected, and whether to have them point forwards or backwards in time, are irrelevant from a modeling perspective (and thus to the Fisher approximation we eventually compute). We will use a directed representation purely for mathematical convenience.

We will assume that our model extends infinitely in both directions, with indices in the range (-∞, ∞), so that the w t 's are all in their stationary distribution (with respect to time). For this to yield a well-defined model we require that Ψ has spectral radius < 1.

The intuition behind this model structure is clear. The correlations between gradient contributions (the w t 's) at two different time-steps should be reasonably well explained by the gradient contributions made at time-steps between them. In other words, they should be approximately Markovian.

We know that the gradient computations are generated by a process, Back-prop Through Time (BPTT), where information flows only between consecutive time-steps (forwards through time during the "forward pass", and backwards during the "backwards pass"). This process involves temporal quantities which are external to the w t 's, such as the inputs x and activations for other layers, which essentially act as "hidden variables". The evolution of these external quantities may be described by their own separate temporal dynamics (e.g. the unknown process which generates the true x's), and thus the w t 's won't be Markovian in general. But insofar as the w t 's (or equivalently the a t 's and g t 's) encode the relevant information contained in these external variables, they should be approximately Markovian. (If they contained all of the information they would be exactly Markovian.)

A similar approximation across consecutive layers was made in the "block-tridiagonal" version of the original K-FAC approach. It was shown by Martens & Grosse (2015) that this approximation was a pretty reasonable one. The linear-Gaussian assumption meanwhile is a more severe one to make, but it seems necessary for there to be any hope that the required expectations remain tractable.

3.5.1. INITIAL COMPUTATIONS:

Define the following "transformed" versions of F T and Ψ:

FT = V 1/2 0 F T V 1/2 0 and Ψ = V1 = V -1/2 0 ΨV 1/2 0 .

As shown in Section A.1 of the appendix we have

EQUATION

where

ζ T (x) = T (1 -x) -x(1 -x T ) (1 -x) 2 .

(Note that rational functions can be evaluated with matrix arguments in this way, as discussed in Section A.1.)

Our goal is to compute F -1 , from which we can recover F -1 via the simple relation

F -1 = V -1/2 0 F -1 V -1/2 0 .

Unfortunately it doesn't appear to be possible to simplify this formula sufficiently enough to allow for the efficient computation of F -1 = E T [ FT ] -1 when Ψ is a Kronecker product (which it will be when V 0 and V 1 are). The difficulty is due to both the appearance of Ψ and its transpose (which are not codiagonalizable/commutative in general), and various higher powers of Ψ.

To proceed from this point and obtain a formula which can be efficiently evaluated when Ψ is a Kronecker product, we will make one of two simplifying assumptions/approximations, which we call "Option 1" and "Option 2" respectively. These are explained in the next two subsections.

3.5.2 OPTION 1: V 1 IS SYMMETRIC If V 1 (the cross-moment over time) is symmetric, this implies that Ψ = V1 = V -1/2 0 V 1 V -1/2 0

is also symmetric. Thus by eqn. 7 we have 

FT = ζ T ( Ψ) + ζ T ( Ψ) -T I = η T ( Ψ), where η T (x) = 2ζ T (x) -T = 2(T (1 -x) -x(1 -x T )) (1 -x) 2 -T = T (1 -x 2 ) -2x(1 -x T ) (1 -x) 2 . Let U diag( ψ)U = Ψ

)] i = f ([b] i ).

We thus have

F = E T [ FT ] = E T [U diag(η T ( ψ))U ] = U diag(E T [η T ( ψ)])U .

Inverting both sides of this yields

EQUATION

where we have defined γ

(x) = 1/E T [η T (x)].

This expression can be efficiently evaluated when Ψ is a Kronecker product since the eigendecomposition of a Kronecker product can be easily obtained from the eigendecomposition of the factors.

Evaluation of γ( ψ) is done component-wise (i.e. [γ( ψ)] i = γ([ ψ] i ))

and is thus easy to perform. See Section 3.5.5 for further details.

V 1 is symmetric if and only if Ψ is symmetric. And as shown in the proof of Proposition 1 (see Appendix A.1) Ψ has the interpretation of being the transition matrix of an LGGM which describes the evolution of "whitened" versions of the w t 's (given by ŵt = V -1/2 0 w t ). Linear dynamical systems with symmetric transition matrices arise frequently in machine learning and related areas (Huang et al., 2016; Hazan et al., 2017) , particularly because of the algorithmic techniques they enable. Intuitively, a symmetric transition matrix allows allows one to model exponential decay of different basis components of the signal over time, but not rotations between these components (which are required to model sinusoidal/oscillating signals).

Note that the observed/measured V 1 may or may not be exactly symmetric up to numerical precision, even if it well approximated as symmetric. For these calculations to make sense it must be exactly symmetric, and so even if it turns out to be approximately symmetric one should ensure that it is exactly so by using the symmetrized version

(V 1 + V 1 )/2.

3.5.3. OPTION 2: COMPUTING THE LIMITING VALUE INSTEAD:

If V 1 is not well approximated as symmetric, another option is to approximate

F = E T [ F (∞) T ] (instead of F = E T [ FT ]

), where we define

F (∞) T ≡ lim T →∞ T T FT .

This is essentially equivalent to the assumption that the training sequences are all infinitely long, which may be a reasonable one to make in practice. We re-scale by the factor T T to achieve the proper scaling characteristics of FT , and to ensure that the limit actually exists.

As shown in Section A.2 of the appendix this yields the following remarkably simple expression for F -1 :

EQUATION

Despite the fact that it includes both Ψ and Ψ , this formula can be efficiently evaluated when Ψ is a Kronecker product due to the existance of decomposition-based techniques for inverting matrices of the form A ⊗ B + C ⊗ D. See Section 3.5.5 for further details.

This approximation can break down if some of the linear components of ŵt have temporal autocorrelations close to 1 (i.e.

[ ψ] i ≈ 1 for some i) and T is relatively small. In such a case we will have that [ ψ] T i is large for some i (despite being raised to the T -th power) so that F (∞) T may essentially "overcount" the amount of temporal correlation that contributes to the sum. This can be made more concrete by noting that the approximation is essentially equivalent to taking

ζ T (x) ≈ lim T →∞ T T ζ T (x) ≡ κ(x) for each x = [ ψ] i .

We can express the error of this as

|κ(x) -ζ T (x)| = T 1 -x - T (1 -x) -x(1 -x T ) (1 -x) 2 = x(1 -x T ) (1 -x) 2 .

It is easy to see how this expression, when evaluated at x = [ ψ] i , might be large when [ ψ] i is close to 1, and T is relatively small.

3.5.4. ESTIMATING Ψ:

The formulae for F -1 from the previous sections depend on the quantity

Ψ = V1 = V -1/2 0 ΨV 1/2 0

, and so it remains to compute Ψ. We observe that

V 1 = V 1,0 = E[w 1 w 0 ] = E[(Ψw 0 + 1 )w 0 ] = ΨE[w 0 w 0 ] + E[ 1 w 0 ] = ΨV 0 + 0 = ΨV 0 . Right-multiplying both sides by V 0 yields Ψ = V 1 V -1

0 . Thus, given estimates of V 0 and V 1 , we may compute an estimate of Ψ as

Ψ = V -1/2 0 V 1 V -1/2 0 .

In practice we estimate V 0 and V 1 by forming estimates of their Kronecker factors and taking the product. The factors themselves are estimated using exponentially decayed averages over mini-batch estimates. And the mini-batch estimates are in turn computed by averaging over cases and summing across time-steps, before divide by the expected number of time-steps.

For example, for A 0 and A 1 these the mini-batch estimates are averages of 1

E T [T ]:

T t=1 a t a t and 1

T -1 t=1 a t+1 a t , respectively. Note that as long as V 0 is computed as the 2nd-order moment of some empirical data, and V 1 computed as the 2nd-order moment between that same data and a temporally shifted version, the spectral radius of

Ψ = V -1/2 0 V 1 V -1/2 0 (and similarly Ψ = V 1 V -1

0 ) will indeed be less than or equal to 1, as we prove in Section B.2 of the appendix. This bound on the spectral radius is a necessary condition for our infinite chain-structured Gaussian graphical model to be well-defined, and for our calculations to make sense.

The sufficient condition that the spectral radius is actually less than 1 will most often be satisfied too, except in the unlikely event that some eigen component remains perfectly constant across time. But even if this somehow happens, the inclusion within the given V 0 of some damping/regularization term such as λI will naturally deal with this problem.

3.5.5 EFFICIENT IMPLEMENTATION ASSUMING V 0 AND V 1 ARE KRONECKER-FACTORED It remains to show that the approximations developed in Section 3.5 can be combined with the Kronecker-factored approximations for V 0 and V 1 from Section 3.2.3 to yield an efficient algorithm for computing F -1 z for an arbitrary vector z = vec(Z). This is a straightforward although very long computation which we leave to Section C of the appendix.

Full pseudo-code for the resulting algorithms is given in Section C.3. As they only involve symmetric eigen-decomposition and matrix-matrix products with matrices the size of A 0 and G 0 they are only several times more expensive to compute than eqn. 6. This extra overhead will often be negligible since the gradient computation via BPTT, whose costs scales with the sequence length T , tends to dominate all the other costs.

4. EXPERIMENTS:

To demonstrate the benefit of our novel curvature matrix approximations for RNNs, we empirically evaluated them within the standard "distributed K-FAC" framework (Ba et al., 2017) on two different RNN training tasks.

The 2nd-order statistics (i.e. the Kronecker factors A 0 , A 1 , G 0 , and G 1 ) are accumulated through an exponential moving average during training. When computing our approximate inverse Fisher, factored Tikhonov damping (Martens & Grosse, 2015) was applied to V 0 = G 0 ⊗ A 0 .

We used a single machine with 16 CPU cores and a Nvidia K40 GPU for all the experiments. The additional computations required to get the approximate Fisher inverse from these statistics (i.e. the "pre-processing steps" described in Section C.3) are performed asynchronously on the CPUs, while the GPU is used for the usual forward evaluation and back-propagation to compute the gradient. Updates are computed using the most recently computed values of these (which are allowed to be stale), so there is minimal per-iteration computational overhead compared SGD. We adopted the step-size selection technique described in Section 5 of Ba et al. (2017) , as we found it let us use larger learning rates without compromising the stability of the optimization. The hyperparameters of our approach, which include the max learning rate and trust-region size for the aforementioned step-size selection procedure, as well as the momentum, damping constants, and the decay-rate for the second-order statistics, as well as the hyper-parameters of the baseline methods, were tuned using a grid search.

Word-level language model: We start by applying our method to a two-layer RNN based on the well-studied Long Short-Term Memory (LSTM) architecture (Hochreiter & Schmidhuber, 1997) for a word-level language modeling task on the Penn-TreeBank (PTB) dataset (Marcus et al., 1993) following the experimental setup in Zaremba et al. (2014) . The gradients are computed using a fixed sequence length truncated back-propagation scheme in which the initial states of the recurrent hidden units are inherited from the final state of the preceding sequence. The truncation length used in the experiments is 35 timesteps. The learning rate is given by a carefully tuned decaying schedule (whose base value we tune along with the other hyperparamters).

In our experiments we simply substitute their optimizer with our modified distributed K-FAC optimizer that uses our proposed RNN Fisher approximations. We performed experiments on two different sizes of the same architecture, which use two-layer 650 and 1024 LSTM units respectively.

LSTMs have 4 groups of internal units: input gates, output gates, forget gates, and update candidates. We treat the 4 weight matrices that compute the pre-activations to each of these as distinct for the purposes of defining Fisher blocks (whereas many LSTM implementations treat them as one big matrix). This results in smaller Kronecker factors that are cheaper to compute and invert. 

4.1. LANGUAGE MODELING WITH LONG SHORT-TERM MEMORY UNITS:

Because the typical vocabulary size used for PTB is 10,000, the Fisher blocks for the input embedding layer and output layer (computing the logits to the softmax) each contain a 10,000 by 10,000 sizes Kronecker factor, which is too large to be inverted with any reasonable frequency. Given that the input vector uses a one-hot encoding it is easy to see its associated factor is actually diagonal, and so we can store and invert it as such. Meanwhile the large factor associated with the output isn't diagonal, but we nonetheless approximate it as such for the sake of efficiency.

In our experiments we found that each parameter update of our method required about 80% more wall-clock time than an SGD update (using mini-batch size of 200) although the updates made more much progress.

In Figure 1 , we plot the training progress as a function of the number of parameter updates. While Adam outperforms SGD in the first few epochs, SGD obtains a lower loss at the end of training. We found the recent layer-normalization technique (Ba et al., 2016) helps speed up Adam considerably, but it hurts the SGD performance. Such an observation is consistent with previous findings. In comparison, our proposed method still significantly outperform both the Adam and the SGD baselines even with the help of layer-normalization.

While optimization performance, not generalization performance, is the focus of this paper, we have included validation performance data in the appendix for the sake of completeness. (See Figure 4 in Appendix D.) Not surprisingly, we found that the 2nd-order methods, including our approach and diagonal ones like Adam, tended to overfit more than SGD on these tasks.

The tendency for SGD w/ early-stopping to self-regularize is well-documented, and there are many compelling theories about why this happens (e.g. Duvenaud et al., 2016; Hardt et al., 2015) . It is also well-known that 2nd-order methods, including K-FAC and diagonal methods like Adam/RMSprop, dont self-regularize nearly as much (e.g. Wilson et al., 2017; Keskar & Socher, 2017) . We feel that this problem can likely be addressed through the careful application of additional explicit regularization (e.g. increased weight decay, drop-out, etc) and/or model modifications, but that exploring this is outside of the scope of this paper.

Character-level model::

To further investigate the optimization performance of our proposed Fisher approximation, we use a small two layer LSTM with 128 units to model the character sequences on the Penn-TreeBank (PTB) dataset (Marcus et al., 1993) . We employ the same data partition in Mikolov et al. (2012) . We plotted the bits-per-character vs the number of parameter updates and the wall-clock times in Figure 2 . The K-FAC updates were roughly twice as time-consuming to compute as the Adam updates in our implementation. Despite this, our results demonstrate that K-FAC has a significant advantage over the Adam baseline in terms of wall-clock time. 

4.2. LEARNING DIFFERENTIABLE NEURAL COMPUTERS:

To further investigate the potential benefits of using our approach over existing methods, we applied it to the Differentiable Neural Computer (DNC) model (Graves et al., 2016) for learning simple algorithmic programs. Recently, there have been several attempts (Weston et al., 2014; Graves et al., 2016) to extend the existing RNN models to incorporate more long-term memory storage devices in order to help solve problems beyond simple sequence prediction tasks. Although these extended RNNs could potentially be more powerful than simple LSTMs, they often require thousands of parameter updates to learn simple copy tasks (Graves et al., 2016) . Both the complexity of these models and the difficulty of the learning tasks have posed a significant challenge to commonly used optimization methods.

The DNC model is designed to solve structured algorithmic tasks by using an LSTM to control an external read-write memory. We applied the Fisher-based precondition to compute the updates for both the weights in the LSTM controller and the read-write weight matrices used to interface with the memory. We trained the model on a simple repeated copy task in which the DNC needs to recreate a series of two random binary sequences after they are presented as inputs. The total length of the sequence is fixed to 22 time-steps. From Figure 3 , we see that our method significantly outperforms the Adam baseline in terms of update count, although only provides a modest improvement in wallclock time.

This gap is explained by the fact that the iterations were significantly more time-consuming to compute relative to the gradient computations than they were in previous two experiments on language models. This is likely due to a different trade-off in terms of the gradient computation vs the overheads specific to our method owing to smallness of the model and dataset. With more careful engineering to reduce the communication costs, and/or a larger model and dataset, we would expect to see a bigger improvement in wall-clock time.

5. CONCLUSION:

We have presented a new family of approximations to the Fisher information matrix of recurrent neural networks (RNNs), extending previous work on Kronecker-factored approximations. With this contribution, recurrent networks can now finally be trained with the K-FAC optimization method. We have demonstrated that our new approximations substantially reduce the required number of iterations for convergence vs standard baseline optimizers on several realistic tasks. And we have also shown that in a modern distributed training setup this results in a substantial savings in wallclock time as well.

A SUPPLEMENTARY COMPUTATIONS A.1 PROOFS FOR SECTION 3.5.1

Proposition 1 Given FT = V 1/2 0 F T V 1/2 0 and Ψ = V1 = V -1/2 0 ΨV 1/2 0 . we have FT = ζ T ( Ψ) + ζ T ( Ψ ) -T I where ζ T (x) = T (1 -x) -x(1 -x T ) (1 -x) 2 .

Proof:

For d > 0 we have that V d = ΨV d-1 , which can be seen as follows:

V d = V d,0 = E[w d w 0 ] = E[Ψw d-1 + d w 0 ] = ΨE[w d-1 w 0 ] + E[ d w 0 ] = ΨV d-1 + 0 = ΨV d-1 . Applying V d = ΨV d-1 recursively yields V d = Ψ d V 0 for d 0.

And using V -d = V d it also follows that

V d = V 0 (Ψ d ) for d 0.

Setting d = 1 and multiplying both sides by V 0 (which is assumed to be invertible) one can also derive the following simple formula for Ψ:

EQUATION

)

To proceed from here we define a "transformed" version of the original chain-structured linear-Gaussian graphical model whose variables are ŵt = V -1/2 0 w t . (Here we assume that V 0 is invertible -it is symmetric by definition.) All quantities related to the original model have their analogues in the transformed model, which we indicate with the hat symbol •.

In the transformed model the 2nd-order moments of the ŵt 's are given by

Vd = E V -1/2 0 w d V -1/2 0 w 0 = V -1/2 0 E w d w 0 V -1/2 0 = V -1/2 0 V d V -1/2 0 .

We observe that V0 = I.

Analogously to the original model, the transformed version obeys

ŵt = Ψ ŵt-1 + ˆ t , with ˆ t = V -1/2 0 t and Ψ = V1 V -1 0 = V1 (using V0 = I).

This can be seen by noting that

ŵt = V -1/2 0 w t = V -1/2 0 (Ψw t-1 + t ) = V -1/2 0 V 1 V -1 0 w t-1 + V -1/2 0 t = (V -1/2 0 V 1 V -1/2 0 )(V -1/2 0 w t-1 ) + ˆ t = V1 ŵt-1 + ˆ t .

It also remains true that the spectral radius of Ψ is less than 1, which can be seen in at least one of two ways: by noticing that the transformed model is well-defined in the infinite limit if and only if the original one is, or that Ψ = V1 = V -1/2 0 ΨV 1/2 0 is a similar matrix to Ψ (in the technical sense) and hence has the same eigenvalues.

As the transformed model is isomorphic to the original one, all of the previously derived relationships which held for it also hold here, simply by replacing each quantity with its transformed version (denoted by the hat symbol •).

Given these relations (included the transformed analogue of equation 5) we can express FT as

FT = T d=0 (T -d) Ψd + T d=0 (T -d) Ψd -T I.

It is a well-known fact that one can evaluate rational functions, and functions that are the limiting values of sequences of rational functions, with matrix arguments. This is done by replacing scalar multiplication with matrix multiplication, division with matrix inversion, and scalar constants with scalar multiples of the identity matrix, etc. Note that because sums of powers and inverses of matrices are co-diagonalizable/commutative when the matrices themselves are, there is no issue of ambiguity caused by mixing commutative and non-commutative algebra in this way.

Moreover, the value of some such function f (x), given a matrix argument B, is 

f (B) = V diag(f (b))V -1 , where V diag(b)V -1 = B is eigendecomposition of B,

] i = f ([b] i ). Note that if [f (b)] i is

undefined from some i, either because of a division by zero, or because the limit which defines f (x) doesn't converge for x = [b] i , then f (B) doesn't exist for that particular B (and otherwise it does).

We observe that our above expression for F T can be rewritten as 

T (x) = T (1 -x) -x(1 -x T ) (1 -x) 2 .

Let U diag( ψ)U -1 = Ψ be the eigendecomposition of Ψ. Because Ψ has a spectral radius less than 1, we have |[ ψ] i | < 1 for each i (so that in particular [ ψ] i = 1), and thus we can evaluate ζ T ( Ψ) and ζ T ( Ψ ) according to the above formula for ζ T (x).

A.2 PROOFS FOR SECTION 3.5.3

Proposition 2 Suppose we approximate

F = E T [ F (∞) T ],

where we have defined

F (∞) T ≡ lim T →∞ T T FT .

Then we have

F -1 = 1 E T [T ] (I -Ψ)(I -Ψ Ψ) -1 (I -Ψ ).

From eqn. 7 we have that

F (∞) T = lim T →∞ T T (ζ T ( Ψ) + ζ T ( Ψ ) -T I) = lim T →∞ T T ζ T ( Ψ) + lim T →∞ T T ζ T ( Ψ ) -lim T →∞ T T T I.

To evaluate this we first term note that lim T →∞ T T T I = T I and lim

T →∞ T T ζ T (A) = κ(A),

where we have defined κ

(x) = lim T →∞ T T ζ T (x).

For |x| < 1 we have that lim T →∞ x T = 0, from which it follows that

κ(x) = lim T →∞ T T T (1 -x) -x(1 -x T ) (1 -x) 2 = lim T →∞ T T T 1 -x -lim T →∞ T T x(1 -x T ) (1 -x) 2 = lim T →∞ T 1 -x -lim T →∞ T T lim T →∞ x(1 -x T ) (1 -x) 2 = T 1 -x -lim T →∞ T T x(1 -0) (1 -x) 2 = T 1 -x .

Let U diag( ψ)U -1 = Ψ be the eigendecomposition of Ψ. Using the fact that |[ ψ] i | < 1 (as established in Section A.1) we can use the above expression to evaluate κ(x) at both x = Ψ and x = Ψ , which yields

F (∞) T = κ( Ψ) + κ( Ψ ) = T ((I -Ψ) -1 + (I -Ψ ) -1 -I).

Pre-multiplying both sides by I -Ψ , and post-multiplying both sides by I -Ψ, we have

(I -Ψ ) F (∞) T (I -Ψ) = T ((I -Ψ ) + (I -Ψ) -(I -Ψ )(I -Ψ)) = T (I -Ψ + I -Ψ -I + Ψ + Ψ -Ψ Ψ) = T (I -Ψ Ψ).

Then applying the reverse operation gives

F (∞) T = T (I -Ψ ) -1 (I -Ψ Ψ)(I -Ψ) -1 .

Taking the expectation over T gives

F = E T [ F (∞) T ] = E T [T ](I -Ψ ) -1 (I -Ψ Ψ)(I -Ψ) -1 .

Finally, inverting both sides yields

F -1 = 1 E T [T ] (I -Ψ)(I -Ψ Ψ) -1 (I -Ψ ).

B ADDITIONAL TECHNICAL PROOFS:



B.1:

Proposition 3 Suppose x ∈ C, x = 0, and T is a non-negative integer. We have

T i=0 (T -i)x i = T (1 -x) -x(1 -x T ) (1 -x) 2 . Proof Observe that x d T i=0 x i dx = x T i=0 dx i dx = x T i=0 ix i-1 = T i=0 ix i .

Another way to express this is to use the geometric series formula

T i=0 x i = 1-x T +1

1-x (which holds for x = 0) before computing the derivative. This gives

x d T i=0 x i dx = x d 1-x T +1 1-x dx = x - (1 + T )x T 1 -x + 1 -x T +1 (1 -x) 2 = x(1 -x T +1 -(1 + T )(1 -x)x T ) (1 -x) 2 .

Thus we have

T i=0 ix i = x(1 -x T +1 -(1 + T )(1 -x)x T ) (1 -x) 2 .

And so

T i=0 (T -i)x i = T T i=0 x i - T i=0 ix i = T (1 -x T ) 1 -x - x(1 -x T +1 -(1 + T )(1 -x)x T ) (1 -x) 2 = T (1 -x)(1 -x T +1 ) -x(1 -x T +1 -(1 + T )(1 -x)x T ) (1 -x) 2 = T -T x -T x T +1 + T x T +2 -x + x T +2 + x T +1 (1 -x + T -T x) (1 -x) 2 = T -T x -T x T +1 + T x T +2 -x + x T +2 + x T +1 -x T +2 + T x T +1 -T x T +2 (1 -x) 2 = T -(T + 1)x + x T +1 (1 -x) 2 = T (1 -x) -x(1 -x T ) (1 -x) 2

where we have again used the geometric series formula

T i=0 x i = 1-x T +1

1-x on the second line.

B.2 SPECTRAL BOUND FOR ESTIMATE OF Ψ:

In what follows all quantities are computed using their defining formulae, starting from the estimated values of A 0 , A 1 , G 0 , and G 1 .

First we observe that since Ψ = V1 = V -1/2 0 ΨV 1/2 0 is similar to Ψ (in the technical sense of the word), they share the same eigenvalues. Thus it suffices bound to the spectral radius of

Ψ = V 1 V -1 0 . Next we observe that V 0 = A 0 ⊗ G 0 and V 1 = A 1 ⊗ G 1 , so that V 1 V -1 0 = (A 1 ⊗ G 1 )(A 0 ⊗ G 0 ) -1 = (A 1 A -1 0 ) ⊗ (G 1 G -1 0 ).

Because the eigendecomposition of a Kronecker product is the Kronecker product of the decompositions of the factors we have that ρ(

V 1 V -1 0 ) = ρ(A 1 A -1 0 )ρ(G 1 G -1 0 )

, where ρ(X) denotes the spectral radius of a matrix X.

Thus it suffices to show that ρ(A 1 A -1 0 ) ≤ 1 and ρ(G 1 G -1 0 ) ≤ 1 for A i and G i as computed by the estimation scheme outlined in Section 3.5.4. Recall that this is the exponentially decayed average of mini-batch averages of estimators of the form 1

T t=1 a t a t and 1

T -1 t=1 a t+1 a t .

In the remainder of this section we will show that ρ(A 1 A -1 0 ) ≤ 1. The argument to show that ρ(G 1 G -1 0 ) ≤ 1 is identical.

where denotes entry-wise multiplication of matrices.

In summary we have that F -1 z can be computed in matrix form as

U G ((U G ZU A ) Y )U A for Y s.t. [Y ] i,j = γ([ ψG ] i [ ψA ] j ).

We note that computing γ([ ψG ] i [ ψA ] j ) is trivial since it is just a scalar evaluation of the rational function γ(x).

C.2 OPTION 2:

For Option 2 we must compute the matrix-vector product

F -1 z = 1 i ω i T i (I -Ψ)(I -Ψ Ψ) -1 (I -Ψ )v.

To do this we will first multiply by I -Ψ , then by (I -Ψ Ψ) -1 , and then by I -Ψ, before finally dividing the result by i ω i T i .

To compute the matrix-vector product (I -Ψ )z we use the identity (C ⊗ B) vec(X) = vec(B XC ) while noting that Ψ = ( ΨA ⊗ ΨG ) = Ψ A ⊗ Ψ G . This gives

(I -Ψ ) vec(Z) = vec(Z) -( Ψ A ⊗ Ψ G ) vec(Z) = vec(Z) -vec( Ψ G Z ΨA ).

The matrix form of this is simply Z -Ψ G Z ΨA .

We may similarly compute (I -Ψ)z in matrix form as Z -ΨG Z Ψ A .

The harder task is to compute (I -Ψ Ψ) -1 z, which is what we tackle next.

We first observe that

I -Ψ Ψ = I ⊗ I -( Ψ A ⊗ Ψ G )( ΨA ⊗ ΨG ) = I ⊗ I -( Ψ A ΨA ) ⊗ ( Ψ G ΨG ) ≡ I ⊗ I -M A ⊗ M G .

Given the eigendecompositions E A diag(m A )E A = M A and E G diag(m G )E G = M G we can thus compute the larger eigendecomposition as

I ⊗ I -M A ⊗ M G = (E A ⊗ E G )(I ⊗ I -diag(m A ) ⊗ diag(m G ))(E A ⊗ E G ) = (E A ⊗ E G ) diag(l ⊗ l -m A ⊗ m G )(E A ⊗ E G ) = (E A ⊗ E G ) diag(vec(ll -m G m A ))(E A ⊗ E G ),

where l is the vector of ones.

Using the eigendecomposition the inverse can then be easily computed as

(I ⊗ I -M A ⊗ M G ) -1 = (E A ⊗ E G ) diag(vec(ll -m G m A )) -1 (E A ⊗ E G ).

Thus (I -Ψ Ψ) -1 z may be computed by first multiplying by (E A ⊗ E G ), then by diag(vec(llm G m A )) -1 (which in matrix form corresponds to element-wise division by llm G m A ), and

then by E A ⊗ E G . The matrix form of this is

E G ((E G ZE A ) (ll -m G m A ))E A ,

where B C denotes element-wise division of the matrix B by the matrix C.

• Multiply the result by I -Ψ :

Z 1 = Z 0 -Ψ G Z 0 ΨA

• Multiply the previous result by (I -Ψ Ψ) -1 using its eigendecomposition:

Z 2 = E G ((E G Z 1 E A ) (ll -m G m A ))E A

where X Y denotes the element-wise division of X by Y .

• Multiply the result by I -Ψ

Z 3 = Z 2 -ΨG Z 2 Ψ A

• Normalize the result by E T [T ] = i ω i T i (for T i and ω i as defined at the bottom of Section 3.4):

Z 4 = 1 i ω i T i Z 3

• Multiply the result by V -1/2 0 and express as a vector:

Z 5 = G -1/2 0 Z 4 A -1/2 0

and F -1 z = vec(Z 5 ) 

:

Roughly speaking, this means that it converges at the asymptotically optimal rate (with optimal constant) for arbitrary statistical estimation procedures (as a function of the amount of samples from Q observed).

Figures and Tables:

be the eigen-decomposition of Ψ. By the above expression for FT we have that FT = U diag(η T ( ψ))U , where f (b) denotes the component-wise evaluation of a function f for each component of the vector b, i.e. [f (b

Figure 1: Optimization performance of our method compared to the baselines in perplexity-perword on length-35 word sequences from Penn-TreeBank. All the methods used a mini-batch size of 200. K-FAC indep. uses the update in eqn. 6, K-FAC option1 uses eqn. 8, and K-FAC option2 uses eqn. 9. (left) Training perplexity v.s. the number of updates. Dashed lines denote the training curves for RNNs with 1024 LSTM units and solid lines denote the training curves for RNNs with 650 LSTM units. (right) Training perplexity v.s. the wall-clock time.

Figure 2: Optimization performance in bit-per-character on length-100 character sequences from Penn-TreeBank. batchsize indicates the mini-batch size used to train the baseline methods (our method always used a mini-batch size of 200). K-FAC indep. uses the update in eqn. 6, K-FAC option1 uses eqn. 8, and K-FAC option2 uses eqn. 9. (left) Training perplexity v.s. the number of updates. (right) Training perplexity v.s. the wall-clock time.

Figure 3: Optimization performance for differentiable Neural Computers (DNC) on a repeated copy task. K-FAC indep. uses the update in eqn. 6, K-FAC option1 uses eqn. 8, and K-FAC option2 uses eqn. 9. (left) Training cross entropy loss v.s. the number of updates. (right) Training cross entropy loss v.s. the wall-clock time.

and where f (b) denotes the componentwise evaluation of f for each component of the vector b, i.e. [f (b)

= ζ T ( Ψ) + ζ T ( Ψ ) -T I, where ζ T (x) = T d=0 (Td)x d . By Proposition 3 in Appendix B.1, we have for x = 1 that ζ

Figure 4: Generalization performance of our method compared to the baselines in perplexity-perword on length-35 word sequences from Penn-TreeBank. All the methods used a mini-batch size of 200. K-FAC indep. uses the update in eqn. 6, K-FAC option1 uses eqn. 8, and K-FAC option2 uses eqn. 9. (left) Validation perplexity v.s. the number of updates. The dashed lines correspond to the experiments that used RNNs with 1024 LSTM units, and the solid lines correspond to experiments that used RNNs with 650 LSTM units. (right) Validation perplexity v.s. the wall-clock time.

Peer review (iclr2018):
 This paper extends the Kronecker-factor Approximate Curvature (K-FAC) optimization method to the setting of recurrent neural networks. The K-FAC method is an approximate 2nd-order optimization method that builds a block diagonal approximation of the Fisher information matrix, where the block diagonal elements are Kronecker products of smaller matrices. In order to approximate the Fisher information matrix for RNNs, the authors assume that the derivative of the loss function with respect to each weight matrix at each time step is independent of the length of the sequence, that these derivatives are temporally homogeneous, that the input and derivatives of the output are independent across every point in time, and that either the one-step cross-covariance of these derivatives is symmetric or that the training sequences are effectively infinite in length. Based on these assumptions, the authors show that the Fisher information can be reduced into a form in which the derivatives of the weight matrices can be approximated by a linear Gaussian graphical model and in which the approximate 2nd order method can be efficiently carried out. The authors compare their method to SGD on two language modeling tasks and against Adam for learning differentiable neural computers. The paper is relatively clear, and the authors do a reasonable job of introducing related work of the original K-FAC algorithm as well as its extension to CNNs before systematically deriving their method for RNNs. The problem of extending the K-FAC algorithm is natural, and the steps taken in this paper seem natural yet also original and non-trivial. The main issue that I have with this paper is the lack of theoretical justification or even intuition for the many approximations carried out in the course of approximating the Fisher information matrix. In many instances, it seemed like these approximations were made purely for convenience and tractability without much regard for (even approximate) correctness. This quality of this paper would be greatly strengthened if it had some bounds on approximation error or even empirical results testing the validity of the assumptions in the paper. Moreover, the experiments do not demonstrate levels of statistical significance in the results, so it is difficult to assert the practical significance of this work. Specific comments and questions Page 2, "r is is". Typo. Page 2, "DV". I found the introduction of V without any explanation to be confusing. Page 2, "P_{y|x}(.theta)". The relation between P_{y|x}(.theta) and f(x,.theta) is never explained. Page 3, "common practice of computing the natural gradient as (F + .lambda I) .nabla h instead of F^{-1} .nabla h". I don*t see how the former can serve as a replacement for the latter. Page 3, "approximate g and a as statistically independent". Even though K-FAC already exists, it would be good to explain why this assumption is reasonable, since similar assumptions are made for the work presented in this paper. Page 4, "This new approximation, called "KFC", is derived by assuming....". Same as previous comment. It would be good to briefly discuss why these assumptions are reasonable. Page 5, Independence of T and w_t*s, temporal homogeneity of w_t*s,, and independence between a_t*s and g_t*s. I can see why these are convenient assumptions, but why are they reasonable? Moreover, why is it further natural to assume that A and G are temporally homogeneous as well? Page 7, "But insofar as the w_t*s ... encode the relevant information contained in these external variables, they should be approximately Markovian". I am not sure what this means. Page 7, "The linear-Gaussian assumption meanwhile is a more severe one to make, but it seems necessary for there to be any hope that the required expectations remain tractable". I am not sure that this is a good enough justification for such an idea, unless there are compelling approximation error bounds. Page 8, Option 1. In what situations is it reasonable to assume that V_1 is symmetric? Pages 8-9, Option 2. What is a good finite sample size in which the assumption that the training sequences are infinitely long is reasonable in practice? Can the error |.kappa(x) - .zeta_T(x)| be translated into a statement on the approximation error? Page 9, "V_1 = V_{1,0} = ...". Typos (that appear to have been caught by the authors already). Page 9, "The 2nd-order statistics ... are accumulated through an exponential moving average during training". How sensitive is the performance of this method to the decay rate of the exponential moving average? Page 10, "The additional computations required to get the approximate Fisher inverse from these statistics ... are performed asynchronously on the CPU*s". I find it a bit unfair to compare SGD to K-FAC in terms of wall clock time without also using the extra CPU*s for SGD as well (e.g. via Hogwild or synchronous parallel SGD). Page 10, "The hyperparameters of our approach...". What is the sensitivity of the experimental results to these hyperparameters? Moreover, how sensitive are the results to initialization? Page 10, "we found that each parameter update of our method required about 80% more wall-clock time than an SGD update". How much of this is attributed to the fact that the statistics are computed asynchronously? Pages 10-12, Experiments. There are no error bars in any of the plots, so it is impossible to ascertain the statistical significance of any of these results. Page 11: Figure 2. Where is the Adam batchsize 50 line in the left plot? Why did the Adam batchsize 200 line disappear halfway through the right plot?