Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks

Introduction

As the size of deep neural network models and the data which they are trained on continues to increase rapidly, the demand for distributed parallel computing is increasing. A common approach for achieving distributed parallelism in deep learning is to use the data-parallel approach, where the data is distributed across different processes while the model is replicated across them. When the mini-batch size per process is kept constant to increase the ratio of computation over communication, the effective mini-batch size over the entire system grows proportional to the number of processes.

When the mini-batch size is increased beyond a certain point, the validation accuracy starts to degrade. This generalization gap caused by large mini-batch sizes have been studied extensively for various models and datasets . Hoffer et al. attribute this generalization gap to the limited number of updates, and suggest to train longer . This has lead to strategies such as scaling the learning rate proportional to the mini-batch size, while using the first few epochs to gradually warmup the learning rate . Such methods have enabled the training for mini-batch sizes of 8K, where ImageNet with ResNet-50 could be trained for 90 epochs to achieve 76.3% top-1 validation accuracy in 60 minutes . Combining this learning rate scaling with other techniques such as RMSprop warm-up, batch normalization without moving averages, and a slow-start learning rate schedule, Akiba et al. were able to train the same dataset and model with a mini-batch size of 32K to achieve 74.9% accuracy in 15 minutes .

More complex approaches for manipulating the learning rate were proposed, such as LARS , where a different learning rate is used for each layer by normalizing them with the ratio between the layer-wise norms of the weights and gradients. This enabled the training with a mini-batch size of 32K without the use of ad hoc modifications, which achieved 74.9% accuracy in 14 minutes (64 epochs) . It has been reported that combining LARS with counter intuitive modifications to the batch normalization, can yield 75.8% accuracy even for a mini-batch size of 65K .

The use of small batch sizes to encourage rapid convergence in early epochs, and then progressively increasing the batch size is yet another successful approach . Using such an adaptive batch size method, Mikami et al. were able to train 122 seconds with an accuracy of 75.3%. The hierarchical synchronization of mini-batches have also been proposed , but such methods have not been tested at scale to the extent of the authors’ knowledge.

In the present work, we take a more mathematically rigorous approach to tackle the large mini-batch problem, by using second-order optimization methods. We focus on the fact that for large mini-batch training, each mini-batch becomes more statistically stable and falls into the realm where second-order optimization methods may show some advantage. Another unique aspect of our approach is the accuracy at which we can approximate the Hessian when compared to other second-order methods. Unlike methods that use very crude approximations of the Hessian, such as the TONGA , Hessian free methods , we adopt the Kronecker-Factored Approximate Curvature (K-FAC) method . The two main characteristics of K-FAC are that it converges faster than first-order stochastic gradient descent (SGD) methods, and that it can tolerate relatively large mini-batch sizes without any ad hoc modifications. K-FAC has been successfully applied to convolutional neural networks , distributed memory training of ImageNet , recurrent neural networks , Bayesian deep learning , and reinforcement learning .

We implement a distributed K-FAC optimizer using a synchronous all-worker scheme. We used half precision floating point numbers for both computation and exploited the symmetry of the Kronecker factor to reduce the overhead.

We were able to show for the first time that second-order optimization methods can achieve similar generalization capability compared to highly optimized SGD, by training ResNet-50 on ImageNet as a benchmark. We converged to 75% top-1 validation accuracy in 35 epochs for mini-batch sizes under 16,384, and achieved 75% even with a mini-batch size of 131,072, which took only 978 iterations (Table 1).

We show that we can reduce the frequency of updating the Fisher matrices for K-FAC after a few hundred iterations. In doing so, we are able to reduce the overhead of K-FAC. We were able to train ResNet-50 on ImageNet in 10 minutes to a top-1 accuracy of 74.9% using 1,024 Tesla V100 GPUs (Table 2).

We show that the Fisher matrices for Batch Normalization layers can be approximated as diagonal matrices, which further reduces the computation and memory consumption.

Related work

