Title: DISTRIBUTED SECOND-ORDER OPTIMIZATION USING KRONECKER-FACTORED APPROXIMATIONS

Abstract: As more computational resources become available, machine learning researchers train ever larger neural networks on millions of data points using stochastic gradient descent (SGD). Although SGD scales well in terms of both the size of dataset and the number of parameters of the model, it has rapidly diminishing returns as parallel computing resources increase. Second-order optimization methods have an affinity for well-estimated gradients and large mini-batches, and can therefore benefit much more from parallel computation in principle. Unfortunately, they often employ severe approximations to the curvature matrix in order to scale to large models with millions of parameters, limiting their effectiveness in practice versus well-tuned SGD with momentum. The recently proposed K-FAC method (Martens and Grosse, 2015) uses a stronger and more sophisticated curvature approximation, and has been shown to make much more per-iteration progress than SGD, while only introducing a modest overhead. In this paper, we develop a version of K-FAC that distributes the computation of gradients and additional quantities required by K-FAC across multiple machines, thereby taking advantage of the method's superior scaling to large mini-batches and mitigating its additional overheads. We provide a Tensorflow implementation of our approach which is easy to use and can be applied to many existing codebases without modification. Additionally, we develop several algorithmic enhancements to K-FAC which can improve its computational performance for very large models. Finally, we show that our distributed K-FAC method speeds up training of various state-of-the-art ImageNet classification models by a factor of two compared to an improved form of Batch Normalization (Ioffe and Szegedy, 2015) .

1. INTRODUCTION:

Current state-of-the-art deep neural networks (Szegedy et al., 2014; Krizhevsky et al., 2012; He et al., 2015) often require days of training time with millions of training cases. The typical strategy to speed-up neural network training is to allocate more parallel resources over many machines and cluster nodes (Dean et al., 2012) . Parallel training also enables researchers to build larger models where different machines compute different splits of the mini-batches. Although we have improved our distributed training setups over the years, neural networks are still trained with various simple first-order stochastic gradient descent (SGD) algorithms. Despite how well SGD scales with the size of the model and the size of the datasets, it does not scale well with the parallel computation resources. Larger mini-batches and more parallel computations exhibit diminishing returns for SGD and related algorithms.

Second-order optimization methods, which use second-order information to construct updates that account for the curvature of objective function, represent a promising alternative. The canonical second-order methods work by inverting a large curvature matrix (traditionally the Hessian), but this doesn't scale well to deep neural networks with millions of parameters. Various approximations to the curvature matrix have been proposed to help alleviate this problem, such as diagonal (LeCun et al., 1998; Duchi et al., 2011; Kingma and Ba, 2014) , block diagonal Le Roux et al. (2008) , and low-rank ones (Schraudolph et al., 2007; Bordes et al., 2009; Wang et al., 2014; Keskar and Berahas, 2015; Moritz et al., 2016; Byrd et al., 2016; Curtis, 2016; Ramamurthy and Duffy) . Another strategy is to use Krylov-subspace methods and efficient matrix-vector product algorthms to avoid the inversion problem entirely (Martens, 2010; Vinyals and Povey, 2012; Kiros, 2013; Cho et al., 2015; He et al., 2016) .

The usual problem with curvature approximations, especially low-rank and diagonal ones, is that they are very crude and only model superficial aspects of the true curvature in the objective function. Krylov-subspace methods on the other hand suffer because they still rely on 1st-order methods to compute their updates.

More recently, several approximations have been proposed based on statistical approximations of the Fisher information matrix (Heskes, 2000; Ollivier, 2013; Grosse and Salakhutdinov, 2015; Povey et al., 2015; Desjardins et al., 2015) . In the K-FAC approach (Martens and Grosse, 2015; Grosse and Martens, 2016) , these approximations result in a block-diagonal approximation to the Fisher information matrix (with blocks corresponding to entire layers) where each block is approximated as a Kronecker product of two much smaller matrices, both of which can be estimated and inverted fairly efficiently. Because the inverse of a Kronecker product of two matrices is the Kronecker product of their inverses, this allows the entire matrix to be inverted efficiently. Martens and Grosse (2015) found that K-FAC scales very favorably to larger mini-batches compared to SGD, enjoying a nearly linear relationship between mini-batch size and per-iteration progress for medium-to-large sized mini-batches. One possible explanation for this phenomenon is that secondorder methods make more rapid progress exploring the error surface and reaching a neighborhood of a local minimum where gradient noise (which is inversely proportional to mini-batch size) becomes the chief limiting factor in convergence1 . This observation implies that K-FAC would benefit in particular from a highly parallel distributed implementation.

In this paper, we propose an asynchronous distributed version of K-FAC that can effectively exploit large amounts of parallel computing resources, and which scales to industrial-scale neural net models with hundreds of millions of parameters. Our method augments the traditional distributed synchronous SGD setup with additional computation nodes that update the approximate Fisher and compute its inverse. The proposed method achieves a comparable per-iteration runtime as a normal SGD using the same mini-batch size on a typical 4 GPU cluster. We also propose a "doubly factored" Kronecker approximation for layers whose inputs are feature maps that are normally too large to handled by the standard Kronecker-factored approximation. Finally, we empirically demonstrate that the proposed method speeds up learning of various state-of-the-art ImageNet models by a factor of two over Batch Normalization (Ioffe and Szegedy, 2015) .

