Amortized Nesterov's Momentum: Robust and Lightweight  Momentum for Deep Learning
Authors: Kaiwen Zhou
Yanghua Jin
Qinghua Ding
James Cheng
Venue:  (2020)
Keywords: momentum
nesterov
optimization
deep learning
neural networks

Abstract: Stochastic Gradient Descent (SGD) with Nesterov's momentum is a widely used optimizer in deep learning, which is observed to have excellent generalization performance. However, due to the large stochasticity, SGD with Nesterov's momentum is not robust, i.e., its performance may deviate significantly from the expectation. In this work, we propose Amortized Nesterov's Momentum, a special variant of Nesterov's momentum which has more robust iterates, faster convergence in the early stage and higher efficiency. Our experimental results show that this new momentum achieves similar (sometimes better) generalization performance with little-to-no tuning. In the convex case, we provide optimal convergence rates for our new methods and discuss how the theorems explain the empirical results. 

Under review as a conference paper at ICLR 2020
AMORTIZED NESTEROV'S s MOMENTUM: ROBUST AND LIGHTWEIGHT MOMENTUM FOR DEEP LEARNING
Anonymous authors
Paper under double-blind review
ABSTRACT
Stochastic Gradient Descent (SGD) with Nesterov's momentum is a widely used optimizer in deep learning, which is observed to have excellent generalization per- formance. In this work, we propose Amortized Nesterov's 's Momentum, which is a special variant of Nesterov's momentum. Compared with Nesterov's momentum. our new momentum has more robust iterates and higher efficiency. Our empir- ical results show that it achieves faster early convergence and comparable final generalization performance with little-to-no tuning. Just like Nesterov's 's method, the new schemes are also proved optimal in general convex setting. Our analysis sheds light on the understanding of the new variant.
INTRODUCTION
In recent years, Gradient Descent (GD) (Cauchy, 1847) and its variants have been widely used to solve large scale machine learning problems. Among them, Stochastic Gradient Descent (SGD) (Robbins & Monro, 1951), which replaces gradient with an unbiased stochastic gradient estimator, is popular choice of optimizer especially for neural network training which requires lower precision. Sutskever et al. (2013) found that using SGD with Nesterov's 's momentum (Nesterov, 1983; 2013b), which was originally designed to accelerate deterministic convex optimization, achieves substantial speedups for training neural networks. This finding essentially turns SGD with Nesterov's 's momen- tum into the benchmarking method of neural network design, especially for classification tasks (He et al., 2016b;a; Zagoruyko & Komodakis, 2016; Huang et al., 2017). It is observed that in these tasks, the momentum technique plays key role in achieving good generalization performance.
Adaptive methods (Duchi et al., 2011; Kingma & Ba, 2015; Tieleman & Hinton, 2012: Reddi et al., 2018), which are also becoming increasingly popular the deep learning community, diagonally scale the gradient to speed up training. However, Wilson et al. (2017) show that these methods always generalize poorly compared with SGD with momentum (both classical momentum (Polyak, 1964) and Nesterov's momentum).
In this work, we introduce Amortized Nesterov's Momentum, which is a special variant of Nes- terov's momentum. From users' perspective, the new momentum has only one additional integer hyper- hyper-parameter m to choose, which we call the amortization length. Learning rate and momentum parameter of this variant strictly aligned with Nesterov's 's momentum and by choosing it recovers Nesterov's 's momentum. This paper conducts an extensive study based on both empirical evaluation and convex analysis to identify the benefits of the new variant (or from users' angle, to set m apart from 1). We list the advantages of Amortized Nesterov's 's Momentum as follows:
1,
Increasing m improves robustness' This an interesting property since the new momentum not only provides acceleration, but also enhances the robustness. We provide an understanding of this property by analyzing the relation between convergence rate and m in the convex setting.
Increasing m reduces (amortized) iteration complexity
A suitably chosen boosts the convergence rate in the early stage of training and produces comparable final generalization performance.
In this work. robusiness refers to the probability of an optimizer significantly deviating from its expected performance, which can be reflected by the deviations of accuracy or loss in the training process over multiple runs that start with the same initial guess.

---

