SCHEDULED RESTART MOMENTUM FOR ACCELERATED STOCHASTIC GRADIENT DESCENT

Abstract

Stochastic gradient descent (SGD) algorithms, with constant momentum and its variants such as Adam, are the optimization methods of choice for training deep neural networks (DNNs). There is great interest in speeding up the convergence of these methods due to their high computational expense. Nesterov accelerated gradient (NAG) with a time-varying momentum, denoted as NAG below, improves the convergence rate of gradient descent (GD) for convex optimization using a specially designed momentum; however, it accumulates error when an inexact gradient is used (such as in SGD), slowing convergence at best and diverging at worst. In this paper, we propose scheduled restart SGD (SRSGD), a new NAG-style scheme for training DNNs. SRSGD replaces the constant momentum in SGD by the increasing momentum in NAG but stabilizes the iterations by resetting the momentum to zero according to a schedule. Using a variety of models and benchmarks for image classification, we demonstrate that, in training DNNs, SRSGD significantly improves convergence and generalization; for instance, in training ResNet-200 for ImageNet classification, SRSGD achieves an error rate of 20.93% vs. the benchmark of 22.13%. These improvements become more significant as the network grows deeper. Furthermore, on both CIFAR and ImageNet, SRSGD reaches similar or even better error rates with significantly fewer training epochs compared to the SGD baseline.

1. INTRODUCTION

Training many machine learning (ML) models reduces to solving the following finite-sum optimization problem min w f (w) := min w 1 N N i=1 f i (w) := min w 1 N N i=1 L(g(x i , w), y i ), w ∈ R d , where {x i , y i } N i=1 are the training samples and L is the loss function, e.g., cross-entropy loss for a classification task, that measure the discrepancy between the ground-truth label y i and the prediction by the model g(•, w), parametrized by w. The problem (1) is known as empirical risk minimization (ERM). In many applications, f (w) is non-convex, and g(•, w) is chosen among deep neural networks (DNNs) due to their preeminent performance across various tasks. These deep models are heavily overparametrized and require large amounts of training data. Thus, both N and the dimension of w can scale up to millions or even billions. These complications pose serious computational challenges. One of the simplest algorithms to solve (1) is gradient descent (GD), which updates w according to: w k+1 = w k -s k 1 N N i=1 ∇f i (w k ), where s k > 0 is the step size at the k-th iteration. Computing ∇f (w k ) on the entire training set is memory intensive and often prohibitive for devices with limited random access memory (RAM) such as graphics processing units (GPUs) used for deep learning (DL). In practice, we sample a subset of the training set, of size m with m N , to approximate ∇f (w k ) by the mini-batch gradient 1/m m j=1 ∇f ij (w k ), resulting in the (mini-batch)-stochastic gradient descent (SGD). SGD and its 1