With respect to large-scale distributed training of deep neural networks, there have been very few studies that use second-order methods. At a smaller scale, there have been previous studies that used K-FAC to train ResNet-50 on ImageNet . However, the SGD they used as reference was not showing state-of-the-art Top-1 validation accuracy (only around 70%), so the advantage of K-FAC over SGD that they claim was not obvious from the results. In the present work, we compare the Top-1 validation accuracy with state-of-the-art SGD methods for large mini-batches mentioned in the introduction (Table 2).

The previous studies that used K-FAC to train ResNet-50 on ImageNet also were not considering large mini-batches and were only training with mini-batch size of 512 on 8 GPUs. In contrast, the present work uses mini-batch sizes up to 131,072, which is equivalent to 32 per GPU on 4096 GPUs, and we are able to achieve a much higher Top-1 validation accuracy of 75%. Note that such large mini-batch sizes can also be achieved by accumulating the gradient over multiple iterations before updating the parameters, which can mimic the behavior of the execution on many GPUs without actually running them on many GPUs.

The previous studies using K-FAC also suffered from large overhead of the communication since they implemented their K-FAC in TensorFlow and used a parameter-server approach. Since the parameter server requires all workers to send the gradients and receive the latest model’s parameters from the parameter server, the parameter server becomes a huge communication bottleneck especially at large scale. Our implementation uses a decentralized approach using MPI/NCCLhttps://developer.nvidia.com/nccl collective communications among the processes. Although, software like Horovodhttps://github.com/horovod/horovod can alleviate the problems with parameter servers, the decentralized approach has been used in high performance computing for a long time, and is known to scale to thousands of GPUs without modification.

Distributed K-FAC

Fisher information matrix. The Fisher information matrix (FIM) of the probabilistic model is estimated by

Strictly speaking, $\mathbf{F}_{\boldsymbol{\theta}}$ is the empirical (stochastic version of) FIM , but we refer to this matrix as FIM throughout this paper for the sake of brevity. In the training of deep neural networks, FIM can be assumed as the curvature matrix in the parameter space .

2 K-FAC

Kronecker-Factored Approximate Curvature (K-FAC) is a second-order optimization method for deep neural networks, which is based on an accurate and mathematically rigorous approximation of the FIM. K-FAC is applied to the training of convolutional neural networks, which minimizes the log likelihood (e.g. a classification task with a loss function (1)).

For the training of the deep neural network with $L$ layers, K-FAC approximates $\mathbf{F}_{\boldsymbol{\theta}}$ as a diagonal block matrix:

The inverse of a Kronecker product is approximated by the Kronecker product of the inverse of each Kronecker factor.

3 Our design

Due to the extra calculation of the inverse FIM, K-FAC has considerable overhead compared to SGD. We designed a distributed parallelization scheme so that this overhead decreases as the number of processes is increased. Furthermore, we introduce a relaxation technique to reduce the computation of the FIM, which is explained in Section 5.4. In doing so, we were able to reduce the overhead of K-FAC to almost a negligible amount.

Figure 1 shows the overview of our design, which shows a single iteration of the training. We use the term stage to refer to each phase of the computation, which is indicated at the top of the figure. The variables in each box illustrates what the process computes during that stage, e.g. at stage 1, each process computes the Kronecker factor $\mathbf{A}$ from the activation.

Once we obtain the pre-conditioned gradients, we switch back to data-parallelism by calling an AllGatherV collective. After stage 6 is finished, all processes can update their parameters using the pre-conditioned gradients for all layers. As we will mention in Section 5.4, we are able to reduce the amount of communication required for the Kronecker factors $\mathbf{A}$ and $\mathbf{G}$. Therefore, the amount of communication is similar to SGD, where the AllReduce is implemented as a ReduceScatter+AllGather. Algorithm 1 shows the pseudo code of our distributed K-FAC design.

4 Further acceleration

Our data-parallel and model-parallel hybrid approach allows us to minimize the overhead of K-FAC in a distributed setting. However, K-FAC still has a large overhead compared to SGD. There are two hotspots in our distributed K-FAC design. The first is the construction of Kronecker factors, which cannot be done in a model-parallel fashion. The second is the extra communication for distributing these Kronecker factors. In this section, we discuss how we accelerated these two hotspots to achieve faster training time.