Under review as a conference paper at ICLR 2020
It is easy to tune The performances of the methods stable for wide range of m and prove that the methods converge for any valid choice of m in the convex setting.
If m is not too large, the methods obtain the optimal convergence rate in general convex setting, just like Nesterov's s method.
The new variant does have some minor drawbacks: it requires one more memory buffer, which is acceptable in most cases, and it shows some undesired behaviors when working with learning rate schedulers, which can be addressed by a small modification. Considering these pros and cons, we believe that the proposed variant can benefit many large-scale deep learning tasks.
Our high level idea is simple: the stochastic Nesterov's momentum can be unreliable since is pro- vided only by the previous stochastic iterate. The iterate potentially has large variance, which may lead to a false momentum that perturbs the training process. We thus propose to use the stochastic Nesterov's 's momentum based on several past iterates, which provides robust acceleration. In other words, instead of immediately using an iterate to provide momentum, we put the iterate into an "amortization plan" and use it later.
PRELIMINARIES: SGD AND NESTEROV'S MOMENTUM
We start with a review of SGD and Nesterov's 's momentum. We discuss some subtlcties in the imple- mentation and evaluation, which contributes to the interpretation of our methods.
Notations In this paper, we use x E Rd denote the vector of model parameters. ||-I1 and (:+ denote the standard Euclidean norm and inner product, respectively. Scalar multiplication for o E denoted as 3.u. f TR.d - R denotes the loss function to be minimized and (xx)) represents the gradient of f evaluated at We denote the unbiased stochastic gradient estimator of V fz) as V fz(z) with the random variable i independent x (c.g.. using mini-batch). We use
Rd and
To Rdd to denote the initial guess.
SGD
SGD has the following simple iterative scheme, where 7 E R denotes the learning rate: Ik+1 = 2b - 7 fig (xx), for k 2 O.
Nesterov's momentum The original Nesterov's accelerated gradient (with constant step) (Nes- terov, 1983; 2013b) has the following scheme2 (y E Rdd n. IR and you zo):
yk+1 = 77 V (xx))
Æa+1 = yk++1 + 3 (yk+1 3k), for t: 2 O,
where call B (yk++1 yk) the momentum. By simply replacing fx( with Vi, (xx), obtain the SGD with Nesterov's momentum, which widely used in deep learning. To make this point clear, recall that the reformulation in Sutskever et al. (2013) (scheme (2), also the Tensorflow (Abadi et al., 2016) version) and the PyTorch (Paszke et al., 2017) version (scheme (3)) have the following schemes (v,, upt E Rd and Up = "o = O): for A 2 0.
Uk+1 = B Uk 7) Vi, (yk + B vx),
r = 3. + V pt
(2)
(3)
|o++1 = yk + UK+1-
#k+1 = 2x - n (B K+1 + V (#x)).
Here the notations are modified based on their equivalence scheme (1). It can be verified that schemes (2) and (3) are equivalent to (1) through 24 8 -ar 11xxaaa yk) and v = 17 2-----1 2 13-1.(yk (yk zk), respectively (see Defazio (2018) for other equivalent forms of scheme (1)).
Interestingly, both PyTorch and Tensorflow track the values { which we refer to as M-SGD. This choice allows a consistent implementation when wrapped in generic optimization layer (De- fazio, 2018). However, the accelerated convergence rate (in the convex case) is built upon {yk) (Nesterov, 2013b) and { may not possess such theoretical improvement. We use OM-SGD to refer to the Original M-SGD that outputs }.
We exchange the notations of x and in Nesterov (2013b).
TTnsooffow tracks the values (ye- 3- vx) {xx)-
2

---

Under review as a conference paper at ICLR 2020
Accuracy
Accuracy
95
3.0
1.0
SGD
SGD (batch)
2.5
0.8
(batch)
90-
2.0
K-5GD (average)
Train Logg
0.6
M-SGD (ful)
655
0.4
0.5
OM SGD
75
0.0
0.0
20
40
60
80
20
80
0
20
60
80
Epoch
Epach
Spoch
(a)
(b)
(c)
Figure 1: ResNet34 on CIFAR-10. Initial learning rate Yo = 0.1, momentum = 0.9, run 5 seeds (same Tn). In (a) (c), plot mean curves with shaded bands indicating #1 standard deviation. (b) shows the standard deviation of test accuracy and its average over 90 epochs. Best viewed in color.
SGD and M-SGD In order to study the features of momentum, in this work, we regard momentum as an add-on to plain SGD, which corresponds to fixing the learning rates* ? = 7). From the interpretation in Allen-Zhu & Orecchia (2017), 7 represents the learning rate for the gradient descent "inside" Nesterov's s method. To introduce the evaluation metries of this paper, we report the results of training ResNet34 (He et al., 2016b) CIFAR-10 10 (Krizhevsky et al., 2009) (our basic case study) using SGD and M-SGD in Figure 1. In this paper, all the multiple runs start with the same initial guess Tn- Figure Ia shows that Nesterov's 's momentum hurts the convergence in the first 60 epochs but accelerates the final convergence, which verifies the importance of momentum for achieving high accuracy. Figure lb depicts the robustness of M-SGD and SGD, which suggests that adding Nesterov's momentum slightly increases the uncertainty in the training process of SGD.
Train-batch loss vs. Full-batch loss In Figure 1c, train-batch loss stands for the average of batch losses forwarded in an epoch, which is commonly used to indicate the training process in deep learning. Full-batch loss is the average loss over the entire training dataset evaluated at the end of each epoch. In terms of optimizer evaluation, full-batch loss is much more informative than train- batch loss it reveals the robustness of an optimizer. However, full-batch loss is too expensive to evaluate and thus we only measure it on small datasets. On the other hand, test accuracy couples optimization and generalization, but since is also evaluated the end of the epoch, its convergence is similar to full-batch loss. Considering the basic usage of momentum in deep learning, we mainly use test accuracy to evaluate optimizers. We provide more discussion on this issue Appendix C.2.
M-SGD vs. OM-SGD We also include OM-SGD in Figure la. In comparison, the final accuracies of M-SGD and OM-SGD are 94.606% + 052% and 94. + 0.111% with average deviations at 1.040% and 0.634%, respectively. This difference can be explained following the interpretation in Hinton (2012) that } are the points after "jump" and {yk} are the points after "correction".
3 AMORTIZED NESTEROV'S MOMENTUM
In this section, we formally introduce SGD with Amortized Nesterov's Momentum (AM1-SGD) in Algorithm with the following remarks:
Options It can be verified that if m 1, AM1-SGD with Option I degenerates to M-SGD and Option corresponds to OM-SGD. Just like the case for M-SGD and OM-SGD, the accelerated convergence rate is built upon Option II while Option is easier to be implemented in a generic optimization layer". Intuitively, Option I SGD with amortized momentum and Option applies an m-iterations tail averaging I.
MMa Yarats (2019) observed that when effective learning rates 7 = =7(1 -1 are fixed, M-SGD and SGD have similar performance. We provide a discussion on this observation in Appendix C.I. 'To implement Option II, we can either maintain another identical network for the shifted point á or tem- porarily change the network parameters in the evaluation phase.
3

---

Under review as a conference paper at ICLR 2020
Final Test Accuracy
Average Accuracy STD
Accuracy
1.0
Option
Accurarys
Optinn
STO%
0.8
80
855
64.4
0.6
5GD
80
M-5GD
SCD
94.2
AM1 SGD
75
10
15
20
25
30
0
10
15
20
25
30
20
60
80
m
m
Epach
(a) Sweeping mn in (3,5, 7, 10, 20,30]. Run 5 seeds.
(b) Fixing m 5. Run 20 seeds.
Figure 2: ResNet34 on CIFAR-10. For all methods, Yo = 0.1, 3- 0.9. Labels of AMI-SGD are Shaded bands (or bars) indicate 1 standard deviation. Best viewed in color.
Algorithm 1 AMI-SGD
Input: Initial guess To, learning rate momentum B, amortization length m, iteration number K. Initialize: x f- 20, ā 2o. _+ (a running average).
1: for k = o..... K do
2: x x x -1- (x).
3: zt xt + m
4:
if (k 11) mod m = 0 then
5:
x+ 3 (a+ i). fadding amortized momentum)
6:
i 3 +- O.
end if 8: end for
7:
Output: Option I: x, Option II: T.
The symbol 't" denotes assignment.
We can improve the efficiency of Algorithm by maintaining a running scaled mo- mentum ū+ 2 m (i+ 3) instead of the running average ā+, by replacing the following steps in Algorithm 1:
Efficiency
Initialize:
3o. xo, ग+ m 2o. Step 3:
ū+ 5 5. +
Step 5:
x x+ + (8/m) 7+
Step 6:
3 ã (1/m) (/tm))++t5 -m
Then, in m-iterations loop, for each of the first addition and scaled vector addition. At the m-th iteration, it requires vector addition, 1 scalar- vector multiplication and scaled vector additions. In comparison, M-SGD (standard PyTorch) requires vector addition, (in-place) scalar-vector multiplication and 2 scaled vector additions per iteration. Thus, as long as m > AM1-SGD has lower amortized cost than M-SGD, For memory complexity, AM1-SGD requires one more auxiliary buffer than M-SGD.
iterations, AMI-SGD requires vector
Tuning m We did parameter sweep for m in our basic case study. We plot the final and the average deviation of test accuracies over 5 runs against vn in Figure 2a. Note that ===1 corresponds to the results of M-SGD and OM-SGD, which are already given in Figure 1. From this empirical result, introduces trade-off between final accuracy and robustness (the convergence behaviors can be found in Appendix A. Figure 2a suggests that m 5 is a good choice for this task. For simplicity, and also recommended setting, fix m= 5 for the rest of experiments in this paper.
momentum that increases robustness To provide a stronger justification, we ran 20 seeds with 5 in Figure 2b and the detailed data given in Figure 3 & Table 1. The results show that the amortized momentum significantly increases the robustness. Intuitively, the gap between Option I and Option II be understood the effect of tail averaging. However, the large gap between Option I and SGD is somewhat mysterious: what Option does is to inject very large momentum into SGD every m iterations. It turns out that this momentum not only provides acceleration, but also helps the algorithm become more robust than SGD. This observation basically differentiates AM1-SGD from simple interpolation in-between M-SGD and SGD.
Amortizzd momentum --(---(t is expected be much large than Nesterov's momentum (-(yk+1 -
4

---

Under review as a conference paper at ICLR 2020
Accuracy
560
METHOD
FINAL ACCURACY
Avg. STD
M-54D
AMI-SGD-
SGD
93.302% + 0.199%
0.928%
AM1-5CD-II
M-SGD
94.710% + 0.169%
0.995%
AMI-SGD-I
94.675% + 0..177%
0.587%
AMI-SGD-II
94.619% + 0.152%
0.314%
20
40
60
80
#poch
Figure & Table Detailed data of the curves in Figure 2b. Best viewed in color.
Learning rate scheduler issue We observed that when factor and the momentum 3 is too large for the task (e.g.. 0.995 for the task of this section), there would be a performance drop after the learning rate reduction. We believe that caused by the different cardinalities of iterates being averaged in क., which leads to a false momentum. This issue resolved by restarting the algorithm after each learning rate reduction inspired by (O' d'nogguu & Candes, 2015). We include more discussion and evidence in Appendix A.4.
schedulers with large decay
3.1 AM2-SGD: A VARIANT WITH IDENTICAL ITERATIONS
Algorithm AM2-SGD
Input:
Initial guess 2o, amortization length a point table [ó1 rate 17. momentum 3, iteration number K.
6m]
Rdxm learning
Initialize: 09 = Fo, vj E [m]]* R. {ix Js E [m]]]<= K-1 is sequence of uniformly random indexes.
If Option II is used, T" zo. {a running average for the point table Φ}
for k
O..... K
do
2: k+1
JN Fh 11 V Si (Ix) and keep other entries unchanged (i.e., oft2 03 for j + jk).
3: #k+1 = k+1 + B [o)+] & (adding amortized momentum}
4: if Option then z++1_ 3h m (ofit) 2
5: end for
Output: Option (not recommended): æK. Option II: ã
[m] denotes the set {l,.... m)
While enjoying an improved efficiency, AMI-SGD does not have identical iterations", which to some extent limits its extensibility to other settings (e.g., asynchronous setting). In this section, propose variant of Amortized Nesterov's 's Momentum (AM2-SGD, Algorithm 2) to address this problem. To show the characteristics of AM2-SGD, we make the following remarks:
Trading memory for extensibility In expectation, the point table é stores the most recent m iterations and thus the output ØK is an m-iterations tail average, which connects AM1-SGD, The relation between AMI-SGD and AM2-SGD resembles that of SVRG (Johnson & Zhang, 2013) and SAGA (Defazio et al., 2014), the most popular methods in finite-sum convex optimization: to reuse the information from several past iterates, can either maintain "snapshot" that aggregates the information or keep the iterates in a table. A side-by-side comparison is given in Section 4.
Options and convergence As in the case of AMI-SGD, m 1, AM2-SGD with Option I corresponds to M-SGD and Option II is OM-SGD. In our preliminary experiments, the convergence of AM2-SGD is similar to AMI-SGD and also has the learning rate scheduler issue. our preliminary experiments (can be found in Appendix A), observed that Option is consistently worse than Option II and it does not seem to benefit from increasing m. Thus, we do not recommend using Option I. We also set = 5 for AM2-SGD for its evaluation due to the similarity.
For AM1-SGD, the workload varies for different iteration le due the if-clause Step 4.
5

---

Under review as a conference paper at ICLR 2020
Additional randomness {js} In our implementation, at each iteration, sample index in [m] as Jk+1 and obtain the stored index jk. We observed that with Option I, AM2-SGD has much larger deviations than AMI-SGD, which we believe is caused by the additional random indexes {-
4 CONVERGENCE RESULTS
The original Nesterov's 's accelerated gradient is famous for its optimal convergence rates for solving convex problems. In this section, we analyze the convergence rates for AM1-SGD and AM2-SGI in the convex case, which explicitly model the effect of amortization (i.e., m). While these rates do not hold for deep learning problems general, they help us understand the observed convergence behaviors of the proposed methods, especially on how they differ from M-SGD (m 1). Moreover, the analysis also provides intuition on tuning m. Since the original Nesterov's 's method is determin- istic (Nesterov, 1983; 2013b), we follow the setting of its stochastic variants (Lan, 2012; Ghadimi & Lan, 2012), in which Nesterov's acceleration also achieves the optimal rates.
We consider the following convex composite problem (Beck & Teboulle, 2009; Nesterov, 2013a):
min zEx
F a f(x) + h(z)).
(4)
where C Rd is a non-empty closed convex set and h is a proper convex function with its proximal operator ,f (-)$ available. We impose the following assumptions on the regularity of f and the stochastic oracle f, (identical to the ones in Ghadimi & Lan (2012) with μ = 0):
Assumptions. For some I. 0, M 0, 2 0,
(a) 0 S f(v) fx) (Vf(x). x) S llv x]| + M fly xl1. Va, X.2
(b) E. [Vi(())] v (xx),, Ve E
(c) E. [V](() - Vf(=) & oo V
The notation E., [.] is E cover several important classes of convex problems. For example, (a) covers the cases of being L-smooth (M = 0) or Lo-Lipschitz continuous (M = 2Lo, L = 0) convex functions and if a = 0 in (c), the assumptions cover several classes of deterministic convex programming problems. We
11) for a random process in, i,,.... These assumptions
denote x² E X as a solution to problem (4) and Zg E X as the initial guess.
Unlike its usage in deep learning, the momentum parameter is always variable in general convex analysis. For the simplicity of analysis, we reformulate AM1-SGD (Algorithm 1) and AM2-SGD (Algorithm 2) into the following 10 E X,a E R):
AM1-SGD (reformulated, proximal)
AM22-GD (reformulated, proximal)
Initialize: ēg 20 = Tg. S K/m.
Initialize: 2n = 03 Fog vjj [m]] do
1: for = ........ do
1: for k = ....., K
2: for 0....,...
do
Sample ja uniformly in [m].
3: k: = sm + j.
3: a (1 Bbx) - zn + Bx 05
4:
Zg (1 Bz) 2z Bg F..
4: 2k+1 prox {2k a V Vf, (a(t)) )}
5:
1 = prox,, h {zz - 0+ fi (xx)).
6:
(xx41 (1 Ba)= Ik++ 8, Tz.)
5: ju
k+1 (1 Ba) 2k+1 + 3 o0
7: end for
6: end for
Output: ok m
1 Σn oK.
8:
Ø.+1
Σソ= Iemtj.
9: end for
Output: øs.
We show in Appendix B.1 that when h and 33 is a constant, the reformulated schemes AMI-
SGD and AM2-SGD are equivalent to Algorithm 1 and Algorithm 2 through a, = 17(1 - (3) and
Vri E IRd, proxx,,, (x) A arg min { ||u æ01² + avh(a)), see Parikh et al. (2014). When M > 0, ff is not necessarily differentiable and we keep using the notation (x(x) to denote an arbitrary subgradient of at x for consistency. For simplicity, we assume K is divisible by m.

---

Under review as a conference paper at ICLR 2020
Full-Batch
Full-Batch
1.0
0.20
0.8
AMI-SGD-
M-54D
560
Fulldaathh Low
0.15
M-SGD
SGD
METHOD
SGD
Avg. STD
0,0329
AM1-SGD-II
AMI-SGD-HI
M-SGD
0.0339
0.4
AMI-SGD-I
0.0142
AMI-SGD-II
0.0056
0.0
0.00 20
20
40
60
80
60
80
Epoch
Epacc
Figure Table ResNet18 with pre-activation on CIFAR-IO. For all methods, 7o 001, 3 run 20 seeds. For AMI-SGD, m 5 and its labels are formatted as 'AMI-SGD-(Option)) Shaded bands indicate 1 standard deviation. Best viewed in color.
On (1 3x) These reformulations are basically how Nesterov's momentum migrated into deep learning (Sutskever et al., 2013). Then we establish the convergence rates for AM1-SGI and AM2-SGD as follows. All the proofs in this paper are given in Appendix B.2.
Theorem 1. For the reformulated AMI-SGD, suppose choose
X1
with A1 min {2
L 11 x²1
B. +
and a,
(5)
L(1 - 3x)
3'
2/m 0² + (S
Then,
(a) The output Fs satisfies
3Lm 11:00 x111
8 811000 - x'1l V0² + M²
E [F - =F(x) S
Ko(m).
(K + m)²
VK + m
(b) the variance has "light tail", iee. Bs( Es[exx [[] ( (1))//1 120l1//222) /02] Eepp(1), Vo
X, and X is compact, denoting Dx A max | x11| flr for any A 2 0, we have
4A (3 |00 x= v6Dx)
Prob F(is) - F(x*)= S Ko(m) +
3v K +
21 (expl-A2 + exp -A))-
Remarks: (a) Regarding (m)m), its minimum obtained at either m = 1 or m = K. Note that for AMI-SGD, is strictly constrained in {l...., .....,K J. It can be verified that when becomes the modified mirror descent SA (Lan, 2012), or under the Euclidean setting, the SGD that outputs the average of the whole history, which is rarely used in practice. In this case, the conver- gence rate in Theorem la becomes the corresponding O(L/K K (a + M)/ K) (cf. Theorem 1 in Lan (2012)). Thus, we can regard AM1-SGD as smooth transition between AC-SA and the mod- ified mirror descent SA. (b) The additional compactness and "light tail" assumptions are similarly required in Nemirovski et al. (2009); Lan (2012); Ghadimi & Lan (2012). Recently, Juditsky et al. (2019) established similar bounds under weaker assumptions by truncating the gradient. However, as indicated by the authors, their technique cannot be used for accelerated algorithms due to the accumulation of bias.
K, AMI-SGD
Understandings: Theorem la gives the expected performance in terms of full-batch loss F(Ø) F(x*), ), from which the trade-off of m clear: Increasing m improves the dependence on variance but deteriorates the (LL K2) term (i.e., the acceleration). Based on this trade-off, we can under- stand the empirical results in Figure 2b: the faster convergence in the early stage could be the result of better control on and the slightly lowered final accuracy is possibly caused by the reduced acceleration effect. Theorem 1b provides the probability of the full-batch loss deviating from its expected performance (i.c., Ko (m)). It is clear that increasing leads to smaller deviations with the same probability, which sheds light on the understanding of the increased robustness observed in Figure 2. Since the theorem is built on the full-batch loss, did experiments based this

---

Under review as a conference paper at ICLR 2020
metric in Figure 4 & Table 2. Here choose training smaller ResNet18 with pre-activation (He et al., 2016a) on CIFAR-10 10 as the case study (the test accuracy is reported in Appendix A.5).
For AM2-SGD, only give the expected convergence results as follows.
Theorem 2. For the reformulated AM2-SGD, if we choose
k/m
A2 T1(1 - Bx)
1|00 - x²1
Bk k/m +
and Ork
with X2 min
3 2m + M) (K-1 + +
the oulput T satisfies
4v2 120 - xx" 10 (o E [F(Ø*]] - F(x*)= I 4(m2-m) (F(x)) 1200 x²1 |=00==11"
(K 2m 1)²
K + 2m -1
In comparison with Theorem la, Theorem 2 has an additional term F(xo) - F(x*) in the upper bound, which is inevitable. This difference comes from different restrictions on the choice of m. For AM2-SGD, improved rate, this additional term is inevitable. As a sanity check, we can let m + co to obtain a point table with almost all xp. and then the upper bound becomes exactly F(x0)) - F(x*). ). In some cases, there exists an optimal choice of m > 1 in Theorem 2. However, the optimal choice could be messy and thus we omit the discussion here.
Remark:
2 1 is the only requirement. Since it is impossible to let > K to obtain an
Comparing the rates, we see that when using the same m. AM2-SGD has slightly better dependence on a, which is related the observation in Figure 5 that AM2-SGD always slightly faster than AM1-SGD. This difference is suggesting that randomly incorporating past iter- ates beyond m iterations helps.
Understanding:
If m O()) Theorems and 2 establish the optimal OLL/K² + (o + M)/VK) rate in the convex setting (see Lan (2012) for optimality), which verifies AMI-SGD and AM2-SGD as variants of the Nesterov's 's method (Nesterov, 1983; 2013b). From the above analysis, the effect of m can be understood as trading acceleration for variance control. However, since both acceleration and variance control boost the convergence speed, the reduced final performance observed the CIFAR experiments may not always be the case as will be shown in Figure 5 and Table 3.
Connections with Katyusha Our original inspiration of AMI-SGD comes from the construction of Katyusha (Allen-Zhu, 2018), the recent breakthrough in finite-sum convex optimization, which uses a previously calculated "snapshot" point to provide momentum, i.e., Katyusha momentum. AMI-SGD also uses aggregated point provide momentum and it shares many structural simi- larities with Katyusha. We refer the interested readers to Appendix B.3.
5 PERFORMANCE EVALUATION
In this section, we evaluate AMI-SGD and AM2-SGD on more deep learning tasks. Our goal is to show their potentials of serving alternatives for M-SGD. Regarding the options: for AMI-SGD, Option I is nice choice, which has slightly better final performance as shown in Table 1; for AM2- SGD, Option is not recommended as mentioned before. Here we choose to evaluate Option for both methods for consistency, which also corresponds to the analysis in Section 4. AM1-SGD and AM2-SGD use exactly the same values for (1, 3) as M-SGD, which was tuned to optimize the performance of M-SGD. We set m 5 for AMI-SG and AM2-SGD.
We trained ResNet50 and ResNet152 (He al., 2016b) on the IISVRC2012 dataset ("ImageNet") (Russakovsky al., 2015) shown in Figure 5b. For this task, used 0.1 initial learning rate and 0.9 momentum for all methods, which is a typical choice. We performed restart after each learning rate reduction discussed in Appendix A.4. We believe that this helps the training process and also does not incur any additional overhead. We report the final accuracy in Table 3.
We also did a language model experiment Penn Treebank dataset (Mareus al., 1993). We used the LSTM (Hochreiter & Schmidhuber, 1997) model defined Merity et al. (2017) and followed the experimental setup in its released code.
We only changed the learning rate and momentum in
8

---

Under review as a conference paper at ICLR 2020
Penn
on ImageNet
ResNet152 imageNet
80
SGD-ASGO
Vil calltinonn perplestty
75
M.SGD
AM1-SDD
70
70
AM2-5GD
65
65
60
55
M
SG0
M-552
Text
360
M-5G0
AM2 5GD
50
AM2 560
45
45
200
400
600
20
40
80
20
40
60
80
Epoch
tpach
Epoch
(a) LSTM on Penn Treebank.
(b) ResNeet on ImageNet. Run seeds.
Figure 5: Convergence of LSTM and ResNet. We plot the curve of validation perplexity and test accuracy, respectively. Shaded bands indicate +1 standard deviation. Best viewed in color.
Table 3: Detailed perplexity and accuracy results for Figure 5.
Penn Treebank (Perplexity)
ImageNet (Final Accuracy)
METHOD
Validation
Test
ResNet50
ResNet152
SGD (+ASGD)
61.283
59.068
72.783% 0.081%
74.361 1 0.293%
M-SGD
60.747
58.355
75.711% + 0.062%
78.065% + 0.100%
AMI-SGD
60.734
57.977
75.779% 0.15%
77.816% 0.287%
AM2-SGD
60.434
58.233
75.845% + 0.073%
78.194% + 0.147%
the setup. The baseline is SGD+ASGI"" (Polyak & Juditsky, 1992) with constant learning rate 30 as used in Merity et al. (2017). For the choice of (77, 3), following Lucas et al. (2019), we chose 0.99 and used the scheduler that reduces the learning rate by half when the validation loss has not decreased for 15 epochs. We swept from 2.5, 1, 0.1, 0.01 and found that 7j = 2.5 resulted in the lowest validation perplexity for M-SGD. We thus ran AMI-SGD and AM2--GD with this (n. 33) and m 5. Due to the small decay factor, we did not restart AMI-SGD and AM2-SGD after learning rate reductions. The validation perplexity curve is plotted in Figure 5a. We report validation perplexity and test perplexity in Table 3. This experiment is directly comparable with the one in Lucas et al. (2019).
Extra results are provided in the appendices for interested readers: the robustness when using large B (Appendix A.2), a CIFAR-100 experiment (Appendix A.6) and comparison with classical momen- tum (Polyak, 1964), AggMo (Lucas al., 2019) and QHM (Ma & Yarats, 2019) (Appendix A.3),
6 CONCLUSIONS
We presented Amortized Nesterov's Momentum, which is special variant of Nesterov's momentum that utilizes several past iterates to provide the momentum. Based on this idea, we designed two different realizations, namely, AM1-SGD and AM2-SGD. Both of them simple to implement with little-to-no additional tuning overhead over M-SGD, Our empirical results demonstrate that switching AM1-SGD and AM2-SGD produces faster early convergence and comparable final generalization performance. AM1-SGD is lightweight and has more robust iterates than M-SGD, and thus can serve favorable alternative M-SGD large-scale deep learning tasks. AM2-SGD could be favorable for more restrictive tasks (e.g., asynchronous training) due to its extensibility and good performance. Both the methods are proved optimal the convex case, just like M-SGD. Based on the intuition from convex analysis, the proposed methods are trading acceleration for variance control, which provides hints for the hyper-parameter tuning.
"SGD+ASGD SGD and switch averaged SGD (ASGD) when threshold is

---

Under review as a conference paper at ICLR 2020
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, al. Tensorflow: system for large- scale machine learning. In /2th {USENIX) Symposium on Operating Systems Design and Imple- mentation ((OSDI) 16), pp. 265-283, 2016.
Zeyuan Allen-Zhu. Katyusha: The first direct acceleration of stochastic gradient methods. Journal of Machine Learning Research, 18(221):1-51, 2018. URL http: jmlr orgppapers/////poooooo html.
v1
Zeyuan Allen-Zhu and Lorenzo Orecchia. Linear coupling: An ultimate unification of gradient and mirror descent. In ITCS, 2017.
Alfred Auslender and Mare Teboulle. Interior gradient and proximal methods for convex and conic optimization. SIAM Journal on Optimization, 16(3):697-725, 2006.
Amir Beck and Mare Teboulle. A fast iterative problems. SIAM journal on imaging sciences, (11):1332022 2009.
algorithm for linear inverse
Augustin Cauchy. Méthode générale pour la résolution des systemes d' d'équations simultanées. Comp. Rend. Sci. Paris, 25((847):536-538, 1847.
Aaron Defazio. On the curved geometry of accelerated optimization. 1811v:1812.044344 2018.
arXiv preprint
Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. Saga: A fast incremental gradient method with support for non-strongly convex composite objectives. Advances neural information processing systems, pp. 1656-1654, 2014.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 120ul):2121-2159) 2011.
Saeed Ghadimi and Guanghui Lan. Optimal stochastic approximation algorithms for strongly
vex stochastic composite optimization i: A generic algorithmic framework. SIAM Journal on Optimization, 22(4):1469-1492 2012.
Saeed Ghadimi and Guanghui Lan. Accelerated gradient methods for nonconvex nonlinear and stochastic programming. Mathematical Programming, 156(1-2):59-99 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference computer vision, pp. 630-645. Springer, 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016b.
Geoffrey Hinton. Lecture 6C The momentum method, 2012. URL httt://www.cs. toronto. edu toronto.edu/-hi htwp://oodeeooonotoooooooreeeeee lecttaeeeeeereeeeeeeee html.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
Chonghai Hu, Weike Pan, and James T Kwok. Accelerated gradient methods for stochastic optimiza- tion and online learning. In Advances in Neural Information Processing Systems, pp. 781-789, 2009.
Gao Huang. Zhuang Liu, Laurens Van Der Maaten, and Kilian Weinberger. Densely connected convolutional networks. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700-4708. 2017.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv: 1:Xxv:550200316: 2015.
10

---

Under review as a conference paper at ICLR 2020
Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In Advances in neural information processing systems, pp. 315-323, 2013.
Anatoli Juditsky, Alexander Nazin, Arkadi Nemirovsky, and Alexandre Tsybakov. Algorithms of robust stochastic optimization based on mirror descent method. arXiv preprint arXiv:1907.02707, 2019.
Diederick P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
Guanghui Lan. An optimal method for stochastic composite optimization. Mathematical Program- ming, 133(1-2):365-397 2012.
Guanghui Lan, Arkadi Nemirovski, and Alexander Shapiro. Validation analysis of mirror descent stochastic approximation method. Mathematical programming, 134(2)425-458, 2012.
Guanghui Lan, Zhize Li, and Yi Zhou, A unified variance-reduced accelerated gradient method for convex optimization. arXiv preprint arXiv: 12412, 2019.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.0300 2016.
James Lucas, Shengyang Sun, Richard Zemel, and Roger Grosse. Aggregated momentum: Stability through passive damping. In International Conference Learning Representations, 2019. URL https: /ttp:////peeeeeeeeeeeemee net forum? Situummddddxxxxx5555555
Jerry Ma and Denis Yarats. Quasi-hyperbolic momentum and adam for deep learning. Interna- n
tional Conference on Learning Representations, 2019. URL
Forum?id-3a S1
Mitchell Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building large annotated corpus of english: The penn treebank. 1993.
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing 1stm lan- guage models. arXiv preprint arXiv: 02182, 2017.
Arkadi Nemirovski, Anatoli Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal optimization, 19(4): 1574 1609, 2009.
Yu Nesterov. Gradient methods for minimizing composite functions. Mathematical Programming, 140(1): 2013a.
Yurii Nesterov. Introductory lectures on convex optimization: A basic course, volume 87, Springer Science & Business Media, 2013b.
Yurii E Nesterov. A method for solving the convex programming problem with convergence rate o11/k)). In Dokl. Akad. Nauk SSSR, volume 269, pp. 543-547, 1983.
Brendan O' d'donoguue and Emmanuel Candes. Adaptive restart for accelerated gradient schemes. Foundations of computational mathematics, 15(3):715-732, 2015.
Neal Parikh, Stephen Boyd, et al. Proximal algorithms. Foundations and Trends in Optinization, 1(3):127-239 2014.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Ditito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Boris Polyak. Some methods of speeding up the convergence of iteration methods. USSR Com- putational Mathematics and Mathematical Physics, 4(5):1-17, 1964.
11