2. BACKGROUND:



2.1. KRONECKER FACTORED APPROXIMATE FISHER:

Let DW be the gradient of the log likelihood L of a neural network w.r.t. some weight matrix W ∈ R Cout×Cin in a layer, where C in , C out are the number of input/output units of the layer. The block of the Fisher information matrix of that layer is given by:

EQUATION

)

where P is the distribution over the input x and the network's distribution over targets y (implied by the log-likelihood objective). Throughout this paper we assume, unless otherwise stated, that expectations are taken with respect to P (and not the training distribution over y).

K-FAC (Martens and Grosse, 2015; Grosse and Martens, 2016) uses a Kronecker-factored approximation to each block which we now describe. Denote the input activation vector to the layer as A ∈ R Cin , the pre-activation inputs as s = W A and the back-propagated loss derivatives as Ds = dL ds ∈ R Cout . Note that the gradient of the weights is the outer product of the input activation and back-propagated derivatives DW = DsA . K-FAC approximates the Fisher block as a Kronecker product of the second-order statistics of the input and the backpropagated derivatives:

F = E vec{DW } vec{DW } = E AA ⊗ DsDs ≈ E AA ⊗ E DsDs F . (2)

This approximation can be interpreted as making the assumption that the second-order statistics of the activations and the backpropagated derivatives are uncorrelated.

2.2. APPROXIMATE NATURAL GRADIENT USING K-FAC:

The natural gradient (Amari, 1998) is defined as the inverse of the Fisher times the gradient. It is traditionally interpreted as the direction in parameter space that achieves the largest (instantaneous) improvement in the objective per unit of change in the output distribution of the network (as measured using the KL-divergence). Under certain conditions, which almost always hold in practice, it can also be interpreted as a second-order update computed by minimizing a local quadratic approximation of the log-likelihood objective, where the Hessian is approximated using the Fisher (Martens, 2014) .

To compute the approximate natural gradient in K-FAC, one multiplies the gradient for the weights of each layer by the inverse of the corresponding approximate Fisher block F for that layer. Denote the gradient of the loss function with respect to the weights W by G W ∈ R Cin×Cout . We will assume the use of the factorized Tikhonov damping approach described by Martens and Grosse (2015) , where the addition of the damping term λI to F is approximated by adding π A λ 1 2 I to E AA and π Ds λ 1 2 I to E DsDs , where π A and π Ds are adjustment factors that are described in detail and generalized in Sec. 4.1. (Note that one can also include the contribution to the curvature from any L2 regularization terms with λ.)

By exploiting the basic identities:

(A ⊗ B) -1 = (A -1 ⊗ B -1 ) and (A ⊗ B) vec(C) = vec(BCA ),

the approximate natural gradient update v can then be computed as:

EQUATION

which amounts to several matrix inversion of multiplication operations involving matrices roughly the same size as the weight matrix W .

3. DISTRIBUTED OPTIMIZATION USING K-FAC:

Stochastic optimization algorithms benefit from low-variance gradient estimates (as might be obtained from larger mini-batches). Prior work suggests that approximate natural gradient algorithms might benefit more than standard SGD from reducing the variance (Martens and Grosse, 2015; Grosse and Martens, 2016) . One way to efficiently obtain low-variance gradient estimates is to parallelize the gradient computation across many machines in a distributed system (thus allowing large mini-batches to be processed efficiently). Because the gradient computation in K-FAC is identical to that of SGD, we parallelize the gradient computation using the standard synchronous SGD model.

However, K-FAC also introduces other forms of overhead not found in SGD -in particular, estimation of second-order statistics and computation of inverses or eigenvalues of the Kronecker factors.

In this section, we describe how these additional computations can be performed asynchronously. While this asynchronous computation introduces an additional source of error into the algorithm, we find that it does not significantly affect the per-iteration progress in practice. All in all, the periteration wall clock time of our distributed K-FAC implementation is only 5-10% higher compared to synchronous SGD with the same mini-batch size.

3.1. ASYNCHRONOUS FISHER BLOCK INVERSION:

Computing the parameter updates as per Eq.3 requires the estimated gradients to be multiplied by the inverse of the smaller Kronecker factors. This requires periodically computing (typically) either inverses or eigendecompositions of each of these factors. While these factors typically have sizes only in the hundreds or low thousands, very deep networks may have hundreds of such matrices (2 or more for each layer). Furthermore, matrix inversion and eigendecomposition see little benefit from GPU computation, so they can be more expensive than standard neural network operations. For these reasons, inverting the approximate Fisher blocks represents a significant computational cost.

It has been observed that refreshing the inverse of the Fisher blocks only occasionally and using stale values otherwise has only a small detrimental effect on average per-iteration progress, perhaps because the curvature changes relatively slowly (Martens and Grosse, 2015) . We push this a step further by computing the inverses asynchronously while the network is still training. Because the required linear algebra operations are CPU-bound while the rest of our computations are GPU-bound, we perform them on the CPU with little effective overhead. Our curvature statistics are somewhat more stale as a result, but this does not appear to significantly affect per-iteration optimization performance. In our experiments, we found that computing the inverses asynchronously usually offered a 40-50% speed-up to the overall wall-clock time of the K-FAC algorithm.

