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 accelerated variants are among the most used optimization algorithms in ML. These gradient-based algorithms have low computational complexity, and they are easy to parallelize, making them suitable for large scale and high dimensional problems (Zinkevich et al., 2010; Zhang et al., 2015) . Nevertheless, GD and SGD have issues with slow convergence, especially when the problem is ill-conditioned. There are two common techniques to accelerate GD and SGD: adaptive step size (Duchi et al., 2011; Hinton et al.; Zeiler, 2012) and momentum (Polyak, 1964) . The integration of both adaptive step size and momentum with SGD leads to Adam (Kingma & Ba, 2014) , one of the most used optimizers for training DNNs. Many recent developments have improved Adam (Reddi et al., 2019; Dozat, 2016; Loshchilov & Hutter, 2018; Liu et al., 2020) . GD with constant momentum leverages the previous step to accelerate GD according to: v k+1 = w k -s k ∇f (w k ); w k+1 = v k+1 + µ(v k+1 -v k ), where µ > 0 is a constant. A similar acceleration can be achieved by the heavy-ball (HB) method (Polyak, 1964) . The momentum update in both (3) and HB have the same convergence rate of O(1/k) as that of GD for convex smooth optimization. A breakthrough due to Nesterov (1983; 2018) replaces µ with (k -1)/(k + 2), which is known as the Nesterov accelerated gradient (NAG) with time-varying momentum. For simplicity, we denote this method as NAG below. NAG accelerates the convergence rate to O(1/k 2 ), which is optimal for convex and smooth loss functions (Nesterov, 1983; 2018) . NAG can also speed up the process of escaping from saddle points (Jin et al., 2017) . In practice, NAG momentum can accelerate GD for nonconvex optimization, especially when the underlying problem is poorly conditioned (Goh, 2017). However, NAG accumulates error and causes instability when the gradient is inexact (Devolder et al., 2014; Assran & Rabbat, 2020) . In many DL applications, constant momentum achieves state-of-the-art result. For instance, training DNNs for image classification. Since NAG momentum achieves a much better convergence rate than constant momentum with exact gradient for general convex optimization, we consider the following question: Can we leverage NAG with a time-varying momentum parameter to accelerate SGD in training DNNs and improve the test accuracy of the trained models? Contributions. We answer the above question by proposing the first algorithm that integrates scheduled restart NAG momentum with plain SGD. Here, we restart the momentum, which is orthogonal to the learning rate restart (Loshchilov & Hutter, 2016) . We name the resulting algorithm scheduled restart SGD (SRSGD). Theoretically, we prove the error accumulation of Nesterov accelerated SGD (NASGD) and the convergence of SRSGD. The major practical benefits of SRSGD are fourfold: • SRSGD remarkably speeds up DNN training. For image classification, SRSGD significantly reduces the number of training epochs while preserving or even improving the network's accuracy. In particular, on CIFAR10/100, the number of training epochs is reduced by half with SRSGD, while on ImageNet the reduction in training epochs is also remarkable. • DNNs trained by SRSGD generalize significantly better than the current benchmark optimizers. The improvement becomes more significant as the network grows deeper as shown in There is also no additional computational or memory overhead. We focus on image classification with DNNs, in which SGD with constant momentum is the choice. & Nesterov, 1985; Nesterov, 2013; Iouditski & Nesterov, 2014; Lin & Xiao, 2014; Renegar, 2014; Freund & Lu, 2018; Roulet et al., 2015; O'donoghue & Candes, 2015; Giselsson & Boyd, 2014; Su et al., 2014) . These studies of restart NAG momentum are for convex optimization with the exact gradient. Restart techniques have also been used for stochastic optimization (Kulunchakov & Mairal, 2019) . In particular, Aybat et al. (2019) developed a multistage variant of NAG with momentum restart between stages. Our work focuses on developing NAG-based optimization for training DNNs. Many efforts have also been



Fig. 1. • SRSGD reduces overfitting in training very deep networks such as ResNet-200 for ImageNet classification, enabling the accuracy to keep increasing with depth. • SRSGD is straightforward to implement and only requires changes in a few lines of the SGD code.

Related Work. Momentum has long been used to accelerate SGD. SGD with scheduled momentum and a good initialization can handle the curvature issues in training DNNs and enable the trained models to generalize well (Sutskever et al., 2013). Kingma & Ba (2014) and Dozat (2016) integrated momentum with adaptive step size to accelerate SGD. In this work, we study the time-varying momentum version of NAG with restart for stochastic optimization. Adaptive and scheduled restart have been used to accelerate NAG with the exact gradient (Nemirovskii