Symmetry-aware communication. The Kronecker factors are all symmetric matrices , so we exploit this property to reduce the volume of communication. To communicate a symmetric matrix of size $N\times N$, we only need to send the upper triangular matrix of size $N(N+1)/2$.

Training schemes

The behavior of K-FAC on large models and datasets has not been studied in length. Also, there are very few studies that use K-FAC for large mini-batches (over 4K) using distributed parallelism at scale . Contrary to SGD, where the hyperparameters have been optimized by many practitioners even for large mini-batches, there is very little insight on how to tune hyperparameters for K-FAC. In this section, we have explored some methods, which we call training schemes, to achieve higher accuracy in our experiments. In this section, we show those training schemes in our large mini-batch training with K-FAC.

We resize the all images in ImageNet to $256\times 256$ ignoring the aspect ratio of original images and compute the mean value ($224\times 224$) of the upper left of the resized images. When reading an image, we randomly crop a $224\times 224$ image from it, randomly flip it horizontally, subtract the mean value, and scale every pixel to $$.

Running mixup. We extend mixup to increase its regularization effect. We synthesize virtual training samples from raw samples and virtual samples from the previous step (while the original mixup method synthesizes only from the raw samples):

where we set $\alpha=\beta=\alpha_{\rm mixup}$.

Random erasing with zero value. We also adopt the Random Erasing . We put zero value on the erasing region of each input instead of a random value as used in the original method. We set the erasing probability $p=0.5$, the erasing area ratio $S_{e}\in[0.02,0.25]$, and the erasing aspect ratio $r_{e}\in[0.3,1]$. We randomly switch the size of the erasing area from $(H_{e},W_{e})$ to $(W_{e},H_{e})$.

2 Warmup damping

The eigenvalue distribution of the Fisher information matrix (FIM) of deep neural networks is known to have an extremely long tail , where most of the eigenvalues are close to zero. This in turn causes the eigenvalues of the inverse FIM to become extremely large, which causes the norm of the preconditioned gradient $\mathcal{G}$ to become huge compared to the parameter $\mathbf{w}$, so the training becomes unstable. To prevent this problem, we add the damping value $\gamma$ to the diagonal of the FIM to get a preconditioned gradient:

We use a modified Tikhonov damping technique for a Kronecker-factored FIM (Equation 5). At early stages of the training, the FIM changes rapidly (Figure 5). Therefore, we start with a large damping rate and gradually decrease it using following rule:

$\gamma^{(t)}$ is the value for the damping in the $t$ th step. $\gamma^{(0)}$ is the initial value, and $t_{\rm warmup}>0$ controls the steps to reach the target value $\gamma_{\rm target}>\gamma^{(0)}$. At each iteration, we use $\gamma_{\rm BN}^{(t)}=\rho_{\rm BN}\cdot\gamma^{(t)}$ ($\rho_{BN}>1$) for the Batch Normalization layers to stabilize the training.

3 Learning rate and momentum

The learning rate used for all of our experiments is schedule by polynomial decay. The learning rate $\eta^{(e)}$ for $e$ th epoch is determined as follows:

$\eta^{(0)}$ is the initial learning rate and $e_{\rm start},e_{\rm end}$ is the epoch when the decay starts and ends. The decay rate $p_{\rm decay}$ guides the speed of the learning rate decay. The learning rate scheduling in our experiments are plotted in Figure 3 .

We use the momentum method for K-FAC updates. Because the learning rate decays rapidly in the final stage of the training with the polynomial decay, the current update can become smaller than the previous update. We adjust the momentum rate $m^{(e)}$ for $e$ th epoch so that the ratio between $m^{(e)}$ and $\eta^{(e)}$ is fixed throughout the training:

where $m^{(0)}$ is the initial momentum rate. The weights are updated as follows:

4 Weights rescaling

To prevent the scale of weights from becoming too large, we adopt the Normalizing Weights technique. We rescale the $\mathbf{w}$ to have a norm $\sqrt{2\cdot d_{\rm out}}$ after (17):

where we use $\epsilon=1\cdot 10^{-9}$ to stabilize the computation. $d_{\rm out}$ is the output dimension or channels of the layer.

Results

