Year: 2019

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. In this work, we propose Amortized Nesterov’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 empirical results show that it achieves faster early convergence and comparable final generalization performance with little-to-no tuning. Just like Nesterov’s method, the new schemes are also proved optimal in general convex setting. Our analysis sheds light on the understanding of the new variant.

heading: 1 INTRODUCTION
text: 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 a popular choice of optimizer especially for neural network training which requires lower precision. Sutskever et al. (2013) found that using SGD with Nesterov’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 momentum 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 a 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 in 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 Nesterov’s momentum. From users’ perspective, the new momentum has only one additional integer hyper-parameter m to choose, which we call the amortization length. Learning rate and momentum parameter of this variant are strictly aligned with Nesterov’s momentum and by choosing m = 1, it recovers Nesterov’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 Momentum as follows:
• Increasing m improves robustness1. This is 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 andm in the convex setting. • Increasing m reduces (amortized) iteration complexity. • A suitably chosen m boosts the convergence rate in the early stage of training and produces
comparable final generalization performance.
1In this work, robustness 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.
• It is easy to tune m. The performances of the methods are stable for a wide range of m and we prove that the methods converge for any valid choice of m in the convex setting. • Ifm is not too large, the methods obtain the optimal convergence rate in general convex setting,
just like Nesterov’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 it is provided 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 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.

