Year: 2021

Abstract:
Adaptive gradient methods, such as ADAM, have achieved tremendous success in machine learning. Scaling gradients by square roots of the running averages of squared past gradients, such methods are able to attain rapid training of modern deep neural networks. Nevertheless, they are observed to generalize worse than stochastic gradient descent (SGD) and tend to be trapped in local minima at an early stage during training. Intriguingly, we discover that substituting the gradient in the second moment estimation term with the momentumized version in ADAM can well solve the issues. The intuition is that gradient with momentum contains more accurate directional information and therefore its second moment estimation is a better choice for scaling than that of the raw gradient. Thereby we propose ADAMOMENTUM as a new optimizer reaching the goal of training fast while generalizing better. We further develop a theory to back up the improvement in optimization and generalization and provide convergence guarantees under both convex and nonconvex settings. Extensive experiments on a wide range of tasks and models demonstrate that ADAMOMENTUM exhibits state-of-the-art performance consistently. The source code is available at https://anonymous. 4open.science/r/AdaMomentum_experiments-6D9B.

heading: 1 INTRODUCTION
text: Prevailing first-order optimization algorithms in modern machine learning can be classified into two categories. One is stochastic gradient descent (SGD) (Robbins & Monro, 1951), which is widely adopted due to its low memory cost and outstanding performance. SGDM (Sutskever et al., 2013) which incorporates the notion of momentum into SGD, has become the best choice for optimizer in computer vision. The drawback of SGD(M) is that it scales the gradient uniformly in all directions, making the training slow especially at the begining and fail to optimize complicated models well beyond Convolutional Neural Networks (CNN). The other type is adaptive gradient methods. Unlike SGD, adaptive gradient optimizers adapt the stepsize (a.k.a. learning rate) elementwise according to the gradient values. Specifically, they scale the gradient by the square roots of some form of the running average of the squared values of the past gradients. Popular examples include AdaGrad (Duchi et al., 2011), RMSprop (Tijmen Tieleman, 2012) and Adam (Kingma & Ba, 2015) etc. Adam, in particular, has become the default choice for many machine learning application areas owing to its rapid optimizing speed and outstanding ability to handle sophisticated loss curvatures.
Despite their fast speed in the early training phase, adaptive gradient methods are found by studies (Wilson et al., 2017; Zhou et al., 2020) to be more likely to exhibit poorer generalization ability than SGD. This is discouraging because the ultimate goal of training in many machine learning tasks is to exhibit high performance during testing phase. In recent years researchers have put many efforts to mitigate the deficiencies of adaptive gradient algorithms. AMSGrad (Reddi et al., 2019) corrects the errors in the convergence analysis of Adam and proposes a faster version. Yogi (Reddi et al., 2018) takes the effect of batch size into consideration. M-SVAG (Balles & Hennig, 2018) transfers the variance adaptation mechanism from Adam to SGD. AdamW (Loshchilov & Hutter, 2017b) firsttime decouples weight decay from gradient descent for Adam-alike algorithms. SWATS (Keskar & Socher, 2017) switches from Adam to SGD throughout the training process via a hard schedule and AdaBound (Luo et al., 2019) switches with a smooth transation by imposing dynamic bounds
on stepsizes. RAdam (Liu et al., 2019) rectifies the variance of the adaptive learning rate through investigating the theory behind warmup heuristic (Vaswani et al., 2017; Popel & Bojar, 2018). AdaBelief (Zhuang et al., 2020) adapts stepsizes by the belief in the observed gradients. Nevertheless, most of the above variants can only surpass (as they claim) Adam or SGD in limited tasks or under specifically and carefully defined scenarios. Till today, SGD and Adam are still the top options in machine learning, especially deep learning (Schmidt et al., 2021). Conventional rules for choosing optimizers are: from task perspective, choose SGDM for vision, and Adam (or AdamW) for language and speech; from model perspective, choose SGDM for Fully Connected Networks and CNNs, and Adam (or AdamW) for Recurrent Neural Networks (RNN) (Cho et al., 2014; Hochreiter & Schmidhuber, 1997b), Transformers (Vaswani et al., 2017) and Generative Adversarual Networks (GAN) (Goodfellow et al., 2014). Based on the above observations, a natural question is:

heading: Is there a computationally efficient adaptive gradient algorithm that can converge fast and meanwhile generalize well?
text: In this work, we are delighted to discover that simply replacing the gradient term in the second moment estimation term of Adam with its momentumized version can achieve this goal. Our idea comes from the origin of Adam optimizer, which is a combination of RMSprop and SGDM. RMSprop scales the current gradient by the square root of the exponential moving average (EMA) of the squared past gradients, and Adam replaces the raw gradient in the numerator of the update term of RMSprop with its EMA form, i.e., with momentum. Since the momentumized gradient is a more accurate estimation of the appropriate direction to descent, we consider putting it in the second moment estimation term as well. We find such operation makes the optimizer more suitable for the general loss curvature and can theoretically converge to minima that generalize better. Extensive experiments on a broad range of tasks and models indicate that: without bells and whistles, our proposed optimizer can be as good as SGDM on vision problems and outperforms all the competitors in other tasks, meanwhile maintaining fast convergence speed. Our algorithm is efficient with no additional memory cost, and applicable to a wide range of scenarios in machine learning, especially deep learning. More importantly, AdaMomentum requires little effort in hyperparameter tuning and the default parameter setting for adaptive gradient method works well consistently in our algorithm.
Notation We use t, T to symbolize the current and total iteration number in the optimization process. θ ∈ Rd denotes the model parameter and f(θ) ∈ R denotes the loss function. We further use θt to denote the parameter at step t and ft to denote the noisy realization of f at time t because of the mini-batch stochastic gradient mechanism. gt denotes the t-th time gradient and α denotes stepsize. mt, vt represent the EMA of the gradient and the second moment estimation term at time t of adaptive gradient methods respectively. is a small constant number added in adaptive gradient methods to refrain the denominator from being too close to zero. β1, β2 are the decaying parameter in the EMA formulation of mt and vt correspondingly. For any vectors a, b ∈ Rd, we employ√ a, a2, |a|, a/b, a ≥ b, a ≤ b for elementwise square root, square, absolute value, division, greater or equal to, less than or equal to respectively. For any 1 ≤ i ≤ d, θt,i denotes the i-th element of θt. Given a vector x ∈ Rd, we use ‖x‖2 to denote its l2-norm and ‖x‖∞ to denote its l∞-norm.
2 ALGORITHM
Table 1: Comparison of AdaMomentum and classic adaptive gradient methods in mt and vt in (1).
Optimizer mt vt SGD gt 1 Rprop gt g2t RMSprop gt (1− β2) ∑t i=1 β t−i 2 g 2 i
Adam (1− β1) ∑t i=1 β t−i 1 gi (1− β2) ∑t i=1 β t−i 2 g 2 i
Ours (1− β1) ∑t i=1 β t−i 1 gi (1− β2) ∑t i=1 β t−i 2 mi 2
Preliminaries & Motivation Omitting the debiasing operation and the damping term , the adaptive gradient methods can be generally written in the following form:
θt+1 = θt − α mt√ vt . (1)
Here mt, vt are called the first and second moment estimation terms. When mt = gt and vt = 1, (1) degenerates to the vanilla SGD. Rprop (Duchi et al., 2011) is the pioneering work using the notion of adaptive learning rate, in which mt = gt and vt = g2t . Actually it is equivalent to only using the sign of gradients for different weight parameters. RMSprop (Tijmen Tieleman, 2012) forces the number divided to be similar for adjacent mini-batches by incorporating momentum acceleration into vt. Adam (Kingma & Ba, 2015) is built upon RMSprop in which it turns
gt into momentumized version. Both RMSprop and Adam boost their performance thanks to the smoothing property of EMA using momentum. Due to the fact that momentumized gradient is a more accurate estimation than raw gradient, we deem that there is no reason to use gt in lieu of mt in second moment estimation term vt. Therefore we propose to replace the gis in vt of Adam with their momentumized versions mis, which further smooths the exponential moving average.
Algorithm 1 AdaMomentum (ours). All mathematical operations are element-wise. 1: Initialization : Parameter initialization θ0, step size α, damping term ,m0 ← 0, v0 ← 0, t← 0 2: while θt not converged do 3: t← t+ 1 . Updating time step 4: gt ← ∇θft(θt−1) . Acquiring stochastic gradient at time t 5: mt ← β1mt−1 + (1− β1)gt . EMA of gradients 6: vt ← β2vt−1 + (1− β2)mt2 + . EMA of squared momentumized gradients 7: m̂t ← mt/(1− βt1) . Bias correction of first moment estimation 8: v̂t ← vt/(1− βt2) . Bias correction of second moment estimation 9: θt ← θt−1 − α · m̂t/ √ v̂t . Updating parameters 10: end while
Detailed Algorithm The detailed procedure of our proposed optimizer is displayed in Algorithm 1. There are two major modifications based on Adam, which are marked in red and blue respectively. One is that we replace the gt in vt of Adam with mt, which is the momentumized gradient. Hence we name our proposed optimzier as AdaMomentum. The other is the location of (in Adam is added after √ · in line 10 of Alg.1). We discover that moving the adding from outside the radical symbol to inside can consistently enhance performance. To the best of our knowledge, our method is the first attempt to put momentumized gradient in the second moment estimation term of adaptive gradient methods. Note that although the modifications seem simple to some degree, they can lead to siginificant changes in the performance of an adaptive gradient optimizer due to the iterative nature of optimization methods, which will also be elaborated in the following section.

heading: 3 WHY ADAMOMENTUM OVER ADAM?

