Year: 2020

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.

heading: 1 INTRODUCTION
text: 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
fi(w) := min w
1
N
N∑
i=1
L(g(xi,w), yi), w ∈ Rd, (1)
where {xi, yi}Ni=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 yi 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:
wk+1 = wk − sk 1
N
N∑
i=1
∇fi(wk), (2)
where sk > 0 is the step size at the k-th iteration. Computing ∇f(wk) 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(wk) by the mini-batch gradient 1/m ∑m j=1∇fij (wk), resulting in the (mini-batch)-stochastic gradient descent (SGD). SGD and its
Under review as a conference paper at ICLR 2021
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:
vk+1 = wk − sk∇f(wk); wk+1 = vk+1 + µ(vk+1 − vk), (3) 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/k2), 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 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.
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.
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 & 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
Under review as a conference paper at ICLR 2021
devoted to studying the non-acceleration issues of SGD with HB and NAG momentum (Kidambi et al., 2018; Liu & Belkin, 2020), as well as accelerating first-order algorithms with noise-corrupted gradients (Cohen et al., 2018; Aybat et al., 2018; Lan, 2012). Ghadimi & Lan (2013; 2016) provides analysis for the general stochastic gradient-based optimization algorithms. .
Organization. In Section 2, we review and discuss momentum for accelerating GD for convex smooth optimization. In Section 3, we present the SRSGD algorithm and its theoretical guarantees. In Section 4, we verify the efficacy of the proposed SRSGD in training DNNs for image classification on CIFAR and ImageNet. In Section 4.3, we perform empirical analysis of SRSGD. We end with some concluding remarks. Technical proofs, some experimental details, and more results in training LSTMs (Hochreiter & Schmidhuber, 1997) and WGANs (Arjovsky et al., 2017; Gulrajani et al., 2017) are provided in the Appendix.
Notation. We denote scalars and vectors by lower case and lower case bold face letters, respectively, and matrices by upper case bold face letters. For a vector x = (x1, · · · , xd) ∈ Rd, we denote its `p norm (p ≥ 1) by ‖x‖p = ( ∑d i=1 |xi|p)1/p. For a matrix A, we use ‖A‖p to denote its induced norm by the vector `p norm. Given two sequences {an} and {bn}, we write an = O(bn) if there exists a positive constant s.t. an ≤ Cbn. We denote the interval a to b (included) as (a, b]. For a function f(w) : Rd → R, we denote its gradient as∇f(w) and its Hessian as∇2f(w).

heading: 2 REVIEW: MOMENTUM IN GRADIENT DESCENT
text: GD. GD (2) is a popular approach to solve (1), which dates back to Cauchy (1847). If f(w) is convex and L-smooth (i.e., ‖∇2f(w)‖2 ≤ L), then GD converges with rate O(1/k) by letting sk ≡ 1/L (we use this sk in all the discussion below), which is independent of the dimension of w. HB. HB (4) (Polyak, 1964) accelerates GD by using the historical information, which gives
wk+1 = wk − sk∇f(wk) + µ(wk −wk−1), µ > 0. (4) We can also accelerate GD by using the Nesterov/lookahead momentum, which leads to (3). Both (3) and (4) have a convergence rate of O(1/k) for convex smooth optimization. Recently, several variants of (3) have been proposed for DL, e.g., (Sutskever et al., 2013) and (Bengio et al., 2013).
NAG. NAG (Nesterov, 1983; 2018; Beck & Teboulle, 2009) replaces µ with (tk − 1)/tk+1, where tk+1 = (1 + √ 1 + 4t2k)/2 with t0 = 1. NAG iterates as following
vk+1 = wk − sk∇f(wk); wk+1 = vk+1 + tk − 1 tk+1 (vk+1 − vk). (5)
NAG achieves a convergence rate O(1/k2) with the step size sk = 1/L. Remark 1. Su et al. (2014) showed that (k− 1)/(k+ 2) is the asymptotic limit of (tk − 1)/tk+1. In the following presentation of NAG with restart, for the ease of notation, we will replace the momentum coefficient (tk − 1)/tk+1 with (k − 1)/(k + 2).
Adaptive Restart NAG (ARNAG). The sequences, {f(wk)− f(w∗)} where w∗ is the minimum of f(w), generated by GD and GD with constant momentum (GD + Momentum, which follows (3)) converge monotonically to zero. However, that sequence generated by NAG oscillates, as illustrated in Fig. 2 (a) when f(w) is a quadratic function. O’donoghue & Candes (2015) proposed ARNAG
Under review as a conference paper at ICLR 2021
(6), which restart the time-varying momentum of NAG according to the change of function values, to alleviate this oscillatory phenomenon. ARNAG iterates as following
vk+1 = wk − sk∇f(wk); wk+1 = vk+1 + m(k)− 1 m(k) + 2 (vk+1 − vk), (6)
where m(1) = 1; m(k + 1) = m(k) + 1 if f(wk+1) ≤ f(wk), and m(k + 1) = 1 otherwise. Scheduled Restart NAG (SRNAG). SR is another strategy to restart the time-varying momentum of NAG. We first divide the total iterations (0, T ] (integers only) into a few intervals {Ii}mi=1 = (Ti−1, Ti], such that (0, T ] = ⋃m i=1 Ii. In each Ii we restart the momentum after every Fi iterations. The update rule is then given by:
vk+1 = wk − sk∇f(wk); wk+1 = vk+1 + (k mod Fi) (k mod Fi) + 3 (vk+1 − vk). (7)
Both AR and SR accelerate NAG to linear convergence for convex problems with the PolyakLojasiewicz (PL) condition (Roulet & d’Aspremont, 2017).
Case Study – Quadratic Function. Consider the following quadratic optimization (Hardt, 2014)
min x f(x) =
1 2 xTLx− xT b, (8)
where L ∈ Rd×d is the Laplacian of a cycle graph, and b is a d-dimensional vector whose first entry is 1 and all the other entries are 0. Note that f(x) is convex with Lipschitz constant 4. In particular, we set d = 1K (1K:= 103). We run T = 50K iterations with step size 1/4. In SRNAG, we restart, i.e., we set the momentum to 0, after every 1K iterations. Fig. 2 (a) shows that GD + Momentum as in (3) converges faster than GD, while NAG speeds up GD + Momentum dramatically and converges to the minimum in an oscillatory fashion. Both AR and SR accelerate NAG significantly.

heading: 3 ALGORITHM PROPOSED: SCHEDULED RESTART SGD (SRSGD)
text: Computing gradient for ERM, (1), can be computational costly and memory intensive, especially when the training set is large. In many applications, such as training DNNs, SGD is used. In this section, we first prove that the error bound of SGD with NAG cannot be bounded by a convergent sequence, then we formulate our new SRSGD as a solution to accelerate the convergence of SGD using the NAG momentum.

heading: 3.1 UNCONTROLLED BOUND OF NESTEROV ACCELERATED SGD (NASGD)
text: Replacing ∇f(wk) := 1/N∑Ni=1∇fi(wk) in (5) with the mini-batch gradient 1/m∑mj=1∇fij (wk) will lead to uncontrolled error bound. Theorem 1 formulates this observation for NASGD. Theorem 1 (Uncontrolled Bound of NASGD). Let f(w) be a convex and L-smooth function with ‖∇f(w)‖ ≤ R, whereR > 0 is a constant. The sequence {wk}k≥0 generated by (5), with stochastic gradient of bounded variance (Bubeck, 2014; Bottou et al., 2018) 1 and using any constant step size sk ≡ s ≤ 1/L, satisfies
E ( f(wk)− f(w∗) ) = O(k), (9)
1We leave the analysis under the other assumptions (Jain et al., 2018) as a future work.
Under review as a conference paper at ICLR 2021
where w∗ is the minimum of f , and the expectation is taken over the generation of the stochastic gradient.
One idea to prove Theorem 1 is by leveraging the established resulting in Lan (2012). We will provide a new proof of Theorem 1 in Appendix A. The proof shows that the uncontrolled error bound is because the time-varying momentum gets close to 1 as iteration increases. To remedy this, we can restart the momentum in order to guarantee that the time-varying momentum with restart is less than a number that is strictly less than 1. Devolder et al. (2014) proved a similar error bound for the δ-inexact gradient, and we provide a brief review of NAG with δ-inexact gradient in Appendix B. As far as we know that there is no lower bound of E(f(wk)− f(w∗)) available even for the δ-inexact gradient, and we leave the lower bound estimation as an open problem.
We consider three different inexact gradients: Gaussian noise with constant and decaying variance corrupted gradients for the quadratic optimization (8), and training logistic regression model for MNIST (LeCun & Cortes, 2010) classification. The detailed settings and discussion are provided in Appendix B. We denote SGD with NAG momentum as NASGD and NASGD with AR and SR as ARSGD and SRSGD, respectively. The results shown in Fig. 2 (b) and (c) (iteration vs. optimal gap for quadratic optimization (8)) and Fig. 3 (a) (iteration vs. loss for training logistic regression) confirm Theorem 1. For these cases, SR improves the performance of NAG with inexact gradients. Moreover, when an inexact gradient is used, ARNAG/ARSGD performs almost the same as GD/SGD asymptotically because ARNAG/ARSGD restarts too often and almost degenerates to GD/SGD.

heading: 3.2 SRSGD AND ITS CONVERGENCE
text: For ERM (1), SRSGD replaces∇f(w) in (7) with stochastic gradient with batch size m and gives
vk+1 = wk − sk 1 m m∑ j=1 ∇fij (wk); wk+1 = vk+1 + (k mod Fi) (k mod Fi) + 3 (vk+1 − vk), (10)
where Fi is the restart frequency used in the interval Ii. We implemented SRSGD, in both PyTorch (Paszke et al., 2019) and Keras (Chollet et al., 2015), by changing just a few lines of code on top of the existing implementation of the SGD optimizer. We provide a snippet of SRSGD code in Appendix J (PyTorch) and K (Keras). We formulate the convergence of SRSGD for general convex and nonconvex problems in Theorem 2 and provide its proof in Appendix C. Theorem 2 (Convergence of SRSGD). Suppose f(w) is L-smooth. Consider the sequence {wk}k≥0 generated by (10) with stochastic gradient that is bounded and has bounded variance, and consider any restart frequency F using any constant step size sk := s ≤ 1/L. Assume that ∑ k∈A ( Ef(wk+1)− Ef(wk) ) = R̄ < +∞ with R̄ being a constant and the set A := {k ∈ Z+|Ef(wk+1) ≥ Ef(wk)}, then we have
min 1≤k≤K
{ E‖∇f(wk)‖22 } = O ( s+ 1
sK
) . (11)
If f(w) is further convex and ∑ k∈B ( Ef(wk+1)− Ef(wk) ) = R̂ < +∞ with R̂ being a constant and the set B := {k ∈ Z+|E‖wk+1 −w∗‖2 ≥ E‖wk −w∗‖2}, then
min 1≤k≤K
{ E ( f(wk)− f(w∗) )} = O ( s+ 1
sK
) , (12)
where w∗ is the minimum of f . To obtain (∀ > 0) error, we set s = O( ) and K = O(1/ 2).
Theorem 2 relies on the assumption that ∑ k∈A or B ( Ef(wk+1)− Ef(wk) ) is bounded, and we provide an empirical verification in Appendix C.1. We leave it open for how to establish the convergence result for SRSGD without this assumption.

