Year: 2020

Abstract:
Adam is a widely used stochastic optimization method for deep learning applications. While practitioners prefer Adam because it requires less parameter tuning, its use is problematic from a theoretical point of view since it may not converge. Variants of Adam have been proposed with provable convergence guarantee, but they tend not be competitive with Adam on the practical performance. In this paper, we propose a new method named Adam (pronounced as Adam-plus). Adam retains some of the key components of Adam but it also has several noticeable differences: (i) it does not maintain the moving average of second moment estimate but instead computes the moving average of first moment estimate at extrapolated data points; (ii) its adaptive step size is formed not by dividing the square root of second moment estimate but instead by dividing the root of the norm of first moment estimate. As a result, Adam requires few parameter tuning, as Adam, but it enjoys a provable convergence guarantee. Our analysis further shows that Adam enjoys adaptive variance reduction, i.e., the variance of the stochastic gradient estimator reduces as the algorithm converges, hence enjoying an adaptive convergence. We also propose a more general variant of Adam with different adaptive step sizes and establish their fast convergence rate. Our empirical studies on various deep learning tasks, including image classification, language modeling, and automatic speech recognition, demonstrate that Adam significantly outperforms Adam and achieves comparable performance with besttuned SGD and momentum SGD.

heading: 1 INTRODUCTION
text: Adaptive gradient methods (Duchi et al., 2011; McMahan & Streeter, 2010; Tieleman & Hinton, 2012; Kingma & Ba, 2014; Reddi et al., 2019) are one of the most important variants of Stochastic Gradient Descent (SGD) in modern machine learning applications. Contrary to SGD, adaptive gradient methods typically require little parameter tuning still retaining the computational efficiency of SGD. One of the most used adaptive methods is Adam (Kingma & Ba, 2014), which is considered by practitioners as the de-facto default optimizer for deep learning frameworks. Adam computes the update for every dimension of the model parameter through a moment estimation, i.e., the estimates of the first and second moments of the gradients. The estimates for first and second moments are updated using exponential moving averages with two different control parameters. These moving averages are the key difference between Adam and previous adaptive gradient methods, such as Adagrad (Duchi et al., 2011).
Although Adam exhibits great empirical performance, there still remain many mysteries about its convergence. First, it has been shown that Adam may not converge for some objective functions (Reddi et al., 2019; Chen et al., 2018b). Second, it is unclear what is the benefit that the moving average brings from theoretical point of view, especially its effect on the convergence rate. Third, it has been empirically observed that adaptive gradient methods can have worse generalization performance than its non-adaptive counterpart (e.g., SGD) on various deep learning tasks due to the coordinate-wise learning rates (Wilson et al., 2017).
The above issues motivate us to design a new algorithm which achieves the best of both worlds, i.e., provable convergence with benefits from the moving average and enjoying good generalization performance in deep learning. Specifically, we focus on the following optimization problem:
min w∈Rd F (w),
where we only have access to stochastic gradients of F . Note that F could possibly be nonconvex in w. Due to the non-convexity, our goal is to design a stochastic first-order algorithm to find the -stationary point, i.e., finding w such that E [‖∇F (w)‖] ≤ , with low iteration complexity. Our key contribution is the design and analysis of a new stochastic method named Adam+. Adam+ retains some of the key components of Adam but it also has several noticeable differences: (i) it does not maintain the moving average of second moment estimate but instead computes the moving average of first moment estimate at extrapolated points; (ii) its adaptive step size is formed not by dividing the square root of coordinate-wise second moment estimate but instead by dividing the root of the norm of first moment estimate. These features allow us to establish the adaptive convergence of Adam+. Different from existing adaptive methods where the adaptive convergence depends on the growth rate of stochastic gradients (Duchi et al., 2011; McMahan & Streeter, 2010; Kingma & Ba, 2014; Luo et al., 2019; Reddi et al., 2019; Chen et al., 2018a;b; Ward et al., 2019; Li & Orabona, 2019; Chen et al., 2019), our adaptive convergence is due to the adaptive variance reduction property of our first order moment estimate. In existing literature, the variance reduction is usually achieved by large mini-batch (Goyal et al., 2017) or recursive variance reduction (Fang et al., 2018; Zhou et al., 2018; Pham et al., 2020; Cutkosky & Orabona, 2019). In contrast, we do not necessarily require large minibatch or computing stochastic gradients at two points per-iteration to achieve the variance reduction. In addition, we also establish a fast rate that matches the state-of-theart complexity under the same conditions of a variant of Adam+. Table 1 provides an overview of our results and a summary of existing results. We refer readers to Section F for a comprehensive survey of other related work. We further corroborate our theoretical results with an extensive empirical study on various deep learning tasks.
Our contributions are summarized below.
• We propose a new algorithm with adaptive step size, namely Adam+, for general nonconvex optimization. We show that it enjoys a new type of data-dependent adaptive convergence that depends on the variance reduction property of first moment estimate. Notably, this data-dependent complexity does not require the presence of sparsity in stochastic gradients to guarantee fast convergence as in previous works (Duchi et al., 2011; Kingma & Ba, 2014; Reddi et al., 2019; Chen et al., 2019; 2018a). To the best of our knowledge, this is the first work establishing such new type of data-dependent complexity.
• We show that a general variant of our algorithm can achieve O( −3.5) worst-case complexity, which matches the state-of-the-art complexity guarantee under the Hessian Lipschitz assumption (Cutkosky & Mehta, 2020).
• We demonstrate the effectiveness of our algorithms on image classification, language modeling, and automatic speech recognition. Our empirical results show that our proposed algorithm consistently outperforms Adam on all tasks, and it achieves comparable performance with the best-tuned SGD and momentum SGD.
Algorithm 1 Adam+: Good default settings for the tested machine learning problems are α = 0.1, a = 1, β = 0.1, 0 = 10−8.
1: Require: α, a ≥ 1: stepsize parameters 2: Require: β ∈ (0, 1): Exponential decay rates for the moment estimate 3: Require: gt(w): unbiased stochastic gradient with parameters w at iteration t 4: Require: w0: Initial parameter vector 5: z0 = g0(w0) 6: for t = 0, . . . , T do 7: Set ηt = αβ a max(‖zt‖1/2, 0) 8: wt+1 = wt − ηtzt 9: ŵt+1 = (1− 1/β)wt + 1/β ·wt+1
10: zt+1 = (1− β)zt + βgt+1(ŵt+1) 11: end for

heading: 2 ALGORITHM AND THEORETICAL ANALYSIS
text: In this section, we introduce our algorithm Adam+ (presented in Algorithm 1) and establish its convergence guarantees. Adam+ resembles Adam in several aspects but also has noticeable differences. Similar to Adam, Adam+ also maintains an exponential moving average of first moment (i.e., stochastic gradient), which is denoted by zt, and uses it for updating the solution in line 8. However, the difference is that the stochastic gradient is evaluated on an extrapolated data point ŵt+1, which is an extrapolation of two previous updates wt and wt+1. Similar to Adam, Adam+ also uses an adaptive step size that is proportional to 1/‖zt‖1/2. Nonetheless, the difference lies at its adaptive step size is directly computed from the square root of the norm of first moment estimate zt. In contrast, Adam uses an adaptive step size that is proportional to 1/ √ vt, where vt is an exponential moving average of second moment estimate. These two key components of Adam+, i.e., extrapolation and adaptive step size from the root norm of the first moment estimate, make it enjoy two noticeable benefits: variance reduction of first moment estimate and adaptive convergence. We shall explain these two benefits later.
Before moving to the theoretical analysis, we would like to make some remarks. First, it is worth mentioning that the moving average estimate with extrapolation is inspired by the literature of stochastic compositional optimization (Wang et al., 2017). Wang et al. (2017) showed that the extrapolation helps balancing the noise in the gradients, reducing the bias in the estimates and giving a faster convergence rate. Here, our focus and analysis techniques are quite different. In fact, Wang et al. (2017) focuses on the compositional optimization while we consider a general nonconvex optimization setting. Moreover, the analysis in (Wang et al., 2017) mainly deals with the error of the gradient estimator caused by the compositional nature of the problem, while our analysis focuses on carefully designing adaptive normalization to obtain an adaptive and fast convergence rates. A similar extrapolation scheme has been also employed in the algorithm NIGT by Cutkosky & Mehta (2020). In later sections, we will also provide a more general variant of Adam+ which subsumes NIGT as a special case.
Another important remark is that the update of Adam+ is very different from the famous Nesterov’s momentum method. In Nesterov’s momentum method, the update of wt+1 uses the stochastic gradient at an extrapolated point ŵt+1 = wt+1 +γ(wt+1−wt) with a momentum parameter γ ∈ (0, 1). In contrast, in Adam+ the update of wt+1 is using the moving average estimate at an extrapolated point ŵt+1 = wt+1 + (1/β − 1)(wt+1 −wt). Finally, Adam+ does not employ coordinate-wise learning rates as in Adam, and hence it is expected to have better generalization performance according to Wilson et al. (2017).