heading: 3.1 ADAMOMENTUM IS MORE SUITABLE FOR GENERAL LOSS CURVATURE
text: In this section, we show that AdaMomentum can converge to (global) minima faster than Adam does via a 1-D example. The left part of Figure 1 is the process of optimization from a plateau to a basin area, where global optimum is assumed to exist. The right part is the zoomed-in version of the situation near the minimum, where we have some peaks and valleys. This phenomenon frequently takes place in optimization since there is only one global minimum with probably a great number of local minima surrounding (Hochreiter & Schmidhuber, 1997a; Keskar et al., 2017).
Benefits of substituting gt with mt. We first explain how substituting mt for gt in the preconditioner vt can improve training via decomposing the trajectory of parameter point along the loss curve. 1) In area A, the parameter point starts to slide down the curve and |gt| begins to enlarge abruptly. So the actual stepsize α/ √ vt is small for Adam. However the absolute value of the momentumized gradient mt is small since it is the EMA of the past gradients, making α/ √ vt still large for AdaMomentum. Hence AdaMomentum can maintain higher training speed than Adam in this changing corner of the loss curve, which is what an optimal optimizer should do. 2) In area B, since the exponential moving average decays the impact of past gradients exponentially w.r.t. t, the magnitude of the elements of mt will gradually becomes as large as gt. 3) In area C, when the parameter approaches the basin, the magnitude of gt decreases, making the stepsizes of Adam increase immediately. In contrast, the stepsize of AdaMomentum is still comparatively small as |mt| is still much larger than |gt|, which is desired for an ideal optimizer. Small stepsize near optimum has benefits for convergence and stability. A more concrete illustration is given in the right part of Figure 1. If the stepsize is too large (e.g. in Adam), the weight parameter θt may rush to θ (2) t+1 and miss the global optimum. In contrast, small stepsize can guarantee the parameter to be close to the global minimum (see θ(1)t+1) even if there may be tiny oscillations within the basin before the final convergence.
Benefits of changing the location of . Next we elaborate why putting under the √ · is beneficial. We denote the debiased second moment estimation in AdaMomentum as v̂t and the second moment estimation term without as v̂′t. By simple calculation, we have
v̂t = ( (1− β2)/(1− βt2) ) · t∑ i=1 βt−i2 m 2 i + 1− β2 , v̂′t = ( (1− β2)/(1− βt2) ) · t∑ i=1 βt−i2 m 2 i .
Hence we have v̂t = v̂′t + /(1 − β2). Then the actual stepsizes are α/( √ v̂′t + /(1− β2)) and
α/( √ v̂′t+ ) respectively. In the final stage of optimization, v̂ ′ t is very close to 0 (because the values of gradients are near 0) and far less than hence the actual stepsizes can be approximately written as√ 1− β2α/ √ and α/ .As usually takes very tiny values ranging from 10−8 to 10−16 and β2 usu-
ally take values that are extremely close to 1 (usually 0.999), we have √ 1− β2α/ √ α/ . Therefore we may reasonably come to the conclusion that after moving term into the radical symbol, AdaMomentum further reduces the stepsizes when the training is near minima, which contributes to enhancing convergence and stability as we have discussed above.

heading: 3.2 ADAMOMENTUM CONVERGES TO MINIMA THAT GENERALIZE BETTER
text: The outline of Adam and our proposed AdaMomentum can be written in the following unified form:
mt = β1mt−1 + (1− β1)gt, vt = β2vt−1 + (1− β2)k2t , θt+1 = θt − αmt /( (1− βt1) √ vt/(1− βt2) ) . (2)
where kt = gt in Adam and kt = mt in AdaMomentum. Inspired by a line of work (Pavlyukevich, 2011; Simsekli et al., 2019; Zhou et al., 2020), we can consider (2) as a discretization of a continuous-time process and reformulate it as its corresponding Lévy-driven stochastic differential equation (SDE). Assuming that the gradient noise ζt = gt−∇f(θt) is independent and centered symmetric α̃-stable (Sα̃S) (Lévy & Lévy, 1954) distributed with covariance matrix Σt possessing a heavy-tailed signature (α̃ ∈ (0, 2]), we are able to derive the Lévy-driven SDE of (2) as:
dθt = −qtR−1t mtdt+ υR−1t ΣtdLt, dmt = β1(∇f(θt)−mt), dvt = β2(k2t − vt), (3) where Rt = diag( √ vt/(1− βt2)), υ = α1−1/α̃, qt = 1/(1 − βt1) and Lt is the α̃-stable Lévy motion with independent components. We are interested in the local stability of the optimizers and therefore we suppose process (3) is initialized in a local basin Ω with a minimum θ∗ (w.l.o.g., we assume θ∗ = 0). To investigate the escaping behavior of θt, we first introduce two technical definitions. Definition 1 (Radon Measure (Simon et al., 1983)). If a measure m(·) defined on the σ-algebra of Borel sets of a Hausdorff topological space X is 1) inner regular on open sets, 2) outer regular on all Borel sets, and 3) finite on all compact sets, then the measure is called a Radon measure. Definition 2 (Escaping Time & Escaping Set). We define escaping time Γ := inf{t ≥ 0 : θt 6∈ Ω−υ γ }, where Ω−υ γ
= {y ∈ Ω : dis(∂Ω, y) ≥ υγ}. Here γ > 0 is a constant. We define escaping set Υ := {y ∈ Rd : R−1θ∗ Σθ∗y 6∈ Ω −υγ}, where Σθ∗ = limθt→θ∗ Σt, Rθ∗ = limθt→θ∗ Rt.
We study the relationship between Γ and Υ and impose some standard assumptions before proceeding.
Assumption 1. f is non-negative with an upper bound, and locally µ-strongly convex in Ω. Assumption 2. There exists some constant L > 0, s.t. ‖∇f(x)−∇f(y)‖2 ≤ L ‖x− y‖2 , ∀x, y.
Assumption 3. We assume that ∫ Γ
0 〈∇f(θt)/(1 + f(θt)), qtR−1t mt〉 dt ≥ 0 a.e., and β1 ≤ β2 ≤ 2β1. There
exist v−, v+ > 0 s.t. each coordinate of √ vt can be uniformly bounded in (v−, v+) and there exist τm, τ > 0 s.t. ‖mt − m̂t‖2 ≤ τm ∥∥∥∫ t−0 (mx − m̂x) dx∥∥∥
2 and ‖m̂t‖2 ≥ τ ∥∥∥∇f(θ̂t)∥∥∥ 2 , where m̂t and θ̂t are calculated
by solving (3) with υ = 0.
Assumption 1 and 2 impose some standard assumptions of stochastic optimization Ghadimi & Lan (2013); Johnson & Zhang (2013). Assumption 3 requires momentumized gradient mt and ∇f(θt) to have similar directions for most of the time, which have been empirically justified to be true in Adam (Zhou et al., 2020). Based on the above assumptions, we can prove that for algorithm of form (2), the expected escaping time is inversely proportional to the Radon measure of the escaping set:
Lemma 1. Under Assumptions 1-3, let υα̃+1 = Θ(α̃) and ln (2∆/(µυ1/3)) ≤ 2µτ(β1−β2/4)/(β1v++µτ), where ∆ = f(θ0)− f(θ∗). Then given any θ0 ∈ Ω−2υ γ , for (3) we have
E(Γ) = Θ(υ/m(Υ)),
where m(·) is a non-zero Radon measure satisfying that m(U) < m(V) if U ⊂ V .
Because larger set has larger volume, i.e., V (U) ≤ V (V) if U ⊂ V , from Lemma 1 we have the escaping time is negatively correlated with the volume of the set Υ. Therefore, we can come to the conclusion that for both Adam and AdaMomentum, if the basin Ω is sharp which is ubiquitous during the early stage of training, Υ has a large Radon measure, which leads to smaller escaping time Γ. This means both Adam and AdaMomentum prefer relatively flat or asymmetric basin He et al. (2019) through the training process.
On the other hand, upon encountering a comparatively flat basin or asymmetric valley Ω, we are able to prove that AdaMomentum will stay longer inside. Before we proceed, we need to impose two mild assumptions.
Assumption 4. There exists a constant H > 0 s.t. ‖∇f(θt)‖2 ≤ H, ‖gt‖2 ≤ H, ∀t ∈ [T ]. Assumption 5. For AdaMomentum, there exists T0 ∈ N s.t., diag(Σt) ≤ β1E(m2t−1)/(2−β1) when t > T0.
Here Assumption 4 is a common assumption in stochastic optimization (Ghadimi & Lan, 2013; Johnson & Zhang, 2013). As β1 is always set as positive number close to 1, Assumption 5 basically requires that the gradient noise variance to be smaller than the second moment of m when t is very large. This assumption is mild as 1) we can select mini-batch size to be large enough to satisfy it as the noise variance is inversely proportional to batch size (Bubeck, 2014). 2) The magnitudes of the variances of the stochastic gradients are usually much lower than that of the gradients (Faghri et al., 2020). Then we can come to the following result.
Proposition 1. Under Assumptions 1-5, upon encountering a comparatively flat basin or asymmetric valley Ω, we have E ( Γ(ADAMOMENTUM) ) ≥ E ( Γ(ADAM) ) .
In other words, when falling into a flat/asymmetric basin, AdaMomentum is more stable than Adam and will not easily escape from it. Combining the aforementioned results and the fact that minima at the flat or asymmetric basins tend to exhibit better generalization performance (as observed in Keskar et al. (2017); He et al. (2019); Hochreiter & Schmidhuber (1997a); Izmailov et al. (2018); Li et al. (2018)), we are able to conclude that AdaMomentum is more likely to converge to minima that generalize better, which may buttress the improvement of AdaMomentum in empirical performance. All the proofs in section 3.2 are provided in Appendix B.

heading: 4 CONVERGENCE ANALYSIS OF ADAMOMENTUM
text: In this section, we establish the convergence theory for AdaMomentum under both convex and non-convex object function conditions. We omit the two bias correction steps in the Algorithm 1 for simplicity and the following analysis can be easily adapted and applied to the de-biased version as well.

heading: 4.1 CONVERGENCE ANALYSIS IN CONVEX OPTIMIZATION
text: We analyze the convergence of AdaMomentum in convex setting utilizing the online learning framework (Zinkevich, 2003). Given a sequence of convex cost functions f1(θ), · · · , fT (θ), the regret is defined as R(T ) = ∑T t=1[ft(θt)− ft(θ ∗)], where θ∗ = argminθ ∑T t=1 ft(θ) is the optimal parameter and ft can be interpreted as the loss function at the t-th step. Then we have:
Theorem 1. Let {θt} and {vt} be the sequences yielded by AdaMomentum. Let αt = α/ √ t, β1,1 = β1, 0 < β1,t ≤ β1 < 1, vt ≤ vt+1 for all t ∈ [T ] and γ = β1/ √ β2 < 1. Assume that the distance between any θt generated by AdaMomentum is bounded, ‖θm − θn‖∞ ≤ D∞ for any m,n ∈ {1, · · · , T}. Then we have the following bound on the regret: R(T ) ≤ D 2 ∞ √ T
2α(1− β1) d∑ i=1 √ vT,i+ D2∞ 2(1− β1) T∑ t=1 d∑ i=1 β1,t √ vt,i αt +
α √ 1 + log T
(1− β1)3(1− γ) √ 1− β2 d∑ i=1 ‖g1:T,i‖2 .
Theorem 1 implies that the regret of AdaMomentum can be bounded by Õ1( √ T ), especially when the data features are sparse as Section 1.3 in Duchi et al. (2011) and then we have ∑d i=1 √ vT,i
√ d and∑d
i=1 ‖g1:T,i‖2 √ dT . When we impose additional assumptions that β1,t decays exponentially and that the gradients of ft are bounded (Kingma & Ba, 2015; Liu et al., 2019), we can obtain the following corollary: Corollary 1. Further Suppose β1,t = β1λt and the function ft has bounded gradients, ‖∇ft(θ)‖∞ ≤ G∞ for all θ ∈ Rd, AdaMomentum achieves the guarantee R(T )/T = Õ(1/ √ T ) for all T ≥ 1:
R(T ) T ≤ dG∞α
√ 1 + log T (1− β1)3(1− γ) √ (1− β2)T + dD2∞G∞ 2α(1− β1) √ T + dD2∞G∞β1 2α(1− β1)(1− λ)2T .
Clearly observed from Corollary 1, the average regret of AdaMomentum converges to zero as T goes to infinity. The proofs of Theorem 1 and Corollary 1 are provided in Appendix C.1.