heading: 4 EXPERIMENTAL RESULTS
text: We evaluate SRSGD on a variety of benchmarks for image classification, including CIFAR10, CIFAR100, and ImageNet. In all experiments, we show the advantage of SRSGD over the widely used and well-calibrated SGD baselines with a constant momentum of 0.9 and decreasing learning
Under review as a conference paper at ICLR 2021
rate at certain epochs, which we denote as SGD. We also compare SRSGD with the well-calibrated SGD in which we switch momentum to the Nesterov momentum of 0.9, and we denote this optimizer as SGD + NM. We fine tune the SGD and SGD + NM baselines to obtain the best validation performance, and we then adopt the same set of parameters for training with SRSGD. In the SRSGD experiments, we tune the restart frequencies on small DNNs for each task based on the validation performance and apply the calibrated restart frequencies to large DNNs for the same task. Note that ARSGD is impractical for training on large-scale datasets since it requires to compute the loss over the whole training set at each iteration, which is very computationally inefficient. Alternatively, ARSGD can estimate loss and restart using mini-batches, but then ARSGD restarts too often and degenerates to SGD without momentum as we mentioned in Section 3. Thus, we do not compare with ARSGD in our CIFAR and ImageNet experiments. The details about hyper-parameters calibration can be found in Appendix D.4. We provide the detailed description of datasets and experimental settings in Appendix D. Additional experimental results in training LSTMs (Hochreiter & Schmidhuber, 1997) and WGANs (Arjovsky et al., 2017; Gulrajani et al., 2017) with SRSGD, as well as the comparison between SRSGD and SGD + NM on ImageNet classification task, are provided in Appendix E. We also note that in all the following experiments, the training loss will blow up if we apply NASGD without restart. These further confirm the stabilizing effect of scheduled restart in training DNNs.

heading: 4.1 CIFAR10 AND CIFAR100
text: We summarize our results for CIFAR in Tables 1 and 2. We also explore two different restarting frequency schedules for SRSGD: linear and exponential schedule. These schedules are governed by two parameters: the initial restarting frequency F1 and the growth rate r. In both scheduling schemes, the restarting frequency at the 1st learning rate stage is set to F1 during training. Then the restarting frequency at the (k + 1)-th learning rate stage is determined by:
Fk+1 = { F1 × rk, exponential schedule F1 × (1 + (r − 1)× k), linear schedule.
We search F1 and r using the method outlined in Appendix D.4. For CIFAR10, (F1 = 40, r = 1.25) and (F1 = 30, r = 2) are good initial restarting frequencies and growth rates for the exponential and linear schedules, respectively. For CIFAR100, those values are (F1 = 45, r = 1.5) for the exponential schedule and (F1 = 50, r = 2) for the linear schedule.
Improvement in Accuracy Increases with Depth. We observe that the linear schedule of restart yields better test error on CIFAR than the exponential schedule for most of the models except for Pre-ResNet-470 and Pre-ResNet-1001 on CIFAR100 (see Tables 1 and 2). SRSGD with either linear or exponential restart schedule outperforms SGD. Furthermore, the advantage of SRSGD over SGD is more significant for deeper networks. This observation holds strictly when using the linear schedule (see Fig. 1) and is generally true when using the exponential schedule with only a few exceptions.
Faster Convergence Reduces the Training Time by Half. SRSGD also converges faster than SGD. This result is consistent with our MNIST case study in Section 3 and indeed expected since SRSGD
Under review as a conference paper at ICLR 2021
can avoid the error accumulation when there is an inexact oracle. For CIFAR, Fig. 3 (b) shows that SRSGD yields smaller training loss than SGD during the training. Interestingly, SRSGD converges quickly to good loss values in the 2nd and 3rd stages. This suggests that the model can be trained with SRSGD in many fewer epochs compared to SGD while achieving a similar error rate.
Results in Table 3 confirm the hypothesis above. We train Pre-ResNet models with SRSGD in only 100 epochs, decreasing the learning rate by a factor of 10 at the 80th, 90th, and 95th epoch while using the same linear schedule for restarting frequency as before with (F1 = 30, r = 2) for CIFAR10 and (F1 = 50, r = 2) for CIFAR100. We compare the test error of the trained models with those trained by the SGD baseline in 200 epochs. We observe that SRSGD training consistently yields lower test errors than SGD except for the case of Pre-ResNet-110 even though the number of training epochs of our method is only half of the number of training epochs required by SGD. For Pre-ResNet-110, SRSGD needs 110 epochs with learning rate decreased at the 80th, 90th, and 100th epoch to achieve the same error rate as the 200-epoch SGD training on CIFAR10. On CIFAR100, SRSGD training for Pre-ResNet-110 needs 140 epochs with learning rate decreased at the 80th, 100th and 120th epoch to outperform the 200-epoch SGD. Comparison with SGD short training is provided in Appendix F.2.
Comparison with Adam and RMSProp. SRSGD outperforms not only SGD with momentum but also other popular optimizers including Adam and RMSProp (Tieleman & Hinton, 2012) for image classification tasks. In fact, for image classification tasks, Adam and RMSProp yield worse performance than the baseline SGD with momentum (Chen & Kyrillidis, 2019). Table 4 compares SRSGD with Adam and RMSprop on CIFAR10.
Under review as a conference paper at ICLR 2021

heading: 4.2 IMAGENET
text: Next, we discuss our experimental results on the 1000-way ImageNet classification task (Russakovsky et al., 2015). We conduct our experiments on ResNet-50, 101, 152, and 200 with 5 different seeds. We use the official PyTorch implementation for all of our ResNet models (Paszke et al., 2019). Following common practice, we train each model for 90 epochs and decrease the learning rate by a factor of 10 at the 30th and 60th epoch. We use an initial learning rate of 0.1, a momentum scaled by 0.9, and a weight decay value of 0.0001. Additional details and comparisons between SRSGD and SGD + NM are given in Appendix E.
We report single crop validation errors of ResNet models trained with SGD and SRSGD on ImageNet in Table 5. In contrast to our CIFAR experiments, we observe that for ResNets trained on ImageNet with SRSGD, linearly decreasing the restarting frequency to 1 at the last stage (i.e., after the 60th epoch) helps improve the generalization of the models. Thus, in our experiments, we use linear scheduling with (F1 = 40, r = 2). From epoch 60 to 90, the restarting frequency decays to 1 linearly.
Advantage of SRSGD continues to grow with depth. Similar to the CIFAR experiments, we observe that SRSGD outperforms the SGD baseline for all ResNet models that we study. As shown in Fig. 1, the advantage of SRSGD over SGD grows with network depth, just as in our CIFAR experiments with Pre-ResNet architectures.
Avoiding Overfitting in ResNet-200. ResNet-200 demonstrates that SRSGD is better than the SGD baseline at avoiding overfitting2. The ResNet-200 trained with SGD has a top-1 error of 22.13%, higher than the ResNet-152 trained with SGD, which achieves a top-1 error of 22.03% (see Table 5). He et al. (2016b) pointed out that ResNet-200 suffers from overfitting. The ResNet-200 trained with our SRSGD has a top-1 error of 20.93%, which is 1.2% lower than the ResNet-200 trained with the SGD and also lower than the ResNet-152 trained with both SRSGD and SGD, an improvement by 0.53% and 1.1%, respectively. We hypothesize that SRSGD with appropriate restart frequency is locally not monotonic (see Fig. 3 (b, c)), and this property allows SRSGD to escape from bad minima in order to reach a better one, which helps avoid overfitting in very deep networks. Theoretical analysis of the observation that SRSGD is less overfitting in training DNNs is under our investigation.
Training ImageNet in Fewer Number of Epochs. As in the CIFAR experiments, we note that when training on ImageNet, SRSGD converges faster than SGD at the first and last learning rate while quickly reaching a good loss value at the second learning rate (see Fig. 3 (c)). This observation suggests that ResNets can be trained with SRSGD in fewer epochs while still achieving comparable error rates to the same models trained by the SGD baseline using all 90 epochs. We summarize the results in Table 6. On ImageNet, we note that SRSGD helps reduce the number of training epochs for very deep networks (ResNet-101, 152, 200). For smaller networks like ResNet-50, training with fewer epochs slightly decreases the accuracy.

heading: 4.3 EMPIRICAL ANALYSIS
text: SRSGD Helps Reduce the Training Time. We find that SRSGD training using fewer epochs yields comparable error rates to both the SGD baseline and the SRSGD full training with 200 epochs on CIFAR. We conduct an ablation study to understand the impact of reducing the number of epochs
2By overfitting, we mean that the model achieves low training error but high test error.
Under review as a conference paper at ICLR 2021
on the final error rate when training with SRSGD on CIFAR10 and ImageNet. In the CIFAR10 experiments, we vary the number of epoch reduction from 15 to 90 while in the ImageNet experiments, we vary the number of epoch reduction from 10 to 30. We summarize our results in Fig. 4, and provide detailed results in Appendix F. For CIFAR10, we can train with 30 fewer epochs while still maintaining a comparable error rate to the full SRSGD training, and with a better error rate than the SGD baseline trained in full 200 epochs. For ImageNet, SRSGD training with fewer epochs decreases the accuracy but still obtains comparable results to the 90-epoch SGD baseline.
Impact of Restarting Frequency. We examine the impact of restarting frequency on the network training. We choose a case study of training a Pre-ResNet-290 on CIFAR10 using SRSGD with a linear schedule scheme for the restarting frequency. We fix the growth rate r = 2 and vary the initial restarting frequency F1 from 1 to 80. As shown in Fig. 5, SRSGD with a large F1, e.g. F1 = 80, approximates NASGD (yellow). We also show the training loss and test accuracy of NASGD in red. As discussed in Section 3, it suffers from error accumulation due to stochastic gradients and converges slowly or even diverges. SRSGD with small F1, e.g. F1 = 1, approximates SGD without momentum (green). It converges faster initially but reaches a worse local minimum (i.e. larger loss). Typical SRSGD (blue) converges faster than NASGD and to a better local minimum than both NASGD and SGD without momentum. It also achieves the best test error. We provide more empirical analysis results in Appendix F, G and H. The impact of the growth rate r is studied in Appendix G.2.

heading: 5 CONCLUSIONS
text: We propose the Scheduled Restart SGD (SRSGD), with two major changes from the widely used SGD with constant momentum. First, we replace the momentum in SGD with the iteration-dependent momentum that used in Nesterov accelerated gradient (NAG). Second, we restart the NAG momentum according to a schedule to prevent error accumulation when the stochastic gradient is used. For image classification, SRSGD can significantly improve the accuracy of the trained DNNs. Also, compared to the SGD baseline, SRSGD requires fewer training epochs to reach the same trained model’s accuracy. There are numerous avenues for future work: 1) deriving the optimal restart scheduling and the corresponding convergence rate of SRSGD and 2) integrating the scheduled restart NAG momentum with adaptive learning rate algorithms, e.g., Adam (Kingma & Ba, 2014).
Under review as a conference paper at ICLR 2021