3.2. ASYNCHRONOUS STATISTICS COMPUTATION:

The other major source of computational overhead in K-FAC is the estimation of the second-order statistics of the activations and derivatives, which are needed for the Kronecker factors. In the standard K-FAC algorithm, these statistics are computed on the same mini-batches as the gradients, allowing the forward pass computations to be shared between the gradient and statistics computations. By computing the gradients and statistics on separate mini-batches, we can enable a higher degree of parallelism, at the expense of slightly more total computational operations. Under this scheme, the statistics estimation is independent of the gradient computation, so it can be done on one or more separate worker nodes with their own independent data shards. These worker nodes receive parameters from the parameter server (just as in synchronous SGD) and communicate statistics back to the parameter server. In our experiments, we assigned at most one worker to computing statistics.

In cases where it is undesirable to devote separate worker nodes to computing statistics, we also introduce a fast approximation to the statistics for convolution layers (see Appendix A).

4. DOUBLY-FACTORED KRONECKER APPROXIMATION FOR LARGE CONVOLUTION LAYERS:

Computing the standard Kronecker factored Fisher approximation for a given layer involves operations on matrices whose dimension is the number of input units or output units. The cost of these operations is reasonable for most fully-connected networks because the number of units in each layer rarely exceeds a couple thousand. Large convolutional neural networks, however, often include a fully-connected layer that "pools" over a large feature map before the final softmax classification.

For instance, the output of the last pooling layer of AlexNet is of size 6 × 6 × 256 = 9216, which then provides inputs to the subsequent fully connected layer of 4096 ReLUs. VGG models also share a similar architecture. For the standard Kronecker-factored approximation one of the factors will be a matrix of size 9216 × 9216, which is too expensive to be explicitly inverted as often as is needed during training.

In this section we propose a "doubly-factored" Kronecker approximation for layers whose input is a large feature map. Specifically, we approximate the second-order statistics matrix of the inputs as itself factoring as a Kronecker product. This gives an approximation which is a Kronecker product of three matrices.

Using the AlexNet example, the 9216 × 4096 weight matrix in the first fully connected layer is equivalent to a filterbank of 4096 filters with kernel size 6 × 6 on 256 input channels. Let A be a matrix of dimension T -by-C in representing the input activations (for a single training case), where T = K w × K h is the feature map height and width, and C in is the number of input channels. The Fisher block for such a layer can be written as:

EQUATION

We begin be making the following rank-1 approximation:

EQUATION

where K ∈ R T , Ψ ∈ R Cin are the factors along the spatial location dimension and the input channel dimension. The optimal solution of a low-rank approximation under the Frobenius norm is given by the singular value decomposition. The activation matrix A is small enough that its SVD can be computed efficiently. Let σ 1 , u 1 , v 1 be the first singular value and its left and right singular vectors of the activation matrix A, respectively. The factors of the rank-1 approximation are then chosen to be K = √ σ 1 u 1 and Ψ = √ σ 1 v 1 . K captures the activation patterns across spatial locations in a feature map and Ψ captures the pattern across the filter responses. Under the rank-1 approximation of A we have:

EQUATION

We further assume the second order statistics are three-way independent between the loss derivatives Ds, the activations along the input channels Ψ, and the activations along spatial locations K:

EQUATION

The final approximated Fisher block is a Kronecker product of three small matrices. And note that although we assumed the feature map activations have low-rank structure, the resulting approximated Fisher is not low-rank.

The approximate natural gradient for this layer can then be computed by multiplying the inverses of each of the smaller matrices against the respective dimensions of the gradient tensor. We define a function R i : R d1×d2×d3 → R dj d k ×di that constructs a matrix from a 3D tensor by "reshaping" it so that the desired target dimension i ∈ {1, 2, 3} maps to columns, while the remaining dimensions (j and k) are "folded together" and map to the rows. Given the gradient of the weights, G W ∈ R T ×Cin×Cout we can compute the matrix-vector product with the inverse double-factored Kronecker approximated Fisher block as:

EQUATION

)

which is a nested application of the reshape function R(•) at each of the dimension of the gradient tensor.

The doubly factored Kronecker approximation provides a computationally feasible alternative to the standard Kronecker-factored approximation for layers that have a number of parameters in the order of hundreds of millions. For example, inverting it for the first fully connected layer of AlexNet takes about 15 seconds on an 8 core Intel Xeon CPU, and such time is amortized in our asynchronous algorithm.

Unfortunately, the homogeneous coordinate formulation is no longer applicable under this new approximation. Instead, we lump the bias parameters together and associate a full Fisher block with them, which can be explicitly computed and inverted since the number of bias parameters per layer is small.

4.1. FACTORED TIKHONOV DAMPING FOR THE DOUBLE-FACTORED KRONECKER APPROXIMATION:

In second-order optimization methods, "damping" performs the crucial task of correcting for the inaccuracies of the local quadratic approximation of the objective that is (perhaps implicitly) optimized when computing the update (Martens and Sutskever, 2012; Martens, 2014, e.g.) . In the well-known Tikhonov damping/regularization approach, one adds a multiple of the identity λI to the Fisher before inverting it (as one also does for L2-regularization / weight-decay), which roughly corresponds to imposing a spherical trust-region on the update.