heading: 4.2 CONVERGENCE ANALYSIS IN NON-CONVEX OPTIMIZATION
text: When f is non-convex and lower-bounded, we derive the non-asymptotic convergence rate of AdaMomentum. Theorem 2. We suppose that Assumptions 2 and 4 hold, and β1,t is chosen such that 0 ≤ β1,t+1 ≤ β1,t < 1, 0 < β2,t < 1,∀t ∈ [T ]. We further assume αt/ √ vt ≥ αt+1/ √ vt+1 for any t ∈ [T ], ∑T t=1 α 2 t ≤ η(T ) . TαT , mint∈[T ],j∈[d] vt,j ≥ c ≥ and mint∈[T ] αt ≤ α, then we have:
min t∈[T ]
E ‖∇f(θt)‖22 ≤ H
TαT
[ C1H 2η(T )
c + C2 dα√ c + C3 dα2 c + C4 ] = 1
TαT (Q1 +Q2η(T )),
for some positive constants Q1, Q2. Here C1, C2, C3 are positive constants independent of d or T , while C4 is a positive constant independent of T .
Note that the conditions in Theorem 2 can be satisfied in most scenarios (Kingma & Ba, 2015; Chen et al., 2019; Reddi et al., 2019). For instance, we can simply employ the common setting αt = α/ √ t, β2,t = 1/t. Corollary 2. When αt is further chosen to be α/ √ t, AdaMomentum satisfies:
min t∈[T ] E ‖∇f(θt)‖22 ≤ H√ Tα
[ C1H 2α2(1 + log(T ))
c + C2 dα√ c + C3 dα2 c + C4 ] =
1√ T (Q∗1 +Q ∗ 2 log(T )),
for some constants Q∗1, Q∗2. C1, C2, C3, C4 are similarly defined in Theorem 2. Corollary 2 manifests theO(log(T )/ √ T ) convergence rate of AdaMomentum in the nonconvex case when we commonly use αt = α/ √ t. We refer readers to the detailed proof in Appendix C.2.

heading: 5 EXPERIMENTS
text: We empirically investigate the performance of AdaMomentum in both optimization and generalization.

heading: 5.1 FASTER & BETTER OPTIMIZATION

heading: 5.1.1 2-D TOY EXAMPLES
text: We compare the optimization process of AdaMomentum with three prevalent optimziers SGDM, RMSprop, Adam and recently proposed AdaBelief , on four classic and representative 2-variable objective functions in
1Õ(·) denotes O(·) with hidden logarithmic factors.
numerical optimization literature: Sphere Function (bowl-shaped) (Dixon, 1978), Three-Hump Camel Function (valley-shaped)2, Beale Function (multimodal)3 and Ackley Function (with numerous local minima) (Adorio & Diliman, 2013). To ensure fair comparison, we use the same hyperparameters in the four adaptive gradient methods and finetune the learning rate of SGDM for the best performance. As illustrated in Figure 2, AdaMomentum achieves the most rapid convergence in all the cases and is stable once reaching the global mininum. Meanwhile, SGDM is highly unstable and inaccurate in the descending directions, and RMSprop, Adam, AdaBelief is much slower. Although these toy examples are simple, they give hints to the behavior of optimizers in complex deep learning tasks as they can be viewed as the local dynamics which occur frequently in deep learning (Zhuang et al., 2020). The details of the loss functions and hyperparameter configurations of the experiment are in Appendix A. The GIFs and the 3D trajectory figures are included in the supplementary.

heading: 5.1.2 GENERATIVE ADVERSARIAL NETWORK
text: Training of GANs is extremely unstable. To further study the optimization ability and stability of AdaMomentum, we experiment on GAN equipped with spectal normalization (Miyato et al., 2018). For the generator and the discriminator network, we adopt ResNets for adequate expression ability. We train the model for 100000 iterations on CIFAR-10 with batch size 64, and the two learning rates are set both as 0.0002. For AdaMomentum all the other hyperparameters are set as default values. Experiments are run 5 times independently and we report the mean and standard deviation of Frechet Inception Distance (FID, the lower the better) Heusel et al. (2017) in Table 2. From Table 2 it is reasonable to draw the conclusion that AdaMomentum outperforms all the best tuned baseline optimizers by a large margin, reaching mean FID score as low as 12.06 with its default hyperparameter values. Here Adam equals AdamW because the optimal weight decay parameter value is 0.

heading: 5.2 SUPERIOR GENERALIZATION
text: We conduct experiments on various modern network architectures for different tasks covering both vision and language processing area: 1) image Classification on CIFAR-10 (Krizhevsky & Hinton, 2009) and ImageNet (Russakovsky et al., 2015) with CNN; 2) language modeling on Penn Treebank (Marcus et al., 1993) dataset using Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997b); 3) neural machine translation on IWSTL’14 DE-EN (Cettolo et al., 2014) dataset employing Transformer. We compare AdaMomentum with seven state-of-the-art optimizers: SGDM (Sutskever et al., 2013), Adam (Kingma & Ba, 2015), AdamW (Loshchilov & Hutter, 2017b), Yogi (Reddi et al., 2018), AdaBound (Luo et al., 2019), RAdam (Liu et al., 2019) and AdaBelief (Zhuang et al., 2020). We perform a careful and extensive hyperparameter tuning for
2https://en.wikipedia.org/wiki/Test_functions_for_optimization 3http://www-optima.amp.i.kyoto-u.ac.jp/member/student/hedar/Hedar_
files/TestGO.htm
all the optimizers compared, and the detailed strategy and configuration are given in Appendix D due to space limit. It is worth mentioning that in experiments we discover that setting α = 0.001, β1 = 0.9, β2 = 0.999 (the default setting for adaptive gradient methods in applied machine learing) works well in most cases. This elucidates that our optimizer is tuning-friendly, which reduces human labor and time cost and is crucial in practice. The mean results with standard deviations over 5 seeds are reported in all the following experiments.

heading: 5.2.1 CNN FOR IMAGE CLASSIFICATION
text: CIFAR-10 We experimented with three prevailing deep CNN architectures: VGG-16 (Simonyan & Zisserman, 2015), ResNet-34 (He et al., 2016) and DenseNet-121 (Huang et al., 2017). The growth rate of DenseNet121 is set as 12 to match CIFAR-10 dataset. In each experiment we train the model for 200 epochs with batch size 128 and decay the learning rate by 0.2 at the 60-th, 120-th and 160-th epoch. We employ label smoothing technique (Szegedy et al., 2016) and the smoothing factor is choosen as 0.1. Figure 3 displays the training and testing results of all the compared optimizers . As indicated, both the training accuracy and the testing accuracy using AdaMomentum can be improved as fast as with other adaptive gradient methods, being much faster than SGDM, especially before the third learning rate annealing. In testing phase, AdaMomentum can exhibit performance as good as SGDM and far exceeds other baseline adaptive gradient methods, including the recently proposed AdaBelief (Zhuang et al., 2020) optimizer. This contradicts the result reported in Zhuang et al. (2020), where they claim AdaBelief can be better than SGDM. This largely stems from the fact that Zhuang et al. (2020) did not take an appropriate stepsize annealing strategy or tune the hyperparameters well. Training 200 epochs with ResNet-34 on CIFAR-10, our experiments show that AdaMomentum and SGDM can reach over 96% accuracy, while in Zhuang et al. (2020) the accuracy of SGDM is only around 94% .
ImageNet To further corroborate the effectiveness of our algorithm on more comprehensive dataset, we perform experiments on ImageNet utilizing ResNet-18 as backbone network. We execute each optimizer for 90 epochs utilizing cosine annealing strategy, which can exhibit better performance results than step-based decay strategy on ImageNet (Loshchilov & Hutter, 2017a; Ma, 2021). As indicated in Table 4, AdaMomentum far exceeds Adam in Top-1 test accuracy and outperforms all the competitors including SGD with momentum.

heading: 5.2.2 LSTM FOR LANGUAGE MODELING
text: We implement LSTMs with layer number from 1 to 3 on Penn Treebank dataset, where adaptive gradient methods are the main-stream choices (much better than SGD). In each experiment we train the model for 200 epochs with batch size of 20 and decay the learning rate by 0.1 at 100-th and 145-th epoch. Test perplexity (the lower the better) against training epochs is plotted in Figure 4 and the best perplexity value is summarized in Table 5. Clealy observed from Figure 4 and Table 5, AdaMomentum achieves the lowest perplexity in all the settings and consistently outperform other competitors by a considerable margin. The training curve is given in Figure 5 in Appendix D due to space limit. Particularly on 2-layer and 3-layer LSTM, AdaMomentum maintains both the fastest convergence and the best performance, which substantiates its superiority.

heading: 5.2.3 TRANSFORMER FOR NEURAL MACHINE TRANSLATION
text: Transformers have been the dominating architecture in NLP and adaptive gradient methods are usually employed to train transformers due to their stronger ability to handle attention-models (Zhang et al., 2019). To test AdaMomentum on transformer, we experiment on IWSTL’14 German-to-English with
the Transformer small model adapting the code from fairseq package.4 We set the length penalty as 1.0, the beam size as 5, warmup initial stepsize as 10−7 and the warmup updates iteration number to be 8000. We train the models for 55 epochs and the results are reported according to the average of the last 5 checkpoints. As shown in Table 12, our optimizer achieves the highest average BLEU score with the lowest variance.

heading: 6 CONCLUSION
text: In this work, we proposed AdaMomentum as a new optimizer for machine learning. We theoretically demonstrate why AdaMomentum outperforms Adam in optimization and generalization. We further validates the superiority of AdaMomentum through both toy examples and large-scale experiments on real-world datasets. Our algorithm is simple and effective with four key advantages: 1) maintaining fast convergence rate; 2) closing the generalization gap between adaptive gradient methods and SGD; 3) applicable to various tasks and models; 4) introducing no additional parameters and easy to tune. Combination of AdaMomentum with other techniques such as Nesterov’s accelerated gradient (Dozat, 2016) may be of independent interest in the future.
4https://github.com/pytorch/fairseq