We train ResNet-50 for ImageNet in all of our experiments. We use the same hyperparameters for the same mini-batch size when comparing the different schemes in Section 4. The training curves for the top-1 validation accuracy shown in Figures 3, 6 are averaged over 2 or 3 executions using the same hyperparameters. The hyperparameters for our results are shown in Table 3. We implement all computation on top of Chainer . We initialize the weights by the HeNormal initializer of Chainer https://docs.chainer.org/en/stable/reference/generated/ chainer.initializers.HeNormal.html with the default parameters.

We conduct all experiments on the ABCI (AI Bridging Cloud Infrastructure) https://abci.ai/ operated by the National Institute of Advanced Industrial Science and Technology (AIST) in Japan. ABCI has 1088 nodes with four NVIDIA Tesla V100 GPUs per node. Due to the additional memory required by K-FAC, all of our experiments use a mini-batch size of 32 images per GPU. For large mini-batch size experiments which cannot be executed directly, we used an accumulation method to mimic the behavior by accumulating over multiple steps. We were only given a 24 hour window to use the full machine so we had to tune the hyperparameters on a smaller number of nodes while mimicking the global mini-batch size of the full node run.

2 Scalability

We measured the scalability of our distributed K-FAC implementation on ResNet-50 with ImageNet dataset. Figure 2 shows the time for one iteration using different number of GPUs. Ideally, this plot should show a flat line parallel to the x-axis, since we expect the time per iteration to be independent of the number of GPUs. From 1 GPU to 64 GPUs, we observed a superlinear scaling, where the 64 GPU case is 131.1% faster compared to 1 GPU, which is the consequence of our hybrid data/model-parallel design. ResNet-50 has 107 layers in total when all the convolution, fully-connected, and batch normalization layers are accounted for. Despite this superlinear scaling, after 256 GPUs we observe performance degradation due to the communication overhead.

3 Large mini-batch training with K-FAC

We trained ResNet-50 for the classification task on ImageNet with extremely large mini-batch size BS={4,096 (4K), 8,192 (8K), 16,384 (16K), 32,768 (32K), 65,536 (65K), 131,072 (131K)} and achieved a competitive top-1 validation accuracy ($\geq 75\%$). The summary of the training is shown in Table 1. The training curves and the learning rate schedules are plotted in Figure 3. When we use BS={4K, 8K, 16K, 32K, 65K}, the training converges in much less than 90 epochs, which is the usual number of epochs required by SGD-based training of ImageNet . For BS={4K,8K,16K}, the required epochs to reach higher than 75% top-1 validation accuracy does not change so much. Even for a relatively large mini-batch size of BS=32K, K-FAC still converges in half the number of epochs compared to SGD. When increasing the mini-batch size to BS=65K, we see a 33% increase in the number of epochs it takes to converge. Note that the calculation time is still decreasing while the number of epochs is less than double when we double the mini-batch size, assuming that doubling the mini-batch corresponds to doubling the number of GPUs (and halving the execution time). At BS=131K, there are less than 10 iterations per epoch since the dataset size of ImageNet is 1,281,167. None of the SGD-based training of ImageNet have sustained the top-1 validation accuracy at this mini-batch size. Furthermore, this is the first work that uses K-FAC for the training with extremely large mini-batch size BS={16K,32K,65K,131K} and achieves a competitive top-1 validation accuracy.

4 Fisher information and large mini-batch training

We analyzed the relationship between the large mini-batch training with K-FAC and the Fisher information matrix of ResNet-50.

Staleness of Fisher information. To achieve faster training with distributed K-FAC, reducing the computation and the communication of the FIM (or the Kronecker factors) is required. In ResNet-50 for ImageNet classification, the data of the Kronecker factors $\mathbf{A},\mathbf{G}$ for the convolutional layers and the FIM $\mathbf{F}$ for the Batch Normalization layers are dominant (Figure 4). Note that we do not factorize the FIM for the Batch Normalization layers into $\mathbf{A}$ and $\mathbf{G}$. Previous work on K-FAC used stale Kronecker factors by only calculating them every few steps . Even though our efficient distributed scheme minimizes the overhead of the Kronecker factor calculation, we thought it was worth investigating how much staleness we can tolerate to further speed up our method. We examine the change rate of the Kronecker factors for the convolutional layers and the FIM for the Batch Normalization layers.