The inverse of a Kronecker product can be computed efficiently as the Kronecker product of the inverse of its factors. Adding a multiple of the identity complicates this computation (although it can still be performed tractably using eigendecompositions). The "factored Tikhonov damping" technique proposed in (Martens and Grosse, 2015) is appealing because it preserves the Kronecker structure of the factorization and thus the inverse can still be computed by inverting each of the smaller matrices (and avoiding the more expensive eigendecomposition operation). And in our experiments with large ImageNet models, we also observe the factored damping seems to perform better in practice. In this subsection we derive a generalized version of factored Tikhonov damping for the double-factored Kronecker approximation.

Suppose we wish to add λI to our approximate Fisher block A ⊗ B ⊗ C. In the factored Tikhonov scheme this is approximated by adding π a λ 

(A + π a λ 1 3 I) ⊗ (B + π b λ 1 3 I) ⊗ (C + π c λ 1 3 I) -(A ⊗ B ⊗ C + λI) (10) =π c λ 1 3 I ⊗ A ⊗ B + π b λ 1 3 I ⊗ A ⊗ C + π a λ 1 3 I ⊗ B ⊗ C + π c λ i 3 I ⊗ π b λ 1 3 I ⊗ A + π c λ 1 3 I ⊗ π a λ 1 3 I ⊗ B + π a λ 1 3 I ⊗ π b λ 1 3 I ⊗ C (11)

Following Martens and Grosse (2015) , we choose π a , π b and π c by taking the nuclear norm in Eq. 11 and minimizing its triangle inequality-derived upper-bound. Note that the nuclear norm of Kronecker products is the product of the nuclear norms of each individual matrices: A ⊗ B * = A * B * . This gives the following formula for the value of π a

EQUATION

)

where the d's are the number of rows (equiv. columns) of the corresponding Kronecker factor matrices. The corresponding formulae for π b and π c are analogous. Intuitively, the Eq. 12 rescales the contribution to each factor matrix according to the geometric mean of the ratio of its norm vs the norms of the other factor matrices. This results in the contribution being upscaled if the factor's norm is larger than averaged norm, for example. Note that this formula generalizes to Kronecker products of arbitrary numbers of matrices as the geometric mean of the norm ratios.

5. STEP SIZE SELECTION:

Although Grosse and Martens (2016) found that Polyak averaging (Polyak and Juditsky, 1992) obviated the need for tuning learning rate schedules on some problems, we observed the choice of learning rate schedules to be an important factor in our ImageNet experiments (perhaps due to higher stochasticity in the updates). On ImageNet, it is common to use a fixed exponential decay schedule (Szegedy et al., 2014; 2015) . As an alternative to learning rate schedules, we instead use curvature information to control the amount by which the predictive distribution is allowed to change after each update. In particular, given a parameter update vector v, the second-order Taylor approximation to the KL divergence between the predictive distributions before and after the update is given by the (squared) Fisher norm:

EQUATION

This quantity can be computed with a curvature-vector product (Schraudolph, 2002) . Observe that choosing a step size of η will produce an update with squared Fisher norm η 2 v F v. Instead of using a learning rate schedule, we choose η in each iteration such that the squared Fisher norm is at most some value c:

EQUATION

Grosse and Martens (2016) used this method to clip updates at the start of training, but we found it useful to use it throughout training. We use an exponential decay schedule c k = c 0 ζ k , where c 0 and ζ are tunable parameters, and k is incremented periodically (every half an epoch in our ImageNet experiments). Shrinking the maximum changes in the model prediction after each update is analogous to shrinking the trust region of the second-order optimization. In practice, computing curvature-vector products after every update introduces significant computational overhead, so we instead used the approximate Fisher F in place of F , which allows the approximate Fisher norm to be computed efficiently as

v F v = v F ( F -1 G W ) = v G W .

The maximum step size η max was set to a large value, and in practice this maximum was reached only at the beginning of training, when F was small in magnitude. We found this outperformed simple exponential learning rate decay on ImageNet experiments (see Appendix B).

6. EXPERIMENTS:

We experimentally evaluated distributed K-FAC on several large convolutional neural network training tasks involving the CIFAR-10 and ImageNet classification datasets.

Due to computational resource constraints, we used a single GPU server with 8 Nvidia K80 GPUs to simulate a large distributed system. The GPUs were used as gradient workers that computed the gradient over a large mini-batch, with the CPUs acting as a parameter server. The Fisher block inversions were performed on the CPUs in parallel, using as many threads as possible. The secondorder statistics required for the various Fisher block approximations were computed either syncronously by the gradient workers after each gradient computation (CIFAR-10 experiments), or asynchronously using a separate dedicated "stats worker" (ImageNet experiments).

Meta-parameters such as learning rates, damping parameters, and the decay-rate for the secondorder statistics, were optimized carefully by hand for each method. The momentum was fixed to 0.9.

Similarly to Martens and Grosse (2015) , we applied an exponentially decayed Polyak averaging scheme to the sequence of output iterates produced by each method. We found this improved their convergence rate in the later stages of optimization, and reduced or eliminated the need to decay the learning rates.