heading: Part

heading: Appendices
text: The appendices are structured as follows. In Section A, we prove Theorem 1. In Section B, we review an error accumulation result of the Nesterov accelerated gradient with δ-inexact gradient. In Section C, we prove Theorem 2. In Section D, we provide some experimental details; in particular, the calibration of restarting hyperparameters. In Section E, we compare SRSGD with benchmark optimization algorithms on some other tasks, including training LSTM and Wasserstein GAN. In Section F, we provide detailed experimental settings in studying the effects of reducing the number of epoch in training deep neural networks with SRSGD, and we provide some more experimental results. In Section G and H, we further study the effects of restarting frequency and training with less epochs by using SRSGD. In Section I, we visualize the optimization trajectory of SRSGD and compare it with benchmark methods. A snippet of our implementation of SRSGD in PyTorch and Keras are available in Section J and K, respectively.

heading: Table of Contents
text: Appendix 14

heading: A Uncontrolled Bound of NASGD 15
text: A.1 Preliminaries . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15 A.2 Uncontrolled Bound of NASGD: Analysis . . . . . . . . . . . . . . . . . . . . 16

heading: B NAG with δ-Inexact Oracle & Experimental Settings in Section 3.1 19

heading: C Convergence of SRSGD 20
text: C.1 Numerical Verification of the assumptions in Theorem 2 . . . . . . . . . . . . . 22

heading: D Datasets and Implementation Details 23
text: D.1 CIFAR . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23 D.2 ImageNet . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23 D.3 Training ImageNet in Fewer Number of Epochs: . . . . . . . . . . . . . . . . . 23 D.4 Details on Restarting Hyper-parameters Search . . . . . . . . . . . . . . . . . . 23

heading: E SRSGD vs. SGD and SGD + NM on ImageNet Classification and Other Tasks 24
text: E.1 Comparing with SGD with Nesterov Momentum on ImageNet Classification . . 24 E.2 Long Short-Term Memory (LSTM) Training for Pixel-by-Pixel MNIST . . . . . 24 E.3 Wasserstein Generative Adversarial Networks (WGAN) Training on MNIST . . 25

heading: F Error Rate vs. Reduction in Training Epochs 27
text: F.1 Implementation Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27 F.2 Short Training on CIFAR10/CIFAR100 Using SGD . . . . . . . . . . . . . . . 28 F.3 Additional Experimental Results . . . . . . . . . . . . . . . . . . . . . . . . . 28
G Impact of Restarting Frequency for ImageNet and CIFAR100 29 G.1 Implementation Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 29 G.2 Impact of the Growth Rate r . . . . . . . . . . . . . . . . . . . . . . . . . . . . 30 G.3 Additional Experimental Results . . . . . . . . . . . . . . . . . . . . . . . . . 31

heading: H Full Training with Less Epochs at the Intermediate Learning Rates 32
text: I Visualization of SRSGD’s trajectory 33

heading: J SRSGD Implementation in Pytorch 35
text: K SRSGD Implementation in Keras 36
Under review as a conference paper at ICLR 2021

heading: A UNCONTROLLED BOUND OF NASGD
text: Consider the following optimization problem
min w f(w), (13)
where f(w) is L-smooth and convex.
Start from wk, GD update, with step size 1r , can be obtained based on the minimization of the function
Qr(v,w k) := 〈v −wk,∇f(wk)〉+ r
2 ‖v −wk‖22. (14)
With direct computation, we can get that
Qr(v k+1,wk)−minQr(v,wk) = ‖gk −∇f(wk)‖2 2r ,
where gk := 1m ∑m j=1∇fij (wk). We assume the variance is bounded, which gives The stochastic gradient rule,Rs, satisfies E[Qr(vk+1,wk)−minQr(v,wk)|χk] ≤ δ, with δ being a constant and χk being the sigma algebra generated by w1,w2, · · · ,wk, i.e.,
χk := σ(w1,w2, · · · ,wk).
NASGD can be reformulated as
vk+1 ≈ arg min v Qr(v,w k) with ruleRs,
wk+1 = vk+1 + tk − 1 tk+1 (vk+1 − vk), (15)
where t0 = 1 and tk+1 = (1 + √ 1 + 4t2k)/2.

heading: A.1 PRELIMINARIES
text: To proceed, we introduce several definitions and some useful properties in variational and convex analysis. More detailed background can be found at Mordukhovich (2006); Nesterov (1998); Rockafellar & Wets (2009); Rockafellar (1970).
Let f be a convex function, we say that f is L-smooth (gradient Lipschitz) if f is differentiable and
‖∇f(v)−∇f(w)‖2 ≤ L‖v −w‖2, and we say f is ν-strongly convex if for any w,v ∈ dom(f)
f(w) ≥ f(v) + 〈∇f(v),w − v〉+ ν 2 ‖w − v‖22.
Below of this subsection, we list several basic but useful lemmas, the proof can be found in Nesterov (1998). Lemma 1. If f is ν-strongly convex, then for any v ∈ dom(f) we have
f(v)− f(v∗) ≥ ν 2 ‖v − v∗‖22, (16)
where v∗ is the minimizer of f . Lemma 2. If f is L-smooth, for any w,v ∈ dom(f),
f(w) ≤ f(v) + 〈∇f(v),w − v〉+ L 2 ‖w − v‖22.
Under review as a conference paper at ICLR 2021