---

Under review as a conference paper at ICLR 2020
Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM Journal on Control and Optimization, 30(4):838-88) 1992.
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference Learning Representations, 2018. URL http::ts: /openreveee net / forum? ie=riddrrrr --
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathemati- cal statisties, pp. 400 407, 1951.
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115(3)2211225 2015.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, llya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-199) 2014.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International conference on machine learning, pp- 1139-1147, 2013.
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26- 31, 2012.
Paul Tseng. On accelerated proximal gradient methods for
optimization. 2008
Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems, pp. 4148-4158, 2017.
Sergey Zagoruyko and Nikos Komodakis. 1rX1v:1605.07146, 2016.
Wide residual networks.
arXiv preprint
Kaiwen Zhou, Fanhua Shang, and James Cheng. A simple stochastic variance reduced algorithm with fast convergence rates. In International conference machine learning, pp. 5980-5989, 2018.
Kaiwen Zhou, Qinghua Ding, Fanhua Shang, James Cheng, Danli Li, Zhi-Quan Luo. Direct Accelcration of SAGA using Sampled Negative Momentum. In AISTATS, pp. 1602-1610, 2019.
12

---

Under review as a conference paper at ICLR 2020
Appendices
A Extra Experimental Results
14
A.1 The effect of on convergence
14
A.2 Robustness on large momentum parameters
15
A.3 Comparison with other momentum
15
A.4 Issues with learning rate schedulers
17
A.5 Test accuracy results of Figure 4 & Table 2
17
A.6 CIFAR-100 experiment
17
A.7 A sanity check
18
Missing parts in Section 4
19
B.1 The reformulations
19
B.2 Proofs of Theorem and Theorem 2
20
B.2.1
Proof of Lemma 1
20
B.2.2
Proof of Theorem la
22
B.2.3
Proof of Theorem 1b
23
B.2.4
Proof of Theorem 2
25
B.3
Connections between AMI-SGD and Katyusha
27
C Miscellanies
28
C.I Comparison of SGD and M-SGD
28
C.2 Training evaluation
29
Experimental Setup
29
D.I Classification Setup
29
D.2 Language Model Setup
30
13