We chose to base our implementation of distributed K-FAC on the TensorFlow framework (Abadi et al., 2016) because it provides well-engineered and scalable primitives for distributed computation. We implement distributed K-FAC in TensorFlow by scanning the gradient-computing graph for groups of parameters whose gradient computations have particular structures. Having identified such groups we compute/approximate their Fisher blocks using a method tailored to the type of structure observed. See Appendix C for details. This type of implementation can be applied to existing modelspecification code without significant modification of said code. And because TensorFlow's parallel primitives were designed with scalability in mind, it should be possible to scale our implementation to a larger distributed system with hundreds of workers.

6.1. CIFAR-10 CLASSIFICATION AND ASYNCHRONOUS FISHER BLOCK INVERSION:

In our first experiment we evaluated the effectiveness of asynchronously computing the approximate Fisher inverses (as described in Section 3.1). We considered the effect that this has both on the quality of the updates, as measured by per-iteration progress on the objective, and on the average per-iteration wall-clock time.

The task is to train a basic convolutional network model on the CIFAR-10 image classification dataset (Krizhevsky and Hinton, 2009) . The model has 3 convolutional layers of 32-32-64 filters, each with a receptive field size of 5x5, followed by a softmax layer that predicts 10 classes. This is a similar but not identical CIFAR-10 model that was used by Grosse and Martens (2016) . All the CIFAR-10 experiments use a mini-batch size of 512.

The baseline method is a simple synchronous version of distributed K-FAC with a fixed learning rate, and up to 4 GPUs acting as gradient and stats workers, which recomputes the inverses of the approximate Fisher blocks once every 20 iterations. This baseline method behaves similarly to the implementation of K-FAC in Grosse and Martens (2016) , while being potentially faster due to its greater use of parallelism. We compare this baseline to a version of distributed K-FAC where the approximate Fisher blocks are inverted asynchronously and in parallel with the rest of the optimization process. Note that under this scheme, inverses are updated about once every 16 iterations for the single GPU condition, and every 30 iterations for the four GPU condition. For networks larger than this relatively small CIFAR-10 net they may get updated (far) less often (e.g. the AlexNet experiments in Section 6.2.2).

The results of this first experiment are plotted in Fig. 2 . We found that the asynchronous version iterated about 1.5 times faster than the synchronous version, while its per-iteration progress remained comparable. The plots show that the asynchronous version is better at taking advantage of parallel computation and displayed an almost linear speed-up as the number of gradient workers increases to 4. In terms of the wall-clock time, using only 4 GPUs the asynchronous version of distributed K-FAC is able to complete 700 iterations in under a minute, where it achieves the minimum test error (19%).

6.2. IMAGENET CLASSIFICATION:

In our second set of experiments we benchmarked distributed K-FAC against several other popular approaches, and considered the effect of mini-batch size on per-iteration progress. To do this we trained various off-the-shelf convnet architectures for image classification on the ImageNet dataset (Krizhevsky et al., 2012) , GoogLeNet InceptionV1 (Szegedy et al., 2014) and the 50-layer Residual network (He et al., 2015) .

Despite having 1.2 million images in the ImageNet training set, a data pre-processing pipeline is almost always used for training ImageNet that includes image jittering and aspect distortion. We used a less extensive dataset augmentation/pre-processing pipeline than is typically used for Ima-geNet, as the purpose of this paper is not to achieve state-of-the-art ImageNet results, but rather to evaluate the optimization performance of distributed K-FAC. In particular, the dataset consists of 224x224 images and during training the original images are first resized to 256x256 and then randomly cropped back down to 224x224 before being fed to the network. Note that while it is typically the case that validation error is higher than training error, this data pre-processing pipeline for ImageNet creates an augmented training set that is more difficult than the undistorted validation set and therefore the validation error is often lower than the training error during the first 90% of training. This observation is consistent with previously published results (He et al., 2015) . In all our ImageNet experiments, we used the cheaper Kronecker factorization from Appendix A, and the KL-based step sized selection method described in Section 5 with parameters c 0 = 0.01 and ζ = 0.96. The SGD baselines use an exponential learning rate decay schedule with a decay rate of 0.96. Decaying is applied after each half-epoch for distributed K-FAC and SGD+Batch Normalization, and after every two epochs for plain SGD, which is consistent with the experimental setup of Ioffe and Szegedy (2015) .

6.2.1. GOOGLELENET AND BATCH NORMALIZATION:

Batch Normalization (Ioffe and Szegedy, 2015) is a reparameterization of neural networks that can make them easier to train with first-order methods, and has been successfully applied to large Ima-geNet models. It can be thought of as a modification of the units of a neural network so that each one centers and normalizes its own raw input over the current mini-batch (or subset thereof), after which it applies a separate shift and scaling operation via its own local "bias" and "gain" parameters (which are optimized). These shift and scaling operations can learn to effectively undo the centering and normalization, thus preserving the class of functions that the network can compute. Batch Normalization (BN) is closely related to centering techniques (Schraudolph, 1998) , and likely helps for the same reason that they do, which is that the alternative parameterization gives rise to loss surfaces with more favorable curvature properties. The main difference between BN and traditional centering is that BN makes the centering and normalization operations part of the model instead of the optimization algorithm (and thus "backprops" through them when computing the gradient), which helps stabilize the optimization.