heading: A.2 UNCONTROLLED BOUND OF NASGD: ANALYSIS
text: In this part, we denote ṽk+1 := arg min
v Qr(v,w
k). (17)
Lemma 3. If the constant r > 0, then
E ( ‖vk+1 − ṽk+1‖22|χk ) ≤ 2δ
r . (18)
Proof. Note that Qr(v,wk) is strongly convex with constant r, and ṽk+1 in (17) is the minimizer of Qr(v,w k). With Lemma 1 we have
Qr(v k+1,wk)−Qr(ṽk+1,wk) ≥
r 2 ‖vk+1 − ṽk+1‖22. (19)
Notice that
E [ Qr(v k+1,wk)−Qr(ṽk+1,wk) ] = E [ Qr(v
k+1,wk)−min v Qr(v,w
k) ] ≤ δ.
The inequality (18) can be established by combining the above two inequalities.
Lemma 4. If the constant satisfy r > L, then we have
E ( f(ṽk+1) + r
2 ‖ṽk+1 −wk‖22 − (f(vk+1) +
r 2 ‖vk+1 −wk‖22)
) (20)
≥ −τδ − r − L 2 E[‖wk − ṽk+1‖22],
where τ = L 2
r(r−L) + 1.
Proof. The convexity of f gives us
0 ≤ 〈∇f(vk+1),vk+1 − ṽk+1〉+ f(ṽk+1)− f(vk+1). (21)
From the definition of the stochastic gradient ruleRs, we have −δ ≤ E ( Qr(ṽ k+1,wk)−Qr(vk+1,wk) )
(22)
= E [ 〈ṽk+1 −wk,∇f(wk)〉+ r
2 ‖ṽk+1 −wk‖22
] −
E [ 〈vk+1 −wk,∇f(wk)〉+ r
2 ‖vk+1 −wk‖22
] .
With (21) and (22), we have
−δ ≤ ( f(ṽk+1) + r
2 ‖ṽk+1 −wk‖22
) − ( f(vk+1) + r
2 ‖vk+1 −wk‖22
) + (23)
E〈∇f(wk)−∇f(ṽk+1), ṽk+1 − vk+1〉.
With the Schwarz inequality 〈a, b〉 ≤ ‖a‖ 2 2
2µ + µ 2 ‖b‖22 with µ = L
2
r−L , a = ∇f(vk+1)−∇f(ṽk+1) and b = wk − ṽk+1,
〈∇f(wk)−∇f(ṽk+1), ṽk+1 − vk+1〉 (24)
≤ (r − L) 2L2 ‖∇f(wk)−∇f(ṽk+1)‖22 + L2 2(r − L)‖v k+1 − ṽk+1‖22
≤ (r − L) 2 ‖wk − ṽk+1‖22 +
L2
2(r − L)‖v k+1 − ṽk+1‖22.
Combining (23) and (24), we have
−δ ≤ E ( f(ṽk+1) + r
2 ‖ṽk+1 −wk‖22
) − E ( f(vk+1) + r
2 ‖vk+1 −wk‖22
) (25)
+ L2
2(r − L)E‖v k+1 − ṽk+1‖22 + r − L 2 E‖wk − ṽk+1‖22.
By rearrangement of the above inequality (25) and using Lemma 3, we obtain the result.
Under review as a conference paper at ICLR 2021
Lemma 5. If the constants satisfy r > L, then we have the following bounds
E ( f(vk)− f(vk+1) ) ≥ r
2 E‖wk − vk+1‖22 + rE〈wk − vk, ṽk+1 −wk〉 − τδ, (26)
E ( f(v∗)− f(vk+1) ) ≥ r
2 E‖wk − vk+1‖22 + rE〈wk − v∗, ṽk+1 −wk〉 − τδ, (27)
where τ := L 2 r(r−L) + 1 and v ∗ is the minimum .
Proof. With Lemma 2, we have
− f(ṽk+1) ≥ −f(wk)− 〈ṽk+1 −wk,∇f(wk)〉 − L 2 ‖ṽk+1 −wk‖22. (28)
Using the convexity of f , we have
f(vk)− f(wk) ≥ 〈vk −wk,∇f(wk)〉, i.e.,
f(vk) ≥ f(wk) + 〈vk −wk,∇f(wk)〉. (29)
According to the definition of ṽk+1 in (14), i.e.,
ṽk+1 = arg min v Qr(v,w k) = arg min v 〈v −wk,∇f(wk)〉+ r 2 ‖v −wk‖22,
and the optimization condition gives
ṽk+1 = wk − 1 r ∇f(wk). (30)
Substituting (30) into (29), we obtain
f(vk) ≥ f(wk) + 〈vk −wk, r(wk − ṽk+1)〉. (31)
Direct summation of (28) and (31) gives
f(vk)− f(ṽk+1) ≥ ( r − L
2
) ‖ṽk+1 −wk‖22 + r〈wk − vk, ṽk+1 −wk〉. (32)
Summing (32) and (20), we obtain the inequality (26)
E [ f(vk)− f(vk+1) ] ≥ r
2 E‖wk − vk+1‖22 + rE〈wk − vk, ṽk+1 −wk〉 − τδ. (33)
On the other hand, with the convexity of f , we have
f(v∗)− f(wk) ≥ 〈v∗ −wk,∇f(wk)〉 = 〈v∗ −wk, r(wk − ṽk+1)〉. (34) The summation of (28) and (34) results in
f(v∗)− f(ṽk+1) ≥ ( r − L
2
) ‖wk − ṽk+1‖22 + r〈wk − v∗, ṽk+1 −wk〉. (35)
Summing (35) and (20), we obtain
E ( f(v∗)− f(vk+1) ) ≥ r
2 E‖wk − vk+1‖22 + rE〈wk − v∗, ṽk+1 −wk〉 − τδ, (36)
which is the same as (27).
Theorem 3 (Uncontrolled Bound of NASGD (Theorem 1 with detailed bounded)). Let the constant r satisfy r < L and the sequence {vk}k≥0 be generated by NASGD with stochastic gradient that has bounded variance. By using any constant step size sk ≡ s ≤ 1/L, then we have
E[f(vk)−min v f(v)] ≤ (2τδ r +R2) 4k 3 . (37)
Under review as a conference paper at ICLR 2021
Proof. We denote F k := E(f(vk)− f(v∗)).
By (26)× (tk − 1) + (27), we have
2[(tk − 1)F k − tkF k+1] r ≥ tkE‖vk+1 −wk‖22 (38)
+ 2E〈ṽk+1 −wk, tkwk − (tk − 1)vk − v∗〉 − 2τtkδ
r .
With t2k−1 = t 2 k − tk, (38)× tk yields
2[t2k−1F k − t2kF k+1] r ≥ E‖tkvk+1 − tkwk‖22 (39)
+ 2tkE〈ṽk+1 −wk, tkwk − (tk − 1)vk − v∗〉 − 2τt2kδ
r
Substituting a = tkvk+1 − (tk − 1)vk − v∗ and b = tkwk − (tk − 1)vk − v∗ into identity
‖a− b‖22 + 2〈a− b, b〉 = ‖a‖22 − ‖b‖22. (40) It follows that
E‖tkvk+1 − tkwk‖22 + 2tkE〈ṽk+1 −wk, tkwk − (tk − 1)vk − v∗〉 (41) = E‖tkvk+1 − tkwk‖22 + 2tkE〈vk+1 −wk, tkwk − (tk − 1)vk − v∗〉
+2tkE〈ṽk+1 − vk+1, tkwk − (tk − 1)vk − v∗〉 =
(40) E‖tkvk+1 − (tk − 1)vk − v∗‖22 − ‖tkwk − (tk − 1)vk − v∗‖22 +2tkE〈ṽk+1 − vk+1, tkwk − (tk − 1)vk − v∗〉 = E‖tkvk+1 − (tk − 1)vk − v∗‖22 − E‖tk−1vk − (tk−1 − 1)vk−1 − v∗‖22 + 2tkE〈ṽk+1 − vk+1, tk−1vk − (tk−1 − 1)vk−1 − v∗〉.
In the third identity, we used the fact tkwk = tkvk + (tk−1 − 1)(vk − vk−1). If we denote uk = E‖tk−1vk − (tk−1 − 1)vk−1 − v∗‖22, (39) can be rewritten as
2t2kF k+1 r + uk+1 ≤ 2t 2 k−1F k r + uk + 2τt2kδ r (42)
+ 2tkE〈vk+1 − ṽk+1, tk−1vk − (tk−1 − 1)vk−1 − v∗〉
≤ 2t 2 kF k
r + uk +
2τt2kδ
r + t2k−1R 2,
where we used
2tkE〈vk+1 − ṽk+1, tk−1vk − (tk−1 − 1)vk−1 − v∗〉 ≤ t2kE‖vk+1 − ṽk+1‖22 + E‖tk−1vk − (tk−1vk − (tk−1 − 1)vk−1 − v∗)‖22 = 2t2kδ/r + t 2 k−1R 2.
Denoting
ξk := 2t2k−1F k
r + uk,
then, we have
ξk+1 ≤ ξ0 + ( 2τδ
r +R2)
k∑
i=1
t2i = ( 2τδ
r +R2)
k3
3 . (43)
With the fact, ξk ≥ 2t 2 k−1F k r ≥ k2F k/4, we then proved the result.
Under review as a conference paper at ICLR 2021
B NAG WITH δ-INEXACT ORACLE & EXPERIMENTAL SETTINGS IN SECTION 3.1
In Devolder et al. (2014), the authors defines δ-inexact gradient oracle for convex smooth optimization as follows: Definition 1 (δ-Inexact Oracle). Devolder et al. (2014) For a convex L-smooth function f : Rd → R. For ∀w ∈ Rd and exact first-order oracle returns a pair (f(w),∇f(w)) ∈ R × Rd so that for ∀v ∈ Rd we have
0 ≤ f(v)− ( f(w) + 〈∇f(w),v −w〉 ) ≤ L
2 ‖w − v‖22.
A δ-inexact oracle returns a pair ( fδ(w),∇fδ(w) ) ∈ R× Rd so that ∀v ∈ Rd we have
0 ≤ f(v)− ( fδ(w) + 〈∇fδ(w),v −w〉 ) ≤ L
2 ‖w − v‖22 + δ.
We have the following convergence results of GD and NAG under a δ-Inexact Oracle for convex smooth optimization. Theorem 4. Devolder et al. (2014)3 Consider
min f(w), w ∈ Rd, where f(w) is convex and L-smooth with w∗ being the minimum. Given access to δ-inexact oracle, GD with step size 1/L returns a point wk after k steps so that
f(wk)− f(w∗) = O ( L
k
) + δ.
On the other hand, NAG, with step size 1/L returns
f(wk)− f(w∗) = O ( L
k2
) +O(kδ).
Theorem 4 says that NAG may not robust to a δ-inexact gradient. In the following, we will study the numerical behavior of a variety of first-order algorithms for convex smooth optimizations with the following different inexact gradients.
Constant Variance Gaussian Noise: We consider the inexact oracle where the true gradient is contaminated with a Gaussian noise N (0, 0.0012). We run 50K iterations of different algorithms. For SRNAG, we restart after every 200 iterations. Fig. 2 (b) shows the iteration vs. optimal gap, f(xk)− f(x∗), with x∗ being the minimum. NAG with the inexact gradient due to constant variance noise does not converge. GD performs almost the same as ARNAG asymptotically, because ARNAG restarts too often and almost degenerates into GD. GD with constant momentum outperforms the three schemes above, and SRNAG slightly outperforms GD with constant momentum.
Decaying Variance Gaussian Noise: Again, consider minimizing (8) with the same experimental setting as before except that ∇f(x) is now contaminated with a decaying Gaussian noise N (0, ( 0.1bt/100c+1 )2). For SRNAG, we restart every 200 iterations in the first 10k iterations, and restart every 400 iterations in the remaining 40K iterations. Fig. 2 (c) shows the iteration vs. optimal gap by different schemes. ARNAG still performs almost the same as GD. The path of NAG is oscillatory. GD with constant momentum again outperforms the previous three schemes. Here SRNAG significantly outperforms all the other schemes.
Logisitic Regression for MNIST Classification: We apply the above schemes with stochastic gradient to train a logistic regression model for MNIST classification LeCun & Cortes (2010). We consider five different schemes, namely, SGD, SGD + (constant) momentum, NASGD, ASGD, and SRSGD. In ARSGD, we perform restart based on the loss value of the mini-batch training data. In SRSGD, we restart the NAG momentum after every 10 iterations. We train the logistic regression model with a `2 weight decay of 10−4 by running 20 epochs using different schemes with batch size of 128. The step sizes for all the schemes are set to 0.01. Fig. 3 (a) plots the training loss vs. iteration. In this case, NASGD does not converge, and SGD with momentum does not speed up SGD. ARSGD’s performance is on par with SGD’s. Again, SRSGD gives the best performance with the smallest training loss among these five schemes.
3We adopt the result from Hardt (2014).
Under review as a conference paper at ICLR 2021

heading: C CONVERGENCE OF SRSGD
text: We prove the convergence of Nesterov accelerated SGD with scheduled restart, i.e., the convergence of SRSGD. We denote that θk := tk−1tk+1 in the Nesterov iteration and θ̂ k is its use in the restart version, i.e., SRSGD. For any restart frequency F (positive integer), we have θ̂k = θk−bk/Fc∗F . In the restart version, we can see that θ̂k ≤ θF =: θ̄ < 1. Lemma 6. Let the constant satisfies r > L and the sequence {vk}k≥0 be generated by the SRSGD with restart frequency F (any positive integer), we have
k∑
i=1
‖vi − vi−1‖22 ≤ r2kR2
(1− θ̄)2 , (44)
where θ̄ := θF < 1 and R := supx{‖∇f(x)‖2}.
Proof. It holds that
‖vk+1 −wk‖2 = ‖vk+1 − vk + vk −wk‖2 (45) ≥ ‖vk+1 − vk‖2 − ‖vk −wk‖2 ≥ ‖vk+1 − vk‖2 − θ̄‖vk − vk−1‖2.
Thus,
‖vk+1 −wk‖22 ≥ ( ‖vk+1 − vk‖2 − θ̄‖vk − vk−1‖2 )2 (46)
= ‖vk+1 − vk‖22 − 2θ̄‖vk − vk−1‖2‖vk − vk−1‖2 + θ̄2‖vk − vk−1‖22 ≥ (1− θ̄)‖vk+1 − vk‖22 − θ̄(1− θ̄)‖vk+1 − vk‖22.
Summing (46) from k = 1 to K, we get
(1− θ̄)2 K∑
k=1
‖vk − vk−1‖22 ≤ K∑
k=1
‖vk+1 −wk‖22 ≤ r2KR2. (47)
In the following, we denote
A := {k ∈ Z+|Ef(vk) ≥ Ef(vk−1)}. Theorem 5 (Convergence of SRSGD). (Theorem 2 with detailed bound) Suppose f(w) is L-smooth. Consider the sequence {wk}k≥0 generated by (10) with stochastic gradient that is bounded and has bound variance. Using any restart frequency F and any constant step size sk := s ≤ 1/L. Assume that ∑ k∈A ( Ef(wk+1)− Ef(wk) ) = R̄ < +∞, then we have
min 1≤k≤K
{ E‖∇f(wk)‖22 } ≤ rR 2
(1− θ̄)2 L(1 + θ̃) 2 + rLR2 2 + θ̃R̃ rK . (48)
If f(w) is further convex and the set B := {k ∈ Z+|E‖wk+1 −w∗‖2 ≥ E‖wk −w∗‖2} obeys∑ k∈B ( Ef(wk+1)− Ef(wk) ) = R̂ < +∞, then
min 1≤k≤K
{ E ( f(wk)− f(w∗) )} ≤ ‖w 0 −w∗‖2 + R̂ 2γk + γR2 2 , (49)
where w∗ is the minimum of f . To obtain (∀ > 0) error, we set s = O( ) and K = O(1/ 2).
Proof. Firstly, we show the convergence of SRSGD for nonconvex optimization. L-smoothness of f , i.e., Lipschitz gradient continuity, gives us
f(vk+1) ≤ f(wk) + 〈∇f(wk),vk+1 −wk〉+ L 2 ‖vk+1 −wk‖22. (50)
Under review as a conference paper at ICLR 2021
Taking expectation, we get
Ef(vk+1) ≤ Ef(wk)− rE‖∇f(wk)‖22 + r2LR2
2 . (51)
On the other hand, we have
f(wk) ≤ f(vk) + θ̂k〈∇f(vk),vk − vk−1〉+ L(θ̂ k)2
2 ‖vk − vk−1‖22. (52)
Then, we have
Ef(vk+1) ≤ Ef(vk) + θ̂kE〈∇f(vk),vk − vk−1〉 (53)
+ L(θ̂k)2
2 E‖vk − vk−1‖22 − rE‖∇f(wk)‖22 +
r2LR2
2 .
We also have
θ̂k〈∇f(vk),vk − vk−1〉 ≤ θ̂k ( f(vk)− f(vk−1) + L
2 ‖vk − vk−1‖22
) . (54)
We then get that
Ef(vk+1) ≤ Ef(vk) + θ̂k ( Ef(vk)− Ef(vk−1) ) − rE‖∇f(wk)‖22 +Ak, (55)
where
Ak := E L 2 ‖vk − vk−1‖22 +
L(θ̂k)2
2 E‖vk − vk−1‖22 +
r2LR2
2 .
Summing the inequality gives us
Ef(vK+1) ≤ Ef(v0) + θ̃ ∑
k∈A
( Ef(vk)− Ef(vk−1) ) (56)
− r K∑
k=1
E‖∇f(wk)‖22 + K∑
k=1
Ak.
It is easy to see that θ̃ ∑
k∈A
( Ef(vk)− Ef(vk−1) ) = θ̃R̃.
We get the result by using Lemma 6
Secondly, we prove the convergence of SRSGD for convex optimization. Let w∗ be the minimizer of f . We have
E‖vk+1 −w∗‖22 = E‖wk − γ∇f(wk)−w∗‖22 (57) = E‖wk −w∗‖22 − 2γE〈∇f(wk),wk −w∗〉+ γ2E‖∇f(wk)‖22 ≤ E‖wk − x∗‖22 − 2γE〈∇f(wk),wk −w∗〉+ γ2R2.
We can also derive
E‖wk −w∗‖2 = E‖vk + θ̂k(vk − vk−1)−w∗‖22 = E‖vk −w∗‖22 + 2θ̂kE〈vk − vk−1,vk −w∗〉+ (θ̂k)2E‖vk − vk−1‖22 = E‖vk −w∗‖22 + θ̂kE ( ‖vk −w∗‖22 + ‖vk−1 − vk‖22 − ‖vk−1 −w∗‖22 )
+ (θ̂)2E‖vk − vk−1‖22 = E‖vk −w∗‖22 + θ̂kE ( ‖vk −w∗‖22 − ‖vk−1 −w∗‖22 ) + 2(θ̂k)2E‖vk − vk−1‖22,
where we used the following identity
(a− b)T (a− b) = 1 2 [‖a− d‖22 − ‖a− c‖22 + ‖b− c‖22 − ‖b− d‖22].
Under review as a conference paper at ICLR 2021
Then, we have
E‖vk+1 −w∗‖22 ≤ E‖vk −w∗‖22 − 2γE〈∇f(wk),wk −w∗〉+ 2(θ̂k)2E‖vk − vk−1‖22 (58) + r2R2 + θ̂kE(‖vk −w∗‖22 − ‖vk−1 −w∗‖22).
We then get that
2γE ( f(wk)− f(w∗) ) ≤ E‖vk −w∗‖22 − E‖vk+1 −w∗‖22 (59) + θ̂k ( E‖vk −w∗‖22 − E‖vk−1 −w∗‖22 ) + r2R2.
Summing the inequality gives us the desired convergence result for convex optimization.

heading: C.1 NUMERICAL VERIFICATION OF THE ASSUMPTIONS IN THEOREM 2
text: In this part, we numerically verify the assumptions in Theorem 2. In particular, we apply SRSGD with learning rate 0.1 to train LeNet 4 for MNIST classification (we test on MNIST due to extremely large computational cost). We conduct numerical verification as follows: starting from a given point w0, we randomly sample 469 mini-batches (note in total we have 469 batches in the training data) with batch size 128 and compute the stochastic gradient using each mini-batch. Next, we advance to the next step with each of these 469 stochastic gradients and get the approximated Ef(w1). We randomly choose one of these 469 positions as the updated weights of our model. By iterating the above procedure, we can get w1,w2, · · · and Ef(w1),Ef(w2), · · · and we use these values to verify our assumptions in Theorem 2. We set restart frequencies to be 20, 40, and 80, respectively. Figure 6 top panels plot k vs. the cardinality of the set A := {k ∈ Z+|Ef(wk+1) ≥ Ef(wk)}, and Figure 6 bottom panels plot k vs. ∑ k∈A ( Ef(wk+1)− Ef(wk) ) . Figure 6 shows that ∑ k∈A ( Ef(wk+1)− Ef(wk) ) converges to a constant R̄ < +∞. We also noticed that when the training gets plateaued, E(f(wk)) still oscillates, but the magnitude of the oscillation diminishes as iterations goes, which is consistent with our plots that the cardinality of A increases linearly, but R̄ converges to a finite number. These numerical results show that our assumption in Theorem 2 is reasonable.
4We used the PyTorch implementation of LeNet at https://github.com/activatedgeek/LeNet-5.
Under review as a conference paper at ICLR 2021

heading: D DATASETS AND IMPLEMENTATION DETAILS

heading: D.1 CIFAR
text: The CIFAR10 and CIFAR100 datasets Krizhevsky et al. (2009) consist of 50K training images and 10K test images from 10 and 100 classes, respectively. Both training and test data are color images of size 32× 32. We run our CIFAR experiments on Pre-ResNet-110, 290, 470, 650, and 1001 with 5 different seeds He et al. (2016b). We train each model for 200 epochs with batch size of 128 and initial learning rate of 0.1, which is decayed by a factor of 10 at the 80th, 120th, and 160th epoch. The weight decay rate is 5× 10−5 and the momentum for the SGD baseline is 0.9. Random cropping and random horizontal flipping are applied to training data. Our code is modified based on the Pytorch classification project Yang (2017),5 which was also used by Liu et al. Liu et al. (2020). We provide the restarting frequencies for the exponential and linear scheme for CIFAR10 and CIFAR100 in Table 7 below. Using the same notation as in the main text, we denote Fi as the restarting frequency at the i-th learning rate.
The ImageNet dataset contains roughly 1.28 million training color images and 50K validation color images from 1000 classes Russakovsky et al. (2015). We run our ImageNet experiments on ResNet50, 101, 152, and 200 with 5 different seeds. Following He et al. (2016a;b), we train each model for 90 epochs with a batch size of 256 and decrease the learning rate by a factor of 10 at the 30th and 60th epoch. The initial learning rate is 0.1, the momentum is 0.9, and the weight decay rate is 1× 10−5. Random 224× 224 cropping and random horizontal flipping are applied to training data. We use the official Pytorch ResNet implementation Paszke et al. (2019),6 and run our experiments on 8 Nvidia V100 GPUs. We report single-crop top-1 and top-5 errors of our models. In our experiments, we set F1 = 40 at the 1st learning rate, F2 = 80 at the 2nd learning rate, and F3 is linearly decayed from 80 to 1 at the 3rd learning rate (see Table 8).

heading: D.3 TRAINING IMAGENET IN FEWER NUMBER OF EPOCHS:
text: Table 9 contains the learning rate and restarting frequency schedule for our experiments on training ImageNet in fewer number of epochs, i.e. the reported results in Table 6 in the main text. Other settings are the same as in the full-training ImageNet experiments described in Section D.2 above.
Additional Implementation Details: Implementation details for the ablation study of error rate vs. reduction in epochs and the ablation study of impact of restarting frequency are provided in Section F and G below.

heading: D.4 DETAILS ON RESTARTING HYPER-PARAMETERS SEARCH
text: In our CIFAR10 and CIFAR100 experiments, for both linear and exponential schedule, we conduct hyperparameter searches over the restarting frequencies using our smallest model, Pre-ResNet-110,
5Implementation available at https://github.com/bearpaw/pytorch-classification 6Implementation available at https://github.com/pytorch/examples/tree/master/imagenet
Under review as a conference paper at ICLR 2021
making choices based on final validation performance. The same chosen restarting frequencies are applied for all models including Pre-ResNet-110, 290, 470, 650, and 1001. In particular, we use 10,000 images from the original training set as a validation set. This validation set contains 1,000 and 100 images from each class for CIFAR10 and CIFAR100, respectively. We first train Pre-ResNet-110 on the remaining 40,000 training images and use the performance on the validation set averaged over 5 random seeds to select the initial restarting frequency F1 and the growth rate r. Both F1 and r are selected using grid search from the sets of {20, 25, 30, 35, 40, 45, 50} and {1, 1.25, 1.5, 1.75, 2}, respectively. We then train all models including Pre-ResNet-110, 290, 470, 650, and 1001 on all 50,000 training images using the selected values of F1 and r and report the results on the test set which contains 10,000 test images. The reported test performance is averaged over 5 random seeds. We also use the same selected values of F1 and r for our short training experiments in Section 4.3.
For ImageNet experiments, we use linear scheduling and sweep over the initial restarting frequency F1 and the growth rate r in the set of {20, 30, 40, 50, 60} and {1, 1.25, 1.5, 1.75, 2}, respectively. We select the values of F1 = 40 and r = 2 which have the highest final validation accuracy averaged over 5 random seeds. Same as in CIFAR10 and CIFAR100 experiments, we select F1 and r using our smallest model, ResNet-50, and apply the same selected hyperparameter values for all models including ResNet-50, 101, 152, and 200. We also use the same selected values of F1 and r for our short training experiments in Section 4.3. However, for ResNet-50, we observe that F1 = 60 and r = 1.75 yields better performance in short training. All reported results are averaged over 5 random seeds.

heading: E SRSGD VS. SGD AND SGD + NM ON IMAGENET CLASSIFICATION AND OTHER TASKS

heading: E.1 COMPARING WITH SGD WITH NESTEROV MOMENTUM ON IMAGENET CLASSIFICATION
text: In this section, we compare SRSGD with SGD with Nesterov constant momentum (SGD + NM) in training ResNets for ImageNet classification. All hyper-parameters of SGD with constant Nesterov momentum used in our experiments are the same as those of SGD described in section D.2. We list the results in Table 10. Again, SRSGD remarkably outperforms SGD + NM in training ResNets for ImageNet classification, and as the network goes deeper the improvement becomes more significant.

heading: E.2 LONG SHORT-TERM MEMORY (LSTM) TRAINING FOR PIXEL-BY-PIXEL MNIST
text: In this task, we examine the advantage of SRSGD over SGD and SGD with Nesterov Momentum in training recurrent neural networks. In our experiments, we use an LSTM with different numbers of hidden units (128, 256, and 512) to classify samples from the well-known MNIST dataset LeCun & Cortes (2010). We follow the implementation of Le et al. (2015) and feed each pixel of the image into the RNN sequentially. In addition, we choose a random permutation of 28× 28 = 784 elements at the beginning of the experiment. This fixed permutation is applied to training and testing sequences. This task is known as permuted MNIST classification, which has become standard to measure the performance of RNNs and their ability to capture long term dependencies.
Under review as a conference paper at ICLR 2021
Implementation and Training Details: For the LSTM model, we initialize the forget bias to 1 and other biases to 0. All weights matrices are initialized orthogonally except for the hidden-to-hidden weight matrices, which are initialized to be identity matrices. We train each model for 350 epochs with the initial learning rate of 0.01. The learning rate was reduced by a factor of 10 at epoch 200 and 300. The momentum is set to 0.9 for SGD with standard and Nesterov constant momentum. The restart schedule for SRSGD is set to 90, 30, 90 . The restart schedule changes at epoch 200 and 300. In all experiments, we use batch size 128 and the gradients are clipped so that their L2 norm are at most 1. Our code is based on the code from the exponential RNN’s Github.7
Results: Our experiments corroborate the superiority of SRSGD over the two baselines. SRSGD yields much smaller test error and converges faster than SGD with standard and Nesterov constant momentum across all settings with different number of LSTM hidden units. We summarize our results in Table 11 and Figure 7.
Table 11: Test errors (%) on Permuted MNIST of trained with SGD, SGD + NM and SRSGD. The LSTM model has 128 hidden units. In all experiments, we use the initial learning rate of 0.01, which is reduced by a factor of 10 at epoch 200 and 300. All models are trained for 350 epochs. The momentum for SGD and SGD + NM is set to 0.9. The restart schedule in SRSGD is set to 90, 30, and 90.
Network No. Hidden Units SGD SGD + NM SRSGD Improvement over SGD/SGD + NM
LSTM 128 10.10± 0.57 9.75± 0.69 8.61± 0.30 1.49/1.14 LSTM 256 10.42± 0.63 10.09± 0.61 9.03± 0.23 1.39/1.06 LSTM 512 10.04± 0.35 9.55± 1.09 8.49± 1.59 1.55/1.06
Beyond DNNs: Training LSTM on PMNIST
Table 10: Test errors (%) on Permuted MNIST of trained with SGD, SGD + NM and SRSGD. The
LSTM model has 128 hidden units. In all experiments, we use the initial learning rate of 0.01, which
is reduced by a factor of 10 at epoch 200 and 300. All models are trained for 350 epochs. The
momentum for SGD and SGD + NM is set to 0.9. The restart schedule in SRSGD is set to 90, 30,
and 90.
Network No. Hidden Units SGD SGD + NM SRSGD Improvement over SGD/SGD + NM LSTM 128 10.10 ± 0.57 9.75 ± 0.69 8.61 ± 0.30 1.49/1.14 LSTM 256 10.42 ± 0.63 10.09 ± 0.61 9.03 ± 0.23 1.39/1.06 LSTM 512 10.04 ± 0.35 9.55 ± 1.09 8.49 ± 1.59 1.55/1.06
SGD + Momentum SRSGDSGD + Nesterov Momentum
Tr ai
n lo
ss
Iteration
Figure 6: Training loss vs. training iterations of LSTM trained with SGD (red), SGD + NM (green), and SRSGD (blue) for PMNIST classification tasks.
We evaluate our models using the discriminator’s loss, i.e. the Earth Moving distance estimate, since572 in WGAN lower discriminator loss and better sample quality are correlated [2].573
Implementation and Training Details: The detailed implementations of our generator and discrim-574 inator are given below. For the generator, we set latent dim to 100 and d to 32. For the discriminator,575 we set d to 32. We train each model for 350 epochs with the initial learning rate of 0.01. The learning576 rate was reduced by a factor of 10 at epoch 200 and 300. The momentum is set to 0.9 for SGD with577 standard and Nesterov constant momentum. The restart schedule for SRSGD is set to 60, 120, 180.578 The restart schedule changes at epoch 200 and 300. In all experiments, we use batch size 64. Our579 code is based on the code from the Pytorch WGAN-GP Github.8580
i m p o r t t o r c h581 i m p o r t t o r c h . nn as nn582
583 c l a s s G e n e r a t o r ( nn . Module ) :584 d e f i n i t ( s e l f , l a t e n t d i m , d =32) :585 s u p e r ( ) . i n i t ( )586 s e l f . n e t = nn . S e q u e n t i a l (587 nn . ConvTranspose2d ( l a t e n t d i m , d ⇤ 8 , 4 , 1 , 0 ) ,588 nn . BatchNorm2d ( d ⇤ 8) ,589 nn . ReLU( True ) ,590
591 nn . ConvTranspose2d ( d ⇤ 8 , d ⇤ 4 , 4 , 2 , 1 ) ,592 nn . BatchNorm2d ( d ⇤ 4) ,593 nn . ReLU( True ) ,594
595 nn . ConvTranspose2d ( d ⇤ 4 , d ⇤ 2 , 4 , 2 , 1 ) ,596 nn . BatchNorm2d ( d ⇤ 2) ,597 nn . ReLU( True ) ,598
599 nn . ConvTranspose2d ( d ⇤ 2 , 1 , 4 , 2 , 1 ) ,600
8Implementation available at https://github.com/arturml/pytorch-wgan-gp
SGD + Momentum SRSGDSGD + esterov Momentum
Tr ai
n lo
ss
Iteration

heading: Test errors (%) on Permuted MNIST
text: Figure 7: Training loss vs. training iterations of LSTM trained with SGD (red), SGD + NM (green), and SRSGD (blue) for PMNIST classification tasks.

heading: E.3 WASSERSTEIN GENERATIVE ADVERSARIAL NETWORKS (WGAN) TRAINING ON MNIST
text: We investigate the advantage of SRSGD over SGD with standard and Nesterov momentum in training deep generative models. In our experiments, we train a WGAN with gradient penalty Gulrajani
7Implementation available at https://github.com/Lezcano/expRNN
25
Under review as a conference paper at ICLR 2021
et al. (2017) on MNIST. We evaluate our models using the discriminator’s loss, i.e. the Earth Moving distance estimate, since in WGAN lower discriminator loss and better sample quality are correlated Arjovsky et al. (2017).
Implementation and Training Details: The detailed implementations of our generator and discriminator are given below. For the generator, we set latent dim to 100 and d to 32. For the discriminator, we set d to 32. We train each model for 350 epochs with the initial learning rate of 0.01. The learning rate was reduced by a factor of 10 at epoch 200 and 300. The momentum is set to 0.9 for SGD with standard and Nesterov constant momentum. The restart schedule for SRSGD is set to 60, 120, 180. The restart schedule changes at epoch 200 and 300. In all experiments, we use batch size 64. Our code is based on the code from the Pytorch WGAN-GP Github.8
i m p o r t t o r c h i m p o r t t o r c h . nn as nn
c l a s s G e n e r a t o r ( nn . Module ) : d e f i n i t ( s e l f , l a t e n t d i m , d =32) :
s u p e r ( ) . i n i t ( ) s e l f . n e t = nn . S e q u e n t i a l (
nn . ConvTranspose2d ( l a t e n t d i m , d ∗ 8 , 4 , 1 , 0 ) , nn . BatchNorm2d ( d ∗ 8) , nn . ReLU( True ) ,
nn . ConvTranspose2d ( d ∗ 8 , d ∗ 4 , 4 , 2 , 1 ) , nn . BatchNorm2d ( d ∗ 4) , nn . ReLU( True ) ,
nn . ConvTranspose2d ( d ∗ 4 , d ∗ 2 , 4 , 2 , 1 ) , nn . BatchNorm2d ( d ∗ 2) , nn . ReLU( True ) ,
nn . ConvTranspose2d ( d ∗ 2 , 1 , 4 , 2 , 1 ) , nn . Tanh ( )
) d e f f o r w a r d ( s e l f , x ) :
r e t u r n s e l f . n e t ( x )
c l a s s D i s c r i m i n a t o r ( nn . Module ) : d e f i n i t ( s e l f , d =32) :
s u p e r ( ) . i n i t ( ) s e l f . n e t = nn . S e q u e n t i a l (
nn . Conv2d ( 1 , d , 4 , 2 , 1 ) , nn . Ins tanceNorm2d ( d ) , nn . LeakyReLU ( 0 . 2 ) ,
nn . Conv2d ( d , d ∗ 2 , 4 , 2 , 1 ) , nn . Ins tanceNorm2d ( d ∗ 2) , nn . LeakyReLU ( 0 . 2 ) ,
nn . Conv2d ( d ∗ 2 , d ∗ 4 , 4 , 2 , 1 ) , nn . Ins tanceNorm2d ( d ∗ 4) , nn . LeakyReLU ( 0 . 2 ) ,
nn . Conv2d ( d ∗ 4 , 1 , 4 , 1 , 0 ) , )
d e f f o r w a r d ( s e l f , x ) : o u t p u t s = s e l f . n e t ( x ) r e t u r n o u t p u t s . s q u e e z e ( )
Results: Our SRSGD is still better than both the baselines. SRSGD achieves smaller discriminator loss, i.e. Earth Moving distance estimate, and converges faster than SGD with standard and Nesterov constant momentum. We summarize our results in Table 12 and Figure 8. We also demonstrate the
8Implementation available at https://github.com/arturml/pytorch-wgan-gp
Under review as a conference paper at ICLR 2021
digits generated by the trained WGAN in Figure 9. By visually evaluation, we observe that samples generated by the WGAN trained with SRSGD look slightly better than those generated by the WGAN trained with SGD with standard and Nesterov constant momentum.
Table 12: Discriminator loss (i.e. Earth Moving distance estimate) of the WGAN with gradient penalty trained on MNIST with SGD, SGD + NM and SRSGD. In all experiments, we use the initial learning rate of 0.01, which is reduced by a factor of 10 at epoch 200 and 300. All models are trained for 350 epochs. The momentum for SGD and SGD + NM is set to 0.9. The restart schedule in SRSGD is set to 60, 120, and 180.
Bey nd DNNs: Traini g WGAN on MNIST
nn . Tanh ( )601 )602 d e f f o r w a r d ( s e l f , x ) :603 r e t u r n s e l f . n e t ( x )604 605 c l a s s D i s c r i m i n a t o r ( nn . Module ) :606 d e f i n i t ( s e l f , d =32) :607 s u p e r ( ) . i n i t ( )608 s e l f . n e t = nn . S e q u e n t i a l (609 nn . Conv2d ( 1 , d , 4 , 2 , 1 ) ,610 nn . Ins tanceNorm2d ( d ) ,611 nn . LeakyReLU ( 0 . 2 ) ,612 613 nn . Conv2d ( d , d ⇤ 2 , 4 , 2 , 1 ) ,614 nn . Ins tanceNorm2d ( d ⇤ 2) ,615 nn . LeakyReLU ( 0 . 2 ) ,616 617 nn . Conv2d ( d ⇤ 2 , d ⇤ 4 , 4 , 2 , 1 ) ,618 nn . Ins tanceNorm2d ( d ⇤ 4) ,619 nn . LeakyReLU ( 0 . 2 ) ,620 621
nn . Conv2d ( d ⇤ 4 , 1 , 4 , 1 , 0 ) ,622 )623 d e f f o r w a r d ( s e l f , x ) :624 o u t p u t s = s e l f . n e t ( x )625 r e t u r n o u t p u t s . s q u e e z e ( )626
Results: Our SRSGD is still better than both the baselines. SRSGD achieves smaller discriminator627 loss, i.e. Earth Moving distance estimate, and converges faster than SGD with standard and Nesterov628 constant momentum. We summarize our results in Table 11 and Figure 7. We also demonstrate the629 digits generated by the trained WGAN in Figure 8. By visually evaluation, we observe that samples630 generated by the WGAN trained with SRSGD look slightly better than those generated by the WGAN631 trained with SGD with standard and Nesterov constant momentum.
epochs reduction), 75 (15 epochs reduction), 70 (20 epochs reduction), 65 (25 epochs reduction), and643
60 (30 epochs reduction).644
F.2 Additional Experimental Results645 Figure 9 shows error rate vs. reduction in epochs for all models trained on CIFAR10 and ImageNet.646 It is a more complete version of Figure 4 in the main text. Table 13 and Table 14 provide detailed test647 errors vs. number of training epoch reduction reported in Figure 4 and Figure 9 . We also conduct an648
21
Figure 8: Earth Moving distance estimate (i.e. discriminator loss) vs. training epochs of WGAN with gradient penalty tra ned with SGD (red), SGD + NM (green), and SRSGD (blue) on MNIST.
SGD SGD + NM SRSGD
Figure 9: MNIST digits generated by WGAN trained with gradient penalty by SGD (left), SGD + NM (middle), and SRSGD (right).

heading: F ERROR RATE VS. REDUCTION IN TRAINING EPOCHS
text: F.1 IMPLEMENTATION DETAILS
CIFAR10 (Figure 4, left, in the main text) and CIFAR100 (Figure 11 in this Appendix): Except for learning rate schedule, we use the same setting described in Section D.1 above and Section 4.1 in the main text. Table 13 contains the learning rate schedule for each number of epoch reduction in Figure 4 (left) in the main text and Figure 11 below.
ImageNet (Figure 4, right, in the main text): Except for the total number of training epochs, other settings are similar to experiments for training ImageNet in fewer number of epochs described in Section D.3. In particular, the learning rate and restarting frequency schedule still follow those in Table 9 above. We examine different numbers of training epochs: 90 (0 epoch reduction), 80 (10 epochs reduction), 75 (15 epochs reduction), 70 (20 epochs reduction), 65 (25 epochs reduction), and 60 (30 epochs reduction).
Under review as a conference paper at ICLR 2021

heading: F.2 SHORT TRAINING ON CIFAR10/CIFAR100 USING SGD
text: For better comparison between SRSGD training using fewer epochs and SGD full training, we also conduct experiments with SGD training using fewer epochs on CIFAR10 and CIFAR100. Table 14 and 15 compares SRSGD short training using 100 epoch, SGD short training using 100 epochs, SRSGD full training using 200 epochs, and SGD full training using 200 epochs for Pre-ResNet-110, 290, and 470 on CIFAR10 and CIFAR100, respectively. The learning rate schedule for SGD short training using 100 epochs is the same as the learning rate schedule for SRSGD short training using 100 epoch given in Section 4 and in Table 13 above. In particular, for both SGD and SRSGD training using 100 epochs, we decrease the learning rate by a factor of 10 at the 80th, 90th, and 95th epoch. We observe that SGD short training has the worst performance compared to the others while SRSGD short training yields either comparable or even better results than SGD full training.

heading: F.3 ADDITIONAL EXPERIMENTAL RESULTS
text: Figure 10 shows error rate vs. reduction in epochs for all models trained on CIFAR10 and ImageNet. It is a more complete version of Figure 4 in the main text. Table 16 and Table 17 provide detailed test errors vs. number of training epoch reduction reported in Figure 4 and Figure 10 . We also conduct an additional ablation study of error rate vs. reduction in epochs for CIFAR100 and include the results in Figure 11 and Table 18 below.
Under review as a conference paper at ICLR 2021
G IMPACT OF RESTARTING FREQUENCY FOR IMAGENET AND CIFAR100
G.1 IMPLEMENTATION DETAILS
For the CIFAR10 experiments on Pre-ResNet-290 in Figure 5 in the main text, as well as the CIFAR100 and ImageNet experiments in Figure 14 and 15 in this Appendix, we vary the initial restarting frequency F1. Other settings are the same as described in Section D above.
Under review as a conference paper at ICLR 2021
CIFAR100
G.2 IMPACT OF THE GROWTH RATE r
We do an ablation study for the growth rate r to understand its impact on the behavior of SRSGD. We choose a case study of training a Pre-ResNet-110 on CIFAR10 using SRSGD with a linear schedule scheme for the restarting frequency. We fix the initial restarting frequency F1 = 30 and vary the growth rate r. We choose r from the set of {0.7, 1.0, 2.0, 10.0}. These values of r represent four different scenarios. When r = 0.7, the restarting frequency decreases every time the learning rate is reduced by a factor of 10. When r = 1.0, the restarting frequency stays constant during the training. When r = 2.0, the restarting frequency increases every time the learning rate is reduced by a factor of 10. Finally, when r = 10.0, it is similar to when r = 2.0, but the restarting frequency increases much faster and to larger values. Figure 12 summarizes the results of our ablation study. We observe that for CIFAR10, decreasing the restarting frequency or keeping it constant during training yields worse results than increasing the restarting frequency. However, increasing the restarting frequency too much also diminishes the performance of SRSGD.
Under review as a conference paper at ICLR 2021

heading: G.3 ADDITIONAL EXPERIMENTAL RESULTS
text: To complete our study on the impact of restarting frequency in Section 5.2 in the main text, we examine the case of CIFAR100 and ImageNet in this section. We summarize our results in Figure 14 and 15 below. Also, Figure 13 is a more detailed version of Figure 5 in the main text.
CIFAR10
CIFAR100
ImageNet
Under review as a conference paper at ICLR 2021

heading: H FULL TRAINING WITH LESS EPOCHS AT THE INTERMEDIATE LEARNING RATES
text: We explore SRSGD full training (200 epochs on CIFAR and 90 epochs on ImageNet) with less number of epochs at the intermediate learning rates and report the results in Table 19, 20, 21 and Figure 16, 17, 18 below. The settings and implementation details here are similar to those in Section F, but using all 200 epochs for CIFAR experiments and 90 epochs for ImageNet experiments.
Under review as a conference paper at ICLR 2021
CIFAR100
I VISUALIZATION OF SRSGD’S TRAJECTORY
Here we visualize the training trajectory through bad minima of SRSGD, SGD with constant momentum, and SGD. In particular, we train a neural net classifier on a swiss roll data as in Huang et al. (2019) and find bad minima along its training. Each red dot in Figure 19 represents the trained model after each 10 epochs in the training. From each red dot, we search for nearby bad local minima,
Under review as a conference paper at ICLR 2021
which are the blue dots. Those bad local minima achieve good training error but bad test error. We plots the trained models and bad local minima using PCA Wold et al. (1987) and t-SNE Maaten & Hinton (2008) embedding. The blue color bar is for the test accuracy of bad local minima; the red color bar is for the number of training epochs.

heading: PCA Embedding of the Trajectory
text: (CONTINUED NEXT PAGE)
Under review as a conference paper at ICLR 2021

heading: J SRSGD IMPLEMENTATION IN PYTORCH
text: i m p o r t t o r c h from . o p t i m i z e r i m p o r t Op t imize r , r e q u i r e d
c l a s s SRSGD( O p t i m i z e r ) : ””” Schedu led R e s t a r t SGD. Args :
params ( i t e r a b l e ) : i t e r a b l e o f p a r a m e t e r s t o o p t i m i z e o r d i c t s d e f i n i n g p a r a m e t e r g ro ups . l r ( f l o a t ) : l e a r n i n g r a t e . w e i g h t d e c a y ( f l o a t , o p t i o n a l ) : w e i gh t decay ( L2 p e n a l t y ) ( d e f a u l t : 0 ) i t e r c o u n t ( i n t e g e r ) : c o u n t t h e i t e r a t i o n s mod 200
Example : >>> o p t i m i z e r = t o r c h . opt im . SRSGD( model . p a r a m e t e r s ( ) , l r = 0 . 1 ,
w e i g h t d e c a y =5e−4, i t e r c o u n t =1) >>> o p t i m i z e r . z e r o g r a d ( ) >>> l o s s f n ( model ( i n p u t ) , t a r g e t ) . backward ( ) >>> o p t i m i z e r . s t e p ( ) >>> i t e r c o u n t = o p t i m i z e r . u p d a t e i t e r ( )
Formula : v { t +1} = p t − l r ∗ g t p { t +1} = v { t +1} + ( i t e r c o u n t ) / ( i t e r c o u n t +3) ∗ ( v { t +1} − v t ) ””” d e f i n i t ( s e l f , params , l r = r e q u i r e d , w e i g h t d e c a y = 0 . ,
i t e r c o u n t =1 , r e s t a r t i n g i t e r =100) : i f l r i s n o t r e q u i r e d and l r < 0 . 0 : r a i s e V a l u e E r r o r ( ” I n v a l i d l e a r n i n g r a t e : {} ” . f o r m a t ( l r ) ) i f w e i g h t d e c a y < 0 . 0 :
r a i s e V a l u e E r r o r ( ” I n v a l i d w e i g h t d e c a y v a l u e : {} ” . f o r m a t ( w e i g h t d e c a y ) )
i f i t e r c o u n t < 1 : r a i s e V a l u e E r r o r ( ” I n v a l i d i t e r c o u n t : {} ” . f o r m a t ( i t e r c o u n t ) ) i f r e s t a r t i n g i t e r < 1 : r a i s e V a l u e E r r o r ( ” I n v a l i d i t e r t o t a l : {} ” . f o r m a t (
r e s t a r t i n g i t e r ) )
d e f a u l t s = d i c t ( l r = l r , w e i g h t d e c a y = w e i g h t d e c a y , i t e r c o u n t = i t e r c o u n t , r e s t a r t i n g i t e r = r e s t a r t i n g i t e r ) s u p e r (SRSGD, s e l f ) . i n i t ( params , d e f a u l t s )
d e f s e t s t a t e ( s e l f , s t a t e ) : s u p e r (SRSGD, s e l f ) . s e t s t a t e ( s t a t e )
d e f u p d a t e i t e r ( s e l f ) : i d x = 1 f o r group i n s e l f . pa r am groups :
i f i d x == 1 : group [ ’ i t e r c o u n t ’ ] += 1 i f group [ ’ i t e r c o u n t ’ ] >= group [ ’ r e s t a r t i n g i t e r ’ ] : g roup [ ’ i t e r c o u n t ’ ] = 1 i d x += 1
r e t u r n group [ ’ i t e r c o u n t ’ ] , g roup [ ’ r e s t a r t i n g i t e r ’ ]
d e f s t e p ( s e l f , c l o s u r e =None ) : ””” Per fo rm a s i n g l e o p t i m i z a t i o n s t e p . Arguments : c l o s u r e ( c a l l a b l e , o p t i o n a l ) : A c l o s u r e t h a t
r e e v a l u a t e s t h e model and r e t u r n s t h e l o s s . ””” l o s s = None i f c l o s u r e i s n o t None :
Under review as a conference paper at ICLR 2021
l o s s = c l o s u r e ( )
f o r group i n s e l f . pa r am groups : w e i g h t d e c a y = group [ ’ w e i g h t d e c a y ’ ] momentum = ( group [ ’ i t e r c o u n t ’ ] − 1 . ) / ( group [ ’ i t e r c o u n t ’ ] +
2 . ) f o r p i n group [ ’ params ’ ] :
i f p . g r ad i s None : c o n t i n u e d p = p . g r ad . d a t a i f w e i g h t d e c a y ! = 0 :
d p . add ( w e i g h t d e c a y , p . d a t a )
p a r a m s t a t e = s e l f . s t a t e [ p ]
i f ’ momentum buffer ’ n o t i n p a r a m s t a t e : buf0 = p a r a m s t a t e [ ’ momentum buffer ’ ] = t o r c h . c l o n e ( p . d a t a ) . d e t a c h ( ) e l s e :
buf0 = p a r a m s t a t e [ ’ momentum buffer ’ ]
buf1 = p . d a t a − group [ ’ l r ’ ]∗ d p p . d a t a = buf1 + momentum∗ ( buf1 − buf0 ) p a r a m s t a t e [ ’ momentum buffer ’ ] = buf1
i t e r c o u n t , i t e r t o t a l = s e l f . u p d a t e i t e r ( )
r e t u r n l o s s

heading: K SRSGD IMPLEMENTATION IN KERAS
text: i m p o r t numpy as np i m p o r t t e n s o r f l o w as t f from k e r a s i m p o r t backend as K from k e r a s . o p t i m i z e r s i m p o r t O p t i m i z e r from k e r a s . l e g a c y i m p o r t i n t e r f a c e s i f K. backend ( ) == ’ t e n s o r f l o w ’ :
i m p o r t t e n s o r f l o w as t f
c l a s s SRSGD( O p t i m i z e r ) : ””” Schedu led R e s t a r t S t o c h a s t i c g r a d i e n t d e s c e n t o p t i m i z e r . I n c l u d e s s u p p o r t f o r N e s t e r o v momentum and l e a r n i n g r a t e decay . # Arguments
l e a r n i n g r a t e : f l o a t >= 0 . L e a r n i n g r a t e . ”””
d e f i n i t ( s e l f , l e a r n i n g r a t e = 0 . 0 1 , i t e r c o u n t =1 , r e s t a r t i n g i t e r =40 , ∗∗ kwargs ) : l e a r n i n g r a t e = kwargs . pop ( ’ l r ’ , l e a r n i n g r a t e ) s e l f . i n i t i a l d e c a y = kwargs . pop ( ’ decay ’ , 0 . 0 ) s u p e r (SRSGD, s e l f ) . i n i t (∗∗ kwargs ) wi th K. name scope ( s e l f . c l a s s . n a m e ) :
s e l f . i t e r a t i o n s = K. v a r i a b l e ( 0 , d t y p e = ’ i n t 6 4 ’ , name= ’ i t e r a t i o n s ’ ) s e l f . l e a r n i n g r a t e = K. v a r i a b l e ( l e a r n i n g r a t e , name= ’ l e a r n i n g r a t e ’ ) s e l f . decay = K. v a r i a b l e ( s e l f . i n i t i a l d e c a y , name= ’ decay ’ ) # f o r s r s g d s e l f . i t e r c o u n t = K. v a r i a b l e ( i t e r c o u n t , d t y p e = ’ i n t 6 4 ’ , name= ’ i t e r c o u n t ’ ) s e l f . r e s t a r t i n g i t e r = K. v a r i a b l e ( r e s t a r t i n g i t e r , d t y p e = ’
i n t 6 4 ’ , name= ’ r e s t a r t i n g i t e r ’ ) s e l f . n e s t e r o v = n e s t e r o v
Under review as a conference paper at ICLR 2021
@ i n t e r f a c e s . l e g a c y g e t u p d a t e s s u p p o r t @K. sy m b o l i c d e f g e t u p d a t e s ( s e l f , l o s s , params ) :
g r a d s = s e l f . g e t g r a d i e n t s ( l o s s , params ) s e l f . u p d a t e s = [K. u p d a t e a d d ( s e l f . i t e r a t i o n s , 1 ) ]
momentum = (K. c a s t ( s e l f . i t e r c o u n t , d t y p e =K. d t y p e ( s e l f . decay ) ) − 1 . ) / ( K. c a s t ( s e l f . i t e r c o u n t , d t y p e =K. d t y p e ( s e l f . decay ) ) + 2 . )
l r = s e l f . l e a r n i n g r a t e i f s e l f . i n i t i a l d e c a y > 0 :
l r = l r ∗ ( 1 . / ( 1 . + s e l f . decay ∗ K. c a s t ( s e l f . i t e r a t i o n s , K. d t y p e ( s e l f . decay )
) ) ) # momentum s h a p e s = [K. i n t s h a p e ( p ) f o r p i n params ]
moments = [K. v a r i a b l e ( v a l u e =K. g e t v a l u e ( p ) , d t y p e =K. d t y p e ( s e l f . decay ) , name= ’ moment ’ + s t r ( i ) )
f o r ( i , p ) i n enumera t e ( params ) ]
s e l f . w e i g h t s = [ s e l f . i t e r a t i o n s ] + moments + [ s e l f . i t e r c o u n t ] f o r p , g , m i n z i p ( params , g rads , moments ) :
v = p − l r ∗ g new p = v + momentum ∗ ( v − m) s e l f . u p d a t e s . append (K. u p d a t e (m, v ) )
# Apply c o n s t r a i n t s . i f g e t a t t r ( p , ’ c o n s t r a i n t ’ , None ) i s n o t None :
new p = p . c o n s t r a i n t ( new p )
s e l f . u p d a t e s . append (K. u p d a t e ( p , new p ) )
c o n d i t i o n = K. a l l (K. l e s s ( s e l f . i t e r c o u n t , s e l f . r e s t a r t i n g i t e r ) ) n e w i t e r c o u n t = K. s w i t c h ( c o n d i t i o n , s e l f . i t e r c o u n t + 1 , s e l f . i t e r c o u n t − s e l f . r e s t a r t i n g i t e r + 1) s e l f . u p d a t e s . append (K. u p d a t e ( s e l f . i t e r c o u n t , n e w i t e r c o u n t ) )
r e t u r n s e l f . u p d a t e s
d e f g e t c o n f i g ( s e l f ) : c o n f i g = { ’ l e a r n i n g r a t e ’ : f l o a t (K. g e t v a l u e ( s e l f . l e a r n i n g r a t e ) )
, ’ decay ’ : f l o a t (K. g e t v a l u e ( s e l f . decay ) ) , ’ i t e r c o u n t ’ : i n t (K. g e t v a l u e ( s e l f . i t e r c o u n t ) ) , ’ r e s t a r t i n g i t e r ’ : i n t (K. g e t v a l u e ( s e l f . r e s t a r t i n g i t e r
) ) } b a s e c o n f i g = s u p e r (SRSGD, s e l f ) . g e t c o n f i g ( ) r e t u r n d i c t ( l i s t ( b a s e c o n f i g . i t e m s ( ) ) + l i s t ( c o n f i g . i t e m s ( ) ) )

Review summaries:
This paper proposes to restart the momentum parameter in SGD (with Nesterov's momentum) according to some carefully chosen schedules in training deep neural network, which is named as SRSGD. Two different restarting schedules are proposed: linear schedule and exponential schedule. The strong point of this paper is its extensive experimental evaluations, which justify that SRSGD significantly improves the convergence speed and generalization over standard momentum SGD. The empirical analysis also sheds some light on the parameter tuning and interpretation of SRSGD.