Scheduled Restart Momentum for Accelerated Stochastic Gradient Descent

TL;DR

This paper introduces Scheduled Restart SGD (SRSGD), a novel optimization scheme that enhances convergence and generalization in deep neural network training by combining NAG-style momentum with periodic resets, outperforming standard SGD.

Contribution

SRSGD is a new NAG-inspired method that stabilizes increasing momentum through scheduled resets, leading to faster convergence and better accuracy in training deep neural networks.

Findings

SRSGD improves convergence speed over standard SGD.

SRSGD achieves lower error rates on ImageNet and CIFAR datasets.

SRSGD requires fewer epochs to reach comparable or better accuracy.

Abstract

Stochastic gradient descent (SGD) with constant momentum and its variants such as Adam are the optimization algorithms of choice for training deep neural networks (DNNs). Since DNN training is incredibly computationally expensive, there is great interest in speeding up the convergence. Nesterov accelerated gradient (NAG) 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 ResNet200 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.