---

Under review as a conference paper at ICLR 2020
A EXTRA EXPERIMENTAL RESULTS
In this appendix, we provide more experimental results further evaluate the Amortized Nesterov's Momentum. Table 4 shows the detailed data of the parameter sweep experiments, where the con- vergence curves of these results are given Appendix A.1. In Appendix A.2, we compare the robustness of AMI-SGD and M-SGD on large momentum parameters. In Appendix A.3, we em- pirically compare the Amortized Nesterov's Momentum with classical momentum (Polyak, 1964), aggregated momentum (Lucas et al., 2019) and quasi-hyperbolic momentum (Ma & Yarats, 2019). We discuss the issues with learning rate schedulers Appendix A.4. We report the test accuracy results of the ResNet18 experiment (in Section 4) in Appendix A.5, A CIFAR 100 experiment is provided in Appendix A.6. We also provide a sanity check for our implementation Appendix A.7.
Table 4: Final test accuracy and average accuracy STD of training ResNet34 on CIFAR- 10 over runs (including the detailed data of the curves in Figure and Figure 2a). For all the methods, Vo = 0.1, 0.9. Multiple runs start with the same Z0--
METHOD
DESCRIPTION
FINAL ACCURACY
Avg. STD
SGD
Standard Pytorch
93.406% + 0.149%
0.991%
M-SGD
Standard Pytorch
94.606% 0.152%
1.040%
AMI-SGD
Option I, m = 1, sanity check
94.672% 0.143%
0.912%
AM1-SGD
Option I, m 3
94.630% + 0.025%
0.643%
AM1-SGD
Option I, m
94.604% 0.104%
0.496%
AMI-SGD
Option I, m 7
94.640% = 0.127%
0,439%
AM1-SGD
Option I, m 10
94.542% 0.133%
0.441% %
AM1-SGD
Option I, m 20
94.378% 0.217%
0.399%
AM1-SGD
Option I, 30
94.300% = 0.145%
0,429%
OM-SGD
AM1-SGD (Opt. II, m 1)
94.728% + 0.111%
0.634%
AM1-SGD
Option m
94.662% 0.139%
0.411%
AMI-SGD
Option II, m = 5
94.602% + 0.083%
0.265%
AM1-SGD
Option II, m
94.514% 0.103%
0.282%
AM1-SGD
Option II, m = 10
94.424% 0.117%
0.293%
AM1-SGD
Option II, m 20
94.358% = 0.181%
0.306%
AM1-SGD
Option m
94.270% 0.134%
0.344%
AM2-SGD
Option I, m 1, sanity check
94.682% 0.212%
0.822%
AM2-SGD
Option I, m 5
94.572% + 0.188%
0.591%
AM2-SGD
Option I, m 10
94.440% 0.138%
0.737%
AM2-SGD
Option I, m = 20
94.312% 0.149%
0.741 %
AM2-SGD
Option 1I, m
94.664% 0.7%
0.263%
AM2-SGD
Option II, m 10
94.496% + 0.211%
0.280%
AM2-SGD
Option m
94.412% 0.140%
0.246%
A.1 THE EFFECT OF m CONVERGENCE
We show in Figure 6 how affects the convergence of test accuracy. The results show that in- creasing m speeds up the convergence in the early stage. While for AM1-SGD the convergences of Option and Option II are similar, AM2-SGD with Option II is consistently better than with Option I in this experiment. seems that AM2-SGD with Option does not benefit from increasing m and the algorithm robust. Thus, do not recommend using Option for AM2-SGD.
14

---