heading: 2 PRELIMINARIES: SGD AND NESTEROV’S MOMENTUM
text: We start with a review of SGD and Nesterov’s momentum. We discuss some subtleties in the implementation and evaluation, which contributes to the interpretation of our methods.
Notations In this paper, we use x ∈ Rd to denote the vector of model parameters. ‖·‖ and 〈·, ·〉 denote the standard Euclidean norm and inner product, respectively. Scalar multiplication for v ∈ Rd and β ∈ R is denoted as β ·v. f : Rd → R denotes the loss function to be minimized and∇f(x) represents the gradient of f evaluated at x. We denote the unbiased stochastic gradient estimator of ∇f(x) as ∇fi(x) with the random variable i independent of x (e.g., using mini-batch). We use x0 ∈ Rd to denote the initial guess.
SGD SGD has the following simple iterative scheme, where γ ∈ R denotes the learning rate:
xk+1 = xk − γ · ∇fik(xk), for k ≥ 0.
Nesterov’s momentum The original Nesterov’s accelerated gradient (with constant step) (Nesterov, 1983; 2013b) has the following scheme2 (y ∈ Rd, η, β ∈ R and y0 = x0):
yk+1 = xk − η · ∇f(xk), xk+1 = yk+1 + β · (yk+1 − yk), for k ≥ 0,
(1)
where we call β · (yk+1 − yk) the momentum. By simply replacing ∇f(xk) with ∇fik(xk), we obtain the SGD with Nesterov’s momentum, which is 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, vpt ∈ Rd and v0 = vpt0 = 0): for k ≥ 0,
(2) { vk+1 = β · vk − η · ∇fik(yk + β · vk), yk+1 = yk + vk+1.
(3) { vptk+1 = β · v pt k +∇fik(xk),
xk+1 = xk − η · (β · vptk+1 +∇fik(xk)).
Here the notations are modified based on their equivalence to scheme (1). It can be verified that schemes (2) and (3) are equivalent to (1) through vk = β−1 ·(xk−yk) and vptk = η−1β−1 ·(yk−xk), respectively (see Defazio (2018) for other equivalent forms of scheme (1)).
Interestingly, both PyTorch and Tensorflow3 track the values {xk}, which we refer to as M-SGD. This choice allows a consistent implementation when wrapped in a generic optimization layer (Defazio, 2018). However, the accelerated convergence rate (in the convex case) is built upon {yk} (Nesterov, 2013b) and {xk} may not possess such a theoretical improvement. We use OM-SGD to refer to the Original M-SGD that outputs {yk}.
2We exchange the notations of x and y in Nesterov (2013b). 3Tensorflow tracks the values {yk + β · vk} = {xk}.
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 rates4 γ = η. From the interpretation in Allen-Zhu & Orecchia (2017), η represents the learning rate for the gradient descent “inside” Nesterov’s method. To introduce the evaluation metrics of this paper, we report the results of training ResNet34 (He et al., 2016b) on CIFAR-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 x0. Figure 1a shows that Nesterov’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 1b 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 trainbatch loss as 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 it is also evaluated at 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 in Appendix C.2.
M-SGD vs. OM-SGD We also include OM-SGD in Figure 1a. In comparison, the final accuracies of M-SGD and OM-SGD are 94.606%± 0.152% and 94.728%± 0.111% with average deviations at 1.040% and 0.634%, respectively. This difference can be explained following the interpretation in Hinton (2012) that {xk} are the points after “jump” and {yk} are the points after “correction”.

heading: 3 AMORTIZED NESTEROV’S MOMENTUM
text: In this section, we formally introduce SGD with Amortized Nesterov’s Momentum (AM1-SGD) in Algorithm 1 with the following remarks:
Options It can be verified that if m = 1, AM1-SGD with Option I degenerates to M-SGD and Option II 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 I is easier to be implemented in a generic optimization layer5. Intuitively, Option I is SGD with amortized momentum and Option II applies an m-iterations tail averaging on Option I.
4Ma & Yarats (2019) observed that when effective learning rates γ = η(1 − β)−1 are fixed, M-SGD and SGD have similar performance. We provide a discussion on this observation in Appendix C.1.
5To implement Option II, we can either maintain another identical network for the shifted point x̃ or temporarily change the network parameters in the evaluation phase.
Algorithm 1 AM1-SGD Input: Initial guess x0, learning rate η, momentum β, amortization length m, iteration number K. Initialize: x← x0, x̃← x0, x̃+ ← 0 {a running average}.
1: for k = 0, . . . ,K − 1 do 2: x← x− η · ∇fik(x). 3: x̃+ ← x̃+ + 1m · x. 4: if (k + 1) mod m = 0 then 5: x← x+ β · (x̃+ − x̃). {adding amortized momentum} 6: x̃← x̃+, x̃+ ← 0. 7: end if 8: end for
Output: Option I: x, Option II: x̃. * The symbol ‘←’ denotes assignment.
Efficiency We can improve the efficiency of Algorithm 1 by maintaining a running scaled momentum ṽ+ , m · (x̃+ − x̃) instead of the running average x̃+, by replacing the following steps in Algorithm 1:
Initialize: x← x0, x̃← x0, ṽ+ ← −m · x0, Step 3: ṽ+ ← ṽ+ + x. Step 5: x← x+ (β/m) · ṽ+. Step 6: x̃← x̃+ (1/m) · ṽ+, ṽ+ ← −m · x̃.
Then, in one m-iterations loop, for each of the first m − 1 iterations, AM1-SGD requires 1 vector addition and 1 scaled vector addition. At the m-th iteration, it requires 1 vector addition, 1 scalarvector multiplication and 3 scaled vector additions. In comparison, M-SGD (standard PyTorch) requires 1 vector addition, 1 (in-place) scalar-vector multiplication and 2 scaled vector additions per iteration. Thus, as long as m > 2, AM1-SGD has lower amortized cost than M-SGD. For memory complexity, AM1-SGD requires one more auxiliary buffer than M-SGD.
Tuning m We did a parameter sweep for m in our basic case study. We plot the final and the average deviation of test accuracies over 5 runs againstm in Figure 2a. Note that m=1 corresponds to the results of M-SGD and OM-SGD, which are already given in Figure 1. From this empirical result, m introduces a trade-off between final accuracy and robustness (the convergence behaviors can be found in Appendix A.1). Figure 2a suggests that m= 5 is a good choice for this task. For simplicity, and also as a recommended setting, we fix m=5 for the rest of experiments in this paper.
A momentum that increases robustness To provide a stronger justification, we ran 20 seeds with m = 5 in Figure 2b and the detailed data are 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 can be understood as the effect of tail averaging. However, the large gap between Option I and SGD is somewhat mysterious: what Option I does is to inject a very large momentum6 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 a simple interpolation in-between M-SGD and SGD.
6Amortized momentum β ·(x̃+−x̃) is expected to be much large than Nesterov’s momentum β ·(yk+1−yk).
Learning rate scheduler issue We observed that when we use schedulers with a large decay factor and the momentum β 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 it is caused by the different cardinalities of iterates being averaged in x̃+, which leads to a false momentum. This issue is resolved by restarting the algorithm after each learning rate reduction inspired by (O’donoghue & Candes, 2015). We include more discussion and evidence in Appendix A.4.

heading: 3.1 AM2-SGD: A VARIANT WITH IDENTICAL ITERATIONS
text: Algorithm 2 AM2-SGD
Input: Initial guess x0, amortization lengthm, a point table φ = [φ1 · · · φm] ∈ Rd×m, learning rate η, momentum β, iteration number K. Initialize: φ0j = x0,∀j ∈ [m]*. {jk | jk ∈ [m]} K−1 k=0 is a sequence of uniformly random indexes.
If Option II is used, φ̄0 = x0. {a running average for the point table φ} 1: for k = 0, . . . ,K − 1 do 2: φk+1jk = xk − η · ∇fik(xk) and keep other entries unchanged (i.e., φ k+1 j = φ k j for j 6= jk). 3: xk+1 = φ k+1 jk + β · (φk+1jk+1 − φ k jk
). {adding amortized momentum} 4: if Option II then φ̄k+1 = φ̄k + 1m · ( φk+1jk − φ k jk ) .
5: end for Output: Option I (not recommended): xK , Option II: φ̄K . * [m] denotes the set {1, . . . ,m}.
While enjoying an improved efficiency, AM1-SGD does not have identical iterations7, which to some extent limits its extensibility to other settings (e.g., asynchronous setting). In this section, we propose a variant of Amortized Nesterov’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 to AM1-SGD. The relation between AM1-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, we can either maintain a “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 AM1-SGD, if 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 AM1-SGD and it also has the learning rate scheduler issue. In our preliminary experiments (can be found in Appendix A), we observed that Option I is consistently worse than Option II and it does not seem to benefit from increasingm. Thus, we do not recommend using Option I. We also set m = 5 for AM2-SGD for its evaluation due to the similarity.
7For AM1-SGD, the workload varies for different iteration k due to the if-clause at Step 4.
Additional randomness {jk} In our implementation, at each iteration, we sample an 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 AM1-SGD, which we believe is caused by the additional random indexes {jk}.

heading: 4 CONVERGENCE RESULTS
text: The original Nesterov’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-SGD in the convex case, which explicitly model the effect of amortization (i.e., m). While these rates do not hold for deep learning problems in 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 method is deterministic (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 x∈X
{ F (x) , f(x) + h(x) } , (4)
whereX ⊆ Rd is a non-empty closed convex set and h is a proper convex function with its proximal operator proxαh(·)8 available. We impose the following assumptions on the regularity of f and the stochastic oracle∇fi (identical to the ones in Ghadimi & Lan (2012) with µ = 0): Assumptions. For some L ≥ 0,M ≥ 0, σ ≥ 0,
(a) 0 ≤ f(y)− f(x)− 〈∇f(x), y − x〉 ≤ L2 ‖y − x‖ 2 +M ‖y − x‖ ,∀x, y ∈ X.9
(b) Ei [∇fi(x)] = ∇f(x),∀x ∈ X. (c) Ei [ ‖∇fi(x)−∇f(x)‖2 ] ≤ σ2,∀x ∈ X.
The notation Eik [ · ] is E [ · | (i0, . . . , ik−1)] for a random process i0, i1, . . .. These assumptions cover several important classes of convex problems. For example, (a) covers the cases of f being L-smooth (M = 0) or L0-Lipschitz continuous (M = 2L0, L = 0) convex functions and if σ = 0 in (c), the assumptions cover several classes of deterministic convex programming problems. We denote x? ∈ X as a solution to problem (4) and x0 ∈ X as the initial guess. Unlike its usage in deep learning, the momentum parameter β is always a variable in general convex analysis. For the simplicity of analysis, we reformulate AM1-SGD (Algorithm 1) and AM2-SGD (Algorithm 2) into the following schemes10(z ∈ X,α ∈ R):
AM1-SGD (reformulated, proximal)
Initialize: x̃0 = z0 = x0, S = K/m. 1: for s = 0, . . . , S − 1 do 2: for j = 0, . . . ,m− 1 do 3: k = sm+ j. 4: xk = (1− βs) · zk + βs · x̃s. 5: zk+1 = proxαsh {zk − αs ·∇fik(xk)}. 6: (xk+1 = (1− βs) · zk+1 + βs · x̃s.) 7: end for 8: x̃s+1 = 1 m ∑m j=1 xsm+j .
9: end for Output: x̃S .
AM2-SGD (reformulated, proximal)
Initialize: z0 = φ0j = x0,∀j ∈ [m]. 1: for k = 0, . . . ,K − 1 do 2: Sample jk uniformly in [m]. 3: xjkk = (1− βk) · zk + βk · φkjk . 4: zk+1 = proxαkh {zk − αk ·∇fik(x jk k )}.
5: φk+1jk = (1− βk) · zk+1 + βk · φ k jk . 6: end for
Output: φ̄K = 1m ∑m j=1 φ K j .
We show in Appendix B.1 that when h ≡ 0 and β is a constant, the reformulated schemes AM1SGD and AM2-SGD are equivalent to Algorithm 1 and Algorithm 2 through αs = η(1−βs)−1 and
8∀x ∈ Rd, proxαh(x) , argminu∈X { 1 2 ‖u− x‖2 + αh(u) } , see Parikh et al. (2014).
9When M > 0, f is not necessarily differentiable and we keep using the notation ∇f(x) to denote an arbitrary subgradient of f at x for consistency.
10For simplicity, we assume K is divisible by m.
αk = η(1 − βk)−1. These reformulations are basically how Nesterov’s momentum was migrated into deep learning (Sutskever et al., 2013). Then we establish the convergence rates for AM1-SGD and AM2-SGD as follows. All the proofs in this paper are given in Appendix B.2. Theorem 1. For the reformulated AM1-SGD, suppose we choose
βs = s
s+ 2 and αs = λ1 L(1− βs)
with λ1 = min { 2
3 , L ‖x0 − x?‖ 2 √ m √ σ2 +M2(S + 1) 3 2
} . (5)

heading: Then,
text: (a) The output x̃S satisfies
E [F (x̃S)]− F (x?) ≤ 3Lm ‖x0 − x?‖2
(K +m)2 +
8 ‖x0 − x?‖ √ σ2 +M2√
K +m , K0(m).
(b) If the variance has a “light tail”, i.e., Ei [ exp { ‖∇fi(x)−∇f(x)‖2/σ2 }] ≤exp{1},∀x ∈
X , and X is compact, denoting DX , maxx∈X ‖x− x?‖, for any Λ ≥ 0, we have
Prob { F (x̃S)− F (x?) ≤ K0(m) + 4Λσ ( 3 ‖x0 − x?‖+ √ 6DX ) 3 √ K +m } ≥ 1− ( exp{−Λ2/3}+ exp{−Λ} ) .
Remarks: (a) Regarding K0(m), its minimum is obtained at either m = 1 or m = K. Note that for AM1-SGD,m is strictly constrained in {1, . . . ,K}. It can be verified that whenm = K, AM1-SGD 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 convergence rate in Theorem 1a becomes the corresponding O(L/K + (σ +M)/ √ K) (cf. Theorem 1 in Lan (2012)). Thus, we can regard AM1-SGD as a smooth transition between AC-SA and the modified 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.
Understandings: Theorem 1a gives the expected performance in terms of full-batch loss F (x̃) − F (x?), from which the trade-off of m is clear: Increasing m improves the dependence on variance σ but deteriorates the O(L/K2) term (i.e., the acceleration). Based on this trade-off, we can understand the empirical results in Figure 2b: the faster convergence in the early stage could be the result of a 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.e., K0(m)). It is clear that increasing m 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, we did an experiments based on this
metric in Figure 4 & Table 2. Here we choose training a smaller ResNet18 with pre-activation (He et al., 2016a) on CIFAR-10 as the case study (the test accuracy is reported in Appendix A.5).
For AM2-SGD, we only give the expected convergence results as follows. Theorem 2. For the reformulated AM2-SGD, if we choose
βk = k/m
k/m+ 2 and αk = λ2 L(1− βk) with λ2 = min 23 , L ‖x0 − x?‖√2m(σ +M) (K−1m + 2) 32 ,
the output φ̄K satisfies E [ F (φ̄K) ] − F (x?)≤ 4(m2−m) ( F (x0)−F (x?) ) +3Lm ‖x0−x?‖2
(K + 2m− 1)2 +
4 √
2 ‖x0−x?‖ (σ+M)√ K + 2m− 1 .
Remark: In comparison with Theorem 1a, Theorem 2 has an additional term F (x0)− F (x?) in the upper bound, which is inevitable. This difference comes from different restrictions on the choice of m. For AM2-SGD,m ≥ 1 is the only requirement. Since it is impossible to letm K to obtain an improved rate, this additional term is inevitable. As a sanity check, we can let m → ∞ to obtain a point table with almost all x0, 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.
Understanding: Comparing the rates, we see that when using the same m, AM2-SGD has slightly better dependence on σ, which is related to the observation in Figure 5 that AM2-SGD is always slightly faster than AM1-SGD. This difference is suggesting that randomly incorporating past iterates beyond m iterations helps. If m = O(1), Theorems 1 and 2 establish the optimal O(L/K2 + (σ + M)/ √ K) rate in the convex setting (see Lan (2012) for optimality), which verifies AM1-SGD and AM2-SGD as variants of the Nesterov’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 in 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 AM1-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. AM1-SGD also uses an aggregated point to provide momentum and it shares many structural similarities with Katyusha. We refer the interested readers to Appendix B.3.

heading: 5 PERFORMANCE EVALUATION
text: In this section, we evaluate AM1-SGD and AM2-SGD on more deep learning tasks. Our goal is to show their potentials of serving as alternatives for M-SGD. Regarding the options: for AM1-SGD, Option I is a nice choice, which has slightly better final performance as shown in Table 1; for AM2SGD, Option I is not recommended as mentioned before. Here we choose to evaluate Option II 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 (η, β) as M-SGD, which was tuned to optimize the performance of M-SGD. We set m = 5 for AM1-SGD and AM2-SGD.
We trained ResNet50 and ResNet152 (He et al., 2016b) on the ILSVRC2012 dataset (“ImageNet”) (Russakovsky et al., 2015) shown in Figure 5b. For this task, we used 0.1 initial learning rate and 0.9 momentum for all methods, which is a typical choice. We performed a restart after each learning rate reduction as 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 on Penn Treebank dataset (Marcus et al., 1993). We used the LSTM (Hochreiter & Schmidhuber, 1997) model defined in Merity et al. (2017) and followed the experimental setup in its released code. We only changed the learning rate and momentum in
the setup. The baseline is SGD+ASGD11 (Polyak & Juditsky, 1992) with constant learning rate 30 as used in Merity et al. (2017). For the choice of (η, β), 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 {5, 2.5, 1, 0.1, 0.01} and found that η = 2.5 resulted in the lowest validation perplexity for M-SGD. We thus ran AM1-SGD and AM2-SGD with this (η, β) and m = 5. Due to the small decay factor, we did not restart AM1-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 β (Appendix A.2), a CIFAR-100 experiment (Appendix A.6) and comparison with classical momentum (Polyak, 1964), AggMo (Lucas et al., 2019) and QHM (Ma & Yarats, 2019) (Appendix A.3).

heading: 6 CONCLUSIONS
text: We presented Amortized Nesterov’s Momentum, which is a 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 are simple to implement with little-to-no additional tuning overhead over M-SGD. Our empirical results demonstrate that switching to 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 as a favorable alternative to M-SGD in 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 in 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.
11SGD+ASGD is to run SGD and switch to averaged SGD (ASGD) when a threshold is met.

heading: Appendices

heading: A Extra Experimental Results 14
text: A.1 The effect of m 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

heading: B Missing parts in Section 4 19
text: B.1 The reformulations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 19
B.2 Proofs of Theorem 1 and Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . 20
B.2.1 Proof of Lemma 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20
B.2.2 Proof of Theorem 1a . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22
B.2.3 Proof of Theorem 1b . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23
B.2.4 Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25
B.3 Connections between AM1-SGD and Katyusha . . . . . . . . . . . . . . . . . . . 27

heading: C Miscellanies 28
text: C.1 Comparison of SGD and M-SGD . . . . . . . . . . . . . . . . . . . . . . . . . . . 28
C.2 Training evaluation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 29

heading: D Experimental Setup 29
text: D.1 Classification Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 29
D.2 Language Model Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 30

heading: A EXTRA EXPERIMENTAL RESULTS
text: In this appendix, we provide more experimental results to further evaluate the Amortized Nesterov’s Momentum. Table 4 shows the detailed data of the parameter sweep experiments, where the convergence curves of these results are given in Appendix A.1. In Appendix A.2, we compare the robustness of AM1-SGD and M-SGD on large momentum parameters. In Appendix A.3, we empirically 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 in 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 in Appendix A.7.

heading: METHOD DESCRIPTION FINAL ACCURACY Avg. STD
text: A.1 THE EFFECT OF m ON CONVERGENCE
We show in Figure 6 how m affects the convergence of test accuracy. The results show that increasing m speeds up the convergence in the early stage. While for AM1-SGD the convergences of Option I and Option II are similar, AM2-SGD with Option II is consistently better than with Option I in this experiment. It seems that AM2-SGD with Option I does not benefit from increasing m and the algorithm is not robust. Thus, we do not recommend using Option I for AM2-SGD.

heading: A.2 ROBUSTNESS ON LARGE MOMENTUM PARAMETERS
text: We compare the robustness of M-SGD and AM1-SGD when β is large in Figure 7 & Table 5. For fair comparison, AM1-SGD uses Option I. As we can see, the STD error of M-SGD scales up significantly when β is larger and the performance is more affected by a large β compared with AM1-SGD.

heading: A.3 COMPARISON WITH OTHER MOMENTUM
text: In this section, we compare AM1-SGD (Option I) with classical momentum (Polyak, 1964), AggMo (Lucas et al., 2019) and QHM (Ma & Yarats, 2019) in our basic case study (training ResNet34 on
CIFAR-10). Since we are not aware of what makes a fair comparison with these methods (e.g., it is not clear what is the effective learning rate for AM1-SGD), we compare them based on the default hyper-parameter settings suggested by their papers.
Classical Momentum The SGD with classical momentum (CM-SGD) that is widely used in deep learning has the following scheme (standard PyTorch) (vcm ∈ Rd, vcm0 = 0):
vcmk+1 = β · vcmk +∇fik(xk), xk+1 = xk − η · vcmk+1, for k ≥ 0.
CM-SGD with its typical hyper-parameter settings (η0 = 0.1, β = 0.9) is observed to achieve similar generalization performance as M-SGD. However, CM-SGD is 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 is inspired by the passive damping from physics literature (Lucas et al., 2019). AggMo uses the following update rules (for t = 1, . . . , T , v(t) ∈ Rd, v(t)0 = 0):
v (t) k+1 = β (t) · v(t)k −∇fik(xk), for t = 1, . . . , T,
xk+1 = xk + η T · T∑ t=1 v (t) k+1, for k ≥ 0.
We used the exponential hyper-parameter setting recommended in the original work with the scalefactor a = 0.1 fixed, β(t) = 1− at−1, for t = 1, . . . , T and choosing T in {2, 3, 4}. We found that T = 2 gave the best performance in this experiment. As shown in Figure 8 & Table 6, with the help of passive damping, AggMo is more stable and robust compared with CM-SGD.
Quasi-hyperbolic Momentum (QHM) Ma & Yarats (2019) introduce the immediate discount factor ν ∈ R for the momentum scheme, which results in the QHM update rules (α ∈ R, vqh ∈ Rd, vqh0 = 0):
vqhk+1 = β · v qh k + (1− β) · ∇fik(xk),
xk+1 = xk − α · (ν · vqhk+1 + (1− ν) · ∇fik(xk)), for k ≥ 0. Here we used the recommended hyper-parameter setting for QHM (α0 = 1.0, β = 0.999, ν = 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 of robustness, huge gaps are observed when comparing AM1-SGD with the remaining methods in Table 6. Note that AM1-SGD is more efficient than both QHM and AggMo, 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 trainbatch loss in this experiment, which is consistent with the results reported in Ma & Yarats (2019); Lucas et al. (2019).
A.4 ISSUES WITH LEARNING RATE SCHEDULERS
We show in Figure 10 that when β is large for the task, using step learning rate scheduler with decay factor 10, a performance drop is 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. We fix this issue by performing a restart after each learning rate reduction (labeled with ‘+’). We plot the train-batch loss here because we find the phenomenon is clearer in this way. If β = 0.9, there is no observable performance drop in this experiment.
For smooth-changing schedulers such as the cosine annealing scheduler (Loshchilov & Hutter, 2016), the amortized momentum works well as shown in Figure 11.

heading: A.5 TEST ACCURACY RESULTS OF FIGURE 4 & TABLE 2
text: We report the test accuracy results of the experiments in Section 4 in Figure 12 & Table 7. These results are reminiscent of the ResNet34 experiments (Figure 3 & Table 1).

heading: A.6 CIFAR-100 EXPERIMENT
text: We report the results of training DenseNet121 (Huang et al., 2017) on CIFAR-100 in Figure 13, which shows that both AM1-SGD and AM2-SGD perform well before the final learning rate reduction. However, the final accuracies are 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
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).

heading: A.7 A SANITY CHECK
text: When m = 1, both AM1-SGD and AM2-SGD are equivalent to M-SGD, we plot their convergence in Figure 14 as a sanity check (the detailed data is given in Table 4).
We observed that when m = 1, both AM1-SGD and AM2-SGD have a 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., vptk = η −1β−1 · (yk − xk)).

heading: B MISSING PARTS IN SECTION 4

heading: B.1 THE REFORMULATIONS
text: When h ≡ 0 and β is a constant, we do the reformulations by eliminating the sequence {zk}. For the reformulated AM2-SGD,
xjkk = (1− β) · zk + β · φ k jk ,
zk+1 = zk − α · ∇fik(x jk k ),
φk+1jk = (1− β) · zk+1 + β · φ k jk ,(
x jk+1 k+1 = (1− β) · zk+1 + β · φ k+1 jk+1
) .
The reformulated AM2-SGD
α(1− β) = η Eliminate {zk} =========⇒
φk+1jk = x jk k − η · ∇fik(x jk k ),
x jk+1 k+1 = φ k+1 jk + β · ( φk+1jk+1 − φ k jk ) .
Algorithm 2
For the reformulated AM1-SGD, when h ≡ 0, the inner loops are basically SGD,
xk = (1− β) · zk + β · x̃s, zk+1 = zk − α · ∇fik(xk),
(xk+1 = (1− β) · zk+1 + β · x̃s.)
α(1− β) = η Eliminate {zk} =========⇒ xk+1 = xk − η · ∇fik(xk).
At the end of each inner loop (i.e., when (k + 1) mod m = 0), we have
x(s+1)m = (1− β) · z(s+1)m + β · x̃s,
while at the beginning of the next inner loop,
x(s+1)m = (1− β) · z(s+1)m + β · x̃s+1,
which means that we need to set xk+1 ← xk+1 + β · (x̃s+1 − x̃s) (reassign the value of xk+1). We also give the reformulation of M-SGD (scheme (1)) to the Auslender & Teboulle (2006) scheme for reference:
xk = (1− β) · zk + β · yk, zk+1 = zk − α · ∇fik(xk), yk+1 = (1− β) · zk+1 + β · yk,( xk+1 = (1− β) · zk+1 + β · yk+1 ) .
Auslender & Teboulle (2006) (AC-SA (Lan, 2012))
α(1− β) = η Eliminate {zk} =========⇒
yk+1 = xk − η · ∇fik(xk), xk+1 = yk+1 + β · (yk+1 − yk).
Nesterov (1983; 2013b)
AC-SA (in the Euclidean case) maps to the Auslender & Teboulle (2006) scheme through (in the original notations)  x = xmd z = x y = xag
1− β = β−1t α = γt
.
Intuition for the Auslender & Teboulle (2006) scheme can be found in Remark 2 in Lan (2012).

heading: B.2 PROOFS OF THEOREM 1 AND THEOREM 2
text: The reformulated schemes are copied here for reference:
AM1-SGD (reformulated, proximal)
Initialize: x̃0 = z0 = x0, S = K/m. 1: for s = 0, . . . , S − 1 do 2: for j = 0, . . . ,m− 1 do 3: k = sm+ j. 4: xk = (1− βs) · zk + βs · x̃s. 5: zk+1 = proxαsh {zk − αs ·∇fik(xk)}. 6: (xk+1 = (1− βs) · zk+1 + βs · x̃s.) 7: end for 8: x̃s+1 = 1 m ∑m j=1 xsm+j .
9: end for Output: x̃S .
AM2-SGD (reformulated, proximal)
Initialize: z0 = φ0j = x0,∀j ∈ [m]. 1: for k = 0, . . . ,K − 1 do 2: Sample jk uniformly in [m]. 3: xjkk = (1− βk) · zk + βk · φkjk . 4: zk+1 = proxαkh {zk − αk ·∇fik(x jk k )}.
5: φk+1jk = (1− βk) · zk+1 + βk · φ k jk . 6: end for
Output: φ̄K = 1m ∑m j=1 φ K j .
Comparing the reformulated schemes, we see that their iterations can be generalized as follows:
x = (1− β) · z + β · y, z+ = proxαh{z − α · ∇fi(x)}, y+ = (1− β) · z+ + β · y.
(6)
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 cornerstone for the convergence proofs of AM1-SGD and AM2-SGD.
Lemma 1. If α(1− β) < 1/L, the update scheme (6) satisfies the following recursion:
1 1− β ( F (y+)− F (x?) ) ≤ β 1− β ( F (y)− F (x?) ) + 1 2α ( ‖z − x?‖2 − ∥∥z+ − x?∥∥2) + (‖∇f(x)−∇fi(x)‖+M)2
2(α−1 − L(1− β)) + 〈∇f(x)−∇fi(x), z − x?〉 .

heading: B.2.1 PROOF OF LEMMA 1
text: This Lemma is similarly provided in Lan (2012); Ghadimi & Lan (2012) under a more general setting that allows non-Euclidean norms in the assumptions, we give a proof here for completeness.
Based on the convexity (Assumption (a)), we have
f(x)− f(x?) ≤ 〈∇f(x), x− z〉︸ ︷︷ ︸ R0 + 〈∇f(x)−∇fi(x), z − x?〉︸ ︷︷ ︸ R1
+ 〈 ∇fi(x), z − z+ 〉︸ ︷︷ ︸ R2
+ 〈 ∇fi(x), z+ − x? 〉︸ ︷︷ ︸ R3 . (7)
We upper bound the terms on the right side one-by-one.
For R0,
R0 (?) =
β 1− β 〈∇f(x), y − x〉 ≤ β 1− β ( f(y)− f(x) ) , (8)
where (?) uses the relation between x and z, i.e., (1− β) · (x− z) = β · (y − x). For R2, based on Assumption (a), we have
f(y+)− f(x) + 〈 ∇f(x), x− y+ 〉 ≤ L
2
∥∥x− y+∥∥2 +M ∥∥x− y+∥∥ .
Then, noting that x− y+ = (1− β) · (z − z+), we can arrange the above inequality as
R2 ≤ L(1− β)
2
∥∥z − z+∥∥2 + 1 1− β ( f(x)− f(y+) ) + 〈 ∇f(x)−∇fi(x), z+ − z 〉 +M
∥∥z − z+∥∥ ≤ L(1− β)
2
∥∥z − z+∥∥2 + 1 1− β ( f(x)− f(y+) ) + ( ‖∇f(x)−∇fi(x)‖+M ) ∥∥z − z+∥∥ . Using Young’s inequality with ζ > 0, we obtain
R2 ≤ L(1− β) + ζ
2
∥∥z − z+∥∥2 + 1 1− β ( f(x)− f(y+) ) + (‖∇f(x)−∇fi(x)‖+M)2 2ζ . (9)
For R3, based on the optimality condition of proxαh{z − α · ∇fi(x)} and denoting ∂h(z+) as a subgradient of h at z+, we have for any u ∈ X ,〈
α · ∂h(z+) + z+ − z + α · ∇fi(x), u− z+ 〉 ≥ 0,〈
∇fi(x), z+ − u 〉 ≤ 〈 ∂h(z+), u− z+ 〉 + 1
α
〈 z+ − z, u− z+ 〉 ≤ h(u)− h(z+) + 1
α
〈 z+ − z, u− z+ 〉 .
Choosing u = x?,
R3 ≤ h(x?)− h(z+) + 1
α
〈 z+ − z, x? − z+ 〉 (?) = h(x?)− h(z+) + 1
2α
( ‖z − x?‖2 − ∥∥z+ − x?∥∥2 − ∥∥z+ − z∥∥2) , (10) where (?) follows from ‖a+ b‖2 = ‖a‖2 + ‖b‖2 + 2 〈a, b〉. Finally, by upper bounding (7) using (8), (9), (10), we conclude that
f(x)− f(x?) ≤ R1 + β 1− β ( f(y)− f(x) ) + L(1− β) + ζ − α−1 2 ∥∥z − z+∥∥2 + 1
1− β ( f(x)− f(y+) ) + h(x?)− h(z+) + (‖∇f(x)−∇fi(x)‖+M) 2 2ζ
+ 1
2α
( ‖z − x?‖2 − ∥∥z+ − x?∥∥2) , After simplification,
1 1− β ( f(y+)− f(x?) ) ≤ β 1− β ( f(y)− f(x?) ) + L(1− β) + ζ − α−1 2 ∥∥z − z+∥∥2 + h(x?)− h(z+) + (‖∇f(x)−∇fi(x)‖+M) 2
2ζ +R1
+ 1
2α
( ‖z − x?‖2 − ∥∥z+ − x?∥∥2) . (11)
Note that with the convexity of h and y+ = (1− β) · z+ + β · y, we have h(y+) ≤ (1− β)h(z+) + βh(y),
h(z+) ≥ 1 1− β h(y+)− β 1− β h(y).
Using the above inequality and choosing ζ = α−1 − L(1 − β) > 0 ⇒ α(1 − β) < 1/L, we can arrange (11) as
1 1− β ( F (y+)− F (x?) ) ≤ β 1− β ( F (y)− F (x?) ) + 1 2α ( ‖z − x?‖2 − ∥∥z+ − x?∥∥2) + (‖∇f(x)−∇fi(x)‖+M)2
2(α−1 − L(1− β)) +R1.

heading: B.2.2 PROOF OF THEOREM 1A
text: Using Assumption (c), Lemma 1 with  x = xk z = zk z+ = zk+1 y = x̃s y+ = xk+1 α = αs β = βs , (12)
and taking expectation, if αs(1− βs) < 1/L, we have 1 1− βs ( Eik [F (xk+1)]− F (x?) ) + 1 2αs Eik [ ‖zk+1 − x?‖2 ] ≤ βs
1− βs ( F (x̃s)− F (x?) ) + 1 2αs ‖zk − x?‖2 +
(σ +M)2
2(α−1s − L(1− βs)) .
Summing the above inequality from k = sm, . . . , sm+m− 1, we obtain
1
(1− βs)m m∑ j=1 ( E [F (xsm+j)]− F (x?) ) + 1 2αsm E [∥∥z(s+1)m − x?∥∥2]
≤ βs 1− βs
( F (x̃s)− F (x?) ) + 1
2αsm ‖zsm − x?‖2 +
(σ +M)2
2(α−1s − L(1− βs)) ,
Using the definition of x̃s+1 and convexity,
αs 1− βs
( E [F (x̃s+1)]− F (x?) ) + 1 2m E [∥∥z(s+1)m − x?∥∥2]
≤ αsβs 1− βs
( F (x̃s)− F (x?) ) + 1
2m ‖zsm − x?‖2 +
αs(σ 2 +M2)
α−1s − L(1− βs) .
(13)
It can be verified that with the choices βs = ss+2 and αs = λ1 L(1−βs) , the following holds for s ≥ 0,
αs+1βs+1 1− βs+1 ≤ αs 1− βs and β0 = 0. (14)
Note that since our analysis aims at providing intuition, we do not refine the choice of αs as in (Hu et al., 2009; Ghadimi & Lan, 2012). Thus, by telescoping (13) from s = S − 1, . . . , 0, we obtain
αS−1 1− βS−1
( E [F (x̃S)]− F (x?) ) + 1 2m E [ ‖zSm − x?‖2 ] ≤ 1
2m ‖x0 − x?‖2 + S−1∑ s=0 αs(σ 2 +M2) α−1s − L(1− βs) ,
and thus,
E [F (x̃S)]− F (x?) ≤ 2L
λ1m(S + 1)2 ‖x0 − x?‖2 +
4L(σ2 +M2)
λ1(S + 1)2 S−1∑ s=0 α2s 1− αs(1− βs)L
(a) ≤ 2L λ1m(S + 1)2 ‖x0 − x?‖2 + 3λ1(σ 2 +M2) L(S + 1)2 S−1∑ s=0 (s+ 2)2
(b) ≤ 2L λ1m(S + 1)2 ‖x0 − x?‖2 + 8λ1(σ 2 +M2)(S + 1) L ,
where (a) follows from λ1 ≤ 23 and (b) holds because 0 ≤ x 7→ (x+2) 2 is non-decreasing and thus
S−1∑ s=0 (s+ 2)2 ≤ ∫ S 0 (x+ 2)2dx ≤ (S + 2) 3 3 ≤ 8(S + 1) 3 3 .
Denoting
λ?1 , L ‖x0 − x?‖
2 √ m √ σ2 +M2(S + 1) 3 2 ,
and based on the choice of λ1 = min { 2 3 , λ ∗ 1 } , if λ∗1 ≤ 23 , we have
E [F (x̃S)]− F (x?) ≤ 8 ‖x0 − x?‖
√ σ2 +M2
m 1 2 (S + 1) 1 2
.
If λ∗1 > 2 3 ,
E [F (x̃S)]− F (x?) ≤ 3L ‖x0 − x?‖2
m(S + 1)2 +
4 ‖x0 − x?‖ √ σ2 +M2
m 1 2 (S + 1) 1 2
.
Thus, we conclude that
E [F (x̃S)]− F (x?) ≤ 3L ‖x0 − x?‖2
m(S + 1)2 +
8 ‖x0 − x?‖ √ σ2 +M2
m 1 2 (S + 1) 1 2
.
Substituting S = K/m completes the proof.

heading: B.2.3 PROOF OF THEOREM 1B
text: In order to prove Theorem 1b, we need the following known result for the martingale difference (cf. Lemma 2 in Lan et al. (2012)):
Lemma 2. With N > 0, let ξ0, ξ1, . . . , ξN−1 be a sequence of i.i.d. random variables, for t = 0, . . . , N − 1, σt > 0 be a deterministic number and ψt = ψt(ξ0, . . . , ξt) be a deterministic measurable function such that Eξt [ψt] = 0 a.s. and Eξt [ exp{ψ2t /σ2t } ] ≤ exp{1} a.s.. Then for any Λ ≥ 0,
Prob  N−1∑ t=0 ψt ≥ Λ √√√√N−1∑ t=0 σ2t  ≤ exp{−Λ2/3}. To start with, using Lemma 1 with the parameter mapping (12), we have
1 1− βs ( F (xk+1)− F (x?) ) + 1 2αs ‖zk+1 − x?‖2
≤ βs 1− βs
( F (x̃s)− F (x?) ) + 1
2αs ‖zk − x?‖2
+ (‖∇f(xk)−∇fik(xk)‖+M)2
2(α−1s − L(1− βs)) + 〈∇f(xk)−∇fik(xk), zk − x?〉
≤ βs 1− βs
( F (x̃s)− F (x?) ) + 1
2αs ‖zk − x?‖2 +
M2
α−1s − L(1− βs)
+ ‖∇f(xk)−∇fik(xk)‖
2
α−1s − L(1− βs) + 〈∇f(xk)−∇fik(xk), zk − x?〉 .
Summing the above inequality from k = sm, . . . , sm+m− 1 and using the choice αs = λ1L(1−βs) with λ1 ≤ 23 , we obtain
αs 1− βs
( F (x̃s+1)− F (x?) ) + 1
2m ∥∥z(s+1)m − x?∥∥2 ≤ αsβs
1− βs ( F (x̃s)− F (x?) ) + 1 2m ‖zsm − x?‖2 + 3α2sM2
+ 3α2s m sm+m−1∑ k=sm ‖∇f(xk)−∇fik(xk)‖ 2 + αs m sm+m−1∑ k=sm 〈∇f(xk)−∇fik(xk), zk − x?〉.
With our parameter choices, the relations in (14) hold and thus we can telescope the above inequality from s = S − 1, . . . , 0,
αS−1 1− βS−1
( F (x̃S)− F (x?) ) ≤ 1
2m ‖x0 − x?‖2 + 3M2 S−1∑ s=0 α2s
+ 3
m K−1∑ k=0 α2bk/mc ‖∇f(xk)−∇fik(xk)‖ 2
︸ ︷︷ ︸ R4
+ 1
m K−1∑ k=0
αbk/mc 〈∇f(xk)−∇fik(xk), zk − x?〉︸ ︷︷ ︸ R5 .
(15)
Denoting V2k , ‖∇f(xk)−∇fik(xk)‖ 2, ᾱ = ∑K−1 k=0 α 2 bk/mc = m ∑S−1 s=0 α 2 s, for R4, by Jensen’s inequality, we have
E [ exp { 1
ᾱ K−1∑ k=0 α2bk/mcV 2 k/σ 2 }] ≤ 1 ᾱ K−1∑ k=0 α2bk/mcE [ exp { V2k/σ2 }] (?) ≤ exp{1},
where (?) uses the additional assumption Eik [ exp { V2k/σ2 }] ≤ exp{1}.
Then, based on Markov’s inequality, we have for any Λ ≥ 0,
Prob { exp { 1
ᾱ K−1∑ k=0 α2bk/mcV 2 k/σ 2
} ≥ exp{Λ + 1} } ≤ exp{−Λ},
Prob { R4 ≥ (Λ + 1)σ2m
S−1∑ s=0 α2s
} ≤ exp{−Λ}. (16)
For R5, since we have Eik [ αbk/mc 〈∇f(xk)−∇fik(xk), zk − x?〉 ] = 0 and
Eik
[ exp { α2bk/mc 〈∇f(xk)−∇fik(xk), zk − x ?〉2
α2bk/mcσ 2D2X
}] ≤ Eik [ exp { V2k/σ2 }] ≤ exp{1},
which is based on the “light tail” assumption, using Lemma 2, we obtain
Prob R5 ≥ ΛσDX √√√√m S−1∑
s=0
α2s  ≤ exp{−Λ2/3}. (17) Combining (15), (16) and (17), based on the parameter setting (cf. (5)) and using the notation
K0(m) , 3Lm ‖x0 − x?‖2
(K +m)2 +
8 ‖x0 − x?‖ √ σ2 +M2√
K +m ,
R6 , 12Lσ2
λ1(S + 1)2 S−1∑ s=0 α2s + 4LσDX λ1(S + 1)2 √ m √√√√S−1∑ s=0 α2s,
we conclude that
Prob {F (x̃S)− F (x?) ≤ K0(m) + ΛR6} ≥ 1− (exp{−Λ2/3}+ exp{−Λ}).
For R6, using the choice of αs and λ1, we obtain
R6 ≤ 4 √ 6σDX
3 √ K +m
+ 8λ1σ
2(S + 1) L ≤ 4 √ 6σDX 3 √ K +m + 4σ2 ‖x0 − x?‖√ K +m √ σ2 +M2
≤ 4σ ( 3 ‖x0 − x?‖+ √ 6DX ) 3 √ K +m ,
which completes the proof.

heading: B.2.4 PROOF OF THEOREM 2
text: Using Assumption (c), Lemma 1 with  x = xjkk z = zk z+ = zk+1 y = φkjk y+ = φk+1jk α = αk β = βk ,
and taking expectation, if αk(1− βk) < 1/L, we have 1
1− βk Eik,jk
[ F (φk+1jk )− F (x ?) ] + 1
2αk Eik,jk
[ ‖zk+1 − x?‖2 ] ≤ βk
1− βk Ejk
[ F (φkjk)− F (x ?) ] + 1
2αk ‖zk − x?‖2 +
(σ +M)2
2(α−1k − L(1− βk)) .
(18)
Note that
Eik,jk [ F (φk+1jk )− F (x ?) ]
= Eik,jk  m∑ j=1 ( F (φk+1j )− F (x ?) ) − m∑ j 6=jk ( F (φkj )− F (x?) ) = Eik,jk  m∑ j=1 ( F (φk+1j )− F (x ?) )− Ejk  m∑ j 6=jk ( F (φkj )− F (x?)
) . Dividing both sides of (18) by m and then adding 1(1−βk)mEjk [∑m j 6=jk ( F (φkj )− F (x?) )] to both sides, we obtain
1
1− βk Eik,jk  1 m m∑ j=1 F (φk+1j )− F (x ?) + 1 2αkm Eik,jk [ ‖zk+1 − x?‖2 ]
≤ − 1 m Ejk [ F (φkjk)− F (x ?) ] + 1 1− βk  1 m m∑ j=1 F (φkj )− F (x?) + 1 2αkm ‖zk − x?‖2
+ (σ +M)2
2m(α−1k − L(1− βk))
= 1− 1−βkm 1− βk  1 m m∑ j=1 F (φkj )− F (x?) + 1 2αkm ‖zk − x?‖2 + (σ +M)2 2m(α−1k − L(1− βk)) . (19)
It can be verified that with our parameters choice: βk = k/m k/m+2 and αk = λ2 L(1−βk) , the following holds for k ≥ 0,
αk+1 1− 1−βk+1m 1− βk+1 ≤ αk 1− βk and β0 = 0.
Note that since our analysis aims at providing intuition, we do not refine the choice of αs as in (Hu et al., 2009; Ghadimi & Lan, 2012). Then, we can telescope (19) from k = K − 1, . . . , 0, which results in
αK−1 1− βK−1 E  1 m m∑ j=1 F (φKj )− F (x?) + 1 2m E [ ‖zK − x?‖2 ]
≤ λ2(m− 1) Lm
( F (x0)− F (x?) ) + 1
2m ‖x0 − x?‖2 + K−1∑ k=0
αk(σ +M) 2
2m(α−1k − L(1− βk)) .
Using the definition of φ̄K and convexity, we obtain
E [ F (φ̄K)− F (x?) ] ≤ 1− βK−1
αK−1
( λ2(m− 1)
Lm
( F (x0)− F (x?) ) + 1
2m ‖x0 − x?‖2 ) +
1− βK−1 αK−1 K−1∑ k=0
αk(σ +M) 2
2m(α−1k − L(1− βk))
(a) =
4(m− 1) ( F (x0)− F (x?) ) m ( K−1 m + 2 )2 + 2L ‖x0 − x?‖2 λ2m ( K−1 m + 2
)2 + 3λ2(σ +M) 2
2Lm ( K−1 m + 2 )2 K−1∑ k=0 ( k m + 2 )2 (b)
≤ 4(m− 1)
( F (x0)− F (x?) ) m ( K−1 m + 2 )2 + 2L ‖x0 − x?‖2 λ2m ( K−1 m + 2
)2 (20) + 4λ2(σ +M) 2 ( K−1 m + 2 ) L ,
where (a) uses λ2 ≤ 23 , (b) follows from simple integration arguments and that K m + 2 ≤ 2 ( K−1 m + 2 ) since K ≥ 1,m ≥ 1.
Based on the choice of
λ2 = min 23 , L ‖x0 − x?‖√2m(σ +M) (K−1m + 2) 32 ,
(20) can be further upper bounded as E [ F (φ̄K)− F (x?) ] ≤ 4(m− 1) ( F (x0)− F (x?) ) m ( K−1 m + 2 )2 + 3L ‖x0 − x?‖2 m ( K−1 m + 2 )2 +4√2 ‖x0 − x?‖ (σ +M) m 1 2 ( K−1 m + 2 ) 1 2 .

heading: B.3 CONNECTIONS BETWEEN AM1-SGD AND KATYUSHA
text: The discussion in this section aims to shed light on the understanding of the experimental results, which also shows some interesting relations between AM1-SGD and Katyusha.
The high level idea of Katyusha momentum is that it works as a “magnet” inside an epoch of SVRG updates, which “stabilizes” the iterates so as to 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 (cf. Lemma 2.4 and its comments in AllenZhu (2018)). In this sense, we can interpret Katyusha momentum as a variance reducer that further reduces the variance of SVRG. Below we show the similarity between the construction of Katyusha and AM1-SGD, based on which we conjecture that the amortized momentum can also reduce the variance of SGD (and thus increase the robustness). However, in theory, following a similar analysis of Katyusha, we cannot guarantee a reduction of σ in the worst case.
Deriving AM1-SGD from Katyusha Katyusha has the following scheme (non-proximal, in the original notations, σ is the strong convexity parameter, cf. Algorithm 1 with Option I in Allen-Zhu (2018))12:
Initialize: x̃0 = y0 = z0 = x0, η = 13L , ω = 1 + ασ. 1: for s = 0, . . . , S − 1 do 2: Compute and store∇f(x̃s). 3: for j = 0, . . . ,m− 1 do 4: k = sm+ j. 5: xk = τ1 · zk + τ2 · x̃s + (1− τ1 − τ2) · yk. 6: ∇̃k = ∇fik(xk)−∇fik(x̃s) +∇f(x̃s). 7: zk+1 = zk − α · ∇̃k. 8: yk+1 = xk − η · ∇̃k. 9: end for
10: x̃s+1 = (∑m−1 j=0 ω j )−1 · ∑m−1 j=0 ω
j · ysm+j+1. 11: end for Output: x̃S .
We can eliminate the sequence {zk} in this scheme. Note that in the parameter setting of Katyusha, we have η = ατ1, and thus
xk+1 = τ1 · zk+1 + τ2 · x̃s + (1− τ1 − τ2) · yk+1 = τ1 · zk − η · ∇̃k + τ2 · x̃s + (1− τ1 − τ2) · yk + (1− τ1 − τ2) · (yk+1 − yk) = xk − η · ∇̃k + (1− τ1 − τ2) · (yk+1 − yk) = yk+1 + (1− τ1 − τ2) · (yk+1 − yk).
Hence, the inner loops can be written as
yk+1 = xk − η · ∇̃k, xk+1 = yk+1 + (1− τ1 − τ2) · (yk+1 − yk),
which is the Nesterov’s scheme (scheme (1)). At the end of each inner loop (when k = sm+m−1),
x(s+1)m = τ1 · z(s+1)m + τ2 · x̃s + (1− τ1 − τ2) · y(s+1)m,
while at the beginning of the next inner loop,
x(s+1)m = τ1 · z(s+1)m + τ2 · x̃s+1 + (1− τ1 − τ2) · y(s+1)m,
which means that we need to set x(s+1)m ← x(s+1)m + τ2 · (x̃s+1 − x̃s) (reassign the value of x(s+1)m).
Then, the following is an equivalent scheme of Katyusha:
12We change the notation xk+1 to xk.
Initialize: x̃0 = y0 = x0, η = 13L , ω = 1 + ασ. 1: for s = 0, . . . , S − 1 do 2: for j = 0, . . . ,m− 1 do 3: k = sm+ j. 4: yk+1 = xk − η · ∇̃k. 5: xk+1 = yk+1 + (1− τ1 − τ2) · (yk+1 − yk). 6: end for 7: x̃s+1 = (∑m−1 j=0 ω j )−1 · ∑m−1 j=0 ω
j · ysm+j+1. 8: x(s+1)m ← x(s+1)m + τ2 · (x̃s+1 − x̃s). 9: end for
Output: x̃S .
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 ∇̃k with ∇fik(xk), set 1− τ1 − τ2 = 0, which is to eliminate Nesterov’s momentum, and use a uniform average for x̃s+1, the above scheme becomes exactly AM1-SGD (Algorithm 1).
If we only replace the SVRG estimator ∇̃k, the scheme can be regarded as adding amortized momentum to M-SGD. This scheme requires tuning the ratio of Nesterov’s momentum and amortized momentum. In our preliminary experiments, after suitable tuning, we observed some performance improvement. However, this scheme increases the complexity, which we do not consider it worthwhile.
A recent work (Zhou et al., 2018) shows that when 1− τ1− τ2 = 0, which is to solely use Katyusha momentum, one can still derive optimal rates and the algorithm is greatly simplified. Their proposed algorithm (i.e., MiG) is structurally more similar to AM1-SGD.

heading: C MISCELLANIES

heading: C.1 COMPARISON OF SGD AND M-SGD
text: Ma & Yarats (2019) normalize the momentum buffer of M-SGD, which results in the following formulation (α ∈ R, vqh ∈ Rd, vqh0 = 0):
vqhk+1 = β · v qh k + (1− β) · ∇fik(xk),
xk+1 = xk − α · (β · vqhk+1 + (1− β) · ∇fik(xk)), for k ≥ 0.
This scheme is equivalent to the PyTorch formulation (scheme (3)) through vptk = (1 − β)−1 · v qh k and η = α(1− β). Based on this formulation, α is understood as 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 α is set to a very large O(k/L), which can be derived from Theorem 1 or Theorem 2 by setting σ = 0,M = 0,m = 1 (then λ1 or λ2 is always 23 since the other term is ∞). If we fix α = 23L 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 βk = kk+2 , it has the convergence rate (applying Lemma 1):
1 1− βk ( F (yk+1)− F (x?) ) + 3L 4 ‖zk+1 − x?‖2 ≤ βk 1− βk ( F (yk)− F (x?) ) + 3L 4 ‖zk − x?‖2 ,
F (yK)− F (x?) ≤ 3L ‖x0 − x?‖2
2(K + 1) .
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 a legal way to adopt a large α that breaks the 1/L limitation. If GD uses the same large α, we would expect it to be unstable and potentially diverge. In this sense, Nesterov’s momentum can be understood as a “stabilizer”. In our 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 5 runs. The significance of QHM (Ma & Yarats, 2019) is that with suitable tuning, it achieves much faster convergence without changing the effective learning rate. Our work uses the convergence behavior of SGD as a reference to reflect and to understand the features of our proposed momentum, which is why we set γ = η.

heading: C.2 TRAINING EVALUATION
text: Due to the mechanism of back-propagation, evaluating train-batch loss basically incurs no overhead. It can be efficiently used to indicate the training progress. However, it 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 et al., 2014) is used in 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 (the experiment in Section 4) in Figure 15.
We also notice that for different optimizers, even if their convergences on train-batch loss are indistinguishable, 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 4 with m = 10.

