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

By Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, Satoshi Matsuoka. 2018

102 citations recorded by the source corpus.

Large-scale distributed training of deep neural networks suffers from the generalization gap caused by the increase in the effective mini-batch size. Previous approaches try to solve this problem by varying the learning rate and batch size over epochs and layers, or some ad hoc modification of the batch normalization. We propose an alternative approach using a second order optimization method that shows similar generalization capability to first order methods, but converges faster and can handle larger mini-batches. To test our method on a benchmark where highly optimized first order methods are available as references, we train ResNet-50 on ImageNet-1K. 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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

In Optimization & Training

Citations in the corpus

Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model (2019, background): “The importance of understanding the benefits of modern parallel hardware has motivated a lot of recent work on training neural networks with larger batch sizes [Goyal et al., 2017, Osawa et al., 2018, McCandlish et al., 2018, Shallue et al., 2018].”
One Epoch Is All You Need (2019, background): “Some promising ideas that tend to overﬁt include mixture-of-experts (Shazeer et al., 2017) and optimization methods exploiting the second order information such as K-FAC (Osawa et al., 2018).”
WoodFisher: Efficient second-order approximations for model compression (2020, methodology): “As a result, K-FAC has found several applications in the last few years in: optimization (Ba et al., 2016; Osawa et al., 2019), pruning (Zeng and Urtasun, 2019; Wang et al.”
M-FAC: Efficient Matrix-Free Approximations of Second-Order Information (2021, methodology): “This approximation has been leveraged for both pruning and for optimization, and it allows efficient computation of the inverse [5, 33, 49, 44, 23]; however, it is known to not always hold [31].”
Augment your batch: better training with larger batches (2019, methodology): “Increasing batch size while mitigating accuracy degradation is actively researched in the ML and systems communities (Goyal et al., 2017; Kurth et al., 2017; Jia et al., 2018; Mikami et al., 2018; Osawa et al., 2018; Ying et al., 2018).”
A Trace-restricted Kronecker-Factored Approximation to Natural Gradient (2020, unknown): “What’s more, KFAC also has been applied to large-scale distributed computing for DNNs and shows excellent experimental performance (Ba et al., 2017; Osawa et al., 2019; Pauloski et al., 2020).”
Adaptive norms for deep learning with regularized Newton methods (2019, background): “The latter observations suggests that advances in distributed second-order algorithms (e.g., Osawa et al. (2018); D¨unner et al. (2018)) constitute a promising direction of research towards the goal of a more widespread use of Newton-type methods in deep learning.”
WoodFisher: Efficient Second-Order Approximation for Neural Network Compression (2020, background): “As a result, K-FAC has found several applications in the last few years in: optimization [38, 39], pruning [21, 22], reinforcement-learning [51], etc.”