Training with stale Fisher information. We found that, regardless of the mini-batch size, the FIM changes rapidly during the first 500 or so iterations. Based on this, we reduce the frequency of updating ($\mathbf{A},\mathbf{G},\mathbf{F}$) after 500 iterations. We apply a heuristic scheduling of the refreshing interval. The refreshing interval (iterations) ${\rm interval}^{(e)}$ for the $e$ th epoch is determined by:

Using 1024 NVIDIA Tesla V100, we achieve 74.9 % top-1 accuracy with ResNet-50 for ImageNet in 10 minutes (45 epochs, including a validation after each epoch). We used the same hyperparameters shown in Table 3. The training time and the validation accuracy are competitive with the results reported by related work that use SGD for training (the comparison is shown in Table 2).

As shown in Figure 4, the FIM for Batch Normalization (BN) layers contribute to a large portion of the memory overhead of K-FAC. To alleviate this overhead, we approximate it with a diagonal matrix. By using the diagonal approximation, we can reduce the memory consumption of the FIM for the all layers of ResNet-50 from 1017MiB to 587MiB (Figure 4). We measure the effect of the diagonal approximation on the accuracy of ResNet-50 for ImageNet with mini-batch size BS=32,768 with/without using stale Fisher information for all layers. In this experiment, we adopt another heuristic for ${\rm interval}^{(e)}$:

The training curves are plotted in Figure 6. Using diagonal FIM does not affect the training curve even with stale FIM. This result suggests that only diagonal values of the FIM is essential for the training of BN layers.

Conclusion

In this work, we proposed a large-scale distributed computational design for the second-order optimization using Kronecker-Factored Approximate Curvature (K-FAC) and showed the advantages of K-FAC over the first-order stochastic gradient descent (SGD) for the training of ResNet-50 with ImageNet classification using extremely large mini-batches. We introduced several schemes for the training using K-FAC with mini-batch sizes up to 131,072 and achieved over 75% top-1 accuracy in much fewer number of epochs/iterations compared to the existing work using SGD with large mini-batch. Contrary to prior claims that second order methods do not generalize as well as SGD, we were able to show that this is not at all the case, even for extremely large mini-batches. Data and model hybrid parallelism introduced in our design allowed us to train on 1024 GPUs and achieved 74.9% in 10 minutes by using K-FAC with the stale Fisher information matrix (FIM). This is the first work which observes the relationship between the FIM of ResNet-50 and its training on large mini-batches ranging from 4K to 131K. There is still room for improvement in our distributed design to overcome the bottleneck of computation/communication for K-FAC -- the Kronecker factors can be approximated more aggressively without loss of accuracy. One interesting observation is that, whenever we coupled our method with a well known technique that improves the convergence of SGD, it allowed us to approximate the FIM more aggressively without any loss of accuracy. This suggests that all these seemingly ad hoc techniques to improve the convergence of SGD, are actually performing an equivalent role to the FIM in some way. The advantage that we have in designing better optimizers by taking this approch is that we are starting from the most mathematically rigorous form, and every improvement that we make is a systematic design decision based on observation of the FIM. Even if we end up having similar performance to the best known first-order methods, at least we will have a better understanding of why it works by starting from second-order methods. Further analysis of the eigenvalues of FIM and its effect on preconditioning the gradient will allow us to further understand the advantage of second-order methods for the training of deep neural networks with extremely large mini-batches.

Acknowledgements

Computational resource of AI Bridging Cloud Infrastructure (ABCI) was awarded by "ABCI Grand Challenge" Program, National Institute of Advanced Industrial Science and Technology (AIST). This work is supported by JST CREST Grant Number JPMJCR19F5, Japan. This work was supported by JSPS KAKENHI Grant Number JP18H03248. (Part of) This work is conducted as research activities of AIST - Tokyo Tech Real World Big-Data Computation Open Innovation Laboratory (RWBC-OIL). This work is supported by "Joint Usage/Research Center for Interdisciplinary Large-scale Information Infrastructures" in Japan (Project ID: jh180012-NAHI).

References