Without any changes to the algorithm, distributed K-FAC can be used to train neural networks that have BN layers. The weight-matrix gradient for such layers has the same structure as it does for standard layers, and so Fisher blocks can be approximated using the same set of techniques. The cross entropy loss and validation error vs the number of updates. Bottom row: cross entropy loss and validation error vs wallclock time (in hours). All methods used 8 GPUs, with distributed K-FAC using the 8-th GPU as a dedicated asynchronous stats worker. per-unit gain and bias parameters cause a minor complication, but because they are relatively few in number, one can compute an exact Fisher block for each of them.

Computing updates for BN networks over large mini-batches is usually done by splitting the minibatch into chunks of size 32, computing the gradients separately for these chunks (using only the data in the chunk to compute the mean and variance statistics), and then summing them together. Using small sample sets to compute the statistics like this introduces additional stochasticity into the BN update that acts as a regularizer, but can also hurt optimization performance. To help decouple the effect of regularization and optimization, we also compared to a BN baseline that uses larger chunks. We found using larger chunks can give a factor of 2 speed-up in optimization performance over the standard BN baseline. In our figures rbz will indicate the chunk size, which defaults 32 if left unspecified.

In Fig. 3 , we compare distributed K-FAC to SGD on GoogLeNet with and without BN. All methods used 4 GPUs, with distributed K-FAC using the 4-th GPU as a dedicated asynchronous stats worker.

We observe that the per-iteration progress made by distributed K-FAC on the training objective is not significantly affected by the use of BN. Moreover, distributed K-FAC is 3.5 times faster than SGD with standard BN baseline (orange line) and 1.5-2 times faster than the enhanced BN baseline (blue line). BN, however, does help distributed K-FAC generalize better, likely due to its aforementioned regularizing effect.

For the simplicity of our discussion, distributed K-FAC is not combined with BN in the the rest of the experiments, as we are chiefly interested in evaluating optimization performance, not regularization, and BN doesn't seem to provide any additional benefit to distributed K-FAC in regards to the former. Note that this is not too surprising, given that K-FAC is provably invariant to the kind of centering and normalization transformations that BN does (Martens and Grosse, 2015) .

6.2.2. ALEXNET AND THE DOUBLY-FACTORED KRONECKER APPROXIMATION:

To demonstrate that distributed K-FAC can efficiently optimize models with very wide layers we train AlexNet using distributed K-FAC and compare to SGD+BN. The doubly-factored Kronecker approximation proposed in Section 4 is applied to the first fully-connected layer of AlexNet, which has 9216 input units and is thus too wide for the standard Kronecker approximation to be feasible. Note that even with this addtional approximation, computing all of the Fisher block inverses for AlexNet is very expensive, and in our experiments they only get updated once every few hundred iterations by our 16 core Xeon 2.2Ghz CPU.

The results from this experiment are plotted in Fig. 4 . They show that Distributed K-FAC still works well despite potentially extreme staleness of the Fisher block inverses, speeding up training by a factor of 1.5 over the improved SGD-BN baseline. 

6.2.3. VERY DEEP ARCHITECTURES (RESNETS):

In recent years very deep convolutional architectures have been successfully applied to ImageNet classification. These networks are particularly challenging to train because the usual difficulties associated with deep learning are especially severe. Fortunately second-order optimization is perhaps ideally suited to addressing these difficulties in a robust and principled way (Martens, 2010) .

To investigate whether distributed K-FAC can scale to such architectures and provide useful acceleration, we compared it to SGD+BN using the 50 layer ResNet architecture (He et al., 2015) . The results from this experiment are plotted in Fig. 5 . They show that distributed K-FAC provides significant speed-up during the early stages of training compared to SGD+BN.

6.2.4. MINI-BATCH SIZE SCALING PROPERTIES:

In our final experiment we explored how well distributed K-FAC scales as additional parallel computing resources become available. To do this we trained GoogLeNet with varying mini-batch sizes of {256, 1024, 2048}, and measured per-training-case progress. Ideally, if extra gradient data is being used efficiently, one should expect the per-training-case progress to remain relatively constant with respect to mini-batch size. The results from this experiment are plotted in Fig. 6 , and show that distributed K-FAC exhibits something close to this ideal behavior, while SGD+BN rapidly loses data efficiency when moving beyond a mini-batch size of 256. These results suggest that distributed K-FAC, more so than the SGD+BN baseline, is capable of speeding up training in proportion to the amount of parallel computational resources used.

We have introduced distributed K-FAC, an asynchronous distributed second-order optimization algorithm which computes Kronecker-factored Fisher approximations and stochastic gradients over larger mini-batches asynchronously and in parallel.

Our experiments show that the extra overhead introduced by distributed K-FAC is mostly mitigated by the use of parallel asynchronous computation, resulting in updates that can be computed in a similar amount of time to those of distributed SGD, while making much more progress on the objective function per iteration. We showed that in practice this can lead to speedups of roughly 3.5x compared to standard SGD + Batch Normalization (BN), and 2x compared to SGD + an improved version of BN on large-scale convolutional network training tasks.

We also proposed a doubly-factored Kronecker approximation that allows distributed K-FAC to scale up to large models with hundreds of millions of parameters, and demonstrated the effectiveness of this approach in experiments.