heading: 2.1 ADAPTIVE VARIANCE REDUCTION AND ADAPTIVE CONVERGENCE
text: In this subsection, we analyze Adam+ by showing its variance reduction property and adaptive convergence. To this end, we make the following assumptions.
Assumption 1. There exists positive constants L,∆, LH , σ and an initial solution w0 such that
(i) F is L-smooth, i.e., ‖∇F (x)−∇F (y)‖ ≤ L ‖x− y‖ , ∀x,y ∈ Rd.
(ii) For ∀x ∈ Rd, we have access to a first-order stochastic oracle at time t gt(x) such that E [gt(x)] = ∇F (x), E ‖gt(x)−∇F (x)‖2 ≤ σ2.
(iii) ∇F is a LH -smooth mapping, i.e., ‖∇2F (x)−∇2F (y)‖ ≤ LH‖x− y‖,∀x,y ∈ Rd.
(iv) F (w0)− F∗ ≤ ∆ <∞, where F∗ = infw∈Rd F (w).
Remark: Assumption 1 (i) and (ii), (iv) are standard assumptions made in literature of stochastic non-convex optimization (Ghadimi & Lan, 2013). Assumption (iii) is the assumption that deviates from typical analysis of stochastic methods. We leverage this assumption to explore the benefit of moving average, extrapolation and adaptive normalization. It is also used in some previous works for establishing fast rate of stochastic first-order methods for nonconvex optimization (Fang et al., 2019; Cutkosky & Mehta, 2020) and this assumption is essential to get fast rate due to the hardness result in (Arjevani et al., 2019). It is also the key assumption for finding a local minimum in previous works (Carmon et al., 2018; Agarwal et al., 2017; Jin et al., 2017).
We might also assume that the stochastic gradient estimator in Algorithm 1 satisfies the following variance property.
Assumption 2. Assume that E[‖g0(w0) − ∇F (w0)‖2] ≤ σ20 and E[‖gt(wt) − ∇F (wt)‖2] ≤ σ2m, t ≥ 1.
Remark: When g0 (resp. gt) is implemented by a mini-batch stochastic gradient with mini-batch size S, then σ20 (resp. σ 2 m) can be set as σ
2/S by Assumption 1 (ii). We differentiate the initial variance and intermediate variance because they contribute differently to the convergence.
We first introduce a lemma to characterize the variance of the moving average gradient estimator zt.
Lemma 1. Suppose Assumption 1 and Assumption 2 hold and a ≥ 1. Then, there exists a sequence of random variables δt satisfying ‖zt −∇F (wt)‖ ≤ δt for ∀t ≥ 0,
E [ δ2t+1 ] ≤ (
1− β 2
) E [ δ2t ] + 2β2σ2m + E [ CL2H‖wt+1 −wt‖4
β3 ] ≤ (
1− β 2
) E [ δ2t ] + 2β2σ2m + E [ CL2Hα 4β4a−3‖zt‖2 ] ,
where C = 1944.
Remark: Note that δt is an upper bound of ‖zt − ∇F (wt)‖, the above lemma can be used to illustrate the variance reduction effect for the gradient estimator zt. To this end, we can bound ‖zt‖2 ≤ 2δ2t + 2‖∇F (wt)‖2, then the term CL2α4β4a−3δ2t can be canceled with −β/4δ2t with small enough α. Hence, we have Eδ2t+1 ≤ (1 − β/4)E[δ2t ] + 2β2σ2m + cE[‖∇F (wt)‖2] with a small constant c. As the algorithm converges with E[‖∇F (wt)‖2] and β decreases to zero, the variance of zt will also decrease. Indeed, the above recursion of zt’s variance resembles that of the recursive variance reduced gradient estimators (e.g., SPIDER (Fang et al., 2018), STORM (Cutkosky & Orabona, 2019)). The benefit of using Adam+ is that we do not need to compute stochastic gradient twice at each iteration.
We can now state our convergence rates for Algorithm 1.
Theorem 1. Suppose Assumption 1 and Assumption 2 hold. Suppose ‖∇F (w)‖ ≤ G for any w ∈ Rd. By choosing the parameters such that α4 ≤ 1
36CL2H , α ≤ 14L , a = 1 and 0 = β a, we have
1
T T∑ t=1 E ‖∇F (wt)‖2 ≤ GE
[∑T t=1 ‖zt‖ ] T + ∆ αT + 18σ20 βT + 30βσ2m . (1)
In addition, suppose the initial batch size is T0 and the intermediate batch size is m, and choose β = T−b with 0 ≤ b ≤ 1, we have
1
T T∑ t=1
E ‖∇F (wt)‖2 ≤ E [ G ∑T t=1 ‖zt‖ ] T + ∆ αT + 18σ2 T 1−bT0 + 30σ2 mT b . (2)
Theorem 2. Suppose Assumption 1 and Assumption 2 hold. By choosing parameters such that 640α3L
3/2 H ≤ 1/120, a = 1, 0 = 0, β = 1/T s with s = 2/3 then it takes T = O
( −4.5 ) number
of iterations to ensure that
1
T T∑ t=1 E [ ‖∇F (wt)‖3/2 ] ≤ 3/2, 1 T E [ T∑ t=1 δ 3/2 t ] ≤ 3/2 .
Remarks: • From Theorem 1, we can observe that the convergence rate of Adam+ crucially depends on
the growth rate of E [∑T t=1 ‖zt‖ ] , which gives a data-dependent adaptive complexity. If
E [∑T t=1 ‖zt‖ ] ≤ Tα with α < 1, then the algorithm converges. Smaller α implies faster
convergence. Our goal is to ensure that 1T ∑T t=1 E‖∇F (wt)‖2 ≤ 2. Choosing b = 1−α,
m = O(1) and T0 = T 1−α = O( −2), and we end up with T = O ( − 2 1−α ) complexity.
• Theorem 2 shows that in the ergodic sense, the Algorithm Adam+ always converges, and the variance gets smaller when the number of iteration gets larger. Theorem 2 rules out the case that the magnitude of zt converges to a constant and the bound (2) in Theorem 1 becomes vacuous.
• To compare with Adam-style algorithms (e.g., Adam, AdaGrad), these algorithms’ convergence depend on the growth rate of stochastic gradient, i.e., ∑d i=1 ‖g1:T,i‖/T , where
g1:T,i = [g1,i, . . . , gT,i] denotes the i-th coordinate of all historical stochastic gradients. Hence, the data determines the growth rate of stochastic gradient. If the stochastic gradients are not sparse, then its growth rate may not be slow and these Adam-style algorithms may suffer from slow convergence. In contrast, for Adam+ the convergence can be accelerated by the variance reduction property. Note that we have E [∑T t=1 ‖zt‖ ] /T ≤
E [∑T t=1(δt + ‖∇F (wt)‖) ] /T . Hence, Adam+’s convergence depends on the variance reduction property of zt.
2.2 A GENERAL VARIANT OF ADAM+: FAST CONVERGENCE WITH LARGE MINI-BATCH
Next, we introduce a more general variant of Adam+ by making a simple change. In particular, we keep all steps the same as in Algorithm 1 except the adaptive step size is now set as ηt = αβa
max(‖zt‖p, 0) , where p ∈ [1/2, 1) is parameter. We refer to this general variant of Adam +
as power normalized Adam+ (Nadam+). This generalization allows us to compare with some existing methods and to establish fast convergence rate. First, we notice that when setting p = 1 and a = 5/4 and β = 1/T 4/7, Nadam+ is almost the same as the stochastic method NIGT (Cutkosky & Mehta, 2020) with only some minor differences. However, we observed that normalizing by ‖zt‖ leads to slow convergence in practice, so we are instead interested in p < 1. Below, we will show that NAdam+ with p < 1 can achieve a fast rate of 1/ 3.5, which is the same as NIGT. Theorem 3. Under the same assumption as in Theorem 1, further assume σ20 = σ2/T0 and σ2m = σ2/m. By using the step size ηt = αβ 4/3
max(‖zt‖2/3, 0) in Algorithm 1 with CL2α4 ≤ 1/14, 0 = 2β4/3,
in order to have E [‖∇F (wτ )‖] ≤ for a randomly selected solution wτ from {w1, . . . ,wT }, it suffice to set β = O( 1/2), T = O( −2), the initial batch size T0 = 1/β = O( −1/2), the intermediate batch size as m = 1/β3 = O( −3/2), which ends up with the total complexity O( −3.5).
Remark: Note that the above theorem establishes the fast convergence rate for Nadam+ with p = 2/3. Indeed, we can also establish a fast rate of Adam+ (where p = 1/2) in the order ofO(1/ 3.625) with details provided in the Appendix E.