Under review as a conference paper at ICLR 2020
Test Accuracy
Test
M-5GD
M-50p
00
AM1-5:0-
AM1-500
80
96D-15
AN1
86
5GD-1
80
40
50
40
60
Epocc
Epoch
Epoch
Figure 6: Convergence of test accuracy from the parameter sweep experiments in Table 4. Labels are formatted as (AN//--SGDD(OOtion------ Best viewed in color.
A.2 ROBUSTNESS ON LARGE MOMENTUM PARAMETERS
We compare the robustness of M-SGD and AMI-SGD when is large in Figure 7 & Table For fair comparison, AMI-SGD uses Option I. As we can see, the STD eπror of M-SGD scales up significantly when is larger and the performance is more affected by large , compared with AM1-SGD.
Test Accuracy
METHOD
FINAL ACCURACY
Avg. STD
M-SGD-0.9
94.606% 0.152%
1.040%
Accuracys
80
M-SGD-0095
94.198% 0. 100
1.200%
M-SGD-0.9
M-SGD-0.99
88.366% 0.357%
2.561% %
70
M-SGD-0.95
M-5GD-0.99
AMI-SGD-0.9
94.604% 0. 104%
0.496%
AM1--GDD-0.9
AM1-SGD-0.95
93.942% +: 0.074%
0.581 %
AM1-SGD-0.95
AM1-5GD-0.99
AM1-SGD-0.99
90.640% 0.375%
0.900%
20
80
Epoch
Figure 7 & Table 5: ResNet34 on CIFAR-1O. Yo = 0.1,8 8 E (0.9. 0.95, 0.99). run seeds (the 0.9 results are copied from Table 4). Labels are formatted as "[Algorithm)-(B} Best viewed in color.
A.3
COMPARISON WITH OTHER MOMENTUM
Test Accuracy
METHOD
ACCURACY
Avg. STD
90
CM-SGD
94.690% + 0.212%
1.107%
AMI-SGD
94.604% + 0.104%
0.496%
85
AggMo
94.556% 0.142%
0.777%
Test
CM-SGD
AM1-SGD
QHM
94.426% 1 0.226%
1.073%
AggMo
M-SGD
94.606% + 0.152%
1.040%
75
0
20
40
60
80
Epoch
Figure Table 6: ResNet34 on CIFAR-I0. Run seeds. The results of AMI-SGD and M-SGD
are copied from Table 4. Best viewed in color.
In this section, we compare AMI-SGD (Option I) with classical momentum (Polyak, 1964), AggMo (Lucas al., 2019) and QHM (Ma & Yarats, 2019) in basic case study (training ResNer34
15

---

Under review as a conference paper at ICLR 2020
CIFAR-10), Since not aware of what makes fair comparison with these methods (e.g.. is not clear what is the effective learning rate for AMI-SGD), we compare them based on the default hyper-parameter settings suggested by their papers.
Classical Momentum leaming has the following scheme (standard Py Torch) (v E R "0 0):
The SGD with classical momentum (CM-SGD) that is widely used in deep
*++ B
uk V fig (I&),,
em k++1 = Ty k++1 for l: 2 0.
CM-SGD with its typical haper-parameter settings (7o = 0.1, = 0.9) is observed to achieve similar generalization performance as M-SGD. However, CM-SGD more unstable and prone to oscillations (Lucas et al., 2019), which makes it less robust than M-SGD as shown in Table 6.
Aggregated Momentum (AggMo) AggMo combines multiple momentum buffers, which in- spired by the passive damping from physics literature (Lucas et al., 2019). AggMo uses the following update rules (for 1, ...., T. v(()
Rd u0 0):
(8)
K41
(0) g(0)
uk
(1)
V fie (IK). fort =1,,,,,... I,
Σ
(0 M++ for k 2.0.
#k+1 2x +
We used the exponential hypereparameter setting recommended in the original work with the scale- factor 0.1 fixed, [ 1 a'=1 for i 1,..., T and choosing T in {2, 4). We found that = 2 gave the best performance in this experiment. As shown Figure Table 6, with the help of passive damping, AggMo is more stable and robust compared with CM-SGD.
Qussi-hyperbolic Momentum (QHM) Ma & Yarats (2019) introduce the immediate discount factor R for the momentum scheme, which results in the QHM update rules (a E IR, w Rdd ghhh =):
ah "k+1 = =. ah + (1-
B) Via (2RR)
Ik+1 = I - a (v u + (1 - v) V (zx)), for 2 0.
Here we used the recommended hyper-parameter setting for QHM (og = 1.0, = 0.999, v 0.7).
Figure 8 shows that AM1-SGD, AggMo and QHM achieve faster convergence in the early stage while CM-SGD has the highest final accuracy. In terms robustness, huge gaps are observed when comparing AMI-SGD with the remaining methods in Table 6. Note that AM1-SGD is more efficient than both QHM and Agg and is as efficient as CM-SGD.
We also plot the convergence of train-batch loss for all the methods in Figure 9. Despite of showing worse generalization performance, both QHM and AggMo perform better on reducing the train- batch loss this experiment, which consistent with the results reported in Ma & Yarats (2019); Lucas et al. (2019).
Train-Batch Loss
Train-Batch Loss
0.6
0.0040
CM.SGD
CM-SGD
0.5
AM1-SGO
0.0035
AM1-SG0
AggMo
Train-Ba Batch Loss
AggMo
0HM
QHM
0.0025
0.2
0.1
0.0020
0.0
0.0015
20
40
60
80
75
80
85
95
Epoch
Epoch
Figure 9: Train-batch loss results. Best viewed in color.
16

---

Under review as a conference paper at ICLR 2020
A.4
ISSUES WITH LEARNING RATE SCHEDULERS
Train-Batch Loss
Train-Batch Loss
Train-Batch Loss
1.0
MG5GD
M-SGD
M: MS0
AM1-500
1.2
AML SG0
Traindattch Less
1.2
AM2
AM2.5GD
0.6
1.0
1.0
0.6
0.6
0.4
0.4
0.2
0.2
0.0
Epoch
Epach
Epoch
(a) [ = 0.95
(b) | = 0.995
(c) 8 = 0.995
Figure 10:
ResNet18 on CIFAR-10. 7ju 0.1, f0.95, 0.995). '+' represents performing
restart after each learning rate reduction.
Train-Batch Loss
We show in Figure 10 that when is large for the task, using step learning rate scheduler with decay factor 10, performance drop observed after each reduction. Both Option I and Option II have this issue and the curves are basically identical. Here we only use Option II. fix this issue by performing restart after each learning rate reduction (labeled with '+'). We plot the train-batch loss here because way. If 3 = 0.9, there no observable performance drop in this experiment.
1.75
1.50
AM1-SG0
Train-Batch
0.75
0.50
find the phenomenon is clearer in this
0.25
Epech
For smooth-changing schedulers such as the cosine an- nealing scheduler (Loshchilov & Hutter, 2016), the amor- tized momentum works well as shown in Figure 11.
Figure 11: ResNet18 on CIFAR-IO.
Cosine annealing scheduler (without restarts), Yo = 0.1, = 0.995.
A.5 TEST ACCURACY RESULTS OF FIGURE & TABLE
We report the test accuracy results of the experiments in Section Figure 12 Table 7. These results reminiscent of the ResNet34 experiments (Figure & Table 1).
Test Accuracy
METHOD
FINAL ACCURACY
Avg. STD
Accurncyh
90
SGD
92.810% + 0.147%
1.005%
8.-
M-SGD
94.057% + 0.170%
1.102%
80-
AMI-SGD-I
93.968% + 0.154%
0.543%
AMI-SGDI
AM1-SCD-II
AMI-SGD-II
93.953% + 0.192%
0.315%
20
40
80
Epoch
Figure 12 & Table 7: ResNet18 with pre-activation CIFAR-10. For all methods, 70 = =0.1,80..3 run 20 seeds. For AMI-SGD, m 5 and its labels are formatted (ELL-SGD-(Option)) Shaded bands indicate 1 standard deviation. Best viewed in color.
A.6 CIFAR- 100 EXPERIMENT
We report the results of training DenseNet121 (Huang et al., 2017) CIFAR-100 in Figure 13, which shows that both AMI-SGD and AM2-SGD perform well before the final learning rate re- duction. However, the final accuracies lowered around 0.6% compared with M-SGD. We also notice that SGD reduces the train-batch loss at an incredibly fast rate and the losses it reaches are consistently lower than other methods in the entire 300 epochs. However, this performance is not
17

---

Under review as a conference paper at ICLR 2020
reflected in the convergence of test accuracy. We believe that this phenomenon suggests that the DenseNet model is actually "overfitting" M-SGD (since in the ResNet experiments, M-SGD always achieves a lower train loss than SGD after the final learning rate reduction).
Test Accuracy
Train-Batch Loss
80
1.0
SGD
75
0.8
M-5GD
Trai Bat--Badhhhhh LosS
70
AM2-SGD
65
60
SGD
M-SGD
55
AM1-SGD
0.2
AM2-5G0
50
0.0
0
50
100
150
200
250
300
0
50
100
150
200
250
300
Epoch
Figure 13: 2enseNet121 on CIFAR-100. For all methods, yo 0.1, 8 0.9, run seeds. AMI- SGD and AM2-SGD use Option II and m 5. Shaded bands indicating #1 standard deviation. Best viewed in color.
A.7 A SANITY CHECK
When 1, both AMI-SGD and AM2-SGD are equivalent to M-SGD, we plot their convergence in Figure 14 as a sanity check (the detailed data given Table 4).
Test Accuracy
Train-Batch Loss
95
M:SGD
AM1-5GD-L-1
Accuracyy66
Train Train-Batch Batch Loss
AM2-5GD-1-1
85
0.4
SGD
AMD-SGD-I-1
0.2
AM2-SGD-I-1
75
0
20
40
60
80
20
40
60
80
Epoch
Epoch
Figure 14: A sanity check. Labels are formatted as 'AM 11M(122--00D-
We observed that when m = 1, both AMI-SGD and AM2-SGD have lower STD error than M-SGD. We believe that it is because they both maintain the iterates without scaling, which is numerically more stable than M-SGD (M-SGD in standard PyTorch maintains a scaled buffer, i.e.,
2f 3-1 (yk zx)).
18

---

Under review as a conference paper at ICLR 2020
B MISSING PARTS IN SECTION 4
B.1 THE REFORMULATIONS
When h = 0 and 3 is constant, we the reformulations by eliminating the sequence {za}.
For the reformulated AM2-SGD,
=x (1 - B) 2k 3 o
-k-1 = 2k a. Vig (ak) )
a(1 - B) = Eliminate {za
k+1 5x 1) (in
k-1 (1 - B) + + 3 Je
Ju+1 k+1
go+ + 3 (0) 032))
(1 = (1 - B) 2k+1 + 3 (111)
The reformulated AM2-SGD
Algorithm
For the reformulated AMI-SGD, when = 0, the inner loops are basically SGD,
Th = (1 - B) = + -.. Zk+1 2k a fi, (Ix),
a(1 3) Eliminate {zz
Fh+1 = Ih 17 V (zx).
(Ie+1 = (1 5) 2 1 + ...)
At the end of each inner loop (i.c., when (I: + 1) mod m = 0), we have
1) (1 B) 2(s+ 13m + B N.. while the beginning of the next inner loop,
f(s-1)m (1 B) [[++1]m + 3 To41s
which means that need to set #4+1 f- Øb+1 + 3 (I++1 Z.) (reassign the value of 2x+1))
We also give the reformulation of M-SGD (scheme (1)) to the Auslender Teboulle (2006) scheme for reference:
Ft (1 - 3) 24 + yk:
ZR++1 =k (z)())))
a(1 - B) =1
Uk+1 (1 (3 2x+1 + 3 yks
Eliminate {2e
Mk+1 T2 1) Vi, (31)) Tk+1 yk+1 + 3 (Vk+1 8k
(IK+1 = (1 - B) 2k+1 + 3 1c+1).
Auslender Teboulle (2006)
Nesterov (1983; 2013b)
(AC-SA (Lan, 2012))
AC-SA (in the Euclidean case) maps the Auslender & Teboulle (2006) scheme through (in the original notations)
x md
29
3 3 a
Intuition for the Auslender & Teboulle (2006) scheme can be found in Remark in Lan (2012).
19

---

Under review as a conference paper at ICLR 2020
B.2 PROOFS OF THEOREM AND THEOREM 2
The reformulated schemes are copied here for reference:
AM1-SGD (reformulated, proximal)
AM2-SGD (reformulated, proximal)
Initialize: Σn- - 20 = 20, S = K/m.
Initialize: 2g = 0 = To. Vj E [m]]
1: for 0, : 5 do
1: for k 0, ..... K do
2: for j = 0, ......
do
2: Sample ja uniformly in [m].
3: k
j.
3: xt (1 5k) 24 + f4 05
4:
2yK = (1 - 33) 2k. + 3g Es.
4: 2k41 = prox {2h Ock V fiz (af*) }}.
5:
2k++1 prox = 0²5" V fi (ax)).
6:
(xx+1 (1 B2) 1k+1 + 8g x..)
5:
k+1 (1 Bax) -h+1 + 3r D9.--
7: end for
6: end for
8: Ta+- 2==11 Jom+jj.
Output: o² m Σド= oj.
1
9: end for
Output: gg--
Comparing the reformulated schemes, we see that their iterations can be generalized as follows:
(1 3) 2 B y,
+ = arox - vf(())),
(6)
y* -)- 8 y.
This type of scheme is first proposed in Auslender & Teboulle (2006), which represents one of the simplest variants of the Nesterov's methods (see Tseng (2008) for other variants). The scheme is then modified into various settings (Hu et al., 2009; Lan, 2012; Ghadimi & Lan, 2012; 2016; Zhou et al., 2019; Lan et al., 2019) to achieve acceleration. The following lemma serves as a comerstone for the convergence proofs of AM1-SGD and AM2-SGD.
Lemma 1. Ifa(1 B) < L/L, the update scheme (6) satisfies the
the following recursion:
1
(F(y*) (F()*)-((x*))
B
(F(w) F(x*)(())))))))))))))))
1
(I=-
15-=x11'
--1--11--- 1-=1==12-111" 11==--x0111
2a
(1V Vf((x)] + M)²
(V)()) Vfi(z),
2(a L(1 B))
B.2.1 PROOF LEMMA
This Lemma similarly provided in Lan (2012); Ghadimi Lan (2012) under a more general setting that allows non-Euclidean norms in the assumptions, give a proof here for completeness.
Based on the convexity (Assumption (a)), we have
- < x VV(() Vf((z), x²)+ VV((()) - 2
Ro
R
+ v fz))) z
We upper bound the terms on the right side one-by-one
For Ro.
R (x) B (Vf(z),y- -
B
1-3 (f(y) - f(x)),
(8)
I-
where (+) uses the relation between and 2, i.e., (1 B) (x 2) 3 (v - x).
For Rg, based on Assumption (a), we have
f(y+) - f(x) + (Vf(=), - yt) - 2 || - y+ $12 M 1 y
20

---

Under review as a conference paper at ICLR 2020
Then, noting that y² = (1 - 3) (z )²),, we can arrange the above inequality
L(1 B) 12- - z 2+||² +
1
R2 -
1- B (f(x) -
2)
2
+ M || 11
L(1 3) 2== - z ||
1
+ - B
+ - -V
+ M)
||=-
2
Using Young's 's inequality with < > 0, we obtain
R₂ L(1 - B) + || x +11 + - (f(x)) f(y))) + ([Vf(z) - vf((2)] + +)² (9)
1
-
2
3
2
For R3. based on the optimality condition of prox an to - Vfi(z)} and denoting 0h(z+) as subgradient of h at 2+, we have for any u E X.
(a ah(22) + 2t - = + a. Vf((z), u --+) 0, ([(()),2z - w) S (0h(z*) +),,,,,, ),u- + (2t - 2, u - 2t)
(2+
S M(u)) - 11tttt))x))=))zz + (2= - 2, u - =t)
Choosing u = x²'
1
Rs S h(x') 112+)
(st - 2, =,,²=2++)
a
@ h(x²) - h(2+)*+++)))))>)))))) + 1
||2 x² - -112t -
- -11= 1||=||=211 - =11²).
(10)
2a
where (+) follows from ||a + b1 = ||a||
2
+ ||6|| + (a. b).
Finally, by upper bounding (7) using (8). (9), (10), we conclude that
1 1- 8 (f(u) - f()) + L(1 B) + Ç a- ||: 1|====- =11
9
f(x) - f(x)) < R1 +
2
1
(f(z) - f(v²)) + +h(x") - -h(=t)+ (1] v(((2)] + +2?
+
1 - 8
2
1
[|2 - 101
1| - x011 =1')))
+ 2a
After simplification.
1 - (d((t(t)) (((())=((x))) - 1 1- (f(() - f(x')) + I(1 - 33 + - -a-1
1
B
1|222|12
2
+ h(xx)= - -1(2*)+ + ([V v(((x)] + +M2
+ R
(1I)
26
1 (I=- x01 112 1|=-- -x=1((²)) =z(1(²)
20
Note that with the convexity of h and yt (1 - (3) xt + y, we have
h(y+) < (1 - 3))((z) ) + Bh(w),
1
3
h(x*) > 1 T- h(y" - B h(y).
Using the above inequality and choosing Ç = a1 - -L(1 - B) > 0= a(1 -3) < 1/L, we can arrange (11) as
1 1- 8 (Flyt) F(x*)) - 1- 3 (F(y) - F(x*)) + 1
1
B
(12- x+1 - || -
x =²1)²)
2a
(l) (z)2 + M)
+ R1.
2(a -
21

---

Under review as a conference paper at ICLR 2020
B.2.2
PROOF OF THEOREM 1A
Using Assumption (c), Lemma I with
T Tk
2 = 2k
zt 2k+1
y F.
(12)
y Fk+11
a Os
3 B,
and taking expectation, if a1 - Ba) < 1/L, we have
1 - 3 (Ei [F(x++1)] - F(x*)) + 20
Ein 2 x²1]3]
3. (F(i.)-((x") -
(o (o M)²
$
+ 20xs
||
x²1l 21
+ 2(as
1 8x
(())
Summing the above inequality from sm,, sm + m we obtain
E [F (xam+ F(x*)) + 20,m
E || [=( -x0 xx11²²)
11 m
8, (F(i.) F(x*))( +
(0 + +M)
20mm 1|āmm - 1²1 + 2(1x5 L(1 - BB)))
- 1 - 8,
Using the definition of 8,41 and convexity,
- F(x*)) + E [[
1 E
[| |
1 - B
2m
(13)
< a 3 (F(i.) - F(x*))( + 2m
1Izam x²11²+ a.(0² +22
a=(
B,
023 L-1-
can be verified that with the choices 3g
and ags M(- the following holds for s 2 0,
842
0'+1 B4+1
ag and Bo 0.
(14)
1 Bs41
1 B
Note that since our analysis aims at providing intuition, we do not refine the choice of a, as in (Hu et al., 2009; Ghadimi & Lan, 2012). Thus, by telescoping (13) from = S I,..., O, we obtain
as
[F(ss) - F(x*)) + 2m E ||.mm =²1]]]
1 - 3s-1
S-1
002 2 + M2)
<
100 2²1 +
Σ
2m
0₂3 L(1 Ba)
==U
and thus,
2L
1:00 x=1 + 4T.(02 + M2)
8-1
a
E [F((ss) - -((=)) 5 A1m(S (S + 1)2
\n + 1)2
0, -3,)L
(a)
2L
3 (02 + M2) S-1
11:00 -
-x*11 2=1
+
Σ
(8 + 2)²
A1 + 1)2
L(S + 1)2
2L
(0² + MS +1)
1:00 - 2=1 -
+
Xim(S + 1)2
22

---

Under review as a conference paper at ICLR 2020
where (a) follows from A I 3 and (b) holds because 0 0z (x 2) is non-decreasing and thus
Σ (s + 2)² < /// 8
S-1
(c+2) - 25
8(S + 1)3
3
==0
Denoting
L 200 x=1
Xi
2 2/ + (S + 1)
and based on the choice of Å = min
( {i if Σ //4 we have
8 -0 x vo E [((8)) F(x)) -((x))
($
If Af 3.
31 || x01
4 11 x' Vo² + M2
[ - F(x))
+
m(S + 1)2
m + 1)
Thus, we conclude that
3L |1100- x=1 mSS 1)2
8 1|000 x' +o++ M²
E [[F(is)) F(x))
+
m (S + 1)
Substituting S = K/m completes the proof.
B.2.3 PROOF OF THEOREM 1B
In order to prove Theorem Ib, need the following known result for the martingale difference (cf. Lemma 2 in Lan et al. (2012)):
Lemma 2. With N 0. let Eo.Ei,..,,EN- be sequence of i...d. random variables, for t O, .,...,N - 1, og > be a deterministic number and t/ = V.......... &t) be a determinis- tic measurable function such that Bg, [bte] = a.s. and Eg, [exp[v//o S exp(1] a.s. Then for any A 2 0,
N-1
N
Prob
Σ
A
Σ 1=
o
<
exp33322333.
t=0
To start with, using Lemma with the parameter mapping (12), we have
1 (F(* -F(x)))
|| - =1
2a.
< 8, (F(5.) F(**)) + 20s
1
|| - =11113
B,
(I (mx) - Vii, (#x) + .M2
+ (Vf(xx))
V.. (IK), 2x x²))
2(a L(1 - (B)))
3, (F(3.) F(**)) + 20s
1
M
<
|| - 1=1
B.
a, L(1 BBa)
[V+( V., (2x)
(Vf(xx))
V (xx), 2x - x²)
3. L( 3.)
23

---

Under review as a conference paper at ICLR 2020
Summing the above inequality from k sm..... .......... + m 1 and using the choice 0₂ -(1-9.)
A1
with A S J. we obtain
a2
(F(f++1) - -F(x*)) + 2m || - x²1
- 8
3,
<
(F(E.)- - F(x*)) + 2m I| x²1² + 30,
1 Bg
+ 3a
sm+mm1
sm
[Z+ - Vaa (Tw))
-(( 24 xx)).
k=
With our parameter choices, the relations in (14) hold and thus we can telescope the above inequality from = S ,,....0.
as-1
(F(ãs) F(x*)) s
1
11
x=12
12 + +3M² Σ 02
38-1
2m
==0
K-1
3
a [k/m] [V](xx))
Vi, (m++))
(15)
R44
K-1
+
[L/m (Vf(z))
Vin (Fx)) x²)
m
f2.
Denoting V A [V[(())) - ViG
(xx))] (xx)))
= Ex= a [4/ml)
m Es00 0g for Rs. by Jensen's
inequality, we have
K-1
}
K-1
E
Σ
2 [k/mm] 02
aik/m E exp (12/021 & exp(1},
K=0
k=
where (+) uses the additional assumption E,g [exp / S exp(1).
Then, based on Markov's incquality, we have for any 2 0,
Prob
exp
alk [k/m] /2/02
2 exp A + 13 S Sexp(-A},
S-1
Prob
R4 2 (A + 1)02m
m
< (xp
(16)
8=0
For Rs, since we have E, [alk/m] (Vf(xx) - Vfc 24 =²))] = 0 and
2 [k/mm]
(V/(s)) V Si, (Mx), 2x - x212 }
E
exp
S S eep (V2/0'}] exp(1).
a² /²//m= D
which is based on the "light tail" assumption, using Lemma 2, we obtain
S-1
Prob Rs 2 AoDx
S Sexp(_A3/33.
(17)
Combining (15), (16) and (17), based on the parameter setting (ef. (5)) and using the notation
Ko(m) A 3Lm ||z0 - 1=
||z0 x= Vo² + M
(K + +m)
VK
m
S-1
12Lo²
ALoDx
Rs
Σ
0g
Σ 0
A1 + 1)
A1 1m mm
24

---

Under review as a conference paper at ICLR 2020
conclude that
Prob (F(ās) F(x") S Ko(m) ARR3 2 (expf-A2/33 + exp(-AA).
For Rg, using the choice of as and A1, we obtain
Rs S 4V6o O
(S+
4V60Dx
40 11:00 x²'
<
3/K m L
3/K +
VK + mV M
4o (3 1|z0 - x + v6Dx)
3v K + m
which completes the proof.
B.2.4 PROOF OF THEOREM 2
Using Assumption (c), Lemma I with
2 Z++1
k++1
y jk
a Ok
8 = 3k
and taking expectation, if ax - BBa) 1/ L, we have
1
Eir.i [F(o)t F(x*)] + 20x Esg.je
|| - x² 2²1
- Bg
(18)
8g
(0 + +M)
<
Eis [F(d - F(x*)] + 200
1 - x 1=1 1122 + 2(ar
Bx
L(1 BBa))
Note that
de
) F
Je
de
Σ (F(o) - F(x*)) ((x))))
Σ (F(o))-E(r)))) F(x'))
j=1
m
m
Σ
(F(ottl) F(x*))
TEjs
(F F(x*))
1=1
Dividing both sides of (18) by and then adding sides, we obtain
1
Eje
EE (F(d)) F F(x*)) both
(1-9k)mm )m
Bic 1 Bk
F(gtt2) - F(x*)
+
Bigci
[| - =01]]]
2a
j=1
< - Ej [F(O) - F(x)))
+ 1 3k
Σ
F( - 1((0))
+ 20gm
|| - x= 112
VTR
Ø=1
(a + +M222
2m
L(1 (x6)
1 - 1-9c
1
F(O)) F((o))
I| x+1
(a + M)2
L(1 - Bx)) (19)
m
2amm
2m(ar
25

---

Under review as a conference paper at ICLR 2020
It can be verified that with our parameters choice: B3 holds for k 0,
k/m k/m+2
Ag and as I(1-B))
the following
1- Pa+1
0
m
and Bg = O.
9k+1
3K
Note that since our analysis aims at providing intuition, we do refine the choice of a, as in (Hu et al., 2009; Ghadimi & Lan, 2012). Then, we can telescope (19) from k = K ..... O, which results in
OK- -1 1 3K-1
1
F(o)) - F(x')
+
||ex - x²1 2²1²]]
m j=1
2m
Am(m -1) 1)
K
ax + +M2
-F(xo) - F(x))) + 2m
1100
x²1
+
Lm
2m(a
-1
L(1 Bx))
k=0
Using the definition of @k and convexity, obtain
1 - -Kx1
(F(x))) (a2(m= (a(m-1))(Fm) (F(xo) -F(x))) + 2m
E [((Ø F(x*)]
100 x² x²1??
ak-1
Lm
1 3K-1
K-1
aulo + M)
0K-1
2m(a - L(1 - Ba))
k=0
(a) 4(m - (F(oo)
F(x*))
2T 1:0 1=1
m
K-1
2)
A2m
K-1
m
m
312(0 + M)²
K-1
Σ
k
2
+
2Lm ( K-1 +2) 2
m
k=0
(b) 4(m (F( -F(x))
2L -00
x111
(20)
m K-1 +2)²
A2m
K-1 + 2) m
m
4A2(0 + M)² ( (K-1 +2)
L
where (a) uses A S 3. (b) follows from simple integration arguments and that K + 2 S
2 K-I +2) since K 2 1, m 2 1.
Based on the choice of
1|:00 - xx11
A2 min
2m M) [K1 +2)
(20) can be further upper bounded
[F(5K) - F(x*)] < 4(m 1)( F(a*))
3L |1:00 - x²1l 112
4v2 11:0 - (1 M)
m K-1 +2)
2
K-1 + 2)
m (K-1 + 2)
m
26

---

Under review as a conference paper at ICLR 2020
B.3 CONNECTIONS BETWEEN AMI-SGD AND KATYUSHA
The discussion in this section aims to shed light on the understanding of the experimental results, which also shows some interesting relations between AMI-SGD and Katyusha.
The high level idea of Katyusha momentum that it works as a "magnet" inside an epoch of SVRG updates, which "stabilizes" the iterates so as make Nesterov's momentum effective (Allen-Zhu, 2018). In theory, the key effect of Katyusha momentum is that it allows the tightest possible variance bound for the stochastic gradient estimator of SVRG (ef. Lemma 2.4 and its comments in Allen- Zhu (2018)). In this sense, we can interpret Katyusha momentum as variance reducer that further reduces the variance of SVRG. Below we show the similarity between the construction of Katyusha and AMI-SGD, based on which we conjecture that the amortized momentum can also reduce the variance of SGD (and thus increase the robustness). However, theory, following a similar analysis of Katyusha, cannot guarantee reduction of in the worst case.
Deriving AMI-SGD from Katyusha Katyusha has the following scheme (non-proxime in the original notations, o is the strong convexity parameter, cf. Algorithm with Option in Allen Zhu (2018))²:
Initialize: ~ = y/ = 20 To. for s ....., S. do 2: Compute and store f(t*).
3L = 1 + 0.0
for j ....., m do
4:
k sm + +j.
Ty = TI 2g T2 x + Ti T2) yk-
6:
V = Vii, (Fk) - Vii, + ((TT).
7:
Z 1 2k Vn.
8:
U/++1 IK 17. Vk.
9:
end for
10: g0+1 E wi) -1
Σタ= w/ tsm+j+1.
11: end for Output: ~.a
We can eliminate the sequence {2 in this scheme. Note that in the parameter setting of Katyusha, we have = (TT, and thus
Ik+1 = Ti . -b+11 + T2 2 + (1 T1 - T2) ) yk+1
Ti - 17 V T2 F=' + (1 T1 T2)
yk + (1 11 T2) (y/+1 yk)
= I 17 / + (1 - Ti - T2) (yk+1 - 344)
yk+1 + (1
Te) (ye+1 yx).
Hence, the inner loops can be written
Vk++1 = 3bk 17 k:
T++1 = Uk+1 + (1 -2) (yk+1 2k),
which is the Nesterov's 's scheme (scheme (1)). At end of each inner loop (when k: = sm m 1),
((++1)m = T1 1)m+))m + T2 H + (1 T1 - -22)= -ist 1)ms
while at the beginning of the next inner loop,
fm = Ti 2(e+1)m + T2 +1 + (1 - TI - -T2) Vee+1)m
which means that we need to set 2(++ + 2(++l)m + T2 (z++1 - F²') (reassign the value of
1))+
Then, the following is an equivalent scheme of Katyusha:
We change the notation J T
27