Finally, we showed that distributed K-FAC enjoys a favorable scaling property with mini-batch size that is seemingly not shared by SGD+BN. In particular, we showed that per-iteration progress tends to be proportional to the mini-batch size up to a much larger threshold than for SGD+BN. This suggests that it will yield even further reductions in total wall-clock training time when implemented in a larger distributed system than the one we considered. A A CHEAPER KRONECKER FACTOR APPROXIMATION FOR CONVOLUTION LAYERS In a convolution layer, the gradient is the sum of the outer product between the receptive field input activation A t and the back-propagated derivatives Ds t at each spatial location t ∈ T . One cannot simply apply the standard Kronecker factored approximation from Martens and Grosse (2015) to each location, sum the results, and then take the inverse, as there is no known efficient algorithm for computing the inverse of such a sum.

In Grosse and Martens (2016) , a Kronecker-factored approximation for convolutional layers called Kronecker Factors for Convolution (KFC) was developed. It works by introducing additional statistical assumptions about how the weight gradients are related across locations. In particular, KFC assumes spatial homogeneity, i.e. that all locations have the same statistics, and spatially uncorrelated derivatives, which (essentially) means that gradients from any two different locations are statistically independent. This yields the following approximation:

EQUATION

In this section we introduce an arguably simpler Kronecker factored approximation for convolutional layers that is cheaper to compute. In practice, it appears to be competitive with the original KFC approximation in terms of per-iteration progress on the objective, working worse in some experiments and better in others, while (often) improving wall-clock time due to its cheaper cost.

It works by approximating the sum of the gradients over spatial locations as the outer product of the averaged receptive field activations over locations Et [At] , and the averaged back-propagated derivatives Et[Dst], multipled by the number of spatial locations |T |. In other words:

EQUATION

Under the approximation assumption that the second-order statistics of the average activations, Et[At], and the second-order statistics of the average derivatives, Et[Dst], are uncorrelated, this becomes: This approximation is cheaper than the original KFC approximation because it is easier to compute a single outer product (after averaging over locations) than it is to compute an outer product at each location and then average. In the synchronous setting, for the large convolutional networks we experimented with, this trick resulted in a 20-30% decrease in overall wall clock time per iteration, with little effect on per-iteration progress.

EQUATION

B EXPERIMENTAL EVALUATION OF THE STEP-SIZE SELECTION METHOD OF SECTION 5:

To compare our proposed step size selection from Sec. 5 with the commonly-used exponential learning rate decay, we performed a simple experiment training GoogLeNet. Both the learning rate and threshold c on the square Fisher norm, is decayed by a factor of 0.96 after every 3200 iterations. The results of this experiment are plotted in Fig. 8 , and indicate that our method outperforms the standard baseline.

C AUTOMATIC CONSTRUCTION OF THE K-FAC COMPUTATION GRAPH:

In recent years, deep learning libraries have moved towards the computational graph abstraction (Bergstra et al., 2010; Abadi et al., 2016) to represent neural network computations. In this section we give a high level description of an algorithm that scans a computational graph for parameters for which one of the various Kronecker-factored approximations can be applied, locates nodes containing the required information to compute the second-order statistics required by the approximations, and then constructs a new graph that computes the approximations and uses them to update the parameters.

For the sake of discussion, we will assume the computation graph is a directed bipartite graph that has a set of operator nodes doing some computation, and some variable nodes that holds intermediate computational results. The trainable parameters are stored in the memory that is loaded or mutated through read/write operator nodes. We also assume that the trainable parameters are grouped layer-wise as a set of weights and biases. Finally, we assume the gradient computation for the trainable parameters is performed by a computation graph (which is usually is generated via automatic differentiation).

In analogy to generating the gradient computation graph through automatic differentiation, given an arbitrary computation graph with a set of the trainable parameters, we would like to use the existing nodes in the given graph to automatically generate a new computation graph, a "K-FAC computation graph", that computes the Kronecker-factored approximate Fisher blocks associated with each group of parameters (typically layers in a neural net), and then uses them to update the parameters.

To compute the Fisher block for a given layer, we want to find all the nodes holding the gradients of the trainable parameters in a computation graph. One simple strategy is to traverse the computation graph from the gradient nodes to their immediate parent nodes.

A set of parameters has a Kronecker-factored approximation to its Fisher block if its corresponding gradient node has a matrix product or convolution operator node as its immediate parent node. For these parameters, the Kronecker factor matrices are the second-order statistics of the inputs to the parent operator node of their gradient nodes (typically the activities A and back-propagated derivatives Ds). For other sets of parameters an exact Fisher block can be computed instead (assuming they have low enough dimension).

In a typical neural network, most of the parameters are concentrated in weight matrices, that are used for matrix product or convolution operations, for which one of the existing Kronecker-factored approximations applies. Homogeneous coordinates can be used if the weights and biases of the same layer are annotated in the computation graph. The rest of the parameters are often gain and bias vectors for each hidden unit, and it is feasible to compute and invert exact Fisher blocks for these.

Kronecker factors can sometimes be shared by approximate Fisher blocks for two or more parameters. This is the case, for example, when a vector of units serves as inputs to two different weightmatrix multiplication operations. In such cases, the computation of the second-order statistics can be reused, which is what we do in our implementation.