heading: 7 ETHICS STATEMENT
text: Our work follows all ethical standards and laws. All the experiments were conducted on publically available datasets, with no new data concerning human or animal subjects generated.

heading: 8 REPRODUCIBILITY STATEMENT
text: We adhere to ICLR reproducibility standards and provide all necessary information to reproduce our experimental and theoretical results. We ensure the reproducibility of our work through several ways, namely
• All the source code and presented figures are available at anonymous link https://anonymous. 4open.science/r/AdaMomentum_experiments-6D9B.
• The detailed descriptions of the classic loss functions used in Toy exmpales in Section 5.1.1 are given in Appendix A.
• All the technical details and proofs in Section 3.2 are included in Appendix B. All the proofs in Section 4 are provided in Appendix C.
• The detailed hyperparameter tuning rule and configurations of the experiments in Section 5 are given in Appendix D.

heading: A DETAILS OF TOY EXAMPLES
text: The detailed description of the four 2-parameter loss functions are given as below:
The Sphere Function is f(x) = x21 + x 2 2, whose global minimum is f(x∗) = 0 at x∗ = (0, · · · , 0).
The Three Hump Camel Function is
f(x) = 2x21 − 1.05x41 + x61 6 + x1x2 + x 2 2,
whose global minimum is f(x∗) = 0 at x∗ = (0, · · · , 0).
The Beale Function is
f(x) = (1.5− x1 + x1x2)2 + (2.2− x1 + x1x22)2 + (2.625− x1 + x1x32)2,
whose global minimum is f(x∗) = 0 at x∗ = (3, 0.5).
The Ackley Function is
f(x) = −20 exp ( −0.2 √ 1
2 (x21 + x 2 2)
) − exp ( 1
2 (cos(2πx1) + cos(2πx2))
) + 20 + exp(1),
whose global minimum is f(x∗) = 0 at x∗ = (0, · · · , 0).
Hyperparameter configuration For all the four toy experiments, we employ the same stepsize for the three adaptive gradient methods: RMSProp, Adam and AdaMomentum. For SGDM We set the momentum parameter as 0.9 and finetune the learning rate for each toy experiment in set {0.0001, 0.001, 0.01, 0.1, 1, 10} as the optimal learning rate varies for different loss functions. After careful parameter tuning, the configuration is: for Sphere Function, the learning rate for adaptive gradient methods and SGDM are both 0.1; for ThreeHump Camel Function, the learning rate for adaptive gradient methods is 0.1 and the learning rate for SGDM is 0.001; for Beale Function, the leanring rate for adaptive, gradient methods is 0.1 and the learning rate for SGDM is 0.1; for Ackley Function, the leanring rate for adaptive, gradient methods is 0.1 and the learning rate for SGDM is 0.1. Combining this with Figure 2, we can see that AdaMomentum with the stepsize 0.1 manifests universally faster and more stable convergence to the optimum than discreetly tuned SGDM and other adaptive gradient approaches.

heading: B TECHNICAL DETAILS OF SUBSECTION 3.2
text: Here we provide more construction details and technical proofs for the Lévy-driven SDE in Adam-alike adaptive gradient algorithm (2). In the beginning we introduce a detailed derivation of the process (3) as well as its corresponding escaping set Υ in definition 2. Then we give some auxiliary theorems and lemmas, and summarize the proof of Lemma 1. Finally we prove the proposition 1 and give a more detailed analysis of the conclusion that the expected escaping time of AdaMomentum is longer than that of Adam in a comparatively flat basin.

heading: B.1 DERIVATION OF THE LÉVY-DRIVEN SDE (3)
text: To derive the SDE of Adam-alike algorithms (2), we firstly define m′t = β1m′t−1 + (1 − β1)∇f(θt) with m′0 = 0. Then by the definition it holds that
m′t −mt = (β1 − 1) t∑ i=0 βt−i1 ζt.
Following Simsekli et al. (2019), the gradient noise ζt has heavy tails in reality and hence we assume that 1
1−β1 (m′t −mt) obeys Sα̃S distribution with time-dependent covariance matrix Σt. Since we can formulate (2) as
θt+1 = θt − α m′t zt + α (m′t −mt) zt where zt = (1− βt1)
√ vt
(1− βt2) , (4)
and we can replace the term (m′t − mt) by α− 1 α̃ (1 − βt1)ΣtS where each coordinate of S is independent and identically distributed as Sα̃S(1) based on the property of centered symmetric α̃-stable distribution. Let
Rt = diag( √
vt (1−βt2) ), and we further assume that the step size α is small, then the continuous-time version of
the process (4) becomes the following SDE:
dθt = −R−1t m′tdt
(1− βt1) + α1− 1 α̃R−1t ΣtdLt, dmt = β1(∇f(θt)−mt), dvt = β2(k2t − vt).
After replacing m′t with mt for brevity, we get the SDE (3) consequently.

heading: B.2 PROOF OF LEMMA 1
text: To prove Lemma 1, we first introduce Theorem 3.
Theorem 3. Suppose Assumptions 1-3 hold. We define κ1 = c1Lv−|τm−1| and κ2 = 2µτ β1v++µτ ( β1 − β24 ) with
a constant c1. Let υα̃+1 = Θ(α̃), ρ0 = 116(1+c2) and ln ( 2∆ µυ1/3 ) ≤ κ2υ−1/3 where ∆ = f(θ0)− f(θ∗) and a constant c2. Then for any θ0∈ Ω−2υ γ
, u >−1, υ ∈ (0, υ0], γ ∈ (0, γ0] and ρ ∈ (0, ρ0] satisfying υγ ≤ ρ0 and limυ→0 ρ = 0, the Adam-alike algorithm in (2) obey
1− ρ 1 + u+ ρ
≤ E [ exp ( −um(Υ)Θ(υ−1)Γ )] ≤ 1 + ρ
1 + u− ρ .
From Theorem 3, by setting υ small, it holds that for any adaptive gradient algorithm the upper and lower bounds of its expected escaping time Γ is at the order of ( υ
m(Υ)
) , which directly implies Lemma 1 conclusively.
Therefore, it suffices to validate Theorem 3.
Theorem 3 is adapted from Theorem 1 in Zhou et al. (2020) and the proof is given in Section B.2.3. Before we proceeed, we first provide some prerequisite notations in Section B.2.1 and list some useful theorems and lemmas in Section B.2.2.

heading: B.2.1 PRELIMINARIES
text: For analyzing the uniform Lévy-driven SDEs in (3), we first introduce the Lévy processLt into two components ξt and εi, namely Lt = ξt + εt, (5) whose characteristic functions are respectively defined as
E [ ei〈λ,ξt〉 ] =e t ∫ Rd\{0}εI { ‖y‖2≤ 1 υδ } ν(dy) ,
E [ ei〈λ,εt〉 ] =e t ∫ Rd\{0}εI { ‖y‖2≤ 1 υδ } ν(dy) ,
where ε = ei〈λ,y〉−1−i〈λ, y〉I { ‖y‖2 ≤ 1 } with υ defined in (3) and a constant δ s.t. υ−δ < 1. Accordingly, the Lévy measure ν of the stochastic processes ξ and ε are
νξ = ν ( A ∩ { ‖y‖2 ≤ 1
υδ
}) , νε = ν ( A ∩ { ‖y‖2 ≥ 1
υδ
}) , where A ∈ B(Rd).
Besides, for analysis, we should consider affects of the Lévy motion Lt to the Lévy-driven SDE of Adam variants. Here we define the Lévy-free SDE accordingly: dθ̂t = −µtQ̂−1t m̂t, dm̂t = β1(∇f(θ̂t)− m̂t), dv̂t = β2(∇(fθ̂t)2 − v̂t). (6) where Q̂t = diag( √ v̂t).

heading: B.2.2 AUXILIARY THEOREMS AND LEMMAS
text: Theorem 4 (Zhou et al. (2020)). Suppose Assumptions 1-3 hold. Assume the sequence {(θ̂t, m̂t, v̂t)} are produced by (6). Let ŝt = htqt (√ ωtv̂t ) with ht = β1, qt = (1 − βt1)−1 and ωt = (1 − βt2)−1. We define
‖x‖2y = ∑ i yix 2 i . Then for Lévy-driven Adam SDEs in (6), its Lyapunov function L(t) = f(θ̂t) − f(θ̂∗) + 1 2 ‖m̂t‖ŝ−1t with the optimum solution θ ∗ in the current local basin Ω obeys
L(t) ≤ ∆ exp ( − 2µτ β1v+ + µτ ( β1 − β2 4 ) t ) ,
where ∆ = f(θ̂0)− f(θ̂∗) due to m̂0 = 0. The sequence {θ̂t} produced by (6) obeys∥∥∥θ̂t − θ∗∥∥∥2 2 ≤ 2∆ µ exp ( − 2µτ β1v+ + µτ ( β1 − β2 4 ) t ) .
Lemma 2 (Zhou et al. (2020)). (1) The process ξ in the Lévy process decomposition can be decomposed into two processes ξ̂ and linear drift, namely,
ξt = ξ̂t + µυt, (7)
where ξ̂ is a zero mean Lévymartingale with bounded jumps. (2) Let δ ∈ (0, 1), µυ = E(ξ1) and Tυ = υ−θ for some θ > 0, ρ0 = ρ0(δ) = 1−δ4 > 0 and θ0 = θ0(δ) = 1−δ
3 > 0. Suppose υ is sufficiently small such that Θ(1) ≤ υ−
1−δ 6 and υ−ρ − 2(C + Θ(1))υ 7 6 (1−δ)+ ρ 2 ≥ 1 with a constant C = | ∫
0<u≤1 u 2dΘ(u)| ∈ (0,+∞). Then for all δ ∈ (0, δ0), θ ∈ (0, θ0) there are p0 =
p0(δ) = δ 2 and υ0 = υ0(δ, ρ) such that the estimates
‖υξTυ‖2 = υ ‖µυ‖2 Tυ < υ 2ρ and P ([υξ]dTυ ≥ υ ρ) ≤ exp(−υ−p),
hold for all p ∈ (0, p0] and υ ∈ (0, υ0] Lemma 3 (Zhou et al. (2020)). Let δ ∈ (0, 1) and gtt≥0 be a bounded adapted cȧdlȧg stochastic process with values in Rd, Tυ = υ−θ , θ > 0. Suppose supt≥0 ‖gt‖ is well bounded. Assume ρ0 = ρ0(δ) = 1−δ16 > 0, θ0 = θ0(δ) = 1−δ 3
> 0, p0 = ρ2 . For ξ̂t in (7), there is δ0 = δ0(δ) > 0 such that for all ρ ∈ (0, ρ0) and θ ∈ (0, θ0), it holds
P ( sup
0≤t≤Tυ υ ∣∣∣∣∣ d∑ i=1 ∫ t 0 gis−dξ̂ i s ∣∣∣∣∣ ≥ υρ ) ≤ 2 exp ( −υ−p ) ,
for all p ∈ (0, p0] and 0 < υ ≤ υ0 with υ0 = υ(ρ), where ξ̂is represents the i-th entry in ξ̂s. Lemma 4 (Zhou et al. (2020)). Under Assumptions 1-3 hold, assume δ ∈ (0, 1), ρ0 = ρ0(δ) = 1−δ16(1+c1κ1) > 0, θ0 = θ0(δ) = 1−δ
3 > 0, p0 = min( ρ̂(1+c1κ1) 2 , p), 1 c2
ln (
2∆ µυρ̂ ) ≤ υ−θ0 where κ1 = c2lv−|τm−1| and
c2 = 2µτ
β1v++µτ
( β1 − β24 ) in Adam-alike adaptive gradient algorithms. For all ρ̂ ∈ (0, ρ0), p ∈ (0, p0],
0 < υ ≤ υ0 with υ0 = υ0(ρ̂), and θ0 = θ̂0, we have
sup θ0∈Ω
P (
sup 0≤t<σ1 ∥∥∥θt − θ̂t∥∥∥ 2 ≥ 2υρ̂ ) ≤ 2 exp(−υ− p 2 ), (8)
where the sequences θt and θ̂t are respectively produced by (3) and (6) in adaptive gradient method .