---

Under review as a conference paper at ICLR 2020
Initialize: 30 = 3/ = 20, 11
37.
1 + ao.
I: for ...... do
2: for j = 0....,
do
3:
8m + j.
4: yk-1 Ik- 17. Vk.
5: Ib+1 = yk+1 + (1 - T1 T2) (yh+1 yk).
end for
7: 3++1
m- wv 1==
-=== w Vem+++1.
mm
8: 1)m (m+ + 72 (I x³).
9: end for Output: Eg--
Now it is clear that the inner loops use Nesterov's momentum and the Katyusha momentum is injected for every m iterations. If we replace the SVRG estimator Va with Vfi, (ox), set T2 0, which eliminate Nesterov's momentum, and use uniform average for #++1, the above scheme becomes exactly AMI-SGD (Algorithm 1).
If we only replace the SVRG estimator Ve. the scheme can be regarded as adding amortized mo- mentum to M-SGD. This scheme requires tuning the ratio of Nesterov's momentum and amortized momentum. In our preliminary experiments, after suitable tuming. we observed some performance improvement. However, this scheme increases the complexity, which we do not consider it worth- while.
recent work (Zhou et al., 2018) shows that when 1 -1 - T2 = O, which is to solely use Katyusha momentum, one still derive optimal rates and the algorithm is greatly simplified. Their proposed algorithm (i.e., MiG) is structurally more similar to AMI-SGD.
C MISCELLANIES
C.1 COMPARISON OF SGD AND M-SGD
Ma & Yarats (2019) normalize the momentum buffer of M-SGD, which results in the following formulation (a E IR, yohh E Rd, nith = 0):
gh u k++ = 3 o + (1 B) fix(),)
F = rx (3 uk1 + (1 B) V., for k 2 0.
This scheme is equivalent to the PyTorch formulation (scheme (3)) through 2h (1 3) 2kk and 7) = a(1 B).
Based on this formulation, is understood the effective learning rate (i.e., the vector it scales has the same cardinality as a gradient) and the experiments in Ma & Yarats (2019) were conducted with fixed = 1. Their results indicate that when using the same effective learning rate, M-SGD and SGD achieve similar performance and thus they suspect that the benefit of momentum basically comes from using sensible learning rates.
Here we provide some intuition on their results based on convex analysis. For simplicity, we consider deterministic smooth convex optimization. In theory, to obtain the optimal convergence rate, the effective learning rate a is set to a very large Okk/L), which can be derived from Theorem 1 or Theorem 2 by setting
0, M = 0, m (then \ or λa is always 3 since the other term is co). If we fix = 3L for both methods, GD has an O(1 K) convergence rate (cf. Theorem 2. 1.13 in Nesterov (2013b)). For the Nesterov's method, if we use 3g A (applying Lemma 1):
it has the convergence rate k=
1
3L (F(J)+1) -((+)) 4
| 10=11-x111=-
3k (F()k) F(a*)) + 3L 1 x+1
3x
1- Bx
3L I x²
F(yx)-F(x7) - F(xx) <
2(K +
28

---

Under review as a conference paper at ICLR 2020
Train-Batch Loss
Full-Batch Train Loss
1.0
1.0
SGD
SGD
0.8
M-SGD
Full-Batch Train Loss
0.8
M-SGD
Train Batch Loss
AM1-SGD-I
AM1-SGD-I
0.6
AM1-SGD-II
0.6
AM1-SGD-II
0.4
0.4
0.2
0.2
0.0
0.0
20
40
60
80
20
40
80
Epoch
Epoch
Figure 15: Train-batch loss full-batch loss. Best viewed in color.
Thus, in this case, both GD and the Nesterov's method yield an O(1/K) rate, and thus we expect them to have similar performance. This analysis suggests that the acceleration effect basically comes from choosing a large effective learning rate, which corresponds to the observations in Ma & Yarats (2019).
However, what is special about the Nesterov's method is that it finds legal way to adopt a large that breaks the I.L. limitation. If GD uses the same large a, we would expect it be unstable and potentially diverge. In this sense, Nesterov's 's momentum can be understood as a "stabilizer". In basic case study (ResNet34 on CIFAR-10), if we align the effective learning rate and set = 1.0 for SGD, the final accuracy is improved but the performance is highly unstable and not robust, which is 2.205% average STD of test accuracy over runs. The significance QHM (Ma & Yarats, 2019) is that with suitable tuning, it achieves much faster convergence without changing the effective learning rate. Our work uses convergence behavior of SGD as a reference to reflect and to understand the features of our proposed momentum, which is why we set
C.2 TRAINING EVALUATION
Due to the mechanism of back-propagation, evaluating train-batch loss basically incurs no overhead. It can be efficiently used indicate the training progress. However, can only be treated as a coarse approximation to the full-batch loss as shown in Figure 1c. If batch normalization (Ioffe & Szegedy, 2015) or dropout (Srivastava al., 2014) is used training, the model changes during the training phase, which makes train-batch loss less accurate. More importantly, train-batch loss is always observed to be statistically stable, which omits many important characteristics of an optimizer such as robustness, oscillations, etc. We include a comparison of train-batch loss and full-batch loss on training ResNet18 with pre-activation on CIFAR-10 10 (the experiment in Section 4) in Figure 15.
We also notice that for different optimizers, even if their convergenees tinguishable, their convergences on test accuracy can vary greatly. We show two examples in Figure 16, where the ImageNet experiment is from Section 5 and the CIFAR- 10 experiment is from Table with m 10.
train-batch loss indis-
D EXPERIMENTAL SETUP
All of our experiments were conducted using PyTorch (Paszke et al., 2017) library.
D.1 CLASSIFICATION SETUP
Our implementation (e.g., ResNet and DenseNet implementations, data pre-processing) generally follows the repository https: / h/gggtttuuuu hub. com/ kuanglivi pytorch-cifar. All the CIFAR experiments in this paper used a single GPU in a mix of RTX2080Ti, TITAN Xp and TITAN V. The batch size is fixed to 128. We used cross-entropy loss with 0.0005 weight decay and used batch normalization (loffe & Szegedy, 2015). Data augmen- tation includes random 32-pixel crops with padding of 4-pixel and random horizontal flips with
CIFAR-10 & CIFAR-100
29

---

Under review as a conference paper at ICLR 2020
ResNet152 on ImageNet
ResNet152 on ImageNet
3.0
80
SGD
M-SGD
Train Batch Loss
AM1-SGD
Test Test Accuracylk
70
AM2-SGD
2.0
65
60
SGD
55
M-SGD
AM1-SGD
50
AM2-5GD
45
20
40
60
80
20
40
60
80
Epoch
Epoch
ResNet34 on CIFAR-10
ResNet34 on CIFAR-10
0.6
AM2-SGD-I
94
0.5
AM2-5GD-II
Train Batch Loss
Test Test Accuracyls
0.4
90
0.3
88
0.2
86
0.1
84
AM2-SGD-1
82
AM2-SGD-II
0.0
20
40
60
80
20
40
60
80
Epoch
Epoch
Figure 16: Train-batch loss test accuracy. Best viewed in color.
0.5 probability. We used step (or multi-step) learning rate scheduler with a decay factor 10. For CIFAR-10 10 experiments, trained 90 epochs and decayed the learning rate every 30 epochs. For the CIFAR-100 experiments, we trained 300 epochs and decayed the learning rate at 150 epoch and 225 epoch following the settings in DenseNet (Huang al., 2017).
ImageNet In the ImageNet experiments, we tried both ResNe150 and ResNet152 (He et al., 2016b). The training strategy is the same as the PyTorch's official repository https://githis com/pytorchs www.rc/////mpppp////p// examplxmeeeeeeeeeeee les / ree reeeraseettee e/masterrrr/rreeeeeeeeeeeeeeeeee /iii//met.nnttt//e which uses batch size of 256. The learning rate starts at 0.1 and decays by factor of 10 every 30 epochs. Also, we applied weight decay with 0.0001 decay rate to the model during the training. For the data augmentation, ap- plied random 224-pixel crops and random horizontal flips with 0.5 probability. Here, we run all experiments across 8 NVIDIA P100 GPUs for 90 epochs.
D.2 LANGUAGE MODEL SETUP
We followed the implementation the repository https: /ttp::.//gtttbb..ooo sa wtsssallefffee// awd-lstm- and trained word level Penn Treebank with LSTM without fine-tuning or con- tinuous cache pointer augmentation for 750 epochs. The experiments were conducted on a single RTX2080Ti. We used the default hyper-pane syper-parameter tuning except for learning rate and momentum: The LSTM has 3 layers containing 1150 hidden units each, embedding size is 400, gradient clipping has maximum norm 0.25, batch size is 80, using variable sequence length, dropout for the layers has probability 0.4, dropout for the RNN layers has probability 0.3, dropout for the input embed- ding layer has probability 0.65, dropout to remove words from embedding layer has probability 0.1, weight drop (Merity et al., 2017) has probability 0.5, the amount of 2 (2-reuularization activation is 2.0, the amount of slowness regularization applied on the RNN activation is 1.0 and all weights receive weight decay of 0.0000012.
the RNN
30


Reviews:
[{"review_id": "r1elVZBh5H", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors proposed Amortized Nesterov’s Momentum, a variant of Nesterov’s momentum that utilizes several past iterates, instead of one iterate, to provide the momentum.  The goal is to have more robust iterates, faster convergence in the early stage and higher efficiency. The authors designed two different realizations, AM1-SGD and AM2-SGD.\n\nComments:\n\nMy major concern for this paper is its unconvincing motivation and experiment results, especially when the approach is designed for deep learning.\n\nThe motivation for the proposed approach is not quite convincing. The authors said that “due to the large stochasticity, SGD with Nesterov’s momentum is not robust...This increased uncertainty slows down its convergence especially in the early stage of training and makes its single run less trustworthy” For image classification, Nesterov momentum is very popular and the final convergence values of different trails seems to be similar. It would be more convincing if the authors can provide practical evidence for supporting this claim.\n\nIt was claimed that Amortized Nesterov’s Momentum has “higher efficiency and faster convergence in the early stage without losing the good generalization performance”. What is the benefit or advantage for having faster early convergence without improving the final generalization performance?\n\nThe authors claim that “M-SGD2 is more robust and achieves slightly better performance”, in Figure 1a, however, it is really hard to tell the difference between M-SGD2 and M-SGD from Figure 1a.\n\nThe efficiency improvement in page 4 is really hard to follow for comparison with Algorithm 1 in page 3. Though m > 2 could reduce the number of operations in step 5 and 6, I don’t think this is a computational bottleneck. I believe these updates should be very fast in comparison with forward and gradient calculation. Making the 1% computation 50% faster does not mean more efficient. I would like to know how much computation cost can be saved with this modification. On the other hand, adding one more auxiliary buffer (scaled momentum) could significantly impact the training as the memory is often the limit. \n\n\nIn section 3.1, what is “identical iteration”? It is hard to compare AM2-SGD with AM1-SGD. It would be easier to follow if the side-by-side algorithm comparison can be shown early.\n\nThe section 4’s theoretical analysis based on the convex composite problem is not quite convincing. I am not sure how these results are related with the deep learning applications.\n\nIn the experiment section, the comparison of AM1/2-SGD with other baselines seems not quite consistent. The authors first state that they use all 0.1 learning rate and 0.9 momentum for all methods, however, the setting for M-SGD is using 0.99 momentum and different learning rate schedule. This makes the comparison not very meaningful, while AM1-SGD and AM2-SGD do not use learning rate restart. With so many differences, the advantage of AM1-SGD and AM2-SGD are not that different with M-SGD.  In the task of ImageNet-152, M-SGD even is better than AM1-SGD. This makes the conclusion that “AM1-SGD has a lightweight design, which can serve as an excellent replacement for M-SGD in large-scale deep learning tasks” not quite valid.\n\nMinor: The author may assume readers maybe familiar Katyusha momentum, I think there may need more background about it. \n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "H1gz-n-oKr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper provides a new simple method to incorporate Nesterov momentum into standard SGD for deep learning, with good empirical and theoretical results. Overall I think this paper should be accepted, some minor comments follow.\n\nAt no point does Polyak's heavy ball method get mentioned, even though the variant of Nesterov acceleration you are considering is very similar to it (since the momentum parameter is fixed, which is not the usual form of Nesterov except in the strongly convex case). It would be beneficial to delineate how this is or isn't related to heavy ball.\n\nThe experiments would benefit from a wall-clock time comparison too, rather than just epochs since these new methods would be slower (but presumably not by much).\n\nThe appendix is huge with most of the technical details relegated there which I did not read fully. I think this impacts the readability significantly, though not grounds for rejection. Perhaps it suggests that a conference with a small page limit is not the best venue?\n\nIt seems that SGD still has better convergence early on. The authors suggest their method fixes this (relative to standard nesterov SGD) but it doesn't seem to be quite as good as SGD. Can you explain or discuss why this is still the case?\n\nThe assumptions require some explanation, they are just listed with no context. What are they and why are they useful?\n\nStep size \"should be constrained to O(1/L)\" is misleading, you should say explicitly that step-size <= 1/L (or whatever it is depending on the algorithm).\n\nSome of the writing is a bit strange / sloppy, e.g.:\n\"AM2-SGD is a bit tricky in randomness\"\n\"However, full-batch loss is way too expensive to evaluate.\"\n\nIn Algorithm 1 AM1-SGD:\n\"xk+1 ← xk+1 + β · (˜x+ − x˜)\"\ndoesn't parse because x_{k+1} appears twice.\n\nMissing references:\n\nAccelerated proximal algorithms:\n\n*) Beck and Teboulle: A Fast Iterative Shrinkage-Thresholding Algorithm\nfor Linear Inverse Problems\n\n*) Nesterov: Gradient Methods for Minimizing Composite Objective Function,\n\nRestarting (slightly different to your approach but still relevant): \n\n*) O'Donoghue: Adaptive Restart for Accelerated Gradient Schemes ", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "H1gUDg7atB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "%% Post Author Response comments %%\nThank you for your detailed response/revision. \n\n1 - Introducing “m-times” larger momentum: Somehow, this is not a particularly intuitive statement or one that reflects clearly in a theoretical bound. Since we are getting to issues surrounding the use of momentum with stochastic optimization, I would like to make a note that the performance of these algorithms more broadly aren't quite sketched out for their use in broader stochastic optimization. In particular, despite broad use in practice, it is unclear if standard variants of Nesterov acceleration/Heavy Ball method achieve \"acceleration\" in stochastic optimization. See for e.g., the work of Kidambi et al ICLR 2018 (“On the insufficiency of existing momentum schemes for stochastic optimization”) - where, the argument was that these methods were designed for deterministic optimization (where we get exact gradients) - in fact, that paper empirically as well as theoretically shows that these schemes do not offer acceleration in a precise sense compared to specialized algorithms for stochastic optimization. It is unclear if the proposed algorithms can offer a similar improvement over SGD in a provable sense, even for the specific examples described in their paper.\n\n2 - The point about theory (just as you mention) is that it doesn’t directly apply towards the simulations, nor, do they improve on already known algorithms - so I am unable to see the point that these results present broader implications that can guide practice.\n\n3 - The response doesn’t address the fact that for the theory bounds presented in the paper to hold (even in the convex settings described), one requires knowledge of parameters that are not known a-priori, and are often fairly difficult to estimate. So the performance of the algorithm in practice may quite significantly be away from the bounds described in the paper.\n\nWhile I appreciate the points and revision made by the authors, I still believe the paper requires some rethinking to present their results (and this includes more detailed comparisons to existing works) in order to make a case towards broader practical applications.\n\n\n%%. %%\n\nThis paper considers robustness issues faced by Nesterov’s Acceleration used with mini-batch stochastic gradients for training Deep Models. In particular, the paper proposes amortized momentum, an algorithm that offers a way to handle these issues. The paper in general is well written and easy to follow. \n\nThe paper proposes algorithms AM-SGD1 and AM-SGD2 and presents extensive results regarding their complexity analysis on convex problems and their performance when training neural networks. The algorithms require storing one more model’s worth of storage compared to standard momentum based methods (which can be viewed as a drawback in certain cases).\n\nComments:\n\n[1] I am concerned about the motivation behind this paper - which, according to the paper is that Nesterov’s accelerated gradient method with stochastic gradients has huge initial fluctuations. The issue with regards to more fluctuations of the initial performance is natural given how aggressive these accelerated methods work. As long as this is not a reason/cause for worse terminal performance (which doesn’t seem to be the case), I am unable to see why large initial fluctuations are concerning.\n\n[2] Theory: The theory bounds for this problem setting do not appear to improve over known bounds in the literature. As a side note, the work of Hu et al. “Accelerated Gradient Methods for Stochastic Optimization and Online Learning” is highly related to this paper’s theoretical aspects, setup and bounds. Furthermore, this bounded variance noise model for stochastic gradients, while being theoretically useful (and important), is often very detached from practice (as this implies that the domain is bounded and we perform projections of iterates whenever they go outside the set - such aspects hardly reflect on practical SGD implementations). Using this as a means to reason about robustness of the proposed algorithm (for e.g. remarks for theorem 1a. and in conclusions) appears to be a big leap that may lead to potentially misleading conclusions.\n\n[3] In order to run the algorithm to achieve the theoretical bounds claimed (in theorems 1 and 2), it appears that the stepsize \\alpha_s depends on unknown quantities such as initial distance to opt, noise variance etc.\n\n[4] The claim in page 2 about comparing SGD and M-SGD says that the stepsize in deterministic and stochastic optimization is constrained to be O(1/L) is rather misleading. In realistic practical implementation of SGD with a multiplicative noise oracle, one really has to use a much smaller stepsize than 1/L. This in a sense leads back to point[2] about the unrealistic nature of bounded variance assumptions for understanding SGD based methods used in the context of Machine Learning. They are better suited for understanding stochastic methods in black-box optimization (as opposed to considering Machine Learning problems).\n\nMy take is that even if the authors justify novelty in terms of theory results (which, to my knowledge is limited compared to existing literature), rewriting the paper by considering its theoretical merit and presenting empirical results (even as considered in this paper) of this algorithm (without attempting to make very strong connections to explain issues experienced in non-convex training of neural networks, since the theory works in vastly different settings under restrictive assumptions) can be appreciated by appropriate sections of audience (both in theory as well as optimization for deep learning communities).", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJg7oswhqS", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n\nThis paper proposes two variants of Nesterov momentum that maintains a buffer of recent updates. The paper proves optimal convergence in the convex setting and makes nice connections to mirror descent and Katyusha.\n\nI vote to reject the submission, with my main concerns being with the experimental results. I would consider raising my score if my concerns are addressed.\n\nConcerns\n-The learning rate schedule on the CIFAR experiments is unconventional. The original ResNet paper trains for 64k iterations (roughly 160 epochs). It’s standard to train for at least 200 epochs (see schedule from Smith et al.). Do the results hold under the standard schedule with careful tuning for the baselines?  \n-The discussion on “Train-batch loss vs. Full-batch loss” in Section 2 is unclear. On smaller datasets, it is feasible to perform evaluation at the end of the epoch on the entire batch. Furthermore, reporting test accuracy couples optimization and generalization, and I am not sure what is meant by the statement ``test accuracy is too informative.”\n-Reporting the peak test accuracy is strange. In general, we do not have access to the test set. It’s more natural to report the final test accuracy or have a holdout set to determine an iteration for evaluation.\n-It’s not clear how significant the results on ImageNet and PTB are. Namely, a comparison to AggMo and/or QHM would be good, since the flavor of these algorithms is quite similar. Experiments in the AggMo paper suggest that AggMo performs better on PTB. In the comparison given in Appendix A3, it seems like AggMo performs slightly better throughout training. SGD should also attain better performance with learning rate tuning on ImageNet.\n-I’m not sure how useful “Test Accuracy STD%” is useful as a metric since it is influenced heavily by the learning rate and its schedule. Tail averaging schemes in general seem like they would increase “robustness.” In addition, there seem to be situations where a higher final variance is beneficial (just run the method for longer and you can find a better solution). It would be nice to expand the discussion on the notion of robustness defined in the paper.\n\nMinor Comments\n-The dashed line in figure 1b) is hard to read. I would recommend removing the grid lines and make the colors more differentiable.\n-Algorithm 1: use of both assignment and equality operator? Whereas other boxes use equality\n-Spacing looks a bit off in parts of the paper. 1) after the first sentence in the introduction 2) “generic optimization layer (Defazio, 2018) . However”)\n-M-SGD and M-SGD2 can be potentially confusing and are not too informative as acronyms.\n-Remark on Theorem 1b: depicts does not seem like the right word\n\nSmith, S. L., Kindermans, P. J., Ying, C., & Le, Q. V. (2017). Don't decay the learning rate, increase the batch size. arXiv preprint arXiv:1711.00489.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a variant of Nesterov momentum which saves computation by only periodically recomputing certain quantities, and which is claimed to be more robust in the stochastic setting. The method seems easy to use, so there's probably no harm in trying it. However, the reviewers and I don't find the benefits persuasive. While there is theoretical analysis, its role is to show that the algorithm maintains the convergence properties while having other benefits. However, the computations saved by amortization seem like a small fraction of the total cost, and I'm having trouble seeing how the increased "robustness" is justified. (It's possible I missed something, but clarity of exposition is another area the paper could use some improvement in.) Overall, this submission seems promising, but probably needs to be cleaned up before publication at ICLR.

Decision: Reject