A neural network can be also instantiated multiple times in a computational graph (with shared parameters) to process different inputs. The gradient of the parameters shared across the instantiations are the sum of the individual gradients from each instantiation. Given such computation graph, the immediate parent operator node from the gradient is a summation whose inputs are computed by the same type of operators. Without additional knowledge about the computation graph, one approximation is to treat the individual gradient contributions in the summation as statistically independent of each other (similarly to how gradient contributions from multiple spatial locations are treated as independent in the KFC approximation (Grosse and Martens, 2016) ). Under this approximation, the Kronecker factors associated with the gradient can be computed by lumping the statistics associated with each of the gradient contributions together.

Our implementation of Distributed K-FAC in TensorFlow applies the above the strategy to automatically generate K-FAC computation graphs without requiring the user to modify their existing model-definition code.

:

Mathematical evidence for this idea can be found inMartens (2014), where it is shown that (convex quadratic) objective functions decompose into noise-dependent and independent terms, and that second-order methods make much more rapid progress optimizing the noise-independent term compared to SGD, while have no effect on the noise-dependent term (which shrinks with the size of the mini-batch)

Figures and Tables:

Figure 1: The diagram illustrates the distributed computation of K-FAC. Gradient workers (blue) compute the gradient w.r.t. the loss function. Stats workers (grey) compute the sampled secondorder statistics. Additional workers (red) compute inverse Fisher blocks. The parameter server (orange) uses gradients and their inverse Fisher blocks to compute parameter updates.

and π c λ 1 3 I to A, B and C respectively, for non-negative scalars π a , π b and π c satisfying π a π b π c = 1. The error associated with this approximation is:

Figure 2: The results from our CIFAR-10 experiment looking at the effectiveness of asynchronously computing the approximate Fisher inverses. gpu indicates the number of gradient workers. Dashed lines denote training curves and solid lines denote test curves. Top row: cross entropy loss and classification error vs the number of updates. Bottom row: cross entropy loss and classification error vs wallclock time.

Figure 3: Optimization performance of distributed K-FAC and SGD training GoogLeNet on Ima-geNet. Dashed lines denote training curves and solid lines denote validation curves. bz indicates the size of mini-batches. rbz indicates the size of chunks used to assemble the BN updates. Top row: cross entropy loss and classification error v.s. the number of updates. Bottom row: cross entropy loss and classification error vs wallclock time (in hours). All methods used 4 GPUs, with distributed K-FAC using the 4-th GPU as a dedicated asynchronous stats worker.

Figure 4: Optimization performance of distributed K-FAC and SGD training AlexNet on ImageNet. Dashed lines denote training curves and solid lines denote validation curves. bz indicates the size of the mini-batches. rbz indicates the size of chunks used to assemble the BN updates. Top row:cross entropy loss and validation error vs the number of updates. Bottom row: cross entropy loss and validation error vs wallclock time (in hours). All methods used 8 GPUs, with distributed K-FAC using the 8-th GPU as a dedicated asynchronous stats worker.

Figure 5: Optimization performance of distributed K-FAC and SGD training ResNet50 on Ima-geNet. The dashed lines are the training curves and solid lines are the validation curves. bz indicates the size of mini-batches. rbz indicates the size of chunks used to assemble the BN updates. Top row: cross entropy loss and classification error v.s. the number of updates. Bottom row: cross entropy loss and classification error v.s. wallclock time (in hours). All methods used 8 GPUs, with distributed K-FAC using the 8-th GPU as a dedicated asynchronous stats worker.

Figure 6: The comparison of distributed K-FAC and SGD on per training case progress on training loss and errors. The experiments were conducted using GoogLeNet with various mini-batch sizes.

Figure 7: Empirical evaluation of the proposed cheaper Kronecker approximation on GoogLeNet. bz indicates the size of the mini-batches. Dashed lines denote training curves and solid lines denote validation curves. Top row: cross entropy loss and classification error vs the number of updates. Bottom row: cross entropy loss and classification error vs wallclock time.

Figure 8: Results from the experiment described in Appendix B. decayKL indicates the proposed step-size selection method and decayLR indicates standard exponential learning rate decay.

Peer review (iclr2017):
The paper proposes an asynchronous distributed K-FAC method for efficient optimization of deep networks. The authors introduce interesting ideas that many computationally demanding parts of the original K-FAC algorithm can be efficiently implemented in distributed fashion. The gradients and the second-order statistics are computed by distributed workers separately and aggregated at the parameter server along with the inversion of the approximate Fisher matrix computed by a separate CPU machine. The experiments are performed in CIFAR-10 and ImageNet classification problems using models such as AlexNet, ResNet, and GoogleReNet. The paper includes many interesting ideas and techniques to derive an asynchronous distributed version from the original K-FAC. And the experiments also show good results on a few interesting cases. However, I think the empirical results are not thorough and convincing enough yet. Particularly, experiments on various and large number of GPU workers (in the same machine, or across multiple workers) are desired. For example, as pointed by the authors in the answer of a comment, Chen et.al. (Revisiting Distributed Synchronous SGD, 2015) used 100 workers to test their distributed deep learning algorithm. Even considering that the authors have a limitation in computing resource under the academic research setting, the maximum number of 4 or 8 GPUs seems too limited as the only test case of demonstrating the efficiency of a distributed learning algorithm.