heading: 3 EXPERIMENTS
text: In this section, we conduct empirical studies to verify the effectiveness of the proposed algorithm on three different tasks: image classification on CIFAR10 and CIFAR100 dataset (Krizhevsky et al.,
2009), language modeling on Wiki-Text2 dataset (Merity, 2016) and automatic speech recognition on SWB-300 dataset (Saon et al., 2017). We choose tasks from different domains to demonstrate the applicability for the real-world deep learning tasks in a broad sense. The detailed description is presented in Table 2. We compare our algorithm Adam+ with SGD, momentum SGD, Adagrad, NIGT and Adam. We choose the same random initialization for each algorithm, and run a fixed number of epochs for every task. For Adam we choose the default setting β1 = 0.9 and β2 = 0.999 as in the original Adam paper.

heading: 3.1 IMAGE CLASSIFICATION
text: CIFAR10 and CIFAR100 In the first experiment, we consider training ResNet18 (He et al., 2016) and VGG19 (Simonyan & Zisserman, 2014) to do image classification task on CIFAR10 and CIFAR100 dataset respectively. For every optimizer, we use batch size 128 and run 350 epochs. For SGD and momentum SGD, we set the initial learning rate to be 0.1 for the first 150 epochs, and the learning rate is decreased by a factor of 10 for every 100 epochs. For Adagrad and Adam, the initial learning rate is tuned from {0.1, 0.01, 0.001} and we choose the one with the best performance. The best initial learning rates for Adagrad and Adam are 0.01 and 0.001 respectively. For NIGT, we tune the their momentum parameter from {0.01, 0.1, 0.9} (the best momentum parameter we found is 0.9) and the learning rate is chosen the same as in SGD. For Adam+, the learning rate is set according to Algorithm 1, in which we choose β = 0.1 and the value of α is the same as the learning rate used in SGD. We report training and test accuracy versus the number of epochs in Figure 1 for CIFAR10 and Figure 2 for CIFAR100. We observe that our algorithm consistently outperforms all other algorithms on both CIFAR10 and CIFAR100, in terms of both training and testing accuracy. Notably, we have some interesting observations for the training of VGG19 on CIFAR100. First, both Adam+ and NIGT significantly outperform SGD, momentum SGD, Adagrad and Adam. Second, Adam+ achieves almost the same final accuracy as NIGT, and Adam+ converges much faster in the early stage of the training.

heading: 3.2 LANGUAGE MODELING
text: Wiki-text2 In the second experiment, we consider the language modeling task on WikiText-2 dataset. We use a 2-layer LSTM (Hochreiter & Schmidhuber, 1997). The size of word embeddings is 650 and the number of hidden units per layer is 650. We run every algorithm for 40 epochs, with batch size 20 and dropout ratio 0.5. For SGD and momentum SGD, we tune the initial learning rate from {0.1, 0.2, 0.5, 5, 10, 20} and decrease the learning rate by factor of 4 when the validation error
saturates. For Adagrad and Adam, we tune the initial learning rate from {0.001, 0.01, 0.1, 1.0}. We report the best performance for these methods across the range of learning rate. The best initial learning rates for Adagrad and Adam are 0.01 and 0.001 respectively. For NIGT, we tune the initial value of learning rate from the same range as in SGD, and tune the momentum parameter β from {0.01, 0.1, 0.9}, and the best parameter choice is β = 0.9. The learning rate and β are both decreased by a factor of 4 when the validation error saturates. For Adam+, we follow the same tuning strategy as NIGT.
We report both training and test perplexity versus the number of epochs in Figure 3. From the Figure, we have the following observations: First, in terms of training perplexity, our algorithm achieves comparable performance with SGD and momentum SGD and outperforms Adagrad and NIGT, and it is worse than Adam. Second, in terms of test perplexity, our algorithm outperforms Adam, Adagrad, NIGT and momentum SGD, and it is comparable to SGD. An interesting observation is that Adam does not generalize well even if it has fast convergence in terms of training error, which is consistent with the observations in (Wilson et al., 2017).

heading: 3.3 AUTOMATIC SPEECH RECOGNITION
text: SWB-300 In the third experiment, we consider the automatic speech recognition task on SWB-300 dataset (Saon et al., 2017). SWB-300 contains roughly 300 hours of training data of over 4 million samples (30GB) and roughly 6 hours of held-out data of over 0.08 million samples (0.6GB). Each training sample is a fusion of FMLLR (40-dim), i-Vector (100-dim), and logmel with its delta and double delta. The acoustic model is a long short-term memory (LSTM) model with 6 bi-directional layers. Each layer contains 1,024 cells (512 cells in each direction). On top of the LSTM layers, there is a linear projection layer with 256 hidden units, followed by a softmax output layer with 32,000 (i.e., 32,000 classes) units corresponding to context-dependent HMM states. The LSTM is unrolled with 21 frames and trained with non-overlapping feature sub-sequences of that length. This model contains over 43 million parameters and is about 165MB large. The training takes about 20 hours on 1 V100 GPU. To compare, we adopt the well-tuned Momentum SGD strategy as described in (Zhang et al., 2019) for this task as the baseline: batch size is 256, learning rate is
0 2 4 6 8 10 12 14
10 4
0
2
4
6
8
10 10
4 ResNet18 training on CIFAR10
Adam+
0 5 10 15
10 4
0
2
4
6
8
10
12 10
4VGG19 training on CIFAR100
Adam+
Figure 5: The growth of quantity
∑t i=1 ‖zi‖ in Adam+
0.1 for the first 10 epochs and then annealed by √
0.5 for another 10 epochs, with momentum 0.9. We grid search the learning rate of Adam and Adagrad from {0.1, 0.01, 0.001}, and report the best configuration we have found (Adam with learning rate 0.001 and Adagrad with learning rate 0.01). For NIGT, we also follow the same learning rate setup (including annealing) as in Momentum SGD baseline. In addition, we fine tuned β in NIGT by exploring β in {0.01, 0.1, 0.9} and reported the best configuration (β=0.9). For Adam+, we follow the same learning rate and annealing strategy as in the Momentum SGD and tuned β in the same way as in NGIT, reporting the best configuration (β=0.01). From Figure 4, Adam+ achieves the indistinguishable training loss and held-out loss w.r.t. well-tuned Momentum SGD baseline and significantly outperforms the other optimizers.
3.4 GROWTH RATE OF ∑t i=1 ‖zi‖
In this subsection, we consider the growth rate of ∑t i=1 ‖zi‖, since they crucially affect the convergence rate as shown in Theorem 1. We report the results of both ResNet18 training on CIFAR10 dataset and VGG19 training on CIFAR100 dataset. From Figure 5, we can observe that it quickly reaches a plateau and then grows at a very slow rate with respect to the number of iterations. This phenomenon verifies the variance reduction effect and also explains the reason why Adam+ enjoys a fast convergence speed in practice.

heading: 4 CONCLUSION
text: In this paper, we design a new algorithm named Adam+ to train deep neural networks efficiently. Different from Adam, Adam+ updates the solution using moving average of stochastic gradients calculated at the extrapolated points and adaptive normalization on only first-order statistics of stochastic gradients. We establish data-dependent adaptive complexity results for Adam+ from the perspective of adaptive variance reduction, and also show that a variant of Adam+ achieves state-of-the-art complexity. Extensive empirical studies on several tasks verify the effectiveness of the proposed algorithm. We also empirically show that the slow growth rate of the new gradient estimator, providing the reason why Adam+ enjoys fast convergence in practice.

heading: A PROOF OF LEMMA 1
text: Proof. The proof is similar to that of Lemma 12 in (Wang et al., 2017). Define
ζ (t) k = { β(1− β)t−k if t ≥ k > 0 (1− β)t−k if t ≥ k = 0 (3)
By the definition of ζ(k)t and the update of Algorithm 1, we have
ζ (t+1) k = (1− β)ζ (t) k , t∑ k=0 ζ (t) k = 1, wt = t∑ k=0 ζ (t) k ŵt+1, zt+1 = t∑ k=0 ζ (t) k ∇f(ŵt+1; ξt+1).
Define mt+1 = ∑t k=0 ζ (t) k ‖wt+1 − ŵk+1‖2, nt+1 = ∑t k=0 ζ (t) k [∇f(ŵk+1; ξk+1)− F (ŵk+1)], where∇f(ŵk+1; ξk+1) is an unbiased stochastic first-order oracle for F (ŵk+1) with bounded variance σ2m. Note that∇F is a LH -smooth mapping (according to Assumption 1 (iii)), then by Lemma 10 of (Wang et al., 2017), we have
‖zt −∇F (wt)‖2 ≤ (LHmt + ‖nt‖)2 ≤ 2L2Hm2t + 2‖nt‖2. Define qt+1 = ∑t k=0 ζ (t) k ‖wt+1 − ŵk+1‖. According to Lemma 11 (a) and (b) of (Wang et al., 2017), we have
mt+1 + 4q 2 t+1 ≤
( 1− β
2
)( mt + 4q 2 t ) + 18
β ‖wt+1 −wt‖2.
Taking squares on both sides of the inequality and using the fact that (a+b)2 ≤ (1+ β2 )a 2+(1+ 2β )b 2 for β > 0, we have( mt+1 + 4q 2 t+1 )2 ≤ (1 + β 2 )( 1− β 2 )2 ( mt + 4q 2 t )2 + ( 1 + 2 β ) 324 β2 ‖wt+1 −wt‖4
≤ (
1− β 2
) (mt + 4q 2 t ) 2 + 972
β3 ‖wt+1 −wt‖4 ,
(4)
where the last inequality holds since 1/β ≥ 1.
Define δ2t = 2L 2 H(mt + 4q 2 t ) 2 + 2 ‖nt‖2, then we have ‖zt −∇F (wt)‖2 ≤ δ2t for all t. Denote Ft+1 by the σ-algebra generated by ξ1, . . . , ξt+1. Taking the summation of (4) and according to the bound of nt derived in Lemma 11 (c) of (Wang et al., 2017), we have
E [ δ2t+1|Ft+1 ] ≤ (
1− β 2
) δ2t + 2β 2σ2m + 1944L2H ‖wt+1 −wt‖ 4
β3 ,
Taking expectation on both sides yields
E [ δ2t+1 ] ≤ (
1− β 2
) E [ δ2t ] + 2β2σ2m + E [ 1944L2H‖wt+1 −wt‖4
β3
] .
Note that ηt = αβ a
max(‖zt‖1/2, 0) , we have E [ δ2t+1 ] ≤ (
1− β 2
) E [ δ2t ] + 2β2σ2m + E [ CL2α4β4a−3‖zt‖2 ] .

heading: B PROOF OF THEOREM 1
text: Proof. By Lemma 1 and the update rule of Algorithm 1, we have
E [ δ2t+1 ] ≤ (
1− β 2
) E [ δ2t ] + 2β2σ2m + E [ CL2Hα 4β4a‖zt‖4
β3(max(‖zt‖1/2, 0))4 ] ≤ (
1− β 2
) E [ δ2t ] + 2β2σ2m + E [ 2CL2Hα 4β4(‖δt‖2 + ‖∇F (wt)‖2)
β3
] ,
(5)
where the second inequality holds since (max(‖zt‖1/2, 0))4 ≥ ‖zt‖2 and ‖zt‖2 ≤ 2‖δt‖2 + 2 ‖∇F (wt)‖2. Note that 2CL2Hα 4 ≤ 1/18. Plugging it in (5), we have
8β 18 E [ δ2t ] ≤ E [ δ2t − δ2t+1 ] + 2β2σ2m + E
[ β
18 ‖∇F (wt)‖2
] . (6)
Summing over t = 1, . . . , T on both sides of (6) and with some simple algebra, we have T∑ t=1 E [ δ2t ] ≤ T∑ t=1 E [ 3 ( δ2t − δ2t+1 ) β ] + T∑ t=1 5βσ2m + T∑ t=1 E [ 1 8 ‖∇F (wt)‖2 ] . (7)
By Assumption 1 (i) and by the property of L-smooth function, we know that
F (wt+1) ≤ F (wt) +∇>F (wt)(wt+1 −wt) + L
2 ‖wt+1 −wt‖2
=F (wt)− ηt∇>F (wt)zt + η2tL
2 ‖zt‖2 ≤ F (wt)− ηt∇>F (wt)(zt −∇F (wt) +∇F (wt)) + η2tL ( ‖zt −∇F (wt)‖2 + ‖∇F (wt)‖2 ) = F (wt)− (ηt − η2tL)‖∇F (wt)‖2 − ηt∇>F (wt)(zt −∇F (wt)) + η2tL ‖zt − F (wt)‖ 2
≤ F (wt)− (ηt
2 − η2tL
) ‖∇F (wt)‖2 + (ηt 2 + η2tL ) ‖zt − F (wt)‖2 .
Noting that ηt = αβ a
max(‖zt‖1/2, 0) , α ≤ 1/4L and 0 = βa, we know that ηtL ≤ 1/4. Hence we
have
‖∇F (wt)‖2 ≤ 4(F (wt)− F (wt+1))
ηt + 3‖zt −∇F (wt)‖2.
Taking summation over t = 1, . . . , T and taking expectation yield T∑ t=1 E ‖∇F (wt)‖2 ≤ E [ T∑ t=1 4(F (wt)− F (wt+1)) ηt ] + 3 T∑ t=1 E‖zt −∇F (wt)‖2
≤ E [ T∑ t=1 4(F (wt)− F (wt+1)) ηt ] + 3 T∑ t=1 E [ δ2t ] .
(8)
Combining (7) and (8) yields T∑ t=1 E ‖∇F (wt)‖2 ≤ E [ T∑ t=1 4(F (wt)− F (wt+1)) ηt ] + T∑ t=1 E [ 9 ( δ2t − δ2t+1 ) β ]
+ T∑ t=1 15βσ2m + T∑ t=1 E [ 3 8 ‖∇F (wt)‖2 ] .
By some simple algebra, we have T∑ t=1 E ‖∇F (wt)‖2 ≤ E [ T∑ t=1 8(F (wt)− F (wt+1)) ηt ] + T∑ t=1 E [ 18 ( δ2t − δ2t+1 ) β ] + T∑ t=1 30βσ2m.
Then we have
1
T T∑ t=1 E ‖∇F (wt)‖2 ≤ E [ T∑ t=1 8 max ( ‖zt‖1/2, 0 ) (F (wt)− F (wt+1)) αβaT ] + 18σ20 βT +30βσ2m. (9)
Noting that |F (wt)− F (wt+1)| ≤ Gηt‖zt‖, we have
1
T T∑ t=1 E ‖∇F (wt)‖2 ≤ 8GE
[∑T t=1 ‖zt‖ ] T + ∆ αT + 18σ20 βT + 30βσ2m. (10)

heading: C PROOF OF THEOREM 2
text: Before introducing the proof, we first introduce several lemmas which are useful for our analysis.
Lemma 2. Adam+ with ηt = αβmax(‖zt‖1/2, 0) and 0 = 0 satisfies
F (wt+1)− F (wt) ≤ αβ ( −‖∇F (wt)‖ 3/2
6 + 9‖zt −∇F (wt)‖3/2
) + 64α4β4L3
3 .
Proof. By the L-smoothness and the update of the algorithm, we have
F (wt+1)− F (wt) ≤ ∇>F (wt)(wt+1 −wt) + L‖wt+1 −wt‖2
2
≤ −αβ · 〈∇F (wt), zt〉 max ( ‖zt‖1/2, 0 ) + α2β2L‖zt‖2( max ( ‖zt‖1/2, 0 ))2 . (11) Define ∆t = zt −∇F (wt). If ‖∇F (wt)‖ ≥ 2‖∆t‖, we have
− 〈zt,∇F (wt)〉 max ( ‖zt‖1/2, 0 ) = − ‖∇F (wt)‖2 + 〈∆t,∇F (wt)〉 max ( ‖∇F (wt) + ∆t‖1/2, 0 ) ≤ − ‖∇F (wt)‖ 2
2‖∇F (wt) + ∆t‖1/2 ≤ −‖∇F (wt)‖ 3/2 3
≤ −‖∇F (wt)‖ 3/2
3 + 8‖∆t‖3/2.
(12)
If ‖∇F (wt)‖ ≤ 2‖∆t‖, we have
− 〈zt,∇F (wt)〉 max ( ‖zt‖1/2, 0 ) = − ‖∇F (wt)‖2 + 〈∆t,∇F (wt)〉 max ( ‖∇F (wt) + ∆t‖1/2, 0 ) ≤ 6‖∆t‖ 2
‖∆t‖1/2 = 6‖∆t‖3/2 ≤ −
‖∇F (wt)‖3/2
3 + 8‖∆t‖3/2.
(13)
By (12) and (13), we have
− 〈zt,∇F (wt)〉 max ( ‖zt‖1/2, 0 ) ≤ −‖∇F (wt)‖3/2 3 + 8‖∆t‖3/2. (14)
By (11) and (14), we have F (wt+1)− F (wt) ≤ αβ ( ‖∇F (wt)‖3/2
3 + 8‖∆t‖3/2
) + α2β2L‖zt‖
= αβ ( −‖∇F (wt)‖ 3/2
3 + 8‖∆t‖3/2
) + α2β2Lmin
x>0
( 2‖zt‖3/2
3x + x2 3 ) ≤ αβ ( −‖∇F (wt)‖ 3/2
3 + 8‖∆t‖3/2
) + α2β2L ( 2‖zt‖3/2
3(8αβL) +
64α2β2L2
3 ) ≤ αβ ( −‖∇F (wt)‖ 3/2
6 + 9‖∆t‖3/2
) + 64α4β4L3
3 ,
where the last inequality holds because ‖zt‖3/2 ≤ 2‖∇F (wt)‖3/2 + 2‖∆t‖3/2.
Lemma 3. For Adam+ with ηt = αβmax(‖zt‖1/2, 0) , there exist random variables δt such that
E [ δ 3/2 t+1 ] ≤ (
1− β 2
) E [ δ 3/2 t ] + 2β3/2σ3/2 + E [ 320L 3/2 H ‖wt+1 −wt‖3
β2
] .
Proof. The proof shares the similar spirit of Lemma 12 in (Wang et al., 2017), but we adapt the proof for our purpose. Define
ζ (t) k = { β(1− β)t−k if t ≥ k > 0 (1− β)t−k if t ≥ k = 0 (15)
By the definition of ζ(k)t and the update of Algorithm 1, we have
ζ (t+1) k = (1− β)ζ (t) k , t∑ k=0 ζ (t) k = 1, wt = t∑ k=0 ζ (t) k ŵt+1, zt+1 = t∑ k=0 ζ (t) k ∇f(ŵt+1; ξt+1).
Define mt+1 = ∑t k=0 ζ (t) k ‖wt+1 − ŵk+1‖2, nt+1 = ∑t k=0 ζ (t) k [∇f(ŵk+1; ξk+1)− F (ŵk+1)], where∇f(ŵk+1; ξk+1) is an unbiased stochastic first-order oracle for F (ŵk+1) with bounded variance σ2. Note that ∇F is a LH -smooth mapping (according to Assumption 1), then by Lemma 10 of (Wang et al., 2017), we have
‖zt −∇F (wt)‖3/2 ≤ (LHmt + ‖nt‖)3/2 ≤ 2L3/2H m 3/2 t + 2‖nt‖3/2.
Define qt+1 = ∑t k=0 ζ (t) k ‖wt+1 − ŵk+1‖. According to Lemma 11 (a) and (b) of (Wang et al., 2017), we have
mt+1 + 4q 2 t+1 ≤
( 1− β
2
)( mt + 4q 2 t ) + 18
β ‖wt+1 −wt‖2.
Taking the power 3/2 on both sides of the inequality and using the fact that (a + b)3/2 ≤√ 1 + β2 a 3/2 + √
1 + 2β b 3/2 for β > 0, we have(
mt+1 + 4q 2 t+1 )3/2 ≤ ( 1 + β
2
)1/2( 1− β
2
)3/2 ( mt + 4q 2 t )3/2 + ( 1 + 2
β
)1/2 80
β3/2 ‖wt+1 −wt‖3 ≤ (
1− β 2
) (mt + 4q 2 t ) 3/2 + 160
β2 ‖wt+1 −wt‖3 ,
(16)
where the last inequality holds since 1/β ≥ 1. By the definition of nt, we have nt+1 = (1 − β)nt + β(∇f(ŵt+1) − F (ŵt+1)). Denote Ft+1 by the σ-algebra generated by ξ1, . . . , ξt+1. Noting that
E [ ‖nt+1‖3/2|Ft+1 ] ≤ ( E [ ‖nt+1‖2|Ft+1 ])3/4 ≤ (1− β/2)3/2‖nt‖3/2 + β3/2σ3/2, (17) where the last inequality holds by invoking Lemma 11(c) of (Wang et al., 2017). Define δ3/2t = 2L
3/2 H (mt + 4q 2 t ) 3/2 + 2‖nt‖3/2, then we have ‖zt −∇F (wt)‖3/2 ≤ δ3/2t for all t. According to (16) and (17), we have
E [ δ 3/2 t+1|Ft+1 ] ≤ (
1− β 2
) ‖δt‖3/2 + 2β3/2σ3/2 + 320L 3/2 H ‖wt+1 −wt‖3
β2 .
Taking expectation on both sides yields
E [ δ 3/2 t+1 ] ≤ (
1− β 2
) E [ δ 3/2 t ] + 2β3/2σ3/2 + E [ 320L 3/2 H ‖wt+1 −wt‖3
β2
] .
Lemma 4. Adam+ with learning rate ηt = αβmax(‖zt‖1/2, 0) and 640α 3L 3/2 H ≤ 1/120 satisfies
1
T T∑ t=1 E [ ‖∇F (wt)‖3/2 ] ≤ 101∆ αβT +
2727E [ δ 3/2 1 ] βT + 4545β1/2σ3/2 + 3β3L3/2 100 .
To ensure that 1T ∑T t=1 E [ ‖∇F (wt)‖3/2 ] ≤ 3/2, we can choose β = 3, T = O( −9/2).
Proof. By Lemma 3 and noting that ηt = αβmax(‖zt‖1/2, 0) , we have
E [ δ 3/2 t+1 ] ≤ (
1− β 2
) E [ δ 3/2 t ] + 2β3/2σ3/2 + E [ 320L 3/2 H α 3β3‖zt‖3/2
β2
]
≤ (
1− β 2
) E [ δ 3/2 t ] + 2β3/2σ3/2 + E [ 640L 3/2 H α 3β ( ‖∇F (wt)‖3/2 + ‖δt‖3/2 )] .
(18)
Note that 640α3L3/2H ≤ 1/120. Plugging it into (18), we have 59β 120 E [ δ 3/2 t ] ≤ E [ δ 3/2 t − δ 3/2 t+1 ] + 2β3/2σ3/2 + E [ β 120 ‖∇F (wt)‖3/2 ] . (19)
Summing over t = 1, . . . , T on both sides of (19) and with some simple algebra, we have T∑ t=1 E [ δ 3/2 t ] ≤ T∑ t=1 E [ 3(δ 3/2 t − δ 3/2 t+1) β ] + T∑ t=1 5β1/2σ3/2 + T∑ t=1 E [ 1 59 ‖∇F (wt)‖2/3 ] .
By Lemma 2, taking expectation on both sides, we have
E [F (wt+1)− F (wt)] ≤ αβ ( − E [ ‖∇F (wt)‖3/2 ] 6 + 9E [ δ 3/2 t ]) + 64α4β4L3 3 . (20)
Summing (20) over t = 1, . . . , T yields
5
504 αβ T∑ t=1 E [ ‖∇F (wt)‖3/2 ] ≤ F (w1)−F∗+αβ
27E [ δ 3/2 1 ] β + T∑ t=1 45β1/2σ3/2 +64α4β4L3T 3 .
Hence, we have
1
T T∑ t=1 E [ ‖∇F (wt)‖3/2 ] ≤ 101∆ αβT +
2727E [ δ 3/2 1 ] β + 4545β1/2σ3/2 + 2155α3β3L3
≤ 101∆ αβT
+ 2727E
[ δ 3/2 1 ] βT + 4545β1/2σ3/2 + 3β3L3/2 100 .
Lemma 5. Under the same setting of Lemma 4, we know that to ensure that 1T ∑T t=1 E [ δ 3/2 t ] ≤
3/2, we need T = O( −9/2) iterations.
Proof. From (19) and Lemma 4, we have T∑ t=1 59β 120 E [ δ 3/2 t ] ≤ E [ δ 3/2 1 ] + 2β3/2σ3/2T + T∑ t=1 E [ β 120 ‖∇F (wt)‖3/2 ] .
Noting that β = T−b with 0 < b < 1, then we know that there exists a universal constant C > 0 such that
1
T T∑ t=1 59 120 E [ δ 3/2 t ] ≤
E [ δ 3/2 1 ] T 1−b + 2σ3/2 T b/2 + 1 T T∑ t=1 E [ 1 120 ‖∇F (wt)‖3/2 ] . (21)
Take b = 23 . From Lemma 4, we know that it takes T = O( −9/2) iterations to ensure that 1 T ∑T t=1 E [ ‖∇F (wt)‖3/2 ] ≤ 3/2. In addition, From (21), we know that it takes T = O( −9/2)
iterations to ensure that 1T ∑T t=1 E [ δ 3/2 t ] ≤ 3/2.
We can easily prove Theorem 2 by incorporating the results in Lemma 4 and Lemma 5. It is also evident to see that if β = 1/T s with 0 < s < 1, then it takes T = O (poly(1/ )) number of iterations to ensure that 1T ∑T t=1 E [ δ 3/2 t ] ≤ 3/2 and 1T ∑T t=1 E [ ‖∇F (wt)‖3/2 ] ≤ 3/2 hold simultaneously.

heading: D PROOF OF THEOREM 3
text: Proof. Define γt = min ( βa
‖zt‖2/3 , β
a
0
) with 0 = 2βa. Then we know that ηt = αγt and γt ≤ 12 .
Note that α ≤ 1L , so we have ηt ≤ 1 2L . By the L-smoothness of F , we have
F (wt+1) ≤ F (wt) +∇>F (wt)(wt+1 −wt) + L
2 ‖wt+1 −wt‖2 ≤ F (wt)− ηt∇>F (wt)zt + ( η2tL
2 + γt 2L
) ‖zt‖2 − 1
2L γt‖zt‖2
= F (wt)− ηt∇>F (wt) (zt −∇F (wt) +∇F (wt)) + ( η2tL
2 + γt 2L
) ‖zt‖2 − 1
2L γt‖zt‖2
(a) ≤ F (wt)− ηt∇>F (wt) (zt −∇F (wt) +∇F (wt)) + ( η2tL+
γt L
)( ‖zt −∇F (wt)‖2 + ‖∇F (wt)‖2 ) − 1
2L γt‖zt‖2
(b) ≤ F (wt)− ηt 2 ‖∇F (wt)‖2 + ηt 2 ‖zt −∇F (wt)‖2 +
( η2tL+
γt L
)( ‖zt −∇F (wt)‖2 + ‖∇F (wt)‖2 ) − 1
2L γt‖zt‖2 = F (wt)− (ηt
2 − η2tL− γt L
) ‖∇F (wt)‖2 + ( η2tL+
γt L + ηt 2
) ‖zt −∇F (wt)‖2 − 1
2L γt‖zt‖2
(c) ≤ F (wt)− 1
2L γt‖zt‖2 +
1 L ‖zt −∇F (wt)‖2 ,
(22) where (a) holds since ‖zt‖2 ≤ 2‖zt−∇F (wt)‖2 +2‖∇F (wt)‖2, (b) holds since−∇>F (wt)zt ≤ 1 2 ( ‖∇F (wt)‖2 + ‖zt −∇F (wt)‖2 ) , (c) holds due to ηt2 − η 2 tL− γt L ≥ 0 (since ηt ≤ 1 2L , we have ηt 2 − η 2 tL ≥ 12L and note that γt L ≤ 1 2L ).
By the definition of γt, we have
γt‖zt‖2 ≥ β2a‖zt‖2/3 min ( ‖zt‖2/3
βa , ‖zt‖4/3 βa 0
) = β2a‖zt‖2/3 min ( ‖zt‖2/3
βa , ‖zt‖4/3 2β2a ) (a)
≥ β2a‖zt‖2/3 ( ‖zt‖2/3
βa − 1 2
) = βa‖zt‖4/3 − β2a‖zt‖2/3
2 ,
(23)
where (a) holds since x ≥ x− 12 , x2 2 ≥ x− 1 2 hold for any x and let x = ‖zt‖2/3 βa .
Combining (22) and (23), we have
βa‖zt‖4/3 ≤ γt‖zt‖2 + β2a‖zt‖2/3
2 ≤ 2L (F (wt)− F (wt+1)) +
β2a‖zt‖2/3
2 + 2‖zt −∇F (wt)‖2
= 2L (F (wt)− F (wt+1)) + βa‖zt‖4/3 · βa
2‖zt‖2/3 + 2‖zt −∇F (wt)‖2.
If β a
2‖zt‖2/3 ≤ 12 , we have β
a‖zt‖4/3 ≤ 4L (F (wt)− F (wt+1))+4‖zt−∇F (wt)‖2. If β a
2‖zt‖2/3 >
1 2 , then β a > ‖zt‖2/3, and hence we have βa‖zt‖4/3 ≤ β3a. As a result, we have
βa‖zt‖4/3 ≤ 4L (F (wt)− F (wt+1)) + 4‖zt −∇F (wt)‖2 + β3a. (24)
Taking summation on both sides of (24) over t = 1, . . . , T yields
T∑ t=1 ‖zt‖4/3 ≤ 4L T∑ t=1 F (wt)− F (wt+1) βa + T∑ t=1 4 βa ‖zt −∇F (wt)‖2 + β2aT. (25)
Define ∆t = zt −∇F (wt), then we have
‖∇F (wt)‖4/3 ≤ 2‖zt‖4/3 + 2‖∆t‖4/3. (26)
Hence,
T∑ t=1 ‖zt‖4/3 + ‖∇F (wt)‖4/3
(a) ≤ 2 T∑ t=1 ‖∆t‖4/3 + 12L T∑ t=1 F (wt)− F (wt+1) βa + T∑ t=1 12 βa ‖zt −∇F (wt)‖2 + 3β2aT
(b) ≤ T∑ t=1 4 3 ( ‖∆t‖2 βa + β2a 2 ) + 12L T∑ t=1 F (wt)− F (wt+1) βa + T∑ t=1 12 βa ‖zt −∇F (wt)‖2 + 3β2aT
≤ 12L T∑ t=1 F (wt)− F (wt+1) βa + T∑ t=1 14 βa ‖zt −∇F (wt)‖2 + 4β2aT,
(27) where (a) holds due to (25) and (26), (b) holds because minx>0 c 2 x + x2 2 = 3c4/3 2 .
By Lemma 1, we know that
E [ δ2t+1 ] ≤ (
1− β 2
) E [ δ2t ] + 2β2σ2 + E [ CL2η4t ‖zt‖4
β3 ] (a)
≤ (
1− β 2
) E [ δ2t ] + 2β2σ2 + E [ CL2α4β4a‖zt‖4
max(‖zt‖8/3, 40)β3 ] ≤ (
1− β 2
) E [ δ2t ] + 2β2σ2 + E [ CL2α4β4a−3‖zt‖4/3 ] .
Note that CL2α4 ≤ 1/14, we have
β 2 E [ δ2t ] ≤ E [ δ2t − δ2t+1 ] + 2β2σ2 + E
[ β4a−3‖zt‖4/3
14
] . (28)
Taking summation on both sides of (28) over t = 1, . . . , T , we have
T∑ t=1 E [ δ2t ] ≤
E [ δ21 ]
β + 2βσ2T + T∑ t=1 E [ β4a−4‖zt‖4/3 14 ]
= E [ δ21 ]
β + 2βσ2T + T∑ t=1 E [ βa‖zt‖4/3 14 ] ,
(29)
where the last equality holds since a = 4/3.
Taking expectation on both sides of (27) and combining (29), we have T∑ t=1 E [ ‖zt‖4/3 + ‖∇F (wt)‖4/3 ] ≤ 12L∆ βa + 14E [ δ21 ] β1+a + 28βσ2T βa + T∑ t=1 E [ ‖zt‖4/3 ] + 4β2aT.
As a result, we have
1
T T∑ t=1 E [ ‖∇F (wt)‖4/3 ] ≤ 12L∆ βaT + 14E [ δ21 ] β1+aT + 28βσ2 βa + 4β2a.
Suppose initial batch size is T0, the intermediate batch size is m, and a = 4/3, then we have
1
T T∑ t=1 E [ ‖∇F (wt)‖4/3 ] ≤ 12L∆ β4/3T + 14σ2 β7/3T0T + 28σ2 β1/3m + 4β8/3. (30)
We can choose β = O( 1/2), T = O( −2), the initial batch size T0 = 1/β = O( −1/2), the intermediate batch size as m = 1/β3 = O( −3/2), which ends up with the total complexity O( −3.5).
E A NEW VARIANT OF ADAM+
Theorem 4. Assume that ‖∇f(w; ξ)‖ ≤ G almost surely for every w ∈ Rd. Choose ηt = αβa max(‖zt‖1/2, 0) with a = 4/3, and we have
1
T T∑ t=1 E [ ‖∇F (wt)‖3/2 ] ≤ 12L∆ βaT + 14E [ δ21 ] β1+aT + 28βσ2 βa + 4β3a.
Denote the initial batch size and the intermediate batch size are T0 andm respectively, then we have
1
T T∑ t=1 E [ ‖∇F (wt)‖3/2 ] ≤ 12L∆ βaT + 14σ2 T0Tβ1+a + 28σ2 βa−1m + 4β3a.
To ensure that 1T ∑T t=1 E [ ‖∇F (wt)‖3/2 ] ≤ 3/2, we choose β = 3/8, T = O(1/ 2), the initial batch size is T0 = 1/ 3/8 andm = 1/ 1.625, then the total computational complexity isO(1/ 3.625). Proof. Define γt = min ( βa
‖zt‖1/2 , β
a
0
) with 0 = 2βa. Then we know that ηt = αγt and γt ≤ 12 .
Note that α ≤ 1L , so we have ηt ≤ 1 2L . By the L-smoothness of F , we have
F (wt+1) ≤ F (wt) +∇>F (wt)(wt+1 −wt) + L
2 ‖wt+1 −wt‖2 ≤ F (wt)− ηt∇>F (wt)zt + ( η2tL
2 + γt 2L
) ‖zt‖2 − 1
2L γt‖zt‖2
= F (wt)− ηt∇>F (wt) (zt −∇F (wt) +∇F (wt)) + ( η2tL
2 + γt 2L
) ‖zt‖2 − 1
2L γt‖zt‖2
(a) ≤ F (wt)− ηt∇>F (wt) (zt −∇F (wt) +∇F (wt)) + ( η2tL+
γt L
)( ‖zt −∇F (wt)‖2 + ‖∇F (wt)‖2 ) − 1
2L γt‖zt‖2
(b) ≤ F (wt)− ηt 2 ‖∇F (wt)‖2 + ηt 2 ‖zt −∇F (wt)‖2 +
( η2tL+
γt L
)( ‖zt −∇F (wt)‖2 + ‖∇F (wt)‖2 ) − 1
2L γt‖zt‖2 = F (wt)− (ηt
2 − η2tL− γt L
) ‖∇F (wt)‖2 + ( η2tL+
γt L + ηt 2
) ‖zt −∇F (wt)‖2 − 1
2L γt‖zt‖2
(c) ≤ F (wt)− 1
2L γt‖zt‖2 +
1 L ‖zt −∇F (wt)‖2 ,
(31) where (a) holds since ‖zt‖2 ≤ 2‖zt−∇F (wt)‖2 +2‖∇F (wt)‖2, (b) holds since−∇>F (wt)zt ≤ 1 2 ( ‖∇F (wt)‖2 + ‖zt −∇F (wt)‖2 ) , (c) holds due to ηt2 − η 2 tL− γt L ≥ 0 (since ηt ≤ 1 2L , we have ηt 2 − η 2 tL ≥ 12L and note that γt L ≤ 1 2L ).
By the definition of γt, we have γt‖zt‖2 ≥ β2a‖zt‖min ( ‖zt‖1/2
βa , ‖zt‖ βa 0
) = β2a‖zt‖min ( ‖zt‖1/2
βa , ‖zt‖ 2β2a ) (a)
≥ β2a‖zt‖ ( ‖zt‖1/2
βa − 1 2
) = βa‖zt‖3/2 −
β2a‖zt‖ 2 ,
(32)
where (a) holds since x ≥ x− 12 , x2 2 ≥ x− 1 2 hold for any x and let x = ‖zt‖1/2 βa .
Combining (31) and (32), we have
βa‖zt‖3/2 ≤ γt‖zt‖2 + β2a‖zt‖
2 ≤ 2L (F (wt)− F (wt+1)) + β2a‖zt‖ 2 + 2‖zt −∇F (wt)‖2
= 2L (F (wt)− F (wt+1)) + βa‖zt‖3/2 · βa
2‖zt‖1/2 + 2‖zt −∇F (wt)‖2.
If β a
2‖zt‖1/2 ≤ 12 , we have β
a‖zt‖4/3 ≤ 4L (F (wt)− F (wt+1))+4‖zt−∇F (wt)‖2. If β a
2‖zt‖1/2 >
1 2 , then β a > ‖zt‖1/2, and hence we have βa‖zt‖3/2 ≤ β4a. As a result, we have
βa‖zt‖3/2 ≤ 4L (F (wt)− F (wt+1)) + 4‖zt −∇F (wt)‖2 + β4a. (33)
Taking summation on both sides of (33) over t = 1, . . . , T yields
T∑ t=1 ‖zt‖3/2 ≤ 4L T∑ t=1 F (wt)− F (wt+1) βa + T∑ t=1 4 βa ‖zt −∇F (wt)‖2 + β3aT. (34)
Define ∆t = zt −∇F (wt), then we have
‖∇F (wt)‖3/2 ≤ 2‖zt‖3/2 + 2‖∆t‖3/2. (35)
Hence, we have
T∑ t=1 ‖zt‖3/2 + ‖∇F (wt)‖3/2
(a) ≤ 2 T∑ t=1 ‖∆t‖3/2 + 12L T∑ t=1 F (wt)− F (wt+1) βa + T∑ t=1 12 βa ‖zt −∇F (wt)‖2 + 3β3aT
(b) ≤ T∑ t=1 3 2 ( ‖∆t‖2 βa + β3a 3 ) + 12L T∑ t=1 F (wt)− F (wt+1) βa + T∑ t=1 12 βa ‖zt −∇F (wt)‖2 + 3β3aT
≤ 12L T∑ t=1 F (wt)− F (wt+1) βa + T∑ t=1 14 βa ‖zt −∇F (wt)‖2 + 4β3aT,
(36) where (a) holds due to (34) and (35), (b) holds because minx>0 c 2 x + x3 3 = 4c3/2 3 .
By Lemma 1, we know that
E [ δ2t+1 ] ≤ (
1− β 2
) E [ δ2t ] + 2β2σ2 + E [ CL2Hη 4 t ‖zt‖4
β3 ] (a)
≤ (
1− β 2
) E [ δ2t ] + 2β2σ2 + E [ CL2Hα 4β4a‖zt‖4
max(‖zt‖2, 40)β3 ] ≤ (
1− β 2
) E [ δ2t ] + 2β2σ2 + E [ CL2Hα 4β4a−3‖zt‖2 ] .
Note that CL2Hα 4 ≤ 1 14G1/2 , we have
β 2 E [ δ2t ] ≤ E [ δ2t − δ2t+1 ] + 2β2σ2 + E
[ β4a−3‖zt‖2
14G1/2
] . (37)
Taking summation on both sides of (37) over t = 1, . . . , T , we have
T∑ t=1 E [ δ2t ] ≤
E [ δ21 ]
β + 2βσ2T + T∑ t=1 E [ β4a−4‖zt‖2 14 ]
= E [ δ21 ]
β + 2βσ2T + T∑ t=1 E [ βa‖zt‖2 14G1/2 ]
≤ E [ δ21 ]
β + 2βσ2T + T∑ t=1 E [ βa‖zt‖3/2 14 ] ,
(38)
where the equality holds since a = 4/3 and last inequality holds since ‖zt‖ ≤ G.
Taking expectation on both sides of (36) and combining (38), we have
T∑ t=1 E [ ‖zt‖3/2 + ‖∇F (wt)‖3/2 ] ≤ 12L(F (w1)− F∗)
βa +
14E [ δ21 ]
β1+a +
28βσ2T
βa + T∑ t=1 E [ ‖zt‖3/2 ] + 4β3aT.
As a result, we have
1
T T∑ t=1 E [ ‖∇F (wt)‖3/2 ] ≤ 12L(F (w1)− F∗) βaT + 14E [ δ21 ] β1+aT + 28βσ2 βa + 4β3a.

heading: F RELATED WORK
text: Adaptive Gradient Methods Adaptive gradient methods were first proposed in the framework of online convex optimization (Duchi et al., 2011; McMahan & Streeter, 2010), which dynamically incorporate knowledge of the geometry of the data to perform more informative gradient-based learning. This type of algorithm was proved to have fast convergence if stochastic gradients are sparse (Duchi et al., 2011). Based on this idea, several other adaptive algorithms were proposed to train deep neural networks, including Adam (Kingma & Ba, 2014), Amsgrad (Reddi et al., 2019), RMSprop (Tieleman & Hinton, 2012). There are many work trying to analyze variants of adaptive gradient methods in both convex and nonconvex case (Chen et al., 2018a; 2019; 2018b; Luo et al., 2019; Chen et al., 2018a;b; Ward et al., 2019; Li & Orabona, 2019; Chen et al., 2019). Notably, all of these works are able to establish faster convergence rate than SGD, based on the assumption that stochastic gradients are sparse. However, this assumption may not hold in deep learning. In contrast, our algorithm can have faster convergence than SGD even if stochastic gradients are not sparse, since our algorithm’s new data-dependent adaptive complexity does not rely on the sparsity of stochastic gradients.
Variance Reduction Methods Variance reduction is a technique to achieve fast rates for finite sum and stochastic optimization problems. It was first proposed for finite-sum convex optimization (Johnson & Zhang, 2013) and then it was extended in finite-sum nonconvex (Allen-Zhu & Hazan, 2016; Reddi et al., 2016; Zhou et al., 2018) and stochastic nonconvex (Lei et al., 2017; Fang et al., 2018; Wang et al., 2019; Pham et al., 2020; Cutkosky & Orabona, 2019) optimization. To prove faster convergence rate than SGD, all these works make the assumption that the objective function is an average of individual functions and each one of them is smooth. In contrast, our analysis does not require such an assumption and to achieve a faster-than-SGD rate.
Other Related Work Arjevani et al. (2019) show that SGD is optimal for stochastic nonconvex smooth optimization, if one does not assume that every component function is smooth. There are recent work trying to establish faster rate than SGD, when the Hessian of the objective function is Lipschitz (Fang et al., 2019; Cutkosky & Mehta, 2020). There are several empirical papers, including LARS (You et al., 2017) and LAMB (You et al., 2019)), which utilize both moving average and normalization for training of deep neural networks with large-batch sizes. Zhang et al. (2020) consider an algorithm for finding stationary point for nonconvex nonsmooth problems. Levy (2017) considers convex optimization setting and design algorithms which adapts to the smoothness parameter. Liu et al. (2019) introduced Rectified Adam to alleviate large variance at the early stage. However, none of them establish data-dependent adaptive complexity as in our paper.
G ADAM+ WITH FIXED LEARNING RATE
We report the results on image classification with CIFAR10 on ResNet18. We further considered the Adam+ with fixed learning rate 0.1 and do not employ any learning rate annealing scheme. We report our result on Figure 6, in which "Adam+ Fixed Stepsize" uses the default setting of Algorithm 1. As we can see from the Figure, Adam+ with fixed stepsize still outperforms Adam,
0 2 4 6 8 10 12 2
4
6
8
10
12
14 ResNet18 training on CIFAR10
Adam+
y=0.8*x+log(20)
Figure 7: log(
∑t i=1 ‖zi‖) versus log(t)
Adagrad and momentum SGD. If we use the same learning rate scheduler of SGD for Adam+, then Adam+ performs the best.
H GROWTH RATE ANALYSIS OF ∑t
i=1 ‖zi‖
We provide the log-log plot (log( ∑t i=1 ‖zi‖) versus log(t)) for the ResNet18 training on CIFAR10 experiment, as illustrated in Figure 7. We can see that the slope is around 0.8. Although the slope is not small, we can see from Figure 5 that the slope becomes almost zero after the iteration 6 × 104 (which corresponds to the epoch number 154). At this particular epoch, the training and test accuracy are not the best so we need to keep the training until epoch 350. Then our algorithm Adam+ is able
to take advantage of the slow growth rate of ∑t i=1 ‖zi‖ for large t and enjoys faster convergence, which is consistent with Figure 1.

Review summaries:
This paper proposes a new optimizer called Adam+, with two main distinctions from standard Adam template: 1) the first order moment estimate is computed using the gradient evaluated at an extrapolated iterate. 2) the step size is scaled with the square root of the norm of the first order moment, rather than the exponential moving average (EMA) in the previous work. Under Lipschitz continuous gradient, Hessian and bounded gradient assumptions, the complexity for finding a point with small gradient norm is $\epsilon^{-3.5}$, which is a SOTA complexity. The practical performance of the algorithm is evaluated in a range of different tasks and the performance is shown to be consistently promising and comparable to SGD/Momentum SGD.

This paper proposes the Adam+ algorithm that maintains an exponential moving average of the first moment and normalizes it by its $p$-th moment for some $p \in (1/2, 1)$. When $p = 2/3$, with appropriate hyperparameters, Adam+ achieves the state-of-the-art complexity $O(1/\epsilon^{3.5})$ to obtain an approximate first-order stationary point for smooth objectives with smooth Hessians. The proof technique is similar to that for SCGD by Wang et al. (2017), which construct a Lyapunov function for algorithms of this kind.