heading: D EXPERIMENTAL SETUP
text: All of our experiments were conducted using PyTorch (Paszke et al., 2017) library.

heading: D.1 CLASSIFICATION SETUP
text: CIFAR-10 & CIFAR-100 Our implementation (e.g., ResNet and DenseNet implementations, data pre-processing) generally follows the repository https://github.com/kuangliu/ 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 (Ioffe & Szegedy, 2015). Data augmentation includes random 32-pixel crops with a padding of 4-pixel and random horizontal flips with
0.5 probability. We used step (or multi-step) learning rate scheduler with a decay factor 10. For the CIFAR-10 experiments, we 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 et al., 2017).
ImageNet In the ImageNet experiments, we tried both ResNet50 and ResNet152 (He et al., 2016b). The training strategy is the same as the PyTorch’s official repository https://github. com/pytorch/examples/tree/master/imagenet, which uses a batch size of 256. The learning rate starts at 0.1 and decays by a 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, we applied 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.

heading: D.2 LANGUAGE MODEL SETUP
text: We followed the implementation in the repository https://github.com/salesforce/ awd-lstm-lm and trained word level Penn Treebank with LSTM without fine-tuning or continuous cache pointer augmentation for 750 epochs. The experiments were conducted on a single RTX2080Ti. We used the default hyper-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 a 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 embedding 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-regularization on the RNN activation is 2.0, the amount of slowness regularization applied on the RNN activation is 1.0 and all weights receive a weight decay of 0.0000012.

Review summaries:
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.