heading: B.2.3 PROOF OF THEOREM 3
text: Proof. The idea of this proof comes from (8) we showed in Lemma 4 where the sequence θt and θ̂t start from the same initialization. Based on Theorem 4, we know that the sequence {θ̂t} from (6) exponentially converges to the minimum θ∗ of the local basin Ω. To escape the local basin Ω, we can either take small steps in the process ζ or large jumps Jk in the process ε. However, (8) suggests that these small jumps might not be helpful for escaping the basin. And for big jumps, the escaping time Γ of the sequence {θt} most likely occurs at the time σ1 if the big jump υJ1 in the process ε is large. The verification of our desired results can be divided into two separate parts, namely establishing upper bound and lower bound of E [ exp ( −um(Υ)Θ(υ−1)Γ )] for any u > −1. Both of them can be established based on the following facts:∣∣∣P(R−1θ ΣθυJk 6∈ Ω±υγ , ‖υJk‖2 ≤ R)− P(R−1θ∗ Σθ∗υJk 6∈ Ω±υγ , ‖υJk‖2 ≤ R)∣∣∣ ≤ δ′4 · Θ(υ−1)Θ(υ−δ) ,∣∣P (R−1θ ΣθυJk 6∈ Ω, ‖υJk‖2 ≤ R)− P (R−1θ∗ Σθ∗υJk 6∈ Ω, ‖υJk‖2 ≤ R)∣∣ ≤ δ′4 · Θ(υ−1)Θ(υ−δ) , P ( R−1θ∗ Σθ∗υJk 6∈ Ω ) − P ( R−1θ∗ Σθ∗υJk 6∈ Ω, ‖υJk‖2 ≤ R ) ≤ δ ′
4 · Θ(υ
−1)
Θ(υ−δ) . (9)
Specifically, for the upper bound of E [ exp ( −um(Υ)Θ(υ−1)Γ )] , we consider both the big jumps in the process ε and small jumps in the process ζ which may escape the local minimum. Instead of estimating the escaping time Γ from Ω, we first estimate the escaping time Ξ̃ from Ω−ρ̄. Here we define the inner part of Ω as Ω−ρ̄ := {y ∈ Ω : dis(∂Ω, y) ≥ ρ̄}. Then by setting ρ̄ → 0, we can use Ξ̃ for a decent estimation of Γ.
We denote ρ̄ = υγ where γ is a constant such that the results of Lemma 2-4 hold. So for the upper bound we mainly focus on Ξ̃ in the beginning and then transfer the results to Γ. In the beginning, we can show that for any u > −1 it holds that,
E [ exp ( −um(Υ)Θ(υ−1)Ξ̃ )] ≤ +∞∑ k=1 E [ e−um(Υ)Θ(υ −1)tkI { Ξ̃ = tk } +Resk ] ,
where
Resk ≤ E [ e−um(Υ)Θ(υ −1)tkI { Ξ̃ ∈ (tk−1, tk) }] , if u ∈ (−1, 0] E [ e−um(Υ)Θ(υ −1)tk−1I { Ξ̃ ∈ (tk−1, tk) }] , if u ∈ (0,+∞).
Then using the strong Markov property we can bound the first term E [ e−um(Υ)Θ(υ −1)tkI { Ξ̃ = tk }] as
R1 = +∞∑ k=1 E [ e−um(Υ)Θ(υ −1)tkI {Γ = tk} ] ≤αυ(1 + ρ/3) 1 + uαυ +∞∑ k=1 ( 1− αυ(1− ρ) 1 + uαυ )k−1
≤αυ(1 + ρ/3) 1 + uαυ +∞∑ k=0 ( 1− αυ(1− ρ) 1 + uαυ )k−1 = 1 + ρ/3
1 + u− ρ .
On the other hand, for the lower bound of E [ exp ( −um(Υ)Θ(υ−1)Γ )] , we only consider the big jumps in the process ε which could escape from the basin, and ignore the probability that the small jumps in the process ζ which may also lead to an escape from the local minimum θ∗. Specifically, we can find a lower bound by discretization:
E [ exp ( −um(Υ)Θ(υ−1)Γ )] ≥ +∞∑ k=1 E [ exp ( −um(Υ)Θ(υ−1)tk ) I{Γ = tk} ] .
Then we can lower bound each term by three equations (9) we just listed here, which implies that for any θ0 ∈ Ω−υ γ ,
E [ e−um(Υ)Θυ −1Γ ] ≥ αυ(1− ρ)
1 + uαυ +∞∑ k=1 ( 1− αυ(1 + ρ) 1 + uαυ )k−1 = 1− ρ 1 + u+ ρ ,
where ρ→ 0 as υ → 0. The proof is completed.

