Nonlinear Conjugate Gradients For Scaling Synchronous Distributed DNN Training
This addresses the challenge of efficient large-scale DNN training for researchers and practitioners, offering a novel optimizer to enhance convergence in distributed settings.
The paper tackles the problem of scaling synchronous distributed deep neural network training by proposing a stochastic preconditioned nonlinear conjugate gradient optimizer, which improves top-1 accuracy by over 10 percentage points on ImageNet and over 15 percentage points on CIFAR-100 at extremely large mini-batch sizes.
Nonlinear conjugate gradient (NLCG) based optimizers have shown superior loss convergence properties compared to gradient descent based optimizers for traditional optimization problems. However, in Deep Neural Network (DNN) training, the dominant optimization algorithm of choice is still Stochastic Gradient Descent (SGD) and its variants. In this work, we propose and evaluate the stochastic preconditioned nonlinear conjugate gradient algorithm for large scale DNN training tasks. We show that a nonlinear conjugate gradient algorithm improves the convergence speed of DNN training, especially in the large mini-batch scenario, which is essential for scaling synchronous distributed DNN training to large number of workers. We show how to efficiently use second order information in the NLCG pre-conditioner for improving DNN training convergence. For the ImageNet classification task, at extremely large mini-batch sizes of greater than 65k, NLCG optimizer is able to improve top-1 accuracy by more than 10 percentage points for standard training of the Resnet-50 model for 90 epochs. For the CIFAR-100 classification task, at extremely large mini-batch sizes of greater than 16k, NLCG optimizer is able to improve top-1 accuracy by more than 15 percentage points for standard training of the Resnet-32 model for 200 epochs.