S CHEDULED R ESTART M OMENTUM FOR A CCELERATED S TOCHASTIC G RADIENT D ESCENT
Anonymous authors
Paper under double-blind review
A BSTRACT
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 classiﬁcation, we demonstrate that, in training DNNs, SRSGD signiﬁcantly improves convergence and generalization; for instance, in training
ResNet-200 for ImageNet classiﬁcation, SRSGD achieves an error rate of 20.93% vs. the benchmark of 22.13%.
These improvements become more signiﬁcant as the network grows deeper.
Furthermore, on both CIFAR and ImageNet, SRSGD reaches similar or even better error rates with signiﬁcantly fewer training epochs compared to the SGD baseline.
I NTRODUCTION
Training many machine learning (ML) models reduces to solving the following ﬁnite-sum optimization problem
[Equation]
where { x i , y i } Ni =1 are the training samples and L is the loss function, e.g., cross-entropy loss for a classiﬁcation 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:
[Equation]
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 (cid:28) N , to approximate ∇ f ( w k ) by the mini-batch gradient 1 /m (cid:80) mj
=1 ∇ f i j ( 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:
[Equation]
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 classiﬁcation.
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 ﬁrst 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 beneﬁts of SRSGD are fourfold:
• SRSGD remarkably speeds up DNN training.
For image classiﬁcation, SRSGD signiﬁcantly 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 signiﬁcantly better than the current benchmark optimizers.
The improvement becomes more signiﬁcant as the network grows deeper as shown in Fig.
1. • SRSGD reduces overﬁtting in training very deep networks such as ResNet-200 for ImageNet classiﬁcation, 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 classiﬁcation 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
[Figure]
devoted to studying the non-acceleration issues of SGD with HB and NAG momentum (Kidambi et al.,
; Liu & Belkin, 2020), as well as accelerating ﬁrst-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 efﬁcacy of the proposed SRSGD in training DNNs for image classiﬁcation 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 = ( x 1 , · · · , x d ) ∈ R d , we denote its (cid:96) p norm ( p ≥ 1 ) by (cid:107) x (cid:107) p
=
( (cid:80) di =1 | x i | p ) 1 /p .
For a matrix A , we use (cid:107) A (cid:107) p to denote its induced norm by the vector (cid:96) p norm.
Given two sequences { a n } and { b n } , we write a n = O ( b n ) if there exists a positive constant s.t.
a n ≤ Cb n .
We denote the interval a to b (included) as ( a, b ] .
For a function f ( w ) : R d → R , we denote its gradient as ∇ f ( w ) and its Hessian as ∇ 2 f ( w ) .
R EVIEW : M OMENTUM IN G RADIENT D ESCENT
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., (cid:107)∇ 2 f ( w ) (cid:107) 2 ≤ L ), then GD converges with rate O (1 /k ) by letting s k ≡ 1 /L
(we use this s k 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
[Equation]
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 ( t k − 1) /t
k +1 , where t k +
=
(1 + (cid:112)
+ 4 t 2 k ) / 2 with t 0 = 1 .
NAG iterates as following
[Equation]
NAG achieves a convergence rate O (1 /k 2 ) with the step size s k = 1 /L .
Remark 1.
Su et al.
(2014) showed that ( k − 1) /
( k + 2) is the asymptotic limit of ( t k − 1) /t
k +1 .
In the following presentation of NAG with restart, for the ease of notation, we will replace the momentum coefﬁcient ( t k − 1) /t
k +1 with ( k − 1) /
( k + 2) .
Adaptive Restart NAG (ARNAG).
The sequences, { f ( w k ) − 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
(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
[Equation]
where m (1) = 1 ; m ( k + 1) = m ( k ) + 1 if f ( w k +1 ) ≤ f ( w k ) , and m ( k + 1) = 1 otherwise.
Scheduled Restart NAG (SRNAG).
SR is another strategy to restart the time-varying momentum of NAG.
We ﬁrst divide the total iterations (0 , T ] (integers only) into a few intervals { I i } mi =
=
( T i − 1 , T i ] , such that (0 , T ] =
(cid:83) mi
=
I i .
In each I i we restart the momentum after every F i iterations.
The update rule is then given by:
[Equation]
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)
[Equation]
where L ∈ R d × d is the Laplacian of a cycle graph, and b is a d -dimensional vector whose ﬁrst 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 = 1 K ( 1 K := 10 3 ).
We run T = 50 K iterations with step size 1 / 4 .
In SRNAG, we restart, i.e., we set the momentum to 0, after every 1 K 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 signiﬁcantly.
[Figure]
A LGORITHM P ROPOSED : S CHEDULED R ESTART SGD (SRSGD)
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 ﬁrst 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.
3. U NCONTROLLED B OUND OF N ESTEROV A CCELERATED SGD (NASGD)
Replacing ∇ f ( w k ) :
= 1 /N (cid:80)
Ni =1 ∇ f i ( w k ) in (5) with the mini-batch gradient 1 /m (cid:80) mj
=1 ∇ f i j ( w k ) 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 (cid:107)∇ f ( w ) (cid:107) ≤ R , where R > 0 is a constant.
The sequence { w k } k ≥ 0 generated by (5) , with stochastic gradient of bounded variance (Bubeck, 2014; Bottou et al., 2018) 1 and using any constant step size s k ≡ s ≤ 1 /L , satisﬁes
[Equation]
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 ( w k ) − 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) classiﬁcation.
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) conﬁrm 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.
3. SRSGD AND I TS C ONVERGENCE
For ERM (1), SRSGD replaces ∇ f ( w ) in (7) with stochastic gradient with batch size m and gives v
[Equation]
where F i is the restart frequency used in the interval
I i .
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 { w k } 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 s k :=
s ≤ 1 /L .
Assume that (cid:80) k ∈A (cid:0) E f ( v k )
− E f ( v k − 1 ) (cid:1) =
¯ R < + ∞ with ¯ R being a constant and the set
A := { k ∈ Z
+ | E f ( w k +1 ) ≥ E f ( w k ) } , then we have
[Equation]
If f ( w ) is further convex and (cid:80) k ∈B (cid:0) E f
( v k ) − E f
( v k − 1 ) (cid:1) =
ˆ R < + ∞ with ˆ R being a constant and the set B := { k ∈ Z + | E (cid:107) x k − x ∗ (cid:107) 2 ≥ E (cid:107) x k − 1 − x ∗ (cid:107) 2 } , then
[Equation]
where w ∗ is the minimum of f .
To obtain (cid:15) ( ∀ (cid:15) > 0 ) error, we set s
= O ( (cid:15) ) and
K
= O (1 /(cid:15) 2 ) .
Theorem 2 relies on the assumption that (cid:80) k ∈A (cid:0) E f ( v k )
− E f ( v k − 1 ) (cid:1) =
¯ R < + ∞ , and we leave it open for how to establish the convergence result for SRSGD without this assumption.
E XPERIMENTAL R ESULTS
We evaluate SRSGD on a variety of benchmarks for image classiﬁcation, 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 rate at certain epochs, which we denote as SGD.
We also compare SRSGD with the well-calibrated
[Figure]
SGD in which we switch momentum to the Nesterov momentum of 0.9, and we denote this optimizer as SGD + NM.
We ﬁne 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 inefﬁcient.
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 classiﬁcation 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 conﬁrm the stabilizing effect of scheduled restart in training DNNs.
4. CIFAR10 AND CIFAR
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 F 1 and the growth rate r .
In both scheduling schemes, the restarting frequency at the 1st learning rate stage is set to F 1 during training.
Then the restarting frequency at the ( k + 1) -th learning rate stage is determined by:
[Equation]
We search F 1 and r using the method outlined in Appendix D.4.
For CIFAR10, ( F 1 = 40 , r = 1 . 25) and ( F 1 = 30 , r = 2) are good initial restarting frequencies and growth rates for the exponential and linear schedules, respectively.
For CIFAR100, those values are ( F 1 = 45 , r = 1 . 5) for the exponential schedule and ( F 1 = 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 signiﬁcant 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 can avoid the error accumulation when there is an inexact oracle.
For CIFAR, Fig.
3 (b) shows that
[Table] [Table]
[Table] [Table]
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 conﬁrm 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 ( F 1 = 30 , r = 2) for CIFAR10 and ( F 1 = 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 classiﬁcation tasks.
In fact, for image classiﬁcation 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.
[Table] [Table]
4. I MAGE N ET
Next, we discuss our experimental results on the 1000-way ImageNet classiﬁcation task (Russakovsky et al., 2015).
We conduct our experiments on ResNet-50, 101, 152, and 200 with 5 different seeds.
We use the ofﬁcial 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 ofat 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 ( F 1 = 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 Overﬁtting in ResNet-200.
ResNet-200 demonstrates that SRSGD is better than the SGD baseline at avoiding overﬁtting 2 .
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 overﬁtting.
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 overﬁtting in very deep networks.
Theoretical analysis of the observation that SRSGD is less overﬁtting 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 ﬁrst 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.
4. E MPIRICAL A NALYSIS
SRSGD Helps Reduce the Training Time.
We ﬁnd 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 on the ﬁnal error rate when training with SRSGD on CIFAR10 and ImageNet.
In the CIFAR 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.
[Figure]
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 ﬁx the growth rate r = 2 and vary the initial restarting frequency F 1 from 1 to 80.
As shown in Fig.
5, SRSGD with a large F 1 , e.g. F 1 = 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 F 1 , e.g. F 1 = 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.
[Figure]
C ONCLUSIONS
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 classiﬁcation, SRSGD can signiﬁcantly 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).
Appendices
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.
Table of Contents Appendix
A Uncontrolled Bound of NASGD
A.1 Preliminaries
A.2 Uncontrolled Bound of NASGD:
Analysis
B NAG with δ -Inexact Oracle & Experimental Settings in Section 3.
C Convergence of SRSGD
D Datasets and Implementation Details
D.1 CIFAR .
D.2 ImageNet . .
D.3 Training ImageNet in Fewer Number of Epochs: . . . . . . . . . . .
Other
r . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A U NCONTROLLED B OUND OF NASGD
Consider the following optimization problem
[Equation]
where f ( w ) is L -smooth and convex.
Start from w k , GD update, with step size 1 r , can be obtained based on the minimization of the function
[Equation]
With direct computation, we can get that
[Equation]
where g k := 1 m (cid:80) m j =1 ∇ f i j ( w k ) .
We assume the variance is bounded, which gives The stochastic gradient rule, R s , satisﬁes E [ Q r ( v k +1 , w k ) − min Q r ( v , w k ) | χ k ] ≤ δ , with δ being a constant and χ k being the sigma algebra generated by w 1 , w 2 , · · · , w k , i.e.,
[Equation]
NASGD can be reformulated as
[Equation] [Equation]
where t 0 = 1 and t k +1 =
(1 + (cid:112)
+ 4 t 2 k ) / 2 .
A. P RELIMINARIES
To proceed, we introduce several deﬁnitions 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
[Equation]
and we say f is ν -strongly convex if for any w , v ∈ dom( f )
[Equation]
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
[Equation]
[Equation]
where v ∗ is the minimizer of f .
Lemma 2.
If f is L -smooth, for any w , v ∈ dom( f ) ,
[Equation]
A. U NCONTROLLED B OUND OF NASGD: A NALYSIS
In this part, we denote
[Equation]
Lemma 3.
If the constant r > 0 , then
[Equation]
Proof.
Note that Q r ( v , w k ) is strongly convex with constant r , and ˜ v k +1 in (17) is the minimizer of
Q r ( v , w k ) .
With Lemma 1 we have
[Equation] [Equation]
Notice that
[Equation]
The inequality (18) can be established by combining the above two inequalities.
Lemma 4.
If the constant satisfy r > L , then we have
[Equation]
where τ =
L 2 r ( r − L ) + 1 .
Proof.
The convexity of f gives us
[Equation]
From the deﬁnition of the stochastic gradient rule R s , we have
[Equation]
With (21) and (22), we have
[Equation]
With the Schwarz inequality (cid:104) a , b
[Equation]
≤ (cid:107) a
[Equation]
22 2 µ + µ 2 (cid:107) b (cid:107) 22 with µ =
L 2 r − L , a = ∇ f ( v k +1 )
− ∇ f
(˜ v k +1 ) and b = w k − ˜ v k +1 , (cid:104)∇
[Equation]
[Equation]
Combining (23) and (24), we have
[Equation]
By rearrangement of the above inequality (25) and using Lemma 3, we obtain the result
[Equation]
Lemma 5.
If the constants satisfy r > L , then we have the following bounds
[Equation]
where τ :=
L 2 r ( r − L ) +
and v ∗ is the minimum .
Proof.
With Lemma
[Equation]
we have
[Equation]
Using the convexity of f , we have
[Equation]
According to the deﬁnition of ˜ v k +1 in (14), i
[Equation]
e.,
[Equation]
and the optimization condition gives
[Equation]
Substituting (30) into (29), we obtain
[Equation]
Direct summation of (28) and (31) gives
[Equation]
[Equation]
Summing (32) and (20), we obtain the inequality (26)
[Equation]
[Equation]
On the other hand
[Equation]
with the convexity of f
[Equation]
we have
[Equation]
The summation of (28) and (34) results in
[Equation]
Summing (35) and (20), we obtain
[Equation]
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 { v k } k ≥ 0 be generated by NASGD with stochastic gradient that has bounded variance.
By using any constant step size s k ≡ s ≤ 1 /L , then we have
[Equation]
Proof.
We denote
F
[Equation]
By (26) × ( t k − 1) + (27), we have
[Equation]
With t 2 k − 1 = t
[Equation] [Equation]
38) × t
[Equation]
yields
[Equation]
Substituting a = t k v k +1 − ( t k − 1) v k − v ∗ and b = t k w k − ( t k − 1) v k − v ∗ into identity
[Equation] [Equation]
It follows that
[Equation]
In the third identity
[Equation]
we used the fact t k w k = t k v k + ( t k − 1 − 1)( v k − v k − 1 ) .
If we denote u k = E
[Equation]
t k
[Equation]
1 v k −
[Equation]
t k
[Equation]
1 − 1) v k − 1 − v ∗ (cid:107) 22 , (39) can be rewritten as
[Equation]
where we used
[Equation]
Denoting
[Equation]
then, we have
[Equation]
With the fact, ξ k ≥ 2 t 2 k − 1 F k r ≥ k 2 F k / 4 , we then proved the result.
B NAG WITH δ -I NEXACT O RACLE & E XPERIMENTAL S ETTINGS IN S ECTION 3.
In Devolder et al.
(2014), the authors deﬁnes δ -inexact gradient oracle for convex smooth optimization as follows:
Deﬁnition 1 ( δ -Inexact Oracle) .
Devolder et al.
For a convex L -smooth function f : R d → R .
For ∀ w ∈ R d and exact ﬁrst-order oracle returns a pair ( f ( w ) , ∇ f ( w ))
∈ R × R d so that for ∀ v ∈ R
d
we have
[Equation]
A δ -inexact oracle returns a pair (cid:0) f δ ( w ) , ∇ f δ ( w ) (cid:1)
∈
R × R d so that ∀ v ∈ R
d
we have
[Equation]
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 ∈ R d , 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 w k after k steps so that
[Equation]
On the other hand, NAG, with step size 1 /L returns
[Equation]
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 ﬁrst-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 .
We run 50 K iterations of different algorithms.
For SRNAG, we restart after every 200 iterations.
Fig.
2 (b) shows the iteration vs. optimal gap, f ( x k ) − 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 . 1 (cid:98) t/ 100 (cid:99) +1 ) 2 ) .
For SRNAG, we restart every 200 iterations in the ﬁrst 10 k iterations, and restart every 400 iterations in the remaining 40 K 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 signiﬁcantly outperforms all the other schemes.
Logisitic Regression for MNIST Classiﬁcation: We apply the above schemes with stochastic gradient to train a logistic regression model for MNIST classiﬁcation LeCun & Cortes (2010).
We consider ﬁve 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 (cid:96) 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 .
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 ﬁve schemes.
C C ONVERGENCE OF SRSGD
We prove the convergence of Nesterov accelerated SGD with scheduled restart, i.e., the convergence of SRSGD.
We denote that θ k :=
t k − 1 t k +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 −(cid:98) k/F (cid:99)∗ F .
In the restart version, we can see that
[Equation]
< 1 .
Lemma 6.
Let the constant satisﬁes r > L and the sequence { v k } k ≥ 0 be generated by the SRSGD with restart frequency F (any positive integer), we have
[Equation]
where ¯ θ := θ F < 1 and R :=
sup x {(cid:107)∇ f ( x ) (cid:107) 2 } .
Proof.
It holds that
[Equation]
Thus,
[Equation]
Summing (46) from k = 1 to K , we get
[Equation]
In the following
[Equation]
we denote
[Equation]
Theorem 5 (Convergence of SRSGD) .
(Theorem 2 with detailed bound) Suppose f ( w ) is L -smooth.
Consider the sequence { w k } 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 s k :
= s ≤ 1 /L .
Assume that (cid:80) k ∈A (cid:0) E f ( v k )
− E f ( v k − 1 ) (cid:1) =
¯ R < + ∞ , then we have
[Equation]
If f ( w ) is further convex and the set B := { k ∈
[Equation]
obeys (cid:80) k ∈B (cid:0) E f ( v k ) − E f ( v k − 1 ) (cid:1)
=
ˆ R < + ∞ , then
[Equation]
where w ∗ is the minimum of f .
To obtain (cid:15) ( ∀ (cid:15) > 0 ) error, we set s
= O ( (cid:15) ) and
K
= O (1 /(cid:15) 2 ) .
Proof.
Firstly, we show the convergence of SRSGD for nonconvex optimization.
L -smoothness of f , i.e., Lipschitz gradient continuity, gives us
[Equation]
[Equation]
Taking expectation, we get
[Equation]
On the other hand
[Equation]
we have
[Equation]
Then, we have
[Equation]
We also have
[Equation]
We then get that
[Equation]
where
[Equation]
Summing the inequality gives us
[Equation]
It is easy to see that
[Equation]
We get the result by using Lemma
Secondly
[Equation]
we prove the convergence of SRSGD for convex optimization.
Let w ∗ be the minimizer of f .
We have
[Equation]
We can also derive
[Equation]
where we used the following identity
[Equation]
Then, we have
[Equation]
We then get that
[Equation]
Summing the inequality gives us the desired convergence result for convex optimization.
D D ATASETS AND I MPLEMENTATION D ETAILS D. CIFAR
The CIFAR10 and CIFAR100 datasets Krizhevsky et al.
(2009) consist of 50 K training images and 10 K 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 ﬂipping are applied to training data.
Our code is modiﬁed based on the Pytorch classiﬁcation project Yang (2017), 4 which was also used by Liu et al.
Liu et al.
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 F i as the restarting frequency at the i -th learning rate.
[Table]
D. I MAGE N ET
The ImageNet dataset contains roughly 1.28 million training color images and 50 K validation color images from 1000 classes Russakovsky et al.
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 ﬂipping are applied to training data.
We use the ofﬁcial Pytorch ResNet implementation Paszke et al.
(2019), 5 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
F 1 = 40 at the 1st learning rate, F 2 = 80 at the 2nd learning rate, and F 3 is linearly decayed fromto 1 at the 3rd learning rate (see Table 8).
[Table]
D. T RAINING I MAGE N ET IN F EWER N UMBER OF E POCHS :
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.
[Table] [Table]
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.
D. D ETAILS ON R ESTARTING H YPER - PARAMETERS S EARCH
In our CIFAR10 and CIFAR100 experiments
[Table]
for both linear and exponential schedule
[Table]
we conduct hyperparameter searches over the restarting frequencies using our smallest model, Pre-ResNet-110, making choices based on ﬁnal validation performance.
The same chosen restarting frequencies are applied for all models including Pre-ResNet-110, 290, 470,
[Table]
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 ﬁrst train Pre-ResNeton
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 F 1 and the growth rate r .
Both F 1 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 F 1 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 F 1 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
F 1 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 F 1 = 40 and r = 2 which have the highest ﬁnal validation accuracy averaged over 5 random seeds.
Same as in CIFAR10 and CIFAR100 experiments, we select F 1 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 F 1 and r for our short training experiments in Section 4.3.
However, for ResNet-50, we observe that F 1 = 60 and r = 1 .
75 yields better performance in short training.
All reported results are averaged over 5 random seeds.
E SRSGD VS . SGD AND SGD + NM ON I MAGE N ET C LASSIFICATION AND O THER T ASKS E. C OMPARING WITH SGD WITH N ESTEROV M OMENTUM ON I MAGE N ET C LASSIFICATION
In this section, we compare SRSGD with SGD with Nesterov constant momentum (SGD + NM) in training ResNets for ImageNet classiﬁcation.
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 classiﬁcation, and as the network goes deeper the improvement becomes more signiﬁcant.
[Table]
In this task
[Table]
we examine the advantage of SRSGD over SGD and SGD with Nesterov Momentum in training recurrent neural networks.
In our experiments
[Table]
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.
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 ﬁxed permutation is applied to training and testing sequences.
This task is known as permuted MNIST classiﬁcation, which has become standard to measure the performance of RNNs and their ability to capture long term dependencies.
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 epochand 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.
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 6. is set to 90, 30, and 90.
[Table]
E. W ASSERSTEIN G ENERATIVE A DVERSARIAL N ETWORKS (WGAN) T RAINING ON MNIST
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 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.
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
Beyond DNNs: Training LSTM on
Network No
Hidden Units
[Figure]
[Table]
[Figure] [Figure]
[Table] [Table]
[Figure]
[Table]
[Figure]
We evaluate our models using the discriminator’s loss, i.e. the Earth Moving distance estimate, sincein WGAN lower discriminator loss and better sample quality are correlated [2].
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 learningrate was reduced by a factor of 10 at epoch 200 and 300.
The momentum is set to 0.9 for SGD withstandard 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.
Ourcode is based on the code from the Pytorch WGAN-GP Github.
[Figure]
nn
[Figure]
nn
[Figure]
( nn . Module ) :
[Figure]
i n i t ( self
[Figure]
dim
[Figure]
d
[Figure] [Figure]
i n i t
[Figure]
s e l f .
[Figure]
nn .
[Figure] [Figure]
[Figure] [Figure]
[Figure]
Implementation available at https://github.com/arturml/pytorch-wgan-gp
[Figure]
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.
[Table]
nn as nn
[Table]
a
[Table]
( nn . Module ) :
def i n i t
[Table]
i n i t (
[Table]
: r e
[Table]
[Table] [Table]
( nn .
[Table]
) : def i n i t ( self ,
[Table]
=32) : super (
[Table]
i n i
[Table]
( ) s
[Table] [Table]
nn .
[Table]
( nn .
[Table] [Table]
, nn .
[Table]
, nn
[Table] [Table]
[Table]
∗
[Equation]
[Table]
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 7.
We also demonstrate the digits generated by the trained WGAN in Figure 8.
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]
Beyond DNNs:
Training WGAN on
MNIST
nn .
Tanh ( ))def forward ( s e l f , x ) :r e t u r n s e l f .
net ( x ) 604c l a s s D
i s c
r
i m i n a t
o r ( nn . Module ) :
def i n i t ( s e l f , d=32) :super ( ) .
i n i t ( )s e l f .
net = nn .
S e q u e n t i a l (nn . Conv2d
(1 , d , 4 , 2 , 1) ,nn .
InstanceNorm2d ( d ) ,nn . LeakyReLU
( 0 . 2 ) , 612nn .
Conv2d
( d , d ⇤ 2 , 4 , 2 ,
nn .
Conv2d
( d ⇤ 2 , d ⇤ 4 , 4 , 2 , 1) ,nn .
InstanceNorm2d ( d ⇤ 4) ,nn . LeakyReLU
( 0 . 2 ) , nn . Conv2d
( d ⇤ 4 , 1 , 4 , 1 , 0) ,)def forward ( s e l f , x ) :outputs =
s e l f .
net ( x )r e t u r n outputs .
squeeze ( )
Results:
Our SRSGD is still better than both the baselines
SRSGD achieves smaller discriminator
loss
Earth distance estimate, and converges faster than SGD with standard and
Nesterov constant momentum.
We summarize our results in Table 11 and Figure
We also demonstrate the digits by the trained WGAN in Figure
By visually evaluation, we observe that samples slightly better
Nesterov
[Table]
F Error Rate vs. Reduction in Training EpochsF. Implementation Details
CIFAR10 (Figure 4, left, in the main text) and CIFAR100 (Figure 10 in this Appendix): Exceptfor learning rate schedule, we use the same setting described in Section D.1 above and Section 4.1 inthe main text.
Table 12 contains the learning rate schedule for each number of epoch reduction in
Figure 4 (left) in the main text and Figure 10 below.
ImageNet (Figure 4, right, in the main text): Except for the total number of training epochs, othersettings 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 8 above.
We examine different numbers of training epochs: 90 (0 epoch reduction), 80 (10epochs reduction), 75 (15 epochs reduction), 70 (20 epochs reduction), 65 (25 epochs reduction), and60 (30 epochs reduction).
F. Additional Experimental Results
Figure 9 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 13 and Table 14 provide detailed testerrors vs. number of training epoch reduction reported in Figure 4 and Figure 9 .
We also conduct an
[Figure]
[Equation]
F E RROR R ATE VS . R EDUCTION IN T RAINING E POCHS F. I MPLEMENTATION D ETAILS
CIFAR10 (Figure 4, left, in the main text) and CIFAR100 (Figure 10 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 10 below.
[Table]
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 ( epochs reduction), 75 (15 epochs reduction), 70 (20 epochs reduction), 65 (25 epochs reduction)
[Table]
and 60 (30 epochs reduction).
F.
S HORT T RAINING ON CIFAR10/CIFAR100 U SING SGD
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.
Tableand 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.
F. A DDITIONAL E XPERIMENTAL R ESULTS
Figure 9 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 9 .
We also conduct an additional ablation study of error rate vs. reduction in epochs for CIFAR100 and include the results in Figure 10 and Table 18 below.
G I MPACT OF R ESTARTING F REQUENCY FOR I MAGE N ET AND CIFAR G. I MPLEMENTATION D ETAILS
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 13 and 14 in this Appendix, we vary the initial restarting frequency F 1 .
Other settings are the same as described in Section D above.
[Table] [Table]
Number of Epoch Reduction
Pre-ResNetPre-ResNetPre-ResNet
[Figure]
r
[Figure]
n
[Figure]
ResNet
[Figure]
ResNet
[Figure] [Figure]
[Table]
G. I MPACT OF THE G ROWTH R ATE 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 ﬁx the initial restarting frequency F 1 = 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.
[Table]
[Figure]
[Table]
[Figure]
[Table]
CIFAR
[Figure]
[Table]
When r = 2 .
[Table]
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 .
[Table]
but the restarting frequency increases much faster and to larger values.
Figure 11 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.
G. A DDITIONAL E XPERIMENTAL R ESULTS
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 Figureand 14 below.
Also, Figure 12 is a more detailed version of Figure 5 in the main text.
H F ULL T RAINING WITH L ESS E POCHS AT THE I NTERMEDIATE L EARNING R ATES
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
T r a i n L o ss
T e s t
E rr o r
[Figure]
r=0.
r=1.
r=2.
r=10.
[Figure]
[Figure] [Figure]
CIFAR
[Figure]
are shown in red and conﬁrm the result of Theorem 1 that NASGD accumulates error due to the stochastic gradients.
Figure 15, 16, 17 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.
[Figure]
ImageNet
T e s t
E rr o r
Number of Epoch Reduction at the 2 nd and 3 rd learning rate
[Figure] [Figure]
[Table]
I V ISUALIZATION 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 classiﬁer on a swiss roll data as in Huang et al.
(2019) and ﬁnd bad minima along its training.
Each red dot in Figure 18 represents the trained model after each 10 epochs in the training.
From each red dot, we search for nearby bad local minima, 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.
and t-SNE Maaten &
[Figure]
Number
[Figure]
Epoch Reduction at the 2 nd and 3 rd learning
[Figure]
Pre-ResNetPre-ResNetPre-ResNetPre-ResNetPre-ResNet-
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
Number of Epoch Reduction at the 2 nd learning rate
S i n g l e C r o p
V a li d a t i o n E rr o r
ResNetResNet- ResNetResNet-
ImageNet
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.
(CONTINUED NEXT PAGE)
[Table]
[Figure]
J SRSGD I MPLEMENTATION IN P YTORCH
import torch from .
optimizer import Optimizer , required c l a s s SRSGD( Optimizer ) : ”””
Scheduled 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 of parameters to optimize or d i c t s defining parameter groups .
l r ( f l o a t ) :
l e a r n i n g r a t e .
weight decay ( float , o p t
i
o
n a l ) : weight decay (L2 penalty )
( d e f a u l t : 0) i t e r c o u n t ( i n t e g e r ) : count the
i t e r a t i o n s mod 200 Example : >>> optimizer = torch .
optim .SRSGD( model .
parameters ( ) , l r =0.1 , weight decay =5e − 4, i t e r c o u n t =1) >>> optimizer .
zero grad ( ) >>> l o s s f n ( model ( input ) , t a r g e t ) .
backward ( )
[Table]
s ) def s e t s t a t e ( self , s t a t e ) : super (SRSGD, s e l f ) .
s e t s t a t e ( s t a t e ) def u p d a t e i t e r ( s e l f ) : idx =for group in s e l f .
param groups : i f idx == 1: group
[ ’ i t e r c o u n t ’ ]
+=
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 ’ ] : group [ ’ i t
e r c o u n t ’ ] = 1 idx +=r e t u r n 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 ’ ]
def step ( self , closure =None ) : ”””Perform a s i n g l e optimization step .
Arguments : closure ( c a l l a b l e , o p t i o n a l ) : A closure t h a t r e e v a l u a t
e s the model and r e t u r n s the l o s s .
”””loss =
None i f closure i s not None : l o s s = closure ( ) for group in s e l f .
param groups : weight decay = group [ ’ weight decay ’ ] momentum =
( group [ ’ i t
e r c
o u n t ’ ] − 1 . )
/ ( group [ ’ i t e r c
o u n t ’ ] + 2 . )
for p in group [ ’ params ’ ]
: i f p .
grad i s
None : continue d p = p .
grad .
data i f weight decay !
=
: d p . add ( weight decay , p . data ) param state =
s e l f .
s t a t e [ p ] i f ’ momentum buffer ’ not in param state : buf0 = param state [ ’ momentum buffer ’ ]
= torch .
clone ( p . data ) .
detach ( ) e l s e : buf
[Table]
K SRSGD I MPLEMENTATION IN K ERAS
[Table]
0 , dtype= ’ int64 ’ , 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 ’ ) # for srsgd 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 , dtype= ’ int64 ’ , 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 , dtype= ’ int64 ’ , name= ’ r e s t a r t i n g i t
e r ’ )
s e l f .
nesterov = nesterov @interfaces .
l e g a c
y g e t u p d a t
e s s u p p
o r t @K. symbolic def get updates ( self , loss , params ) :
grads =
s e l f .
g e t g r a d i e n t s ( loss , params )
s e l f . updates =
[K. update add ( 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 , dtype=K. dtype ( s e l f . decay ) ) − 1 . )
/ (K. c a s t ( s e l f .
i t e r c o u n t , dtype=K. dtype ( 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
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Table]