heading: B.3 PROOF OF PROPOSITION 1
text: Proof. Since we assumed the minimizer θ∗ = 0 in the basin Ω which is usually small,we can employ secondorder Taylor expansion to approximate Ω as a quadratic basin whose center is θ∗. In other words, we can write
Ω = { y ∈ Rd ∣∣∣∣ f(θ∗) + 12y>H(θ∗)y ≤ h(θ∗) } ,
where H(θ∗) is the Hessian matrix at θ∗ of function f and h(θ∗) is the basin height. Then according to Definition 2, we have
Υ = { y ∈ Rd ∣∣∣ y>Σθ∗R−1θ∗ H(θ∗)R−1θ∗ Σθ∗y ≥ h∗f} . Here Rθ∗ = limθt→θ∗ diag( √ vt/(1− βt2)) is a matrix depending on the algorithm, h∗f = 2(h(θ∗)− f(θ∗)) and Σθ∗ is independent of the alogorithm, i.e. the same for Adam and AdaMomentum. Firstly, we will prove that v(ADAMOMENTUM)t ≥ v (ADAM) t when t → ∞. To clarify the notation, we use θt,mt, vt, gt to denote the symbols for Adam and θ̃t, m̃t, ṽt, g̃t for AdaMomentum, and ζt is the gradient noise. By using Lemma 1 and above results, we have θt ≈ θ̃t ≈ θ∗ before escaping when t is large, and thus vt = limθt→θ∗ [∇f(θt) + ζt]2 and ṽt = limθt→θ∗ [β1m̃t−1 + (1 − β1)(∇f(θ̃t) + ζt)]2. We will firstly show that E(ṽt) ≥ E(vt) when t is large.
E(vt) = E( lim θt→θ∗ [∇f(θt) + ζt]2) (i) = lim θt→θ∗ E([∇f(θt) + ζt]2)
= lim θt→θ∗
( E(∇f(θt)2) + E(2∇f(θt)ζt) + E(ζ2t ) ) (ii) =E( lim
θt→θ∗ ∇f(θt)2) + lim θt→θ∗ E(2∇f(θt)ζt) + lim θt→θ∗ E(ζ2t )
(iii) = lim θt→θ∗ E(ζ2t ),
where (i) and (ii) are due to the dominated convergence theorem (DCT) since we have that we know both ‖∇f(θt)‖2 and ‖∇f(θt) + ζ2‖2 could be bounded by H in Assumption 4. And (iii) is due to the fact that ∇f(θ∗) = 0 since function f attains its minimum point at θ∗, and ζt has zero mean, i.e.
lim θt→θ∗ E(∇f(θt)ζt) = lim θt→θ∗ E(∇f(θt))E(ζt) = 0.
And similarly we can prove that, E(ṽt) = E (
lim θt→θ∗
[β1m̃t−1 + (1− β1)(∇f(θ̃t) + ζt)]2 )
= lim θt→θ∗
( E(β21m̃2t−1) + E((1− β1)2(∇f(θ̃t) + ζt)2) + E(2β1(1− β1)m̃t−1∇(f(θ̃t) + ζt)) ) (i) =β21 lim
θt→θ∗ E(m̃2t−1) + (1− β1)2 lim θt→θ∗ E(ζ2t ),
where we can get the equality (i) simply by the same argument with dominated convergence theorem we just used:
lim θ̃t→θ∗ E(∇(f(θ̃t)2) = E( lim θ̃t→θ∗ ∇(f(θ̃t)2) (i) = 0,
lim θ̃t→θ∗ E(∇(f(θ̃t)ζt) = E( lim θ̃t→θ∗
∇(f(θ̃t)ζt) (ii) = 0,
lim θ̃t→θ∗ E(m̃t−1(∇f(θ̃t) + ζt)) = E( lim θ̃t→θ∗ m̃t−1∇f(θ̃t)) + lim θ̃t→θ∗
E(m̃t−1)E(ζt) (iii) = 0,
where we get the equality (i) and (ii) since the function f(θ̃t)2 and f(θ̃t)ζt could be absolutely bounded by H2. And the first term in equality (iii) is 0 since we have ‖m̃t−1‖2 ≤ H by its definition and ∇f(θ
∗) = 0, and the second term vanishes since the noise ζt has zero mean. Based on the Assumption 5, we have
E(m̃2t−1) ≥ 2− β1 β1 E(ζ2t ),
which implies that E(ṽt) ≥ E(vt) when t is large. It further indicates that R(ADAMOMENTUM)θ∗ ≥ R (ADAM) θ∗ . We consider the volume of the complementary set
Υc = { y ∈ Rd ∣∣∣ y>Σθ∗R−1θ∗ H(θ∗)R−1θ∗ Σθ∗y < h∗f} , which can be viewed as a d-dimensional ellipsoid. We can further decompose the symmetric matrix M := Σθ∗R −1 θ∗ H(θ ∗)R−1θ∗ Σθ∗ by SVD decomposition
M = U>AU,
where U is an orthogonal matrix and A is a diagonal matrix with nonnegative elements. Hence the transformation y → Uy is an orthogonal transformation which means the volume of Υc equals the volume of set{
y′ ∈ Rd ∣∣∣ y′>Ay′ < h∗f} .
Considering the fact that the volume of a d-dimensional ellipsoid centered at 0 Ed(r) = {(x1, x2, · · · , xn) :∑d i=1
x2i R2i ≤ 1} is
V (Ed(r)) = π n 2
Γ(n 2
+ 1) Πni=1Ri,
and the fact we just proved that R(ADAMOMENTUM)θ∗ ≥ R (ADAM) θ∗ . Therefore we deduce the volume of Υ(ADAMOMENTUM) is smaller than that of Υ(ADAM), which indicates that for Radon measure m(·) we have m(Υ(ADAMOMENTUM)) ≥ m(Υ(ADAM)). Based on Lemma 1, we consequently have E(Γ(ADAMOMENTUM)) ≥ E(Γ(ADAM)).

heading: C PROOFS IN SECTION 4
text: C.1 PROOF OF THE CONVERGENCE RESULTS FOR THE CONVEX CASE

heading: C.1.1 PROOF OF THEOREM 1
text: Proof. Firstly, according to the definition of AdaMomentum in Algorithm 1, by algebraic shrinking we have T∑ t=1 m2t,i√ tvt,i = T−1∑ t=1 m2t,i√ tvt,i + (∑T j=1(1− β1,j)Π T−j k=1 β1,T−k+1gj,i )2 √ T ∑T j=1(1− β2)β T−j 2 m 2 j,i
≤ T−1∑ t=1 m2t,i√ tvt,i + ( ∑T j=1 Π T−j k=1 β1,T−k+1)( ∑T j=1 Π T−j k=1 β1,T−k+1g 2 j,i)√ T ∑T j=1(1− β2)β T−j 2 m 2 j,i (i) ≤ T−1∑ t=1 m2t,i√ tvt,i + ( ∑T j=1 β T−j 1 )( ∑T j=1 β T−j 1 g 2 j,i)√ T (1− β2) ∑T j=1 β T−j 2 m 2 j,i ≤ T−1∑ t=1 m2t,i√ tvt,i + 1 1− β1 ∑T j=1 β T−j 1 g 2 j,i√ T (1− β2) ∑T j=1 β T−j 2 m 2 j,i
= T−1∑ t=1 m2t,i√ tvt,i + 1 (1− β1) √ T (1− β2) T∑ j=1 βT−j1 g 2 j,i√∑T
j=1 β T−j 2 (∑j l=1(1− β1,l)Π j−l k=1β1,j−k+1gl,i )2 ≤ T−1∑ t=1 m2t,i√ tvt,i + 1 (1− β1) √ T (1− β2) T∑ j=1 βT−j1 g 2 j,i√∑T
j=1 β T−j 2 ((1− β1,j)gj,l) 2
≤ T−1∑ t=1 m2t,i√ tvt,i + 1 (1− β1) √ T (1− β2) T∑ j=1 βT−j1 g 2 j,i√
βT−j2 (1− β1,j)2g2j,i (ii) ≤ T−1∑ t=1 m2t,i√ tvt,i + 1 (1− β1)2 √ T (1− β2) T∑ j=1 γT−jgj,i,
where (i) arises from β1,t ≤ β1, and (ii) comes from the definition that γ = β1√β2 . Then by induction, we have T∑ t=1 m2t,i√ tvt,i ≤ T∑ t=1 1 (1− β1)2 √ t(1− β2) t∑ j=1 γt−jgj,i
≤ 1 (1− β1)2 √ 1− β2 T∑ t=1 1√ t t∑ j=1 γt−jgj,i
(i) ≤ 1 (1− β1)2 √ 1− β2 T∑ t=1 gt,i T∑ j=t γj−t√ j
≤ 1 (1− β1)2 √ 1− β2 T∑ t=1 gt,i T∑ j=t γj−t√ t
≤ 1 (1− β1)2 √ 1− β2 T∑ t=1 gt,i · 1 (1− γ) √ t
≤ 1 (1− β1)2(1− γ) √ 1− β2 T∑ t=1 gt,i√ t
(ii) ≤ 1
(1− β1)2(1− γ) √ 1− β2 ‖g1:T,i‖2 √√√√ T∑ t=1 1 t
(iii) ≤
√ 1 + log T
(1− β1)2(1− γ) √ 1− β2 ‖g1:T,i‖2 ,
where (i) exchangings the indices of summing, (ii) employs Cauchy-Schwarz Inequality and (iii) comes from the following bound on harmonic sum:
T∑ t=1 1 t ≤ 1 + log T.
Due to convexity of ft, we get
ft(θt)− ft(θ∗) ≤ g>t (θt − θ∗)
= d∑ i=1 gt,i(θt,i − θ∗,i). (10)
According to the updating rule, we have
θt+1 = θt − αt mt√ vt
= θt − αt ( β1,t√ vt mt−1 + 1− β1,t√ vt gt ) . (11)
Substracting θ∗, squaring both sides and considering only the i-th element in vectors, we obtain
(θt+1,i − θ∗,i)2 = (θt,i − θ∗,i)2 − 2αt ( β1,t√ vt,i mt−1,i + 1− β1,t√ vt,i gt,i ) (θt,i − θ∗,i) + α2t ( mt,i√ vt,i )2 .
By rearranging the terms, we have
2αt 1− β1,t√
vt,i gt,i(θt,i−θ∗,i) = (θt,i−θ∗,i)2−(θt+1,i−θ∗,i)2−2αt · β1,t√ vt,i ·mt−1,i(θt,i−θ∗,i)+α2t ( mt,i√ vt,i )2 .
Further we have
gt,i(θt,i − θ∗,i) = √ vt,i
2αt(1− β1,t) [(θt,i − θ∗,i)2 − (θt+1,i − θ∗,i)2] +
αt √ vt,i
2(1− β1,t) ( mt,i√ vt,i )2 +
β1,t 1− β1,t (θ∗,i − θt,i)mt−1,i
=
√ vt,i
2αt(1− β1,t) [(θt,i − θ∗,i)2 − (θt+1,i − θ∗,i)2] +
αt √ vt,i
2(1− β1,t) ( mt,i√ vt,i )2
+ β1,t 1− β1,t · v 1 4 t,i√ αt · (θ∗,i − θt,i) · √ αt · mt−1,i
v 1 4 t,i
≤ √ vt,i
2αt(1− β1) [(θt,i − θ∗,i)2 − (θt+1,i − θ∗,i)2] +
α 2(1− β1) · m2t,i√ tvt,i
(12)
+ β1,t 2αt(1− β1,t) (θ∗,i − θt,i)2 √ vt,i + β1α 2(1− β1) · m2t−1,i√ tvt,i , (13)
where (13) bounds the last term of (12) by Cauchy-Schwarz Inequality and plugs in the value of αt. Plugging (13) into (11) and summing from t = 1 to T , we obtain
R(T ) = T∑ t=1 d∑ i=1 gt,i(θt,i − θ∗,i)
≤ T∑ t=1 d∑ i=1
√ vt,i
2αt(1− β1) [(θt,i − θ∗,i)2 − (θt+1,i − θ∗,i)2] + T∑ t=1 d∑ i=1
α 2(1− β1) · m2t,i√ tvt,i
+ T∑ t=1 d∑ i=1 β1,t 2αt(1− β1,t) (θ∗,i − θt,i)2 √ vt,i + T∑ t=1 d∑ i=1 β1α 2(1− β1) · m2t−1,i√ tvt,i
≤ d∑ i=1
√ v1,i
2α1(1− β1) (θ1,i − θ∗,i)2 +
1
2(1− β1) T∑ t=2 d∑ i=1 (θt,i − θ∗,i)2 (√ vt,i αt − √ vt−1,i αt−1 ) (14)
+ T∑ t=1 d∑ i=1 β1,t 2αt(1− β1) (θ∗,i − θt,i)2 √ vt,i + T∑ t=1 d∑ i=1 α 1− β1 · m2t,i√ tvt,i , (15)
where (15) rearranges the first term of (14). Finally utilizing the assumptions in Theorem 1, we get
R(T ) ≤ d∑ i=1
√ v1,i
2α1(1− β1) D2∞ +
1
2(1− β1) T∑ t=2 d∑ i=1 D2∞ (√ vt,i αt − √ vt−1,i αt−1 )
+ D2∞
2(1− β1) T∑ t=1 d∑ i=1 β1,tv 1 2 t,i αt + d∑ i=1
α √ 1 + log T
(1− β1)3(1− γ) √ 1− β2 ‖g1:T,i‖2
= d∑ i=1
√ vT,i
2αT (1− β1) D2∞ + D2∞ 2(1− β1) T∑ t=1 d∑ i=1 β1,tv 1 2 t,i αt
+ d∑ i=1
α √ 1 + log T
(1− β1)3(1− γ) √ 1− β2 ‖g1:T,i‖2 , (16)
which is our desired result.

heading: C.1.2 PROOF OF COROLLARY 1
text: Proof. Plugging αt = α√t and β1,t = β1λ t into (16), we get
R(T ) ≤ D 2 ∞ √ T
2α(1− β1) d∑ i=1 √ vT,i + D2∞ 2α(1− β1) T∑ t=1 d∑ i=1 β1λ t √ tvt,i
+ d∑ i=1
α √ 1 + log T
(1− β1)3(1− γ) √ 1− β2 ‖g1:T,i‖2 . (17)
Next, we employ Mathematical Induction to prove that vt, i ≤ G∞ for any 0 ≤ t ≤ T, 1 ≤ i ≤ d. ∀i, we have m20,i = 0 ≤ G2∞. Suppose mt−1,i ≤ G∞, we have
m2t,i = (β1,tmt−1,i + (1− β1,t)gt,i)2
(i) ≤β1,tm2t−1,i + (1− β1,t)g2t,i ≤ β1,tG2∞ + (1− β1,t)G2∞ = G2∞,
where (i) comes from the convexity of function f = x2. Hence by induction, we have m2t,i ≤ G2∞ for all 0 ≤ t ≤ T . Furthermore, ∀i, we have v0,i = 0 ≤ G2∞. Suppose vt−1,i ≤ G2∞, we have
vt,i = β2vt−1,i + (1− β2)m2t,i ≤ β2G2∞ + (1− β2)G2∞ = G2∞.
Therefore, by induction, we have vt,i ≤ G2∞, ∀i, t. Combining this with the fact that ∑d i=1 ‖g1:T,i‖2 ≤
dG∞ √ T and (17), we obtain
R(T ) ≤ dG∞D 2 ∞ √ T
2α(1− β1) + dG∞D
2 ∞β1
2α(1− β1) T∑ t=1 λt √ t+
dG∞α √ 1 + log T (1− β1)3(1− γ) √ (1− β2)T . (18)
For ∑T t=1 λ t √ t, we apply arithmetic geometric series upper bound:
T∑ t=1 λt √ t ≤ T∑ t=1 tλt ≤ 1 (1− λ)2 . (19)
Plugging (19) into (18) and dividing both sides by T , we obtain
R(T ) T ≤ dG∞α
√ 1 + log T (1− β1)3(1− γ) √ (1− β2)T + dD2∞G∞ 2α(1− β1) √ T + dD2∞G∞β1 2α(1− β1)(1− λ)2T ,
which concludes the proof.
C.2 PROOF OF THE CONVERGENCE RESULTS FOR THE NON-CONVEX CASE

heading: C.2.1 USEFUL THEOREM
text: Theorem 5. (Chen et al. (2019)) Suppose Assumptions 2 and 4 are satisfied, β1,t is chosen such that 0 ≤ β1,t+1 ≤ β1,t < 1, 0 < β2 < 1,∀t > 0. There exists some constant G such that ∥∥∥αt · mt√vt ∥∥∥2 ≤ G,∀t. Then Adam-type algorithms yield
E [ T∑ t=1 αt 〈∇f(θt),∇f(θt)/ √ vt〉 ]
≤ E [ C1
T∑ t=1 ‖αtgt/ √ vt‖22 + C2 T∑ t=1 ∥∥∥∥ αt√vt − αt−1√vt−1 ∥∥∥∥ 1 + C3 T∑ t=1 ∥∥∥∥ αt√vt − αt−1√vt−1 ∥∥∥∥2 2 ] + C4, (20)
where C1, C2 and C3 are constants independent of d and T , C4 is a constant independent of T , the expectation is taken w.r.t all randomness corresponding to {gt}. Furthermore, let γt := minj∈[d] min{gi}ti=1 αi√ vi,j denotes the minimum possible value of effective stepsize at time t over all possible coordinate and past gradients {gi}ti=1. The convergence rate of Adam-type algorithm is given by
min t∈[T ]
E [ ‖∇f(θt)‖22 ] = O
( s1(T )
s2(T )
) ,
where s1(T ) is defined through the upper bound of RHS of (20), and ∑T t=1 γt = Ω(s2(T )).
We present the proof of this Theorem in subsection C.2.4 and C.2.5 for completeness and reader’s convenience.

heading: C.2.2 PROOF OF THEOREM 2
text: Proof. We will first bound each term on RHS of Equation (20). Given all conditions in Theorem 2 hold, we have that
E [ T∑ t=1 ∥∥∥∥αt gt√vt ∥∥∥∥2 2 ] (i) ≤ 1 c E [ T∑ t=1 ‖αtgt‖22 ]
≤ 1 c E [ T∑ t=1 α2t ‖gt‖22 ] (ii) ≤ H 2
c T∑ t=1 α2t , (21)
where (i) arises from the fact that 0 < c ≤ vt,∀t ∈ [T ] and inequality (ii) is based on ‖gt‖2 ≤ H, ∀t ∈ [T ]. Then,
E [ T∑ t=1 ∥∥∥∥ αt√vt − αt−1√vt−1 ∥∥∥∥ 1 ] (i) =E [ d∑ i=1 T∑ t=1 αt−1√ vt−1,i − αt√ vt,i ]
= E [ d∑ i=1 α1√ v1,i − αT√ vT,i ]
≤ E [ d∑ i=1 α1√ v1,i ] (ii) ≤ dα√
c , (22)
and here (i) holds since we assume that αt√ vt ≥ αt+1√ vt+1 , ∀t ∈ [T ], and (ii) comes from the fact that 0 < c ≤ vt, 0 < αt ≤ α,∀t ∈ [T ]. Next,
E [ T∑ t=1 ∥∥∥∥ αt√vt − αt−1√vt−1 ∥∥∥∥2 2 ] = E [ T∑ t=1 d∑ i=1 ∥∥∥∥ αt√vt,i − αt−1√vt−1,i ∥∥∥∥2 2 ] (i) ≤E
[ T∑ t=1 d∑ i=1 ∥∥∥∥ αt√vt,i − αt−1√vt−1,i ∥∥∥∥ 2 α√ c ]
≤ dα 2
c , (23)
where we have (i) because ∥∥∥ αt√vt,i − αt−1√vt−1,i ∥∥∥2 = αt−1√vt−1,i − αt√vt,i ≤ α√c . On the other hand, we can obtain a lower bound of the LHS of Equation (20):
E [ T∑ t=1 αt 〈 ∇f(θt), ∇f(θt)√ vt 〉] ≥ 1 H E [ T∑ t=1 αt ‖∇f(θt)‖22 ] ≥ TαT H min t∈[T ] E ‖∇f(θt)‖22 , (24)
where the last equality comes from the fact that vt is weighted average ofm2t and ‖mt‖2 ≤ H sincemt is also an exponential moving average of gt. By combining the results in (21), (22), (23) and (24) to (20), we obatain
TαT H min t∈[T ] E ‖∇f(θt)‖22
≤ E [ T∑ t=1 αt 〈 ∇f(θt), ∇f(θt)√ vt 〉]
≤ E [ C1
T∑ t=1 ∥∥∥∥αt gt√vt ∥∥∥∥2 2 + C2 T∑ t=1 ∥∥∥∥ αt√vt − αt−1√vt−1 ∥∥∥∥ 1 + C3 T∑ t=1 ∥∥∥∥ αt√vt − αt−1√vt−1 ∥∥∥∥2 2 ] + C4
≤ C1H 2
c T∑ t=1 α2t + C2dα√ c + C3dα 2 c + C4 = C1H 2η(T ) c + C2 dα√ c + C3 dα2 c + C4.
After rearrangement, we can easily deduce that
min t∈[T ]
E ‖∇f(θt)‖22 ≤ H
TαT
[ C1H 2η(T )
c + C2 dα√ c + C3 dα2 c + C4 ] = 1
TαT (Q1 +Q2η(T )), (25)
where
Q1 = H ( C2
dα√ c + C3 dα2 c + C4
) , Q2 = C1H 3
c .

heading: C.2.3 PROOF OF COROLLARY 2
text: By choosing αt = α/ √ t ≤ α,∀t ∈ [T ], we have
TαT = α √ T , η(T ) = T∑ t=1 α2t = α 2 T∑ t=1 1 t ≤ α2(1 + log(T )).
Combining this with (25) and making some rearrangement, we have:
min t∈[T ]
E ‖∇f(θt)‖22 ≤ H
α √ T
[ C1 α2H2(1 + log(T ))
c + C2 dα√ c + C3 dα2 c + C4 ] =
1√ T (Q∗1 +Q ∗ 2 log(T )).
where
Q∗1 = H ( C1 H2α
c + C2 dα√ c + C3 dα2 c + C4
) , Q∗2 = C1H 3α
c .

heading: C.2.4 TECHNICAL LEMMAS FOR PROOF OF THEOREM 5
text: In this section, we introduce seven useful lemmas.
Lemma 5 (Chen et al. (2019)). Let θ0 , θ1 in the Algorithm, consider the sequence
zt = θt + β1,t
1− β1,t (θt − θt−1),∀t ≥ 2.
The following holds true: zt+1 − zt =− (
β1,t+1 1− β1,t+1 − β1,t 1− β1,t ) αtmt√ vt
− β1,t 1− β1,t ( αt√ vt − αt−1√ vt−1 ) mt−1 − αtgt√ vt ,∀t > 1,
and
z2 − z1 = − (
β1,2 1− β1,2 − β1,1 1− β1,1 ) α1m1√ v1 − α1g1√ v1 .
Lemma 6 (Chen et al. (2019)). Suppose that the conditions in Theorem 5 hold, then we have
E [f(zt+1 − f(zt))] ≤ 6∑ i=1 Ti,
where
T1 = −E [ t∑ i=1 〈 ∇f(zi), β1,i 1− β1,i ( αi√ vi − αi−1√ vi−1 ) mi−1 〉] , (26)
T2 = −E [ t∑ i=1 αi 〈 ∇f(zi), gi√ vi 〉] , (27)
T3 = −E [ t∑ i=1 〈 ∇f(zi), ( β1,i+1 1− β1,i+1 − βi 1− βi ) αimi√ vi 〉] , (28)
T4 = E [ t∑ i=1 3L 2 ∥∥∥∥( β1,i+11− β1,i+1 − β1,i1− β1,i ) αimi√ vi ∥∥∥∥2 2 ] , (29)
T5 = E [ t∑ i=1 3L 2 ∥∥∥∥ β1,i1− β1,i ( αi√ vi − αi−1√ vi−1 ) mi−1 ∥∥∥∥2 2 ] , (30)
T6 = E [ t∑ i=1 3L 2 ∥∥∥∥αigi√vi ∥∥∥∥2 2 ] . (31)
Lemma 7 (Chen et al. (2019)). Suppose that the condition in Theorem 5 hold, then for T1 in (26) it holds that
T1 = −E [ t∑ i=1 〈 ∇f(zi), β1,i 1− β1,i ( αi√ vi − αi−1√ vi−1 ) mi−1 〉]
≤ H2 β1 1− β1 E [ t∑ i=2 d∑ j=1 ∣∣∣∣∣ ( αi√ vi − αi−1√ vi−1 ) j ∣∣∣∣∣ ] .
Lemma 8 (Chen et al. (2019)). Suppose the conditions in Theorem 5 are satisfied, then T3 in (28) can be bounded as:
T3 = −E [ t∑ i=1 〈 ∇f(zi), ( β1,i+1 1− β1,i+1 − βi 1− βi ) αimi√ vi 〉]
≤ (
β1 1− β1 − β1,t+1 1− β1,t+1
) (H2 +G2).
Lemma 9 (Chen et al. (2019)). Suppose assumptions in Theorem 5 are satisfied, then for T4 in (29), it holds that
T4 = E [ t∑ i=1 3L 2 ∥∥∥∥( β1,i+11− β1,i+1 − β1,i1− β1,i ) αimi√ vi ∥∥∥∥2 2 ]
≤ 3L 2
( β1
1− β1 − β1,t+1 1− β1,t+1
)2 G2.
Lemma 10 (Chen et al. (2019)). Suppose the assumptions in Theorem 5 are satisfied, then for T5 in (30), we have
T5 = E [ t∑ i=1 3L 2 ∥∥∥∥ β1,i1− β1,i ( αi√ vi − αi−1√ vi−1 ) mi−1 ∥∥∥∥2 2 ]
≤ 3L 2
( β1
1− β1
)2 H2E [ t∑ i=2 d∑ j=1 ( αi√ vi − αi−1√ vi−1 )2 j ] .
Lemma 11 (Chen et al. (2019)). Suppose the assumptions in Theorem 5 are satisfied, then T2 in (27) can be bounded as:
T2 =− E [ t∑ i=1 αi 〈 ∇f(zi), gi√ vi 〉]
≤E t∑ i=2 1 2 ∥∥∥∥αigi√vi ∥∥∥∥2 2 + L2 ( β1 1− β1 )2( 1 1− β1 )2 E [ d∑ j=1 t−1∑ i=2 ( αigi√ vi )2 j ]
+ L2H2 (
β1 1− β1
)4( 1
1− β1
)2 E [ d∑ j=1 t−1∑ i=2 ( αi√ vi − αi−1√ vi−1 )2 j ]
+ 2H2E [ d∑ j=1 t∑ i=2 ∣∣∣∣∣ ( αi√ vi − αi−1√ vi−1 ) j ∣∣∣∣∣ ]
+ 2H2E [ d∑ j=1 ( α1√ v1 ) j ] − E [ t∑ i=1 αi 〈∇f(xi),∇f(xi)/ √ vi〉 ] .

heading: C.2.5 PROOF OF THEOREM 5
text: Proof. We can prove Theorem 5 after combining Lemma 5, 6, 7, 8, 9, 10 and 11. Specifically, firstly based on Lemma 5 it holds that
E [f(zt+1)− f(z1)]
≤ 6∑ i=1 Ti
=E [ t∑ i=1 3 2 L ∥∥∥∥( β1,t+11− β1,t+1 − β1,t1− β1,t ) αtmt/ √ vt ∥∥∥∥2 2 ] + E [ t∑ i=1 3 2 L ‖αigi/ √ vi‖22 ]
+ E [ t∑ i=1 3 2 L ∥∥∥∥ β1,i1− β1,i ( αt√ vi − αi−1√ vi−1 ) mi−1 ∥∥∥∥2 2 ]
− E [ t∑ i=1 〈 ∇f(zi), β1,i 1− β1,i ( αi√ vi − αi−1√ vi−1 ) mi−1 〉] − E [ t∑ i=1 αi 〈∇f(zi), gi/ √ vi〉 ]
− E [ t∑ i=1 〈 ∇f(zi), ( β1,i+1 1− β1,i+1 − β1,i 1− β1,i ) αimi/ √ vi 〉] .
Then we can combine Lemma 6, 7, 8, 9, 10 and 11 and further merge similar terms. We have that
E [f(zt+1)− f(zt)] ≤ 6∑ i=1 Ti
≤H2 β1 1− β1 E [ t∑ i=2 d∑ j=1 ∣∣∣∣∣ ( αi√ vi − αi−1√ vi−1 ) j ∣∣∣∣∣ ]
+
( β1
1− β1 − β1,t+1 1− β1,t+1
) (H2 +G2) + 3L
2
( β1
1− β1 − β1,t 1− β1,t
)2 G2
+ 3L
2
( β1
1− β1
)2 H2E [ t∑ i=2 d∑ j=1 ( αi√ vi − αi−1√ vi−1 )2 j ]
+ E t∑ i=2 1 2 ∥∥∥∥αigi√vi ∥∥∥∥2 + L2( β11− β1 )2( 1 1− β1 )2 E [ d∑ j=1 t−1∑ i=2 ( αigi√ vi )2 j ]
+ L2H2 (
β1 1− β1
)4( 1
1− β1
)2 E [ d∑ j=1 t−1∑ i=2 ( αi√ vi − αi−1√ vi−1 )2 j ]
+ 2H2E [ d∑ j=1 t∑ i=2 ∣∣∣∣∣ ( αi√ vi − αi−1√ vi−1 ) j ∣∣∣∣∣ ] + 2H2E [ d∑ j=1 ( α1√ v1 ) j ]
− E [ t∑ i=1 αi 〈∇f(xi),∇f(xi)/ √ vi〉 ] .
Finally, after rearrangement and some calculation, it can be verified that
E [ t∑ i=1 αi 〈∇f(xi),∇f(xi)/ √ vi〉 ]
≤ E [ C1
T∑ t=1 ‖αtgt/ √ vt‖22 + C2 T∑ t=1 ∥∥∥∥ αt√vt − αt−1√vt−1 ∥∥∥∥ 1
+ C3 T∑ t=1 ∥∥∥∥ αt√vt − αt−1√vt−1 ∥∥∥∥2 2 ] + C4,
where we have
C1 , 3
2 L+
1 2 + L2 β1 1− β1
( 1
1− β1
)2 ,
C2 , H 2 β1
1− β1 + 2H2,
C3 , [ 1 + L2 ( 1
1− β1
)2( β1
1− β1
)] H2 (
β1 1− β1
)2 ,
C4 ,
( β1
1− β1
) (H2 +G2) + ( β1
1− β1
)2 G2 + 2H2E [ ‖α1/ √ v1‖1 ] + E[f(z1)− f(z∗)],
and z∗ is an optimal where f(·) takes its minimum. This result directly implies Theorem 5.

heading: D ADDITIONAL EXPERIMENTAL DETAILS
text: Hyperparameter tuning rule For hyperparameter tuning, we perform extensive and careful grid search to choose the best hyperparameters for all the baseline algorithms.
For SGDM, we set the momentum parameter as 0.9 and search the optimal learning rate α among set {30.0, 1.0, 0.1, 0.01, 0.0015, 0.001, 0.0005, 0.0002} for all the experiments.
For adaptive gradient method baselines (Adam, AdamW, Yogi, Adabound, RAdam and AdaBelief), we search for β1 among {0.5, 0.6, 0.7, 0.8, 0.9}, β2 among {0.9, 0.98, 0.99, 0.999, 0.9999}, α among {0.1, 0.01, 0.0015, 0.001, 0.0005}, weight decay parameter among {1.2×10−6, 10−4, 5×10−4} and among {10−8, 10−12, 10−16} for all the experiments except ImageNet. All the additional method-specific parameters
are carefully chosen according to the original paper setting. Due to extremely large computing workload on ImageNet dataset, we set β1 = 0.9, β2 = 0.999, α = 0.001 and = 1 × 10−8 for all the adaptive gradient methods, and finetune weight decay parameter from set {0, 1×10−4, 5×10−4, 10−2, 5×10−2} for the values we reported as we run.
For our AdaMomentum, we employ the default parameters as β1 = 0.9 and choose learning rate and weight decay parameter using same parameter searching scheme as the adaptive gradient baseline methods. β2 is set as 0.999 for all the tasks. We Choose = 10−8 for image classification tasks and = 1 × 10−16 for other tasks. We find this setting is universally ample for satisfactory performance, which further demonstrates the superiority of AdaMomentum that little tuning effort is needed.
All the experiments reported are performed on NVIDIA GeForce RTX 2080Ti GPUs with Intel Core i7-8700K 3.70GHz CPUs. We provide some additional information concerning the empirical experiments for completeness.
D.1 IMAGE CLASSIFICATION
CIFAR datasets The values of the hyperparameters after careful tuning of the reported results of the adaptive gradient methods on CIFAR-10 in the main paper is summarized in Table 7. For SGDM, the optimal hyperparameter setting is: the learning rate is 0.1, the momentum parameter is 0.9, the weight decay parameter is 5 × 10−4. For Adabound, the final learning rate is set as 0.1 (matching SGDM) and the value of the hyperparameter gamma is 10−3.
ImageNet The values of the hyperparameters after careful tuning of the reported results of the adaptive gradient methods on CIFAR-10 in the main paper is summarized in Table 8. For SGDM, the stepsize is 0.1, the momentum parameter is 0.9 and the weight decay is 5× 10−4.

heading: D.2 LSTM ON LANGUAGE MODELING
text: The training perplexity curve is illustrated in Figure 5. We can clearly see that AdaMomentum is able to make the perplexity descent faster than SGDM and most other adaptive gradient methods. In experimental settings, the size of the word embeddings is 400 and the number of hidden units per layer is 1150. We employ dropout in training and the dropout rate for RNN layers is 0.25 and the dropout rate for input embedding layers is 0.4.
The optimal hyperparameters of adaptive gradient methods for 1-layer, 2-layer and 3-layer LSTM are listed in Tables 9, 10 and 11 respectively. For SGDM, the Well tuned stepsize is 30.0 and the momentum parameter is 0.9. For Adabound, the final learning rate is set as 30.0 (matching SGDM) and the value of the hyperparameter gamma is 10−3.

heading: D.3 TRANSFORMER ON NEURAL MACHINE TRANSLATION
text: For transformer on NMT task, the well tuned hyperparameter values are summarized in Table 12. The stepsize of SGDM is 0.1 and the momentum parameter of SGDM is 0.9. Initial learning rate is 10−7 and the minimum learning rate threshold is set as 10−9 in the warm-up process for all the optimizers.

Review summaries:
In this work, the authors proposed an improved version of Adam by using a momentumized version of the second moment. The authors justify the proposed modification and give some standard convergence analysis. The authors test the proposed methods in several deep learning tasks and show the improvement over some existing methods. 

This paper proposed AdaMomentum, which uses the EMA (exponential moving average) of the square of EMA of gradient as sthe denominator, while Adam uses the EMA of the square of the gradient. The authors conducted experiments in CNN, Transformer, LSTM and GAN to show the superior performance of AdaMomentum. The authors also tried to provide theoretical analysis.

This paper study the  adaptive type gradient methods, such as ADAM which scaling gradients by square roots of the running averages of squared past gradients. Nevertheless,  they discover that substituting the gradient in the second moment estimation term with the momentumized version in ADAM can well solve the weaknesses of ADAM which  generalize worse than stochastic gradient descent (SGD) and tend to be trapped in local minima at an early stage during training.Their intuition is that gradient with momentum contains more accurate directional information and therefore its second moment estimation is a better choice for scaling than that of the raw gradient. Thereby they propose ADAMOMENTUM as a new optimizer reaching the goal of training fast while generalizing better. Convergence guarantee is provided   and   extensive experiments on a wide range of tasks and models is also conducted.