Original: Adam Can Converge With Proper β 1 and β
Revised: Adam Can Converge with Proper β 1 and β

Original: Abstract
Revised: Abstract

Original: Ever since Reddi et al. (2018) pointed out the divergence issue of Adam, many new variantshave been designed to obtain convergence. However, vanilla Adam remains exceptionallypopular and it works well after proper tuning.
Revised: Adam is one of the most popular algorithms in deep learning.

Original: Why is the divergence not observed? Wepoint out there is a gap between theory and practice: Reddi et al. (2018) pick the samplesize n after picking the hyperparameters of Adam, i.e., ( β 1 , β 2 ) ; while practical applicationsoften fix n first and then tune ( β 1 , β 2 ) . Since Adam often works well in the latter setting, weconjecture it can still converge.
Revised: Despite its practical popularity,the convergence of Adam is far from being understood.

Original: In this work, we confirm this conjecture under fixed n . Weprove that, when the 2nd-order momentum parameter β 2 is large and 1st-order momentumparameter β 1 < √ β 2 < 1 ,  Adam converges to the neighborhood of critical points.
Revised: In this work, we prove that whenthe 2nd-order momentum parameter β 2 is large and 1st-order momentum parameter β 1 <√ β 2 < 1 , randomly shuffled Adam converges to the neighborhood of stationary points.

Original: The sizeof the neighborhood is propositional to the  variance of stochastic gradients.
Revised: Thesize of the neighborhood is propositional to the noise variance of stochastic gradients.

Original: Under an extracondition (strong growth condition), Adam converges to critical points.
Revised: Withextra strong growth condition, we prove that Adam can converge to stationary points.

Original: As β 2 increases, ourconvergence result can cover any β 1 ∈ [0 , 1) including the default setting β 1 = 0 .
Revised: Asβ 2 increases, our convergence result can cover any β 1 ∈ [0 , 1) including the default settingβ 1 = 0 .

Original: To ourbest knowledge, we are the first to prove that Adam with arbitrarily large β 1 can convergewithout any form of bounded gradient assumption.
Revised: To our best knowledge, we are the first to prove that Adam with any large β 1 canconverge without any form of bounded gradient assumption.

Original: This result shows that vanilla Adam,without any modification, is still theoretically justified to converge.
Revised: We also investigate the divergence behavior of Adam.

Original: When β 2 is small, wefurther point out a large region of ( β 1 , β 2 ) combinations where Adam can diverge to infinity.
Revised: When β 2 is small, we point out a largeregion of ( β 1 , β 2 ) combinations where Adam can diverge to infinity.

Original: Our divergence result considers the same setting (fixing n ahead) as our convergence result,indicating that there is a phase transition from  divergence to convergence when increasing β 2 .
Revised: We emphasize that bothour results above consider a quite different setting from the divergence result in Reddi et al.(2018): they consider picking ( β 1 , β 2 ) before picking the sample size n ; while we consider amore practical scenario with n fixed in advance. Our analysis shows that Adam’s behaviorchanges significantly with hyperparameters.

Original: These results may provide  guidance towards better hyperparameter tuning of Adam.
Revised: These results may provide new guidance towardsbetter hyperparameter tuning of Adam.

Original:  Introduction
Revised:  Introduction

Original: Modern machine learning tasks often aim to solve the following finite-sum problem.
Revised: Modern machine learning tasks often aim to solve the following finite-sum problem.

Original: [Equation]
Revised: [Equation]

Original: where n is the number of samples or mini-batches and x denotes the trainable parameters.
Revised: where n often stands for the number of samples or mini-batches and x denotes the trainable parameters.

Original: In deeplearning, Adam (Kingma & Ba, 2014) is one of the most popular algorithms for solving (1).
Revised: When combined with deep neural nets, Adam (Kingma & Ba, 2014) is one of the most dominantalgorithms for solving (1).

Original: Ithas been applied to various machine learning domains such as natural language processing (NLP)(Vaswani et al., 2017;Brown et al., 2020;
Revised: In the past five years, Adam has reached tremendous success in differentmachine learning applications including natural language processing (NLP) (Vaswani et al., 2017;

Original:  Devlin et al., 2018), generative adversarial networks (GANs)(Radford et al., 2015; Isola et al., 2017; Zhu et al., 2017) and computer vision (CV) (Dosovitskiyet al., 2021).
Revised: Brown et al., 2020; Devlin et al., 2018), generative adversarial networks (GANs) (Isola et al., 2017; Zhu et al., 2017; Radford et al., 2015) and computer vision (CV) (Dosovitskiy et al., 2021), etc..

Original: Despite its prevalence, Reddi et al.
Revised: Despite its prevalence, Reddi et al.

Original: (2018) point out that Adam can diverge with a widerange of hyperparameters.
Revised: (2018) point out that Adam can diverge under a wide range ofhyperparameters.

Original: A main result in (Reddi et al., 2018) states that 1 :
Revised: A main result in (Reddi et al., 2018) conveys the message that 1 :

Original: For any β 1 , β 2 s.t.[-
 ≤-] β 1 < √ β 2 < 1 , there exists a problemsuch that Adam diverges.
Revised: For any β 1 , β 2 s.t.[-
 ≤-] β 1 < √ β 2 < 1 , there exists a problem that Adam will diverge.

Original: Here, β 1 and β 2 are the hyperparameter to control Adam’s 1st-order  and 2nd-order momentum. More description of Adam can be seen in Algorithm 1 (presented later in Section 2.1).
Revised: Here, β 1 is the hyperparameter to control Adam’s 1st-order momentum and β 2 control the weightedaverage of squared past gradients (a.k.a.

Original: Ever since(Reddi et al., 2018) pointed out the divergence issue, many new variants have been designed. Forinstance, AMSGrad (Reddi et al., 2018) enforced the adaptor v t (defined later in Algorithm 1) to Appendix D.2. be non-decreasing; AdaBound (Luo et al., 2019) imposed constraint v t ∈ [ C l , C u ] to ensure theboundedness on effective stepsize.
Revised: 2nd-order momentum).

Original: Although these variants can provably obtain convergence, theyoften introduce new issues. For instance, AMSGrad is often reported to slow down the trainingprocess (Zhou et al., 2018b); AdaBound introduces two extra hyperparameters C l and C u . On the other hand, counter-intuitively, vanilla Adam remains exceptionally popular (see evidenceat (Scholar)).
Revised: More description can be seenin Algorithm 1 (presented later in Section 2.1).

Original: Without any modification on its update rules, Adam works well after proper tuning.
Revised: Despite the divergence issue, Adam works wellin practice after proper tuning.

Original: Even more mysteriously, we find that the commonly reported hyperparameters actually satisfythe divergence condition stated earlier.
Revised: Even more mysteriously, we find that the commonly reportedhyperparameters actually satisfy the divergence condition above.

Original: For instance, Kingma & Ba (2014) claimed that ( β 1 , β 2 ) =(0 .
Revised: For instance, Kingma & Ba (2014) claim that ( β 1 , β 2 ) =

Original:  9 , 0 . 999) is a “good choice for the tested machine learning problems" and it is indeed the defaultsetting in PyTorch.
Revised: (0 . 9 , 0 . 999) is a “good choice for the tested machine learning problems" and itis indeed the default setting in practice.

Original: In  super-large models GPT-3 and Megatron (Brown et al., 2020; Smith et al.,2022), ( β 1 , β 2 ) is chosen to be (0 . 9 , 0 . 95) .
Revised: In recent super-large models GPT-3 and Megatron (Brownet al., 2020; Smith et al., 2022), ( β 1 , β 2 ) = (0 . 9 , 0 . 95) is chosen as their final hyperparameters.

Original:  GAN researchers (e.g. Radford et al. ; Isolaet al. (2017)) use ( β 1 , β 2 ) =
Revised: For GAN researchers (e.g. (Isola et al., 2017)), ( β 1 , β 2 ) =

Revised: (0 . 5 , 0 . 999) is often preferred.

Original: All these hyperparameters live in the divergence regionβ 1 < √ β 2 .
Revised: All thesehyperparameters satisfy β 1 < √ β 2 < 1 , and thus live in the “divergence region" by (Reddi et al.,2018).

Original: Surprisingly, instead of observing the divergence issue, these hyperparameters achievegood performances andthey actually show the sign of convergence.
Revised: However, they all reach good performances and the divergence issue is not observed.

Original: Why does Adam work well despite itstheoretical divergence issue?
Revised: Why does Adam work well despite its divergence issue?

Original: Is there any mismatch betweendeep learning problems and the divergent example?
Revised: One might argue that neural nets trainingmight be somewhat “simpler".

Original: We take a closer look into the divergenceexample andfind out the mismatch does exist. In particular, we notice an important (but oftenignored) characteristic of  the divergence example: (Reddi et al., 2018) picks  ( β 1 , β 2 )before pickingthe sample size n . 
Revised: We would like to point out it is not the crucial reason. Instead, wenotice a non-negligible gap between practical applications and the above divergence result: practicalapplications often fix the number of samples or mini-batches n in advance; while the divergenceresult picks n after picking ( β 1 , β 2 ) . This difference turns out to be crucial.

Original: Put in another way, to construct the divergence example, they change n fordifferent ( β 1 , β 2 ) .
Revised: We explain as follows. To construct the divergence example, one needs to change n for different ( β 1 , β 2 ) .

Revised: This is shown inthe proof of (Reddi et al., 2018) and is also restated in Appendix C.2.

Original: For instance, for ( β 1 , β 2 ) = (0 , 0 . 99) , we use one n to construct the divergentexample; for ( β 1 , β 2 ) 
Revised: For instance, for ( β 1 , β 2 ) =(0 , 0 . 99) , we need one n to construct the divergent example; for ( β 1 , β 2 ) =

Original: = (0 , 0 . 9999) , we use another n to construct another divergent example.
Revised: (0 , 0 . 9999) , we needanother n to construct another divergent example.

Original: This understanding reveals a gap between the divergence result and practice: In practical applicationsof Adam listed above, practitioners tune the hyperparameters ( β 1 , β 2 ) after the sample size n is given,while the divergence example considers the opposite order. Considering the good performance of Adam under fixed n , we conjecture that “the order of pickingn or ( β 1 , β 2 ) " might be  the major gapbetween theory and practice.
Revised: However, in practical applications, sample size nis often given before picking ( β 1 , β 2 ) so the divergence result cannot be applied. Considering thedifferent settings, the above divergence result might not provide proper guidance for Adam users.

Original: So could Adam possibly converge when  n is fixed?
Revised: So how would Adam behave when the sample size n is fixed?

Original: Unfortunately,the behavior of vanilla Adam is far less studied than its variants (perhaps due to the criticism ofdivergence).
Revised: Can Adam possibly converge?

Original: To explore this question, we run experiments for different choices of ( β 1 , β 2 ) on afew tasks.
Revised: If so,what ( β 1 , β 2 ) pairs can ensure convergence? With these questions in mind, we run experiments ondifferent tasks.

Original: First, we run Adam for a convex function (2) with fixed n (see the definition in Section3.2).
Revised: First, we run Adam under different choices of ( β 1 , β 2 ) on a convex function (2) withfixed n (see the definition in Section 3.2).

Original: Second, we run Adamfor the classification problem on data MNIST and CIFAR-10 with fixedbatchsize.
Revised: Then, we run Adam on real dataset MNIST and CIFAR-10with fixed batchsize.

Original: We observe some interesting phenomena in Figure 1 (a), (b) and (c).
Revised: We observe some interesting phenomena in Figure 1 (a), (b) and (c).

Original: First, when β 2 is large, the optimization error is small for almost all values of β 1 .
Revised: First, when β 2 is large, the optimization error is small for almost all values of β 1 .

Original: Second, whenβ 1 , β 2 are both small, there is a red region with relatively large error.
Revised: Second, whenβ 1 , β 2 are both small, there is a red region with relatively large error.

Original: On MNIST, CIFAR-10, theerror in the red region is around 0.4 times higher than that in the blue region.
Revised: On MNIST, CIFAR-10, theerror in the red region is around 0.4 times higher than that in the blue region.

Original: The situation is a lotworse on function (2): the error in the red region is 70 times higher than that in the blue region.[Table] [Table] a
Revised: The situation is a lotworse on function (2): the error in the red region is 70 times higher than that in the blue region.

Original: [Table]
Revised: [Figure]

Original: Function

Original: [Table] [Table] [Figure] While Adam’s performances seem unstable in the red region, we find that Adam always performswell in the top blue region in Figure 1.
Revised: While Adam’s performances seem unstable in the red region, we find that Adam always performswell in the top blue region in Figure 1.

Original: This seems to suggest that Adam can converge without anyalgorithmic modification, as long as β 1 and β 2 are chosen properly.

Original: We ask the following question:
Revised: We are led to the following question:

Original: When n is fixed, can vanilla Adam provably converge without any algorithmic modification?
Revised: When n is fixed ahead, does Adam provably converge in the top region of Figure 1?

Original: In this work, we theoretically explore this question.
Revised: In this work, we theoretically explore this question.

Original: Our contributions are summarized below andalso visualized in Figure 1 (d).
Revised: When n is

Original: We prove thefollowing results when n is fixed.
Revised: fixed in advance, we prove the results below.

Original:  • We prove that when β 2 is largeenough and β 1 < √ β 2 ,  Adam converges to the neighborhoodof critical points.
Revised: Our contributions are also visualized in Figure 1 (d) • When β 2 is large and β 1 < √ β 2 , we prove that Adam converges to the neighborhood of stationarypoints.

Original: The size of the neighborhood is propositional to the variance of stochasticgradients.
Revised: The size of the neighborhood is propositional to noise variance of stochastic gradients.

Original: Withan extracondition (so-called strong growth condition), we prove that Adam canconverge to critical points.
Revised: With extra strong growth condition, we prove that Adam can converge to stationary points.

Original: As β 2 increases, these results can cover any momentum parameterβ 1 ∈ [0 , 1) including the default setting β 1 = 0 .
Revised: Asβ 2 increases, these results can cover any momentum parameter β 1 ∈ [0 , 1) including the defaultsetting β 1 = 0 .

Original: In particular, our analysis does not requirebounded gradient assumption.•
Revised: In particular, our analysis does not require bounded gradient assumption.•

Original: We study the  divergence issue of Adam.
Revised: We point out the algorithmic divergence in the red region in Figure 1 (d).

Original: We prove that: for any fixed n , there exists a functionsuch that, Adam diverges to infinity when ( β 1 , β 2 ) is picked in the red region in Figure 1 (d).The size of the red region increases with n .
Revised: We prove that: thereexists a function such that, Adam diverges to infinity when ( β 1 , β 2 ) is picked in the red region in Figure 1 (d). The size of the red region depends on n .

Original: The shape of the region follows the solution to ouranalytic conditions.•
Revised: The shape of the region follows thesolution to our analytic conditions.

Original: We emphasize a few characteristics of ourresults. (1) phase transition . The divergence resultconsiders the same setting (fixing n and then pick ( β 1 , β 2 ) ) as our convergence result, indicatingthat there is a phase transition from divergence to convergence. (2) problem-dependent bounds .Our
Revised: We emphasize that our divergence result consider “fix n beforechoosing ( β 1 , β 2 ) ".

Original: convergence and divergence regions of ( β 1 , β 2 ) are problem-dependent, which is drasticallydifferent from (Reddi et al., 2018)  which established the problem-independent worst-case choiceof ( β 1 , β 2 ) . (3) non-asymptotic characterization . the “divergence region” of ( β 1 , β 2 ) expandsas n increases and converges to the whole region [0 , 1) 2 as n goes to infinity, which recovers(actually stronger than) the problem-independent divergence result of (Reddi et al., 2018) thatrequires β 1 < √ β 2 . In this sense, we can view the divergence result of (Reddi et al., 2018) as anasymptotic characterization of the divergence region (as n → ∞ ) and our divergence result as anon-asymptotic characterization (for any fixed n ).
Revised: This setting is quite different from (Reddi et al., 2018) mentioned in Section1 which considers a reverse ordering.

Original: We provide more discussion in Section 4.We believe

Original:  our results can boost new understandings for Adam.
Revised: We believe our results can boost new understandings for Adam.

Original: While Reddi et al.
Revised: While Reddi et al.

Original: (2018) reveal that“Adam can diverge", our results show the other side of the coin: When n is fixed, Adam can still workwithout any algorithmic modification.
Revised: (2018) revealthat “Adam can diverge", our results show the other side of the coin: when n is fixed, Adam can stillwork if hyperparameters are chosen properly.

Original: Since “fixing n " is closer to practical scenarios, our resultssuggests that Adam is still a theoretically justified algorithm and practitioners can use it confidently.
Revised: Since “fixing n " is closer to practical scenarios, ourconvergence and divergence results may provide better guidance for hyperparameter tuning.

Original: Wefurther emphasize two features in this work.
Revised: We emphasize two features in this work.

Original: First, our convergence results can cover anyβ 1 ∈ [0 , 1) , which can involve massive momentum.
Revised: First, our convergence results can cover any β 1 ∈ [0 , 1) ,which can involve massive momentum.

Original: It turns out that large-momentum Adam is noteasy to analyze.
Revised: It turns out that large-momentum Adam is not easy to analyze.

Original: Even with stronger assumptions like bounded gradient ( ∥∇ f ( x ) ∥ < C , ∀ x ), itsconvergence is not well understood (see related works in Section 2.2).
Revised: Even with stronger assumptions like bounded gradient ( ∥∇ f ( x ) ∥ < C , ∀ x ), its convergence is notwell understood (see related works in Section 2.2).

Original: To our best knowledge, this isthe first result that proves Adam with any large β 1 can converge without any form of bounded gradientassumption.
Revised: To our best knowledge, this is the first result thatproves Adam with any large β 1 can converge without any form of bounded gradient assumption.

Original: The proof contains a new method to handle unbounded momentum in the stochasticnon-linear dynamics system.

Original: We will highlight our technical novelties in Section 5.
Revised: Wewill highlight our technical novelty in Section 5.

Original: Second, when n is fixed, we point out the behaviorof Adam significantly changes under different( β 1 , β 2 ) : It converges in certain region but it may diverge in others.
Revised: Second, when n is fixed, we point out Adam’s behavior significantly changes under different ( β 1 , β 2 ) :it converges in certain region but it may diverge in others.

Original: Pointing out these behaviorsin the ( β 1 , β 2 ) diagram may provide better guidance for tuning β 1 and β 2 .

Original: Removing boundednesscondition is important to characterize such qualitative behaviors.
Revised: Removing boundedness condition isimportant to characterize such qualitative behavior.

Original: With assumed boundedness ongradient or iterates, Adam cannot diverge to infinity, yet we have theoretical and numerical evidenceshowing that it can.
Revised: With assumed boundedness on gradient oriterates, Adam never diverges, yet we have theoretical and numerical evidence showing that it can.

Original:  Preliminaries
Revised:  Preliminaries

Original: 2.1 Problem Settings
Revised: 2.1 Problem Settings

Original: We consider finite-sum problem (1).
Revised: We consider finite-sum problem (1).

Original: We use x to denote the optimization variable.
Revised: We use x to denote the optimization variable.

Original: We denote ∇ f j asthe gradient of f j and let ◦ be the component-wise product.
Revised: We denote ∇ f j asthe gradient of f j and let ◦ be the component-wise product.

Original: The division and square-root operatorare component-wise as well.
Revised: The division and square-root operatorare component-wise as well.

Original: We present randomly shuffled Adam in Algorithm 1. 
Revised: We present randamly shuffled Adam in Algorithm 1. Algorithm

Revised: Initialize x 1 , 0 = x 0 , m 1 , − 1 = ∇ f ( x 0 ) and v 1 , − 1 = ∇ f ( x 0 ) ◦ ∇ f ( x 0 ) .

Revised: for k = 1 → ∞ do

Original: {+Sample { τ k, 0 , τ k, 1 , · · · , τ k,n − 1 } as a random permutation of { 0 , 1 , 2 , · · · , n − 1 } for i = 
+}
Revised: {+Sample { τ k, 0 , τ k, 1 , · · · , τ k,n − 1 } as a random permutation of { 0 , 1 , 2 , · · · , n − 1 } for i = 
+}

Original: In Algorithm 1, m denotes the 1st-order momentum and v denotes the 2nd-order momentum.
Revised: → n − 1 do m k,i = β 1 m k,i − 1 + (1 − β 1 ) ∇ f τ k,i ( x k,i ) v k,i = β 2 v k,i − 1 + (1 − β 2 ) ∇ f τ k,i ( x k,i ) ◦ ∇ f τ k,i ( x k,i ) x k,i +1 =

Original: {+x k,i − η k √ v k,i + ϵ ◦ m k,i end for x k +1 , 
+}
Revised: {+x k,i − η k √ v k,i + ϵ ◦ m k,i end for x k +1 , 
+}

Revised: = x k,n v k

Revised: [Equation]

Revised: m

Revised: [Equation]

Revised: Break if certain exit condition is satisfied.

Revised: end for

Revised: In Algorithm 1, m denotes the 1st-order momentum and v denotes the 2nd-order momentum.

Original: theyare weighted averaged by hyperparameter β 1 , β 2 , respectively.
Revised: theyare weighted averaged by hyperparameter β 1 , β 2 , respectively.

Original: Larger β 1 , β 2 will adopt more historyinformation.
Revised: Larger β 1 , β 2 will adopt more historyinformation.

Original: We denote x k,i , m k,i , v k,i ∈
Revised: We denote x k,i , m k,i , v k,i ∈

Original: R d as the value of x, m, v at the k -th outer loop (epoch)and i -th inner loop (batch), respectively.
Revised: R d as the value of x, m, v at the k -th outer loop (epoch)and i -th inner loop (batch), respectively.

Original: We choose η k = η 1 √ nk as the stepsize.
Revised: We choose η k = η 1 √ nk as the stepsize.

Original: In practice, ϵ isadopted for numerical stability and it is often chosen to be 10 − 8 .
Revised: In practice, ϵ isadopted for numerical stability and it is often chosen to be 10 − 8 .

Original: In our theory, we allow ϵ to be anarbitrary non-negative constant including 0 .
Revised: In our theory, we allow ϵ to be anarbitrary non-negative constant including 0 .

Original: In the original version of Adam in (Kingma & Ba, 2014), it has an additional “bias correction” step.
Revised: In the vanilla version of Adam, it has an additional “bias correction” step (see (Kingma & Ba, 2014)).

Original: This “bias correction” step can be implemented by changing the stepsize η k into ˆ η k = √ 1 − β k1 − β k 1 η kand using zero initialization.
Revised: This “bias correction” step can be implemented by changing the stepsize η k into ˆ η k = √ 1 − β k1 − β k 1 η kand using zero initialization.

Original: In Algorithm 1, the “bias correction” step is replaced by a specialinitialization, which also corrects the bias.
Revised: In Algorithm 1, the “bias correction” step is replaced by a specialinitialization, which also corrects the bias.

Original: Note that ˆ η k ∈ [ √ 1 − β 2 η k , 1 1 − β 1 η k ] is well-bounded Algorithm Initialize x 1 , 0 = x 0 , m 1 , − 1 = ∇ f ( x 0 ) and v 1 , − 1 = max i ∇ f i ( x 0 ) ◦ ∇ f i ( x 0 ) . for k = 1 → ∞ do
Revised: Note that ˆ η k ∈ [ √ 1 − β 2 η k , 1 1 − β 1 η k ] is well-boundednear η k , so η k and ˆ η k brings the same convergence rate.

Original: [-Sample { τ k, 0 , τ k, 1 , · · · , τ k,n − 1 } as a random permutation of { 0 , 1 , 2 , · · · , n − 1 } for i = 
-]
Revised: [-Sample { τ k, 0 , τ k, 1 , · · · , τ k,n − 1 } as a random permutation of { 0 , 1 , 2 , · · · , n − 1 } for i = 
-]

Original: → n − 1 do m k,i =

Original: β 1 m k,i − 1 + (1 − β 1 ) ∇ f τ k,i ( x k,i ) v k,i =

Original: β 2 v k,i − 1 + (1 − β 2 ) ∇ f τ k,i ( x k,i ) ◦ ∇ f τ k,i ( x k,i ) x

Original: k,i +1 =

Original: [-x k,i − η k √ v k,i + ϵ ◦ m k,i end for x k +1 , 
-]
Revised: [-x k,i − η k √ v k,i + ϵ ◦ m k,i end for x k +1 , 
-]

Original: = x k,n ; v k +1 , − 1 = v k,n − 1 ; m k +1 , − 1 = m k,n − 1 end for near η k , so η k and ˆ η k brings the same convergence rate.

Original: In addition, as the effect of initializationbecomes negligible when the training progresses, Adam with zero & our initialization will have thesame asymptotic behavior.
Revised: In addition, as the effect of initializationbecomes negligible when the training progresses, Adam with zero & our initialization will have thesame asymptotic behavior.

Original: We follow the form of Algorithm 1, which makes results cleaner.
Revised: We follow the form of Algorithm 1, which makes results cleaner.

Original: In ouranalysis, we make the assumptions below.
Revised: In ouranalysis, we make the assumptions below.

Original: Assumption 2.1.
Revised: Assumption 2.1.

Original: we consider x ∈ R d and f i ( x ) satisfies gradient Lipschitz continuous with constant
Revised: we consider x ∈ R d and f i ( x ) satisfies gradient Lipschitz continuous with constant

Original: L .
Revised: L .

Original: We assume f ( x ) is lower bounded by a finite constant f ∗ .
Revised: We assume f ( x ) is lower bounded by a finite constant f ∗ .

Original: Assumption 2.2.
Revised: Assumption 2.2.

Original: f i ( x ) and f ( x ) satisfy: (cid:80) n
Revised: f i ( x ) and f ( x ) satisfy: (cid:80) n

Original: [Equation]
Revised: [Equation]

Original:  Assumption 2.2 is quite general.
Revised: [Equation] Assumption 2.2 is quite general, it covers some stronger conditions.

Original: When D 1 = 1 /n , it becomes the “constant variance” with constant D 0 /n .
Revised: When D 1 = n , it becomes the“constant variance” with constant D 0 /n .

Original: “constant variance" condition is commonly used in both SGD and Adam analysis (e.g.(Ghadimi et al., 2016;
Revised: “constant variance" condition is commonly used in both SGDand Adam analysis (e.g. (Ghadimi et al., 2016;

Original: Zaheer et al., 2018; Huang et al., 2021)).
Revised: Zaheer et al., 2018; Huang et al., 2021)).

Original: Assumption 2.2 allows more flexiblechoices of D 1 ̸ = n and thus it is weaker than “constant variance”.
Revised: Assumption2.2 allows more flexible choices of D 1 ̸ = n and thus it is weaker than “constant variance”.

Original: When D 0 > 0 , the problem instance is sometimes called “non-realizable" (Shi et al., 2020).
Revised: When D 0 > 0 , the problem instance is sometimes called “non-realizable" (Shi et al., 2020).

Original: In thiscase, adaptive gradient methods are not guaranteed to reach the exact critical point.
Revised: In thiscase, adaptive gradient methods are not guaranteed to reach the exact stationary point.

Original: Instead, theyonly converge to a bounded region (near critical points)
Revised: Instead, theyonly converge to a bounded region (near stationary points)

Original: (Zaheer et al., 2018; Shi et al., 2020).
Revised: (Zaheer et al., 2018; Shi et al., 2020).

Original: Thisphenomenon indeed occurs for Adam in experiments, even with diminishing stepsize (see Appendix B.1).
Revised: Thisphenomenon indeed occurs for Adam in practice, even with diminishing stepsize (see Appendix D.2).

Original: The behavior of SGD is similar: constant stepsize SGD converges to a bounded region with itssize propositional to the noise level D 

Revised: The behavior of SGD is similar: constant stepsize SGD converges to a bounded region with its sizepropositional to the noise level D

Original: (Yan et al., 2018; Yu et al., 2019; Liu et al., 2020b).
Revised: (Yan et al., 2018; Yu et al., 2019; Liu et al., 2020b).

Original: When D 0 = 0 , Assumption 2.2 is often called “strong growth condition" (SGC) (Vaswani et al.,2019).
Revised: When D 0 = 0 , Assumption 2.2 is often called “strong growth condition" (SGC) (Vaswani et al.,2019).

Original: When ∥∇ f ( x ) ∥ = 0 , under SGC we have ∥∇ f j ( x ) ∥ = 0 forall j .
Revised: When ∥∇ f ( x ) ∥ = 0 , under SGC we have ∥∇ f j ( x ) ∥ = 0 for

Original:  SGC is increasinglypopular recently e.g.(Schmidt & Roux, 2013; Vaswani et al., 2019).
Revised: all j . SGC is sometimes usedin SGD analysis e.g. (Schmidt & Roux, 2013; Vaswani et al., 2019).

Original: This condition is reasonablein the overparameterized regime where neural networks are capable of interpolating all data points: When every data point reaches zero training loss, the gradient of each point also converges to zero(Vaswani et al., 2019).
Revised: This condition is reasonablein the overparameterized regime where neural networks are capable of interpolating all data points(Vaswani et al., 2019).

Original: We will show that Adam can converge to critical points if SGC holds.
Revised: We will show that Adam can converge stationary points if SGC holds.

Original: When n, f ∗ , L, D 0 , D 1 are fixed a priori, we use F n,f ∗ L,D 0 ,D 1 ( R d ) to denote the (sub) function classcontaining f ( x ) satisfying Assumption 2.1 and 2.2 with constant n, f ∗ , etc..
Revised: When n, f ∗ , L, D 0 , D 1 are fixed a priori, we use F n,f ∗ L,D 0 ,D 1 ( R d ) to denote the (sub) function classcontaining f ( x ) satisfying Assumption 2.1 and 2.2 with constant n, f ∗ , etc..

Original: We introduce thisnotation to clearly present the divergence result in Proposition 3.3.
Revised: We introduce thisnotation to clearly present the divergence result in Proposition 3.3.

Original: Without this pre-defined functionclass, the claim of divergence might be  confusing.
Revised: Without this pre-defined functionclass, the claim of divergence might be misleading or confusing.

Original: 2.2 Related Works
Revised: 2.2 Related Works

Original: Ever since Reddi et al.
Revised: Ever since Reddi et al.

Original: (2018) pointed out the divergence issue, there are many attempts on designingnew variants of Adam.
Revised: (2018) pointed out the divergence issue, there are many attempts on designingnew variants of Adam.

Revised: For instance, Zou et al.

Original: As introduced in Section 1, these methods can address the divergence of theoriginal Adam.
Revised: and Chen et al. (2018) replace the constant( β 1 , β 2 ) by iterate-dependent ones e.g. β 1 t or β 2 t . These methods can address the divergence of theoriginal Adam.

Original: But they might brings extra concerns (e.g. extra hyperparameters).
Revised: Many variants also achieve satisfactory performance in practice.

Original: Since we focus onunderstanding Adam without modification , we introduce more variants later in Appendix D.1.
Revised: Since we focus onunderstanding Adam without modification, we introduce more variants later in Appendix C.1.

Original: Compared with proposing new variants, the convergence of vanilla Adam is far less studied thanits variants (perhaps due to the criticism of divergence). We summarize a few works here. Zhouet al. (2018b) analyze the counter-example in (Reddi et al., 2018) and find certain hyperparametercan work.
Revised: There are several works trying to understand Adam’s convergence behavior, but they often requireextra conditions.

Original: However, their analysis is restricted to the counter-example. Zaheer et al.
Revised: For instance, the analysis in Zaheer et al.

Original: (2018) revealsrelation between mini-batch sizes and (non-)convergence of Adam.
Revised: (2018) reveals relation between mini-batchsizes and (non-)convergence of Adam.

Revised: Gadat & Gavra (2020) point out that Adam only converges tominimizers.

Original: However, this work requireβ 1 = 0 and Adam is reduced to RMSProp (Hinton et al., 2012).
Revised: However, these two works require β 1 = 0 and Adam is reduced to RMSProp (Hintonet al., 2012).

Original: De et al.
Revised: De et al.

Original: (2018)  analyze RMSPropand non-zero- β 1 Adam, but they assume the sign of all stochastic gradients to keep the same.
Revised: (2018) further analyze Adam with non-zero β 1 , but they assume the signof all stochastic gradients to keep the same.

Original: Itseems unclear how to check this condition a priori.
Revised: It seems unclear how to check this condition a priori.

Original: Additionally, they require β 1 to be inverselyrelated to the upper bound of gradient, which forces β 1 to be small (as a side note, this result only applies to full-batch Adam).
Revised: Additionally, they require β 1 inversely related to the upper bound of gradient, which forces β 1 to besmall.

Original: Défossez et al.
Revised: Défossez et al.

Original: (2020) analyze Adamwith β 1 < β 2 and provide someinsights on the momentum mechanisms.
Revised: (2020) analyze Adam β 1 < β 2 and provide some insights on the momentummechanisms.

Original: However, their bound is inversely proportional to ϵ (thehyperparameter for numerical stability)and the bound goes to infinity when ϵ goes to 0.
Revised: However, they require ϵ (the hyperparameter for numerical stability) to be large, inwhich case ϵ will be much larger than √ v .

Original: This isdifferent from practical application since small ϵ such as 10 − 8 often works well.

Original: Further, using largeϵ is against the nature of adaptive gradient methods because √ v no longer dominates in the choiceof stepsize.
Revised: This is against the nature of adaptive gradient methodsbecause √ v no longer dominates in the choice of stepsize.

Original: In this case, Adam is essentially transformed back to SGD.

Original: Two  recent works (Huanget al., 2021) and (Guo et al., 2021) propose novel and simple frameworks to analyze Adam-familywith large β 1 .
Revised: Two very recent works (Huang et al.,2021) and (Guo et al., 2021) propose novel and simple frameworks to analyze Adam-family withlarge β 1 .

Original: Yet, they require the effective stepsize of Adam to be bounded in certain interval, i.e.,1 √ v t + ϵ ∈ [ C l , C u ] 2 .
Revised: Yet, they require the effective stepsize of Adam will be bounded in certain interval, i.e.,1 √ v t + ϵ ∈ [ C l , C u ] 2 .

Original: This boundedness condition changes Adam into AdaBound (Luo et al.
Revised: This boundedness condition changes Adam into AdaBound (Luo et al.

Original: , 2019)and thus they cannot explain the observations on original Adam in Section 1.
Revised: , 2019)and thus they cannot explain the observations on original Adam in Section 1.

Original: Additionally, all ofthese works (including those for new variants) require bounded gradient assumptions.
Revised: Additionally, all ofthese works (including those for new variants) require bounded gradient assumptions.

Original: Among all these works, (Shi et al., 2020) is the first work that proves RMSProp (a non-momentumversion of Adam with β 1 = 0 ) can converge without bounded gradient assumption.
Revised: Among all these works, (Shi et al., 2020) is the first work that proves RMSProp (a non-momentumversion of Adam with β 1 = 0 ) can converge without bounded gradient assumption.

Original: When the2nd-order parameter β 2 is large, they show that RMSProp converges to a bounded region in general,and to critical points under SGC.
Revised: When the2nd-order parameter β 2 is large, they show that RMSProp converges to a bounded region in general,and to stationary points under SGC.

Original: They also provide a convergence result for Adam with β 1 closeenough to 0.
Revised: They also provide a convergence result for Adam with β 1 closeenough to 0.

Original: However, a simple calculation by Zhang et al.
Revised: However, a simple calculation by Zhang et al.

Original: (2022) shows that they require β 1 < 10 − 7on real datasets.
Revised: (2022) shows that they require β 1 < 10 − 7on real datasets.

Original: Thus their result does not provide much extra information other than RMSProp.
Revised: Thus their result does not provide much extra information other than RMSProp.

Original: Itis still unclear how β 1 affects the (non-)convergence of Adam.
Revised: It is still unclear how β 1 affects the (non-)convergence of Adam.

Original: Proving convergence under large β 1is non-trivial.
Revised: Proving convergence under largeβ 1 is non-trivial.

Original: We point out there are at least two challenges.
Revised: We point out there are at least two issues.

Original: First, it seems unclear how to controlthe massive momentum m t .
Revised: First, it seems unclear how to control they this Appendix the massive momentum m t .

Original: Second, the momentum m t in Adam is multiplied by 1 / √ v t , causingnon-linear perturbation.
Revised: Second, the momentum m t in Adam is multiplied by 1 / √ v t , causingnon-linear perturbation.

Original: We use much different techniques to resolve these two issues.
Revised: It turns out that we requires much different techniques to resolve these twoissues.

Original: We highlightour technical novelties in Section 5.
Revised: We highlight our technical novelties in Section 5.

Original: 2.3 The Importance and Difficulties of Removing Bounded Gradient Assumptions
Revised: 2.3 The Importance and Difficulties of Removing Bounded Gradient Assumptions

Original: Here, we emphasize the importance to remove bounded gradient assumption.
Revised: Here, we emphasize the importance to remove bounded gradient assumption.

Original: First, unlike theassumptions in Section 2.1, bounded gradient is not common in SGD analysis.
Revised: First, unlike theassumptions in Section 2.1, bounded gradient is not common in SGD analysis.

Original: So it is of theoreticalinterests to remove this condition for Adam.
Revised: So it is of theoreticalinterests to remove this condition for Adam.

Original: Second, bounded gradient condition rules out thechances of gradient divergence a priori.
Revised: Second, bounded gradient condition rules out thechances of gradient divergence a priori.

Original: However, there are numerical evidences showing that Adam’sgradient can diverge (see Appendix B and Appendix B.1).
Revised: However, there are numerical evidences showing that Adam’sgradient can diverge (see Section 6 and Appendix D.2).

Original: Removing the boundedness assumptionhelps us point out the divergence and convergence behavior of Adam in the ( β 1 , β 2 ) diagram.
Revised: Removing the boundedness assumption helpsus point out the divergence and convergence behavior of Adam in the ( β 1 , β 2 ) diagram.

Original: However, it is often difficult to remove the boundedness condition, especially for stochastic algorithmswith momentum.
Revised: However, it is often difficult to remove the boundedness condition, especially for stochastic algorithmswith momentum.

Original: For instance, An early paper Bertsekas & Tsitsiklis (2000) analyzed SGD-typemethods without any boundedness condition.
Revised: For instance, An early paper Bertsekas & Tsitsiklis (2000) analyzed SGD-typemethods without any boundedness condition.

Original: But it is not until recently that Yu et al.
Revised: But it is not until recently that Yu et al.

Original: Liuet al. (2020b) prove SGDM (SGD with momentum) converges without bounded gradient assumption.
Revised: Liuet al. (2020b) prove SGDM (SGD with momentum) converges without bounded gradient assumption.

Original: Such attempts of removing boundedness assumptions are often appreciated for general optimizationproblems where “bounded-assumption-free" is considered as a major contribution.
Revised: Such attempts of removing boundedness assumptions are often appreciated for general optimizationproblems where “bounded-assumption-free" is considered as a major contribution.

Original: For Adam, the role of momentumm t is even more intricate since it is multiplied by 1 / √ v t .
Revised: For Adam, the role of momentum is even more intricate since it is multiplied by 1 / √ v t .

Original: Combinedwith v t , the impact of previous signals not only affect the update direction, but also change the stepsizefor each component.
Revised: Combinedwith v t , the impact of previous signals not only affect the update direction, but also change the stepsizefor each component.

Original: Further, both momentum m t and stepsize 1 / √ v t are random variables andthey are highly correlated.

Original: Such statistical dependency causes trouble for analysis.

Original: In summary,the role of momentum in Adam could be much different from that in SGDM or GDM.
Revised: As such, the role of momentum in Adam could be much different from that in SGDM or GDM.

Original: Even withboundedness conditions, the convergence of large- β 1 Adam is still not well understood (see relatedworks in Section 2.2).
Revised: Even with boundedness conditions, the convergence of large- β 1 Adam is still notwell understood (see related works in Section 2.2).

Original: In this work, we propose new techniques to handle Adam’s momentum underany large β 1 , regardless of the gradient magnitude.
Revised: In this work, we propose new techniques to handle Adam’s momentum under any large β 1 , regardless of the gradient magnitude.

Original: These techniques are not revealed in any existingworks.
Revised: These techniques arenot revealed in any existing works.

Original: We introduce our technical contribution in Section 5.
Revised: We introduce our technical contribution in Section 5.

Original:  Main Results
Revised:  Main Results

Original: 3.1 Convergence Results
Revised: 3.1 Convergence Results

Revised: [Figure]

Original: Here, we give the convergence results under large β 2 .
Revised: Here, we give the convergence results under large β 2 .

Original: Theorem 3.1.
Revised: Theorem 3.1.

Original: For any f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) , we assume the hyperparameters in Algorithm 1satisfy: β 1 < √ β 2 < 1 ; β 2 is greater or equal to a threshold γ 1 ( n ) ; and η k = η 1 √ nk .
Revised: For any f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) , when the hyperparameters in Algorithm 1 satisfies:β 1 < √ β 2 < 1 ; β 2 is greater or equal to a threshold γ 1 ( n ) ; and η k = η 1 √ nk , we will have: [Equation]

Original: Let k m ∈
Revised: where k m ∈

Original: Nsatisfies k m ≥ 4 and β ( k m − 1) 
Revised: N satisfies k m ≥ 4 and β ( k m − 1) n

Original: n≤ β n 1 √ k m − 1 , 4 we have the following results for any T > k m :

Original: [Table]
Revised: [Equation]

Original: Remark 1: the choice of β 2 .
Revised: Remark 1: the choice of β 2 .

Original: Our theory suggests that large β 2 should be used to ensure convergence.
Revised: The threshold γ 1 ( n ) satisfies β 2 ≥ 1 − O (cid:16) 1 −

Original: This message matches our experimental findings in Figure 1. We would like to point outthat the requirement of “large β 2 " is neccessary, because small β 2 will indeed lead to divergence(shown later in Section 3.2). We here comment a bit on the the threshold γ 1 ( n ) .
Revised: [Equation]

Original: γ 1 ( n ) satisfiesβ 2 ≥ 1 − O (cid:16) 1 − β n 1 n 2 ρ (cid:17) (see inequality (33) and Remark G.7), where ρ is a constant that depends on thetraining trajectory.
Revised: n 2 ρ (cid:17) (see inequality(32) and Remark G.8), where ρ is a constant that depends on the training trajectory.

Original: In worst cases, ρ is upper bounded by n 2 . 5 , but we find the practical ρ to be muchsmaller.
Revised: In worst cases,ρ is upper bounded by n 2 . 5 , but we find the practical ρ to be much smaller.

Original: In Appendix B.1, we estimate ρ on MNIST and CIFAR-10.
Revised: In Appendix D.2, weestimate ρ on MNIST and CIFAR-10.

Original: In practical training process, weempirically observe that ρ ≈ O ( n ) , thus the required γ 1 ( n ) ≈ 1 − O (cid:0) n − 3 (cid:1) .
Revised: In practical training process, we empirically observe thatρ ≈ O ( n ) , thus the required γ 1 ( n ) ≈ 1 − O (cid:0) n − 3 (cid:1) .

Original: Note that our thresholdof β 2 is a sufficient condition for convergence, so there may be a gap between the practical choicesand the theoretical bound of β 2 .
Revised: Note that our threshold of β 2 is a sufficientcondition for convergence, so there may be a gap between the practical choices and the theoreticalbound of β 2 .

Original: Closing the gap will be an interesting future direction.
Revised: Closing the gap will be an interesting future direction.

Original: We find that γ 1 ( n ) increases with n .
Revised: We further remark that γ 1 ( n ) increases with n .

Original: This property suggests that larger β 2 should be used when n islarge.
Revised: This condition suggests that larger β 2 should be usedwhen the batchsize is small.

Original: This phenomenon is also verified by our experiments in Figure 8 (d), Appendix B. We alsoremark that γ 1 ( n ) slowly increases with β 1 .
Revised: This phenomenon is also verified by our experiments in Figure 4 (d), Section 6.

Original: This property is visualized in Figure 1 (d) where thelower boundary of blue region slightly lifts up when β 1 increases.
Revised: Intuitively, this phenomenon is understandable since larger n brings larger noise in thestochastic gradient, and thus larger β 2 is required.

Original: Remark 2: the choice of β 1 .
Revised: Remark 2: the choice of β 1 .

Original: β 1 < √ β 2 is also a requirement in (Kingma & Ba, 2014), but theirconvergence analysis requires bounded gradient and Reddi et al.
Revised: β 1 < √ β 2 is also a requirement in (Kingma & Ba, 2014), buttheir convergence analysis requires bounded gradient and Reddi et al.

Original: (2018) later pointed out that theproof has flaws.
Revised: (2018) later pointed out thatthe proof has flaws.

Original: Theorem 3.1 is the first result that proves the convergence of Adam with the samechoice of β 1 , but with less assumptions.
Revised: Theorem 3.1 is the first result that proves the convergence of Adam with thesame choice of β 1 , but with less assumptions.

Original: Since β 2 is suggested to be large, our convergence result cancover flexible choice of β 1 ∈ [0 , 1) .
Revised: Since β 2 is suggested to be large, our convergenceresult can cover flexible choice of β 1 ∈ [0 , 1) .

Original: For instance, β 2 = 0 . 999 brings the threshold of β 1 < 0 .
Revised: For instance, β 2 = 0 . 999 brings the threshold ofβ 1 < 0 .

Original: 9995 ,which covers  basically all practical choices of β 1 reported in the literature (seeSection 1),
Revised: 9995 , which covers and basically all practical choices of β 1 reported in the literature (see

Original: includingthe default setting β 1 = 0 .9 .
Revised: Section 1), including the default setting β 1 = 0 .

Original: Remark 3: convergence to a bounded region.
Revised: Remark 3: convergence to a bounded region.

Original: When D 0 > 0 , Adam converges to a boundedregion near critical points.
Revised: When D 0 > 0 , Adam converges to a boundedregion near stationary points.

Original: As discussed in Section 2.1, converging to bounded region is commonfor stochastic methods including constant-stepsize SGD (Yan et al., 2018;
Revised: As discussed in Section 2.1, converging to bounded region is commonfor stochastic methods including constant-stepsize SGD (Yan et al., 2018;

Original: Yu et al., 2019;
Revised: Yu et al., 2019;

Original: Liu et al.,2020b) and diminishing-stepsize RMSProp (Zaheer et al., 2018; Shi et al., 2020).
Revised: Liu et al.,2020b) and diminishing-stepsize RMSProp (Zaheer et al., 2018; Shi et al., 2020).

Original: This phenomenonis also observed in practice: even for convex quadratic function with D 0 > 0 , Adam with diminishingstepsize cannot reach exactly zero gradient (see Appendix B.1).
Revised: This phenomenonis also observed in practice: even for convex quadratic function with D 0 > 0 , Adam with diminishingstepsize cannot reach exactly zero gradient (see Appendix D.2).

Original: The good news is that, the constant
Revised: The good news is that, the constant

Original: O ( √ D 0 ) vanishes to 0 as β 2 goes to 1.
Revised: O ( √ D 0 ) vanishes to 0 as β 2 goes to 1.

Original: This can be seen in Remark G.14 in Appendix G.9.
Revised: This can be seen in Remark G.17 in Appendix G.9.

Original: As a corollary of Theorem 3.1, we have the following result under SGC (i.e., D 0 = 0 ).
Revised: As a corollary of Theorem 3.1, we have the following result under SGC (i.e., D 0 = 0 ).

Original: Corollary 3.2.
Revised: Corollary 3.2.

Original: Under the setting in Theorem 3.1.
Revised: Under the setting in Theorem 3.1.

Original: When D 0 = 0 for Assumption 2.2, we have
Revised: When D 0 = 0 for Assumption 2.2, we have

Original: [Equation]
Revised: [Equation]

Revised: where k m ∈

Revised: N satisfies k m ≥ 4 and β ( k m −

Original:  Under SGC (i.e. D 0 = 0 ), Corollary 3.2 states that Adam can converge to critical points.
Revised: [Equation] n≤ [Equation] Under SGC (i.e. D 0 = 0 ), Corollary 3.2 states that Adam can converge to stationary points.

Original: This isindeed the case in practice.
Revised: This isindeed the case in practice.

Original: For instance, function (2) satisfies SGC and we observe 0 gradient normafter Adam converges (see Appendix B and Appendix B.1).
Revised: For instance, function (2) satisfies SGC and we observe 0 gradient normafter Adam converges (see Appendix D.2).

Original: The convergence rate in Corollary 3.2 iscomparable to that of SGD under the same condition in (Vaswani et al., 2019).
Revised: The convergence rate in Corollary 3.2 is comparable tothat of SGD under the same condition in (Vaswani et al., 2019).

Original: 3.2 Divergence Results
Revised: 3.2 Divergence Results

Original: Theorem 3.1 shows that when β 2 is large, any β 1 < √ β 2 ensures convergence.
Revised: Theorem 3.1 shows that when β 2 is large, any β 1 < √ β 2 ensures convergence.

Original: Now we consider254 the  case where β 2 is small.
Revised: Now we consider the 222 case where β 2 is small.

Original: We will show that in this case, a wide range of β 1 is facing the risk255 of diverging to infinity.
Revised: We will show that in this case, a wide range of β 1 is facing the risk of divergingto infinity.

Original: The divergence of small- β 2 Adam suggests that “large β 2 " is necessary inthe convergence result Theorem 3.1.

Original: We construct a counter-example in F n,f ∗ L,D 0 ,D 1 ( R d ) .
Revised: We construct a counter-example in F n,f ∗ L,D 0 ,D 1 ( R d ) .

Original: Considerf ( x ) = (cid:80) n − 1 i =

Revised: Consider f ( x ) = (cid:80) n − 1 i =

Original: f i ( x ) for x ∈ R , we define f i ( x ) as:
Revised: f i ( x ) forx ∈ R , we define f i ( x ) as:

Original: [Equation]
Revised: [Equation]

Original: Summing up all the f i ( x ) , we can see that
Revised: Summing up all the f i ( x ) , we can see that

Original: [Equation]
Revised: [Equation]

Original: is a lower bounded convex smooth function with optimal solution x ∗ = − 2 .
Revised: is a lower bounded convex smooth function with optimal solution x ∗ = − 2 .

Original: Function (2) allowsboth iterates and gradients to diverge to infinity.
Revised: Function (2) allows bothiterates and gradients to diverge to infinity.

Original: As shown in Figure 1 (a), when running Adam on(2), there exists a red large-error region.
Revised: As shown in Figure 1 (a), when running Adam on (2), there exists a red large-error region.

Original: This shows the sign of divergence.
Revised: Thisshows the sign of divergence.

Original: We further theoreticallyverify the conjecture in Proposition 3.3.
Revised: We further theoretically verify the conjecture.

Original: Proposition 3.3.
Revised: Proposition 3.3.

Original: For any function class F n,f ∗ L,D 0 ,D 1 ( R d ) , there exists a f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) , s.t.when ( β 1 , β 2 ) satisfies analytic condition (11) , (12) , (13) in Appendix E, Adam’s iterates and functionvalues diverge to infinity.
Revised: For any function class F n,f ∗ L,D 0 ,D 1 ( R d ) , there exists a f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) , s.t.when ( β 1 , β 2 ) satisfies analytic condition (9) , (10) , (11) in Appendix E, Adam’s iterates and functionvalues diverge to infinity.

Original: By solving these conditions in NumPy , we plot the orange region in Figure2.
Revised: By solving these conditions in NumPy , we plot the orange region in 2.

Original: The size of the region depends on n and it expands to the whole region when n goes to infinity.0.
Revised: Thesize of the region depends on n and it expands to the whole region when n goes to infinity. [Equation] b e t a

Original: The proofcan be seen in Appendix E.
Revised: We present the proof in Appendix E.

Original: We find the “divergenceregion" always stays below the “convergence threshold"γ 1 ( n ) in
Revised: We find the “divergence region"always stays below the “convergence threshold"

Original: Theorem 3.1, so the two results are self-consistent (see the remarkin Appendix E).
Revised: γ 1 ( n ) in Theorem3.1, so the two results are self-consistent (see the remark in Appendix E).

Original: Proposition 3.3 states the divergence of iteratesand function values.
Revised: Proposition 3.3 states the divergence of iterates and functionvalues.

Original: Consistently, our experiments also show thedivergence of gradient (see Appendix B and Appendix B.1).
Revised: Our experiments also show the divergence of gradient (see Section 6 and Appendix D.2).

Original: Theseresults characterize Adam’s divergence behavior both numericallyand theoretically.
Revised: These results characterize Adam’sdivergence behavior both numerically and theoretically.

Original: We emphasize that the orange region is not discussed in (Reddi et al.,2018) because we consider n fixed while they allow n changing.
Revised: We emphasize that the orange region is not discussed in (Reddi et al.,2018) because we consider n fixed while they allow n changing.

Original: When n is allowed to increase, our orange region will expand to thewhole region and thus we can derive a similar (actually stronger)result as (Reddi et al., 2018).
Revised: When n is allowed to increase, our orange region will expand tothe whole region and thus we can derive a similar result as (Reddiet al., 2018).

Original: We provide more explanation in Section 4.
Revised: We provide more explanation in Section 4.

Original: Combining Theorem 3.1 and Proposition 3.3, we establisha clearer image on the relation between ( β 1 , β 2 ) and qualitativebehavior of Adam.
Revised: Combining Theorem 3.1 and Proposition 3.3, we establish a clearer image onthe relation between ( β 1 , β 2 ) and qualitative behavior of Adam.

Original:  Reconciling Our Results with (Reddi et al., 2018)
Revised:  Reconciling Our Results with (Reddi et al., 2018)

Original: We discuss more on the relation between (Reddi et al., 2018) and our results.
Revised: We discuss more on the relation between (Reddi et al., 2018) and our results.

Original: Recall Reddi et al.(2018) show that “for any β 1 < √ β 2 < 1 , Adam can diverge".
Revised: Recall Reddi et al.(2018) show that “for any β 1 < √ β 2 < 1 , Adam can diverge".

Original: This does not contradict with ourconvergence results in Theorem 3.1.
Revised: This does not contradict with ourconvergence results in Theorem 3.1.

Original: Further, it is different from our divergence result in Proposition3.3.
Revised: Further, it is different from our divergence result in Proposition3.3.

Original: As discussed in Section 1, the key difference lies in whether ( β 1 , β 2 ) is picked before or afterpicking n , which often stands for sample size or the number of mini-batch.
Revised: As discussed in Section 1, the key difference lies in whether ( β 1 , β 2 ) is picked before or afterpicking n , which often stands for sample size or the number of mini-batch.

Original: Since n is fixed when the (sub-)function class F n,f ∗ L,D 0 ,D 1 ( R d ) is given, the above key difference canalso be expressed as whether ( β 1 , β 2 ) is picked before or after picking F n,f ∗ L,D 0 ,D 1 ( R d ) .
Revised: Since n is fixed when the (sub-)function class F n,f ∗ L,D 0 ,D 1 ( R d ) is given, this key difference can also beexpressed as whether ( β 1 , β 2 ) is picked before or after picking F n,f ∗ L,D 0 ,D 1 ( R d ) .

Original: Now wediscuss the following two cases.
Revised: Now we discuss thefollowing two cases.

Original: Case I: When ( β 1 , β 2 ) is picked before picking F n,f ∗ L,D 0 ,D 1 ( R d ) .
Revised: Case I: When ( β 1 , β 2 ) is picked before picking F n,f ∗ L,D 0 ,D 1 ( R d ) .

Original: As discussed in Section 1, the divergence result requires differentn for different ( β 1 , β 2 ) .
Revised: As discussed in Section 1, the divergence result requires differentn for different ( β 1 , β 2 ) .

Original: In this sense, the considered function classis constantly changing.
Revised: In this sense, the considered function classis constantly changing.

Original: It does not contradict with our Theorem 3.1which considers a fixed function class with fixed n .
Revised: It does not contradict with our Theorem3.1 which considers a fixed function class with fixed n .

Original: For Case I , weillustrate Adam’s behavior in Figure 3.
Revised: For Case I , we illustrate Adam’s behavior in Figure 3.

Original: The red region is proved by(Reddi et al., 2018).
Revised: The red region withis proved by (Reddi et al., 2018).

Original: For completeness, we remove the condition“ β 1 < √ β 2 " and further prove that Adam will diverge to infinityfor any ( β 1 , β 2 ) ∈ [0 , 1) 2 .
Revised: For completeness, we remove thecondition “ β 1 < √ β 2 " and further prove that Adam will divergeto infinity for any ( β 1 , β 2 ) ∈ [0 , 1) 2 .

Original: The result is shown in the following Corollary 4.1.
Revised: The result is shown in thefollowing Corollary 4.1.

Original: Corollary 4.1.
Revised: Corollary 4.1.

Original: For any ( β 1 , β 2 ) ∈ [0 , 1) 2 , there exists a functionsatisfying Assumption 2.1 and 2.2 that the Adam’s iterates andfunction values diverge to infinity.
Revised: For any ( β 1 , β 2 ) ∈ [0 , 1) 2 , there exists a functionsatisfying Assumption 2.1 and 2.2 that the Adam’s iterates andfunction values diverge to infinity.

Revised: Rigorously speaking, Reddi et al.

Revised: (2018) consider constrained problems.

Revised: So their claimed “divergence" actually means the iterates willstay in a huge region with the size of whole feasible set.

Revised: We callit “non-convergence" to distinguish from our result of “diverge toinfinity" in Figure 3.

Original: Proof of Corollary 4.1 can be seen in the final paragraph in Appendix
Revised: Proof of Corollary 4.1 can be seen in the final paragraph in Appendix

Original: E.
Revised: E.

Original: In the proof, we also require different n to cause divergence for different ( β 1 , β 2 ) .
Revised: In the proof, we also require different n to cause divergence for different ( β 1 , β 2 ) .

Original: So the functionclass is constantly changing.
Revised: So the functionclass is constantly changing.

Original: As a result, in Case I , we cannot prove any convergence result.
Revised: As a result, in Case I , we cannot prove any convergence result.

Original: [Table]
Revised: [Table]

Revised: The orange region in Figure

Original:  Case II: When ( β 1 , β 2 ) is picked after picking F n,f ∗ L,D 0 ,D 1 ( R d ) .
Revised: [Table] Case II: When ( β 1 , β 2 ) is picked after picking F n,f ∗ L,D 0 ,D 1 ( R d ) .

Original: When the function class is pickedin advance, sample size n will also be fixed.
Revised: When the function class is pickedin advance, sample size n will also be fixed.

Original: This case is closer to most practical applications.
Revised: This case is closer to most practical applications.

Original: Inthis case, we find that Adam’s behavior changes significantly in the different region of Figure 3.
Revised: Inthis case, we find that Adam’s behavior changes significantly in the different region of Figure 3.

Original: First, ∀ f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) will converge when β 1 < √ β 2 and β 2 is large.
Revised: First, ∀ f ( x ) ∈ F n,f ∗ L,D 0 ,D 1 ( R d ) will converge when β 1 < √ β 2 and β 2 is large.

Original: Second, ∃ f ( x ) ∈
Revised: Second, ∃ f ( x ) ∈

Original: F n,f ∗ L,D 0 ,D 1 ( R d ) will diverge to infinity when ( β 1 , β 2 ) are in the orange region in Figure 2.
Revised: F n,f [Equation] L,D 0 ,D 1 ( R d ) will diverge to infinity when ( β 1 , β 2 ) are in the orange region in Figure 2.

Original: Since Case
Revised: Since Case

Original: II is closer to practical scenarios, these results can provide better guidance for hyperparameter tuningfor Adam users.
Revised: II is closer to practical scenarios, these results can provide better guidance for hyperparameter tuningfor Adam users.

Original: We provide some suggestionsfor practitioners in Appendix C.
Revised: We provide some suggestions in Appendix B.

Original: For Case II , we summarize the possible behaviors of Adam in Table 1.
Revised: For Case II , we summarize the possible behaviors of Adam in Table 1.

Original: We also illustrate ourconvergence and divergence results in Figure 1 (d).
Revised: We also illustrate ourconvergence and divergence results in Figure 1 (d).

Original: Note that there are some blanket areas where
Revised: Note that there are some blanket areas where

Original: Adam’s behavior remains unknown, this part will be left as interesting future work.
Revised: Adam’s behavior remains unknown, this part will be left as interesting future work.

Original:  Proof Ideas for the Convergence Result
Revised:  Proof Ideas for the Convergence Result

Original: We now  introduce our proof ideas for the convergence result in Theorem 3.1.
Revised: We now briefly introduce our proof ideas for the convergence result in Theorem 3.1.

Original: Simply put, we want to319 control the update direction m k,i / √ v k,i inside the dual cone of gradient direction.
Revised: Simply put, we 289 need to control the update direction m k,i / √ v k,i near the negative gradient direction.

Original: Namely: E ⟨∇ f ( x k, 0 ) ,
Revised: Namely, we want to show

Original: [Equation]
Revised: [Equation]

Original: However, directly proving (3) could be difficult because both m k,i and v k,i distort the trajectory.
Revised: However, directly proving (3) could be difficult because both m k,i and v k,i distort the trajectory.

Original: To start with, we try to control the movement of v k,i .
Revised: To start with, we try to control the movement of v k,i .

Original: We use a standard approach for Adam-type algorithms: use large or increasing β 2 (e.g. (Shi et al., 2020;
Revised: We use a standard approach for Adam-type algorithms: use large or increasing β 2 (e.g. (Shi et al., 2020;

Original: Zou et al., 2019; Chen et al., 2018b)).
Revised: Zou et al., 2019; Chen et al., 2018)).

Original: Recall v k,i = (1 − β 2 ) (cid:80) ij =1 β i − j 2 ∇ f τ k,j ( x k,j ) ◦ ∇ f τ k,j ( x k,j ) +
Revised: Recall v k,i = (1 − β 2 ) (cid:80) ij =1 β i − j 2 ∇ f τ k,j ( x k,j ) ◦ ∇ f τ k,j ( x k,j ) +

Original: β
Revised: β

Original: i 2 v k, 0 , we have v k,i ≈ v k, 0 when β 2 is large.
Revised: i 2 v k, 0 , we have v k,i ≈ v k, 0 when β 2 is large.

Original: In this case, we have:
Revised: In this case, we have:

Original: [Equation]
Revised: [Equation]

Original: where the first “ ≈ " is due to the large β 2 and the second “ ≈ " is our goal.

Original: Now we need to show:
Revised: Now we need to show:

Original: E

Original: [Equation]
Revised: [Equation]

Original: f

Original: [Equation] x k, 0 ) √ [Equation] where ∂ l f ( x k, 0 ) is the l -th component of ∇ f ( x k, 0 ) , similarly for m l,k, 

Revised: where ∂ l f ( x k, 0 ) is the l -th component of ∇ f ( x k, 0 ) , similarly for m l,k,

Original: and v l,k, 0 .
Revised: and v l,k, 0 .

Original: ( ∗ ) is due to thefinite-sum structure.
Revised: ( ∗ ) is due to thefinite-sum structure.

Original: However, it is not easy to prove (4).
Revised: However, it is not easy to prove (4).

Original: We point out some technical issues below.
Revised: We point out some technical issues below.

Original: Issue I: massive momentum.
Revised: Issue I: massive momentum.

Original: Directly proving (4) is still not easy.
Revised: Directly proving (4) is still not easy.

Original: We need to first consider asimplified problem: for every l ∈ [ d ] , assume we treat ∂ l f ( x k, 0 ) / √ v l,k, 0 as a constant, how to bound
Revised: We need to first consider asimplified problem: for every l ∈ [ d ] , assume we treat ∂ l f ( x k, 0 ) / √ v l,k, 0 as a constant, how to bound

Original: E (cid:80) n − 1 i =0 (cid:0) m l,k,i − ∂ l f τ k,i ( x k, 0 ) (cid:1) ?
Revised: E (cid:80) n − 1 i =0 (cid:0) m l,k,i − ∂ l f τ k,i ( x k, 0 ) (cid:1) ?

Original: It turns out that this simplified problem is still non-trivial.
Revised: It turns out that this simplified problem is still non-trivial.

Original: When β 1 is large, m l,k,i contains heavyhistorical signals which significantly distort the trajectory from gradient direction.
Revised: When β 1 is large, m l,k,i contains heavyhistorical signals which significantly distort the trajectory from gradient direction.

Original: Existing literature(Zaheer et al., 2018; De et al., 2018; Shi et al., 2020) take a naive approach: they set β 1 ≈ 

Revised: Existing literatures(Zaheer et al., 2018; De et al., 2018; Shi et al., 2020) take a naive approach: they set β 1 ≈

Original: so thatm l,k,i ≈ ∂
Revised: so thatm l,k,i ≈ ∂

Original: l f τ k,i ( x k,i ) .
Revised: l f τ k,i ( x k,i ) .

Original: Then we get (4) ≈ 0 .
Revised: Then we get (4) ≈ 0 .

Original: However, this method cannot be applied here since weare interested in practical cases where β 1 is large in [0 , 1) .
Revised: However, this method cannot be applied here since weare interested in practical cases where β 1 is large in [0 , 1) .

Original: Issue II:stochastic non-linear dynamics.
Revised: Issue II: non-linear dynamics.

Original: Even if we solve Issue I , it is still unclear how to prove(4).
Revised: Even if we solve Issue I , it is still unclear how to prove (4).

Original: This is because: for every l ∈ [ d ] , ∂ l f ( x k, 0 ) / √ v l,k, 0 is a r.v. instead of a constant.
Revised: This isbecause: for every l ∈ [ d ] , ∂ l f ( x k, 0 ) / √ v l,k, 0 is a r.v. instead of a constant.

Original: With this terminvolved, we are facing with a stochastic non-linear dynamics, which could be difficult to analyze.
Revised: With this term involved,we are facing with a stochastic non-linear dynamics, which could be difficult to analyze.

Original: Further, ∂ l f ( x k, 0 ) / √ v l,k, 0 is statistically dependent with (cid:0) m l,k,i − ∂ l f τ k,i ( x k, 0 ) (cid:1) , so we are notallowed to handle the expectation E ( ∂ l f ( x k, 0 ) / √ v l,k, 0 ) separately .

Original: Unfortunately, even with additional assumptions like bounded gradient, there is no general approachto tackle the above issues.
Revised: Unfortunately, even with additional assumption like bounded gradient, there is no general approachto tackle the above issues.

Original: In this work, we propose solutions regardless of gradient magnitude.
Revised: In this work, we propose a solution regardless of gradient magnitude.

Original: Solution to Issue I.
Revised: Step 1.

Original: We prove the following Lemma to resolve Issue I .
Revised: We prove the following Lemma to resolve Issue I .

Original: Lemma 5.1.
Revised: Lemma 5.1.

Original: (Informal) Consider Algorithm 1.
Revised: (Informal) Consider Algorithm 1.

Original: For every l ∈ [ d ] and any β 1 ∈ [0 , 1) , we have the 340 following result under Assumption 2.1.
Revised: For every l ∈ [ d ] and any β 1 ∈ [0 , 1) , we have the 309 following result under Assumption 2.1.δ [Equation] β 1 ) : =

Revised: E

Original: [Equation]
Revised: [Equation]

Original: where ∂ l f ( x k, 0 ) is the l -th component of ∇ f ( x k, 0 ) ; m l,k,i
Revised: where ∂ l f ( x k, 0 ) is the l -th component of ∇ f ( x k, 0 ) ; m l,k,i

Original: = (1 − β 1 ) ∂ l fτ k,i ( x k,i ) + β 1 m l,k,i − 1 .
Revised: = (1 − β 1 ) ∂ l f ( x k,i ) + β 1 m l,k,i − 1 .

Original: We present the proof idea in Appendix A. Simply put, we construct a simple toy example called“color-ball" model (of the 1st kind). This toy model shows a special property of δ ( β 1 ) .
Revised: We present the proof idea in Appendix A. Simply put, we construct a simple toy example called“color-ball" model, which shows a special property of δ ( β 1 ) .

Original: We find out:for Algorithm 1, error terms from successive epochs can be canceled, which keeps the momentumroughly in the descent direction.
Revised: We find out: for Algorithm 1, error termsfrom successive epochs can be canceled, which keeps the momentum roughly in the descent direction.

Original: This important property is not revealed in any existing work.
Revised: This important property is not revealed in any existing work.

Revised: In Appendix A, we demonstrate the“color-ball” model in simple language where only basic probability knowledge is required.

Original: Remark 4: When assuming bounded gradient ∥∇ f ( x ) ∥ ≤ G , a naive upper bound would beδ ( β 1 ) =
Revised: Remark 4: When assuming bounded gradient ∥∇ f ( x ) ∥ ≤ G , a naive upper bound would beδ ( β 1 ) =

Original: O ( G ) .
Revised: O ( G ) .

Original: However, such constant upper bound does not imply δ ( β 1 ) is close to 0.
Revised: However, such constant upper bound does not imply δ ( β 1 ) is close to 0.

Original: It will nothelp prove the convergence.
Revised: It will nothelp prove the convergence.

Original: This might be partially the reason why large- β 1 Adam is hard to analyzeeven under bounded gradient (see related works in Section 2.2).
Revised: This might be partially the reason why large- β 1 Adam is hard to analyzeeven under bounded gradient (see related works in Section 2.2).

Original: We emphasize Lemma 5.1 holds trueregardless of gradient norm, so it could be deployed in both bounded or unbounded gradient analysis.
Revised: We emphasize Lemma 5.1 holds trueregardless of gradient norm, so it could be deployed in both bounded or unbounded gradient analysis.

Original: Solution to Issue II.
Revised: Step 2.

Original: We try to show (4) by adopting Lemma 5.1.
Revised: We try to show (4) by adopting Lemma 5.1.

Original: However, the direct applicationcannot work since ∂ l f ( x k, 0 ) √ v l,k, 0 is not constant.
Revised: However, the direct application cannotwork since ∂ l f ( x k, 0 ) √ v l,k, 0 is not constant.

Original: Despite its randomness, we find out that when β 2 islarge, (cid:110) ∂
Revised: Despite its randomness, we find out that when β 2 is large, (cid:110) [Equation]

Original: l f ( x k, 0 ) √ v l,k, 0 (cid:111) ∞
Revised: l [Equation] ( x [Equation]

Original:  k =1 behaves like Cauchy sequence which has unique limit point.
Revised: [Equation] ,k, [Equation] ∞ k =1 behaves like Cauchy sequence which has unique limit point.

Original: As such, although∂ l f ( x k, 0 ) √ v l,k, 0 brings extra perturbation, the quantity in (4) share the similar asymptotic behavior as δ ( β 1 ) .
Revised: As such, although∂ l f ( x k, 0 ) √ v l,k, 0 brings extra perturbation, the quantity in (4) share the similar asymptotic behavior as δ ( β 1 ) .

Original: We prove the following Lemma 5.2.
Revised: Therefore, with the help of Step 1 , we prove the following Lemma 5.2.

Original: Lemma 5.2.
Revised: Lemma 5.2.

Original: (Informal) Under Assumption 2.1 and 2.2, consider Algorithm 1 with large β 2 andβ 1 < √ β 2 .
Revised: (Informal) Under Assumption 2.1 and 2.2, consider Algorithm 1 with large β 2 and β 1 < √ β 2 .

Original: For those l with gradient component larger than certain threshold, we have:
Revised: For those l with gradient component larger than certain threshold, we have:

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: In Appendix A, we introduce how to derive (6) from (5).

Original: To do so, we introduce a new type of“color-ball" model (we call it color-ball of the 2nd kind) which adopts the random perturbation of∂ l f ( x k, 0 ) √ v l,k, 0 .

Original: Understanding color-ball model of the 2nd kind is crucial for proving Lemma 5.2.

Original: We conclude the proof of (4) by some additional analysis on “those l with small gradient component".
Revised: We conclude the proof of (4) by some additional analysis on “those l with small gradient component".

Original: This case is a bit easier since it reduces to bounded gradient case.
Revised: This case is a bit easier since it reduce to bounded gradient case.

Original: For readers who wants to learn moreabout the idea of tackling Issue I and II , please refer to Appendix A where we formally introduce the1st and 2nd kind of color-ball models.

Original: Since the whole proof is quite long, we provide a proof sketchin Appendix G.1.
Revised: Since the whole proof is quite long,we provide a proof sketch in Appendix G.1.

Original: The whole proof is presented in Appendix G.
Revised: The whole proof is presented in Appendix G.

Original:  Conclusions
Revised:  Experiments

Revised: To support our theory, we provide more simulations and real-data experiments.

Revised: All the experimentalsettings and hyperparameters are presented in Appendix D. We aim to show:(I).

Revised: When β 2 is large, a large range of β 1 gives good performance, including all β 1 < √ β 2 .(II).

Revised: When β 2 is small, a large range of β 1 performs relatively badly

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: (d) MNIST(batchsizev.s.

Revised: β 2 )

Revised: Counter-example (2) .

Revised: We run more experiments on (2) with initialization x = − 5 and n =5 , 10 , 15 , 20 .

Revised: We show the result of n = 20 in Figure 4 (a), (b); the rest are shown in Appendix D.2.

Revised: We find that: first, when β 2 is small, gradient norms increase dramatically.

Revised: This emphasizes theimportance of discarding bounded gradient assumptions.

Revised: Second, when β 2 is large, Adam convergesto stationary points for β 1 < √ β 2 .

Revised: This matches our theory since (2) satisfies D 0 = 0 .

Revised: For functionswith D 0 > 0 , we further show that Adam can only converge to a bounded region in Appendix D.2.

Revised: MNIST and CIFAR-10.

Revised: As shown in Figure 1 (b)& (c) in Section 1, the training results matchboth claim (I) and (II).

Revised: In addition, there is a convex-shaped boundary on the transition from low lossto higher loss, this boundary roughly matches the condition in Theorem 3.1.

Revised: NLP.

Revised: We use Adam to train Transformer XL (Dai et al., 2019) on the WikiText-103 dataset (Merityet al., 2016).

Revised: This architecture and dataset is widely used in NLP tasks (e.g. (Howard & Ruder, 2018;

Revised: See et al., 2017)).

Revised: As shown in Figure 4 (c), the training results match both claim (I) and (II).

Revised: Batchsize and β 2 .

Revised: As shown in Figure 4 (d): on MNIST, smaller batchsize requires larger β 2 toreach small loss.

Revised: This matches the message by Theorem 3.1: the threshold of β 2 increases with n .

Revised:  Conclusions

Original: In this work, we explore the (non-)convergence of Adam.
Revised: In this work, we explore the (non-)convergence of Adam.

Original: When β 2 is large, we prove that Adamcan converge with any β 1 < √ β 2 .
Revised: When β 2 is large, we prove that Adamcan converge with any β 1 < √ β 2 .

Original: When β 2 is small, we further show that Adam might diverge toinfinity for a wide range of β 1 .
Revised: When β 2 is small, we further show that Adam might diverge toinfinity for a wide range of β 1 .

Original: One interesting question is to verify the advantage of Adam over
Revised: One interesting question is to verify the advantage of Adam over

Original: SGD.
Revised: SGD.

Original: In this work, we focus on the fundamental issue of convergence.
Revised: In this work, we focus on the fundamental issue of convergence.

Original: Proving faster convergence of
Revised: Proving faster convergence of

Original: Adam would be our future work.
Revised: Adam would be our future work.

Revised: Checklist

Revised: For all authors...

Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: Appendix Organization
Revised: contributions and scope? [Yes] (b) Did you describe the limitations of your work? [Yes]

Revised: (c) Did you discuss any potential negative societal impacts of your work?

Revised: [Yes]

Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Revised: them?

Revised: [Yes]

Revised: If you are including theoretical results...

Revised: (a) Did you state the full set of assumptions of all theoretical results?

Revised: [Yes]

Revised: (b) Did you include complete proofs of all theoretical results?

Revised: [Yes]

Revised: If you ran experiments...

Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Revised: mental results (either in the supplemental material or as a URL)?

Revised: [No]

Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Revised: how they

Revised: were chosen)?

Revised: [Yes]

Revised: (c) Did you report error bars (e

Revised: g., with respect to the random seed after running experi-

Revised: ments multiple times)?

Revised: [No]

Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., type

Revised: of GPUs, internal cluster, or cloud provider)?

Revised: [No]

Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Revised: (a) If your work uses existing assets, did you cite the creators?

Revised: [Yes]

Revised: (b) Did you mention the license of the assets?

Revised: [No]

Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Revised: [No]

Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Revised: using/curating?

Revised: [No]

Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Revised: information or offensive content?

Revised: [No]

Revised: If you used crowdsourcing or conducted research with human subjects...

Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Revised: applicable?

Revised: [N/A]

Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Revised: Board (IRB) approvals, if applicable?

Revised: [N/A]

Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Revised: spent on participant compensation?

Revised: [N/A]

Revised: Negative Social Impact

Revised: This script may provide better guidance for deep neural network training.

Revised: It would have certainnegative social impact if neural nets are deployed for illegal usage.

Revised: Appendix Organization

Original: The Appendix is organized as follows.
Revised: The Appendix is organized as follows.

Original: • Appendix A introduces two “color-ball" models and their applications in proving [-Lemma 5.
 andLemma 5.2.-]
Revised: • Appendix A discusses the ideas of “color-ball" model and its application in proving [-Lemma 5.
 andLemma 5.2.-]Lemma5.1.

Original: This part is important for proving Theorem 3.1.• Appendix Bintroduces experiments to support our theory.• Appendix C provide some suggestions for hyperparameter tuning of Adam.• Appendix D.1 provide more discussions on some recent related works.
Revised: This part is important for proving Theorem 3.1.• Appendix B provide some suggestions for hyperparameter tuning of Adam.• Appendix C.1 provide more discussions on some recent related works.

Original: Appendix D.2 re-state thenon-convergence results in (Reddi et al., 2018).• Appendix  E provides detailed proof for Proposition 3.3.• Appendix F provides some more notations and technical lemmas that serve for the proof of
Revised: Appendix C.2re-state the non-convergence results in (Reddi et al., 2018).• Appendix D provides more detailed setting of the experiments, as well as some additionalexperiments mentioned in the full script.• Appendix E provides detailed proof for Proposition 3.3.

Original: Theorem 3.1. • Appendix  G provides detailed proof for Theorem 3.1.
Revised: • Appendix F provides more some notations and technical lemmas that serve for the proof ofTheorem 3.1.• Appendix G provides detailed proof for Theorem 3.1.

Original: Especially, Appendix G.1 provide a proof roadmap.
Revised: Especially, Appendix G.1 provide aproof roadmap.

Original: A  Introduction to the “Color-Ball” Models: the Key Ingredients to ProveTheorem 3.1
Revised: A An Introduction to the “Color-Ball” Model

Original: We now introduce two “color-ball" models and their applications in tackling Issue I and Issue IImentioned in Section 5.
Revised: We now introduce ideas of the “color-ball" model and its application in bounding δ ( β 1 ) in Lemma5.1.

Original: These two color-ball models are important for proving Theorem 3.1.Solution to Issue I.
Revised: This part is important for proving Theorem 3.1.

Original: As discussed in Section 5, we wish to show that δ ( β 1 ) vanishes with k .
Revised: As discussed in Section 5, we wish to show that δ ( β 1 ) vanishes with k .

Original: Formally, we wish to get the following equation (7) for every l ∈ [ d ] .
Revised: Formally, we wish to get thefollowing equation (5) for every l ∈ [ d ] .

Original: [Equation]
Revised: [Equation]

Original: When k is large, O ( β nk 1 ) vanishes faster than O ( 1 √ k ) and thus Lemma 5.1 can be proved.
Revised: When k is large, O ( β k 1 ) vanishes faster than O ( 1 √ k ) and thus Lemma 5.1 can be proved.

Original: In thefollowing context, we will carefully quantify the mismatch between m l,k,i and the ∂ l f ( x k, 0 ) .
Revised: In thefollowing context, we will carefully quantify the mismatch between m l,k,i and the ∂ l f ( x k, 0 ) .

Original: Wefind out that the error terms from successive epochs can be cancelled, which keeps the momentumroughly in the descent direction.
Revised: Wefind out that the error terms from successive epochs can be cancelled, which keeps the momentumroughly in the descent direction.

Original: To help readers understand our idea, we introduce the “color-ball "model (of the 1st kind) as follows.
Revised: To help readers understand our idea, we consider a toy model calledthe “color-ball " model as follows.

Original: The color-ball model of the 1st kind. Consider a box containing two balls labeled with constantc 0 , c 1 ∈ R , respectively.
Revised: The “color-ball " model. Consider a box containing two balls labeled with constant c 0 , c 1 ∈ R ,respectively.

Original: In each round (epoch), we randomly sample balls from the box withoutreplacement, then we put them back.
Revised: In each round (epoch), we randomly sample balls from the box without replacement,then we put them back.

Original: We denote the 1st sampled label in the k -th epoch as a k and the2nd sampled one as b k ; a k , b k ∈ { c 0 , c 1 } .
Revised: We denote the 1st sampled label in the k -th epoch as a k and the 2nd sampledone as b k ; a k , b k ∈ { c 0 , c 1 } .

Original: We define two random variables m 0 and m 1 as follows(assume β ∈[0 , 1) ):
Revised: We define two random variables m 0 and m 1 as follows (assumeβ ∈

Original: [Equation]
Revised: [Equation]

Original: where m 0 ,k denotes the summand of m 0 in k -th epoch, similarly for m 1 ,k .
Revised: where m 0 ,k denotes the summand of m 0 in k -th epoch, similarly for m 1 ,k .

Original: Note that in each epoch,m 0 and m 1 share the same sample order.
Revised: Note that in each epoch,m 0 and m 1 share the same sample order.

Original: Further, we introduce the following deterministic constants.
Revised: Further, we introduce the following deterministic constants.

Revised: f 1 =

Revised: [Equation]

Revised: f 1 ,k + β 2 + β(cid:124) (cid:123)(cid:122) (cid:125) f 1 ,k −· · · + β 2( k − 1) + β 2( k − 1)+

Revised: cid:124) (cid:123)(cid:122) (cid:125) f 1 ,);

Original: [Equation]
Revised: [Equation]

Revised: f

Original: [Equation]
Revised: [Equation]

Revised: k

Revised: [Equation]

Revised: f

Revised: [Equation]

Revised: k −· · ·

Revised: [Equation]

Original:  where f 0 ,k denotes the summand of f 0 in k -th epoch, similarly for f 1 ,k .
Revised: [Equation] ( k [Equation] 1)+ (cid:124) (cid:123)(cid:122) (cid:125) f 0 ,); where f 0 ,k denotes the summand of f 0 in k -th epoch, similarly for f 1 ,k .

Original: Now we prove the following
Revised: Now we prove the following

Original: Lemma A.1. Lemma A.1.
Revised: Lemma A.1. Lemma A.1.

Original: In the color-ball model of the 1st kind, we have
Revised: In this color-ball example, we have

Original: [Equation]
Revised: [Equation]

Original: E
Revised: where the expectation is taken on all the possible draws. For the color-ball example with n ≥ 2 balls,we have

Original: [Equation]
Revised: [Equation]

Original: (cid:88) i=
Revised: i

Original: [Equation]
Revised: [Equation]

Original: where the expectation is taken on all the possible draws.
Revised: m

Original: For the color-ball example with n ≥ 2 balls,we have(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) E (cid:34) n − 1 (cid:88) i
Revised: [Equation]

Original: = m i − n − 1 (cid:88)
Revised: i = [Equation] (cid:35)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)

Original: [Equation]
Revised: [Equation]

Original: How is Lemma A.1 related to (7) ?
Revised: How is Lemma A.1 related to (5) ?

Original: In this color-ball toy example, m i mimics the possiblerealization of momentum up to the i -th inner loop in k -th epoch.
Revised: In this color-ball toy example, m i mimics the possiblerealization of momentum up to the i -th inner loop in k -th epoch.

Original: f i mimics the stochastic gradient∇
Revised: f i mimics the stochastic gradient∇

Original: f i ( x k, 0 ) .
Revised: f i ( x k, 0 ) .

Original: This is because we can expand ∇ f i ( x k, 0 ) into an infinite-sum sequence ∇ f i ( x k, 0 ) =(1 − β 1 ) ∇ f i ( x k, 0 ) (cid:80) ∞
Revised: This is because we can expand ∇ f i ( x k, 0 ) into an infinite-sum sequence ∇ f i ( x k, 0 ) =(1 − β 1 ) ∇ f i ( x k, 0 ) (cid:80) ∞

Original: j =0 β
Revised: j =0 β

Original: j 1 , which shares a similar structure as f i .
Revised: j 1 , which shares a similar structure as f i .

Original: In this sense, Lemma A.1 mayprovide ideas to prove (7).
Revised: In this sense, Lemma A.1 mayprovide ideas to prove (5).

Original: Nevertheless, there are still gap between these two, we will explain thegap later.
Revised: Nevertheless, there are still gap between these two, we will explain thegap later.

Original: Proof.
Revised: Proof.

Original: We use E k [ · ] to denote the conditional expectation given all the history up to the beginningof k -th epoch.
Revised: We use E k [ · ] to denote the conditional expectation given all the history up to the beginningof k -th epoch.

Original: Since E [ · ] = E [ E k [ · ]] , we first calculate E k (cid:104)(cid:80) 1 i =

Revised: Since E [ · ] = E [ E k [ · ]] , we first calculate E k (cid:104)(cid:80) 1 i =

Original: m i (cid:105) .
Revised: m i (cid:105) .

Original: Since all the history beforek -th epoch is fixed, we relegate this part to later discussion and first focus on the expectation of 529 (cid:80)i =

Revised: Since all the history beforek -th epoch is fixed, we relegate this part to later discussion and first focus on the expectation of 536 (cid:80)i =

Original: m i,k .
Revised: m i,k .

Original: As shown in Figure 4 (upper part), there are 2 possible realization.
Revised: As shown in Figure 5 (upper part), there are 2 possible realization.

Original: With the help of Figure 4 (upper part), we have the following result.E k [Equation] (cid:88)
Revised: With the help of Figure 5 (upper part), we have the following result.

Original: i

Original: [Equation]
Revised: [Equation]

Original: m

Original: [Equation]

Original: −

Original: [Equation]

Original: (cid:88) i = f

Original: [Equation] (cid:35) [Equation] Now we move one step further to calculate E k − 1 E k (cid:104)(cid:80) 1 i =

Revised: Now we move one step further to calculate E k − 1 E k (cid:104)(cid:80) 1 i =

Original: m i − (cid:80) 1 i =

Revised: m i − (cid:80) 1 i =

Original: f i (cid:105) .
Revised: f i (cid:105) .

Original: Using the similarstrategy as (8), we have the following result.
Revised: Using the similarstrategy as (6), we have the following result.

Original: The calculation is illustrated in Figure 4 (lower part)and Figure 5. [Figure]
Revised: The calculation is illustrated in Figure 5 (lower part)and Figure 6.

Original: [Equation]
Revised: [Equation]

Original: [Figure] [Equation] [Figure] [Equation]
Revised: i =

Original: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure] [Equation] [Figure]
Revised: i

Original: [Equation]
Revised: [Equation]

Original: [Figure] [Equation] [Figure] [Equation]
Revised: i = f i

Original: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure] (cid:88) 
Revised: k + 1 (cid:88) i

Original: [Figure] (cid:88) 
Revised: = m i,k − 1 − 1 (cid:88) i

Original: [Figure]
Revised: = f

Original: (cid:35)(cid:41)
Revised: i,k −(cid:35)(cid:41)

Original: [Equation]
Revised: [Equation]

Original: [Figure] [Figure]
Revised: (cid:40)

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: i = m i,k −(cid:41) − 1 (cid:88) i = f i,k −

Original: [Equation]
Revised: [Equation]

Original: We observe that only the highest order term remains in the calculation.
Revised: We observe that only the highest order term remains in the calculation.

Original: Repeat this process untilk = 1 , we will get the results in Lemma A.1.
Revised: Repeat this process untilk = 1 , we will get the results in Lemma A.1.

Original: The above analysis also holds for general n ≥ 2 .
Revised: The above analysis also holds for general n ≥ 2 .

Revised: Formore evidence, please refer to Appendix G.1 where we further provide a derivation for n = 3 .

Original: The gap between Lemma A.1 and Equation (7) .
Revised: The gap between Lemma A.1 and Equation (5) .

Original: Lemma A.1 shows the key idea of provingequation (7).
Revised: Lemma A.1 shows the key idea of provingequation (5).

Original: However, due to its idealized setting, the color-ball toy example is still far from our realgoal (
Revised: However, due to its idealized setting, the color-ball toy example is still far from our realgoal (5).

Original: We list some of the gaps here.
Revised: We list some of the gaps here.

Original: • In each possible trajectory: x k,i is changing with k and i , while the balls are fixed in the color-ballexample.• When taking expectation: x is changing in different trajectory, while the balls is fixed in thecolor-ball example.• x is a vector in R d while the label of the balls are constant in R .

Original: It requires extra technical lemmas to handle these differences.

Original: We provide more discussions in

Original: Appendix G

Original: Solution to Issue II.

Original: We now discuss how to resolve Issue II mentioned in Section 5.

Original: The solutioncontains two parts.

Original: First, we need to prove (5). Second, we derive (6) from (5). Due to the limitedspace, we relegate the first part to Appendix G.1 (related to Lemma G.2). Now, we discuss the secondpart: Assume we have (5), how do we use it to prove (6)?
Revised: (i)

Original: To answer this question, we introduce thecolor-ball model of the 2nd kind.

Original: The color-ball model of the 2nd kind. Consider the same setting as the color-ball model in Step 1 .

Original: We define a sequence of real random variable { r j } kj =1 with the following relation.

Original: [Equation]

Original: Further, we assume r j is fixed when fixing the history up to j -th round. The sequence { r j } kj =1mimics the sequence { ∂ l f ( x k, 0 ) √ v l,k, 0 } ∞ k =1 in (5). Now [Equation]
Revised: In each possible trajectory: x k,i is changing with k and i , while the balls are fixed in the

Original: we consider the following quantities.

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: k −

Original: [Equation]

Original: [Equation]

Original: ( k −

Original: [Equation]

Original: + β 2( k − 1)+ (cid:124) (cid:123)(cid:122) (cid:125) f 1 ,

Original: [Equation]

Original:   

Original: [Equation] [Equation] (cid:125) [Equation]
Revised: color-ball example.

Original: [Equation]

Original: ( k −

Original: [Equation]

Original: ( k − 1)

Original: [Equation]

Original: (cid:124) (cid:123)(cid:122) (cid:125) f

Original: [Equation]

Original: )   ; . We now prove the following Lemma A.2. Lemma A.2. Consider the color-ball model of the 2nd kind, we have
Revised: (ii) When taking expectation: x is changing in different trajectory, while the balls is fixed in thecolor-ball example.

Original: [Equation] For general n = 1 , 2 , 3 , · · · , we have [Equation] Proof.
Revised: (iii) x is a vector in R d while the label of the balls are constant in R .

Original: We only describe the proof idea here. The proof contains the following 4 steps. Step 2.1.
Revised: It requires extra technical lemmas to handle these differences.

Original: We firstly take E k [ · ] and thus r k can be viewed as a constant.

Original: We use the color-ballprocedure as in Figure 6 to calculate E k (cid:104)

Original: [Figure]

Original: [-Step 2.2. We change E k (cid:104) r k (cid:80) 1 i =
 m i,k − 1 (cid:105) into E k (cid:104) r k − 1 (cid:80) 1 i =
 m i,k − 1 (cid:105) + Error; where Error =-]
Revised: [-Step 2.2. We change E k (cid:104) r k (cid:80) 1 i =
 m i,k − 1 (cid:105) into E k (cid:104) r k − 1 (cid:80) 1 i =
 m i,k − 1 (cid:105) + Error; where Error =-]Briefly speaking, the error caused by(i) and (ii) can be neutralized by diminishing stepsize.

Original: O (1 / √ k ) .[-Step 2.3. We calculate E k − 1 E k (cid:104) r k − 1 (cid:80) 1 i =
 m i,k − 1 (cid:105) =-]
Revised: Issue (iii) can be handled by component-wiseanalysis of x .[-Step 2.3. We calculate E k − 1 E k (cid:104) r k − 1 (cid:80) 1 i =
 m i,k − 1 (cid:105) =-]

Original: [-E k − 1 (cid:104) r k − 1 (cid:80) 1 i =
-]
Revised: [-E k − 1 (cid:104) r k − 1 (cid:80) 1 i =
-]

Original: m i,k − 1 (cid:105) .

Original: [-This part isillustrated in the top row in Figure 
-]
Revised: [-This part isillustrated in the top row in Figure 
-]

Original: Step 2.4.

Original: For the leftovers in Step 2.1 , we change all r k into r k − 1 .

Original: Then we do the cancellation tocalculate E k − 1 E k (cid:104)

Original: [-r k − 1 (cid:80) 1 i =
-]
Revised: [-r k − 1 (cid:80) 1 i =
-]

Original: [-m i,k − 1 + r k − 1 (cid:80) 1 i =
-]
Revised: [-m i,k − 1 + r k − 1 (cid:80) 1 i =
-]

Original: [-m i,k − r k − 1 (cid:80) 1 i =
-]
Revised: [-m i,k − r k − 1 (cid:80) 1 i =
-]

Original: f i (cid:105) .

Original: This step is shown inthe second and third row in Figure 7.

Original: Repeat this process to the 1st epoch.

Original: We can prove the Lemma A.2.

Original: We emphasize that here are still some gap between Lemma A.2 and our goal in (6) in Lemma 5.2.

Original: First, we have the similar gap as discussed at the end of the Solution to Issue I .

Original: Second, the conditionin Lemma 5.2 has requirement on the gradient norm, while this requirement is temporarily ignored inthe color-ball method of the 2nd kind.

Original: It requires some technical lemmas to handle these gaps.

Original: For more details, please refer to the complete proof in Appendix G.
Revised: For more details, please refer to the complete proof in Appendix G.

Original: B Experiments

Original: To support our theory, we provide more simulations and real-data experiments.

Original: All the experimentalsettings and hyperparameters are presented in Appendix B.2.

Original: We aim to show:(I

Original: When β 2 is large, a large range of β 1 gives good performance, including all β 1 < √ β 2 .

Original: II)

Original: When β 2 is small, a large range of β 1 performs relatively badly.

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: g r a d i e n t n o r m

Original: [Figure]

Original: Counter-example (2) .

Original: We run more experiments on (2) with initialization x = − 5 and n =5 , 10 , 15 , 20 .

Original: We show the result of n = 20 in Figure 8 (a), (b); the rest are shown in Appendix B.1.

Original: We find that: first, when β 2 is small, gradient norms increase dramatically.

Original: This emphasizes the importance of discarding bounded gradient assumptions.

Original: Second, when β 2 is large, Adam convergesto critical points for β 1 < √ β 2 .

Original: This matches our theory since (2) satisfies D 0 = 0 .

Original: For functionswith D 0 > 0 , we further show that Adam can only converge to a bounded region in Appendix B.1.

Original: MNIST and CIFAR-10.

Original: As shown in Figure 1 (b)& (c) in Section 1, the training results matchboth claim (I) and (II).

Original: In addition, there is a convex-shaped boundary on the transition from low lossto higher loss, this boundary roughly matches the condition in Theorem 3.1.

Original: NLP.

Original: We use Adam to train Transformer XL (Dai et al., 2019) on the WikiText-103 dataset (Merityet al., 2016).

Original: This architecture and dataset is widely used in NLP tasks (e.g. (Howard & Ruder, 2018;

Original: See et al., 2017)).

Original: As shown in Figure 8 (c), the training results match both claim (I) and (II).

Original: Batchsize and β 2 .

Original: As shown in Figure 8 (d): on MNIST, smaller batchsize requires larger β 2 toreach small loss.

Original: Since batchsize equals to (number of total sample)/(number of batches).

Original: In thecontext of finite-sum setting with n summand, n usually stands for the number of batches (e.g., Inthe extreme case when batchsize = 1 , n equals to the number of total samples).

Original: Therefore, smallerbatchsize brings larger n .

Original: As such, Figure 8 (d) matches the message by Theorem 3.1: the thresholdof β 2 increases with n .

Original: B.1 More experiments

Original: Estimation on ρ in Theorem 3.1.

Original: To ensure convergence, Theorem 3.1 requires β 2 ≥ γ 1 ( n ) =1 − O ((1 − β n 1 ) /n 2 ρ ) .

Original: Now we estimate the constant ρ .

Original: According to our definition in Remark G.7in Appendix G.1, ρ = ρ 1 ρ 2 ρ 3 , where ρ 1 , ρ 2 , ρ 3 are defined as follows.

Original: [Equation]

Original: These definitions can also be seen in Appendix F.1.

Original: In worst case, we have 0 ≤ ρ 3 ≤ √ nρ 1 ≤ n .

Original: However, ρ is highly dependent on the problem instance f ( x ) and training process.

Original: We now estimatehow ρ changes with Adam’s trajectory on MNIST and CIFAR-10.

Original: We use β 1 = 0 .

Original: 9 , β 2 = 0 .

Original: On both datasets, we set batchsize to be 64, which brings n = 937 on MNIST and n = 781 on

Original: CIFAR-10.

Original: We collect ρ 1 , ρ 2 , ρ 3 along the training process and estimate their distribution density.

Original: The results are shown in Figure 9 and 10.

Original: [Figure]

Original: [-On both CIFAR-10 and MNIST, we observe that the maximal ρ 1 < 25 ≈ O ( √ n ) , ρ 2 < 
-]
Revised: [-On both CIFAR-10 and MNIST, we observe that the maximal ρ 1 < 25 ≈ O ( √ n ) , ρ 2 < 
-]

Original: ≈ O ( n ) ,ρ 3 < 0 .

Original: 1 ≈ O (1 / √ n ) .

Original: Therefore, ρ = ρ 1 ρ 2 ρ 3 ≈ O ( n ) .

Original: [Figure]

Original: More experiments on function (2) .

Original: Figure 1 (a) shows the optimality gap after 50k iterationswhen n = 10 and initialization x = 1 .

Original: Here, we provide more relevant experiments.

Original: First, wheninitialized at x = 1 , we run experiments with n = 5 , 10 , 15 , 20 .

Original: The results are shown in Figure 11.

Original: We observe that the blue region shrinks as n increases.

Original: This matches our conclusion in Theorem 3.1:when n increases, the convergent threshold of β 2 increases, which means we need larger β 2 to ensureconvergence.

Original: It also matches the conclusion in Theorem 3.3: when n increases, the divergence regionwill expand (more evidence can be seen in Appendix E).

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: When initializing at x = − 5 , we further demonstrate that the gradient norm of f ( x ) can dramaticallyincrease.

Original: All the setting is the same as that in Figure 11 except for the change of initialization.

Original: Starting at x = − 5 , the algorithm can touch the “quadratic" side of function (2) where the gradient isunbounded.

Original: The results are shown in Figure 12.

Original: We observe a similar pattern as that in Figure 11.

Original: Asa result, the gradient norm of f ( x ) is large in the left bottom corner.

Original: We further plot the change of gradient norm along the iterations.

Original: We pick β 1 = 0 .

Original: 1 and β 2 =0 , 0 . 1 , 0 . 3 , 0 . 5 , 0 .

Original: 99 to see the phase transition when increasing β 2 .

Original: The result is shown in

Original: Figure 13.

Original: When β 2 is small, the gradient norm of f ( x ) increases rapidly along the iteration.

Original: Mostof them are even much larger than the upper bound of color bar in Figure 12.

Original: As a result, there is aphase transition from diverge to converge when increasing β 2 from 0 to 1.

Original: [Figure]

Original: Adam converges to a exact critical point under SGC where D 0 = 0 .

Original: Since function (2) satisfies

Original: SGC, we further provide empirical evidence that Adam converges to a exact critical point under SGC.

Original: We use function (2) n = 20 and initialization x = − 5 .

Original: We choose some large enough β 2 to ensure

Original: [Figure]

Original: As shown in Figure

Original: All thehyperparameter setting is the same as before.

Original: [Figure]

Original: Adam converges to a bounded region when D 0 > 0 .

Original: Now we show that Adam may not convergeto an exact critical point when D 0 > 0 .

Original: Instead, it converges to a bounded region near the criticalpoint.

Original: For this part, we re-state the example from (Shi et al., 2020).

Original: Consider the following function.

Original: [Equation]

Original: Summing up f j ( x ) we get

Original: [Equation]

Original: This f ( x ) is lower bounded by − a 2 9 with optimal solution x ∗ = 0 .

Original: [-At the optimal point x ∗ = 0 ,∇ f j ( x ∗ ) ̸ = 
-]
Revised: [-At the optimal point x ∗ = 0 ,∇ f j ( x ∗ ) ̸ = 
-]

Original: so we have D 0 > 0 .

Original: When running Algorithm 1 on this function, we observe that Adamwith diminishing stepsize does not converge to 0 gradient norm.

Original: Instead, it converge to a boundedregion.

Original: Further, the size of the region shrinks when β 2 increases.

Original: These phenomena matches ourdiscussion in Remark 3.

Original: The result is shown in Figure 15.

Original: In the experiment, we choose β 1 = 0 .

Original: [-9 ,a = 3 , x 
-]
Revised: [-9 ,a = 3 , x 
-]

Original: = − 2 and diminishing stepsize η k = 0 .

Original: 1 √ k .

Original: B.2 Experimental Settings

Original: Here, we introduce our experimental settings.

Original: • Experiments on function (2) .

Original: We use Algorithm 1 with cyclic order f 0 , f 1 , f 2 and so on.

Original: We

Original: report the optimality gap x − x ∗ after 50k iteration, or equivalently 50000 /n epochs.

Original: We use

Original: ϵ = 10 − 8 for numerical stability.

Original: We use diminishing stepsize η k = 0 .

Original: 1 / √ k , where k is the indexof epoch.

Original: Unless otherwise stated, this setting applies to all the other experiments on function (2).

Original: In Figure 1 (a), we use n = 10 and initialization x = 1 .

Original: We will report more results with differentn and different initialization in Appendix B.1.

Original: Instead, it only converge to a bounded region near the critical

Original: • MNIST (Deng, 2012).

Original: We use one-hidden-layer neural network with width =16.

Original: We set batchsize

Original: =1, weight decay =0, stepsize =0.0001 and train for 20 epochs.

Original: We use ϵ = 10 − 8 for numericalstability.• CIFAR-10 (Krizhevsky et al., 2009).

Original: We use ResNet-18 (He et al., 2016) as the architecture.

Original: Wechoose batchsize =16, weight decay =5e-4 , initial stepsize=1e-3.

Original: We use a stage-wise constantlearning rate scheduling with a multiplicative factor of 0.1 on epoch 30, 60 and 90.

Original: We useϵ =

Original: 10 − 8 for numerical stability.

Original: For MNIST and CIFAR-10, larger batchsize will bring similar pattern as that in Figure 1, but thephase transition will occur at some smaller β 2 .

Original: • NLP.

Original: The WikiText-103 dataset is a collection of over 100 million tokens extracted from the setof verified ‘Good’ and ‘Featured’ articles on Wikipedia.

Original: The base model of Transformer XLcontains 16 self-attention layers.

Original: In each self-attention layer, there are 10 heads.

Original: The encodingdimension of each head is set to be 41.

Original: We set batchsize = 60, number of iteration =

Original: 200k, andinitial stepsize = 0.00025.

Original: We use cosine learning rate scheduler, which is a popular choice fortraining Transformer.

Original: We use ϵ = 10 − 8 for numerical stability.

Original: C Some Potential Implications for Practical Users
Revised: B Some Potential Implications for Practical Users

Original: Theorem 3.1 and Proposition 3.3 establish a clearer image on the relation between ( β 1 , β 2 ) and qualitative behavior of Adam, which may have certain implication for practical users.
Revised: Theorem 3.1 and Proposition 3.3 establish a clearer image on the relation between ( β 1 , β 2 ) and qualitative behavior of Adam, which may have certain implication for practical users.

Original: Many practitionersare still doing grid or random search over ( β 1 , β 2 ) , which could be costly.
Revised: Many practitioners are still doing grid or random search over ( β 1 , β 2 ) , which could be costly.

Original: The following advice maycut down a large portion of search space.
Revised: The following advice maycut down a large portion of search space.

Original: Suppose we start with some random point ( β ∗ 1 , β ∗ 2 ) , weprovide the following suggestions for tuning β 1 and β 2 .
Revised: Suppose we start with some random point ( β ∗ 1 , β ∗ 2 ) , weprovide the following suggestions for tuning β 1 and β 2 .

Original: Case 1: If ( β ∗ 1 , β ∗ 2 ) lies above the blue curve in Figure 2.
Revised: Case 1: If ( β ∗ 1 , β ∗ 2 ) lies above the blue curve in Figure 2.

Original: We point out two sub-cases.
Revised: We point out two sub-cases.

Original: First, if
Revised: First, if

Original: Adam with ( β ∗ 1 , β ∗ 2 ) diverges, then any points with “ β 1 < √ β 2 " and “ β 2 < β ∗ 2 " shall not be trieddue to the risk of divergence.
Revised: Adam with ( β ∗ 1 , β ∗ 2 ) diverges, then any points with “ β 1 < √ β 2 " and “ β 2 < β ∗ 2 " shall not be trieddue to the risk of divergence.

Original: Second, if Adam with ( β ∗ 1 , β ∗ 2 ) converges, then any points above thispoint will converge.
Revised: Second, if Adam with ( β ∗ 1 , β ∗ 2 ) converges, then any points above thispoint will converge.

Original: You may either fix β 1 = β ∗ 1 and increase β 2 , or fix β 2 = β ∗ 2 and try smaller β 1 .
Revised: You may either fix β 1 = β ∗ 1 and increase β 2 , or fix β 2 = β ∗ 2 and try smaller β 1 .

Original: Both ways have convergence guarantee.
Revised: Both ways have convergence guarantee.

Original: Case 2: If ( β ∗ 1 , β ∗ 2 ) lies below the blue curve in Figure 2.
Revised: Case 2: If ( β ∗ 1 , β ∗ 2 ) lies below the blue curve in Figure 2.

Original: We still discuss two sub-cases.
Revised: We still discuss two sub-cases.

Original: First, ifyou observe divergence at ( β ∗ 1 , β ∗ 2 ) , then do not further try any point on the left.
Revised: First, ifyou observe divergence at ( β ∗ 1 , β ∗ 2 ) , then do not further try any point on the left.

Original: We do not suggestexploring the points in below either, since the majority of them will still face the risk of divergence.
Revised: We do not suggestexploring the points in below either, since the majority of them will still face the risk of divergence.

Original: Instead, we suggest fix β 1 = β ∗ 1 and increase β 2 .
Revised: Instead, we suggest fix β 1 = β ∗ 1 and increase β 2 .

Original: Since our Theorem 3.1 applies for any β 1 < √ β 2 ,algorithm will converge when β 2 is large enough.
Revised: Since our Theorem 3.1 applies for any β 1 < √ β 2 ,algorithm will converge when β 2 is large enough.

Original: Second, if you observe convergence at ( β ∗ 1 , β ∗ 2 ) (which is also possible according to Figure 1 (a)),then we suggest: (i) either fix β 1 = β ∗ 1 and increase β 2 , or (ii) fix β 2 = β ∗ 2 and try smaller β 1 .
Revised: Second, if you observe convergence at ( β ∗ 1 , β ∗ 2 ) (which is also possible according to Figure 1 (a)),then we suggest: (i) either fix β 1 = β ∗ 1 and increase β 2 , or (ii) fix β 2 = β ∗ 2 and try smaller β 1 .

Original: Bothways push ( β 1 , β 2 ) into the region of Theorem 3.1 with convergence guarantee.
Revised: Bothways push ( β 1 , β 2 ) into the region of Theorem 3.1 with convergence guarantee.

Original: D More Discussion on Related Works.D.1 More Related Works on the Convergence Analysis of Adam Family
Revised: C More Discussion on Related Works.

Revised: C.1 More Related Works on the Convergence Analysis of Adam Family

Original: New variants of Adam.
Revised: New variants of Adam.

Original: Ever since Reddi et al.
Revised: Ever since Reddi et al.

Original: (2018) pointed out the non-convergence issue of Adam, one active line of work has tried to design new variants of Adam that can be proved to converge.
Revised: (2018) pointed out the non-convergence issueof Adam, one active line of work has tried to design new variants of Adam that can be proved toconverge.

Original: For instance, Zou et al.
Revised: For instance, Zou et al.

Original: (2019); Gadat & Gavra (2020)
Revised: (2019); Gadat & Gavra (2020)

Original: ; Chen et al.
Revised: ; Chen et al.

Original: (2018b, 2021) replace the constanthyperparameters by iterate-dependent ones e.g. β 1 t or β 2 t .
Revised: replacethe constant hyperparameters by iterate-dependent ones e.g. β 1 t or β 2 t .

Original: AMSGrad (Reddi et al., 2018) and AdaFom (Chen et al., 2018b) modify { v t } to be an non-decreasing sequence.
Revised: AMSGrad (Reddi et al.,2018) and AdaFom (Chen et al., 2018) modify { v t } to be an non-decreasing sequence.

Original: Similarly, AdaBound(Luo et al., 2019) impose lower and upper bounds on { v t } to prevent the effective stepsize fromvanishing or exploding.
Revised: Similarly, AdaBound (Luo et al., 2019) impose lower and upper bounds on { v t } to prevent the effective stepsizefrom vanishing or exploding.

Original: Zhou et al.
Revised: Zhou et al.

Original: (2018b) also adopt a new estimate of v t to correct the bias.
Revised: (2018) also adopt a new estimate of v t to correct thebias.

Original: Thereare also attempts to combine Adam with Nesterov momentum (Dozat, 2016)as well as warm-uptechniques (Liu et al., 2020a). Padam (Chen et al., 2018a) also introduce a partial adaptive parameterto improve the generalization performance. There are also some works providing theoretical analysison the variants of Adam. For instance, Zhou et al.
Revised: There are also attempts to combine Adam with Nesterov momentum (Dozat, 2016)

Original: (2018a) study the convergence of AdaGrad and AMSGrad under bounded gradient condition. Gadat & Gavra (2020) study the asymptotic behaviorof a subclass of adaptive gradient methods from landscape point of view. Their analysis appliesto Adam-variants with β 1 = 0 and β 2 increasing along the iterates (it could also be understood as 
Revised: as well aswarm-up techniques (Liu et al., 2020a).

Original: RMSProp with increasing β 2 ). When this script is under review, a new work (Iiduka, 2022) appearon line. Iiduka (2022) analyze the convergence of AMSGrad by relaxing the Lipschitz-gradientcondition. However, their analysis requires extra conditions on both bounded gradient and boundeddomain.
Revised: As reported in the literature, these variants often reach goodperformance in many applications.

Original: Some more discussions on (Guo et al., 2021) and (Huang et al., 2021).
Revised: Some more discussions on (Guo et al., 2021) and (Huang et al., 2021).

Original: Here, we discuss more ontwo recent works Guo et al.
Revised: Here, we discuss more ontwo recent works Guo et al.

Original: and Huang et al.
Revised: and Huang et al.

Original: As mentioned in Section 2.2, they bothrequire some extra conditions.
Revised: As mentioned in Section 2.2, they bothrequires some extra conditions.

Original: First, both Guo et al.
Revised: First, both Guo et al.

Original: and Huang et al.
Revised: and Huang et al.

Original: (2021) requires boundedgradient assumption.
Revised: (2021) requires boundedgradient assumption.

Original: This can be seen in Assumption 2 in (Guo et al., 2021).
Revised: This can be seen in Assumption 2 in (Guo et al., 2021).

Original: In (Huang et al.
Revised: In (Huang et al.

Original: , 2021),they require bounded iterates ( their Theorem 1) or change Adam into AdaBound (Luo et al.
Revised: , 2021),they require bounded iterates ( their Theorem 1) or change Adam into AdaBound (Luo et al.

Original: , 2019)by clipping (their Remark 2, Corollary 1).
Revised: , 2019)by clipping (their Remark 2, Corollary 1).

Original: Both settings inherent boundedness on gradient.
Revised: Both settings inherent boundedness on gradient.

Original: Besides bounded gradient, both (Huang et al., 2021) and (Guo et al., 2021) requires 1 / ( √ v t + ϵ ) ≤ C u .
Revised: Besides bounded gradient, both (Huang et al., 2021) and (Guo et al., 2021) requires 1 / ( √ v t + ϵ ) ≤ C u .

Original: This condition is stated in Assumption 2 in (Guo et al., 2021) and Assumption 3 in (Huang et al.,2021)
Revised: This condition is stated in Assumption 2 in (Guo et al., 2021) and Assumption 3 in (Huang et al.,2021)

Original: (they presented it as H t ⪰ ρI d ≻ 0 , where matrix H t = diag ( √ v t + ϵ ) ).
Revised: (they presented it as H t ⪰ ρI d ≻ 0 , where matrix H t = diag ( √ v t + ϵ ) ).

Original: Combining these twoconditions, the effective stepsize of Adam will be bounded in certain interval 1 √ v t + ϵ ∈
Revised: Combining these twoconditions, the effective stepsize of Adam will be bounded in certain interval 1 √ v t + ϵ ∈

Original: [ C l , C u ] .
Revised: [ C l , C u ] .

Original: Suchboundedness condition changes Adam into AdaBound (Luo et al., 2019) and thus cannot explain theobservations of Adam in Section 1.
Revised: Suchboundedness condition changes Adam into AdaBound (Luo et al., 2019) and thus cannot explain theobservations of Adam in Section 1.

Original: D.2 A Brief Introduction to (Reddi et al., 2018) 
Revised: C.

Revised: An Brief Introduction to (Reddi et al., 2018)

Original: Here, we re-state two counter examples by (Reddi et al., 2018).
Revised: Here, we re-state two counter examples by (Reddi et al., 2018).

Original: For the consistence of notation, wewill re-state their results under our notation in the full script.
Revised: For the consistence of notation, wewill re-state their results under our notation in the full script.

Original: They consider the convex problem((Reddi et al., 2018)): min (cid:80) n − 1 i =

Revised: They consider the convex problem((Reddi et al., 2018)): min (cid:80) n − 1 i =

Original: f i ( x ) where x ∈ [ − 1 , 1] , n ≥ 3 :f i ( x ) = (cid:26) nx, for i =− x, otherwise, (9)
Revised: f i ( x ) where x ∈ [ − 1 , 1] , n ≥ 3 :

Original: Note that (9) satisfy both Assumption 2.1 and 2.2 (with D 1 = n 2 + n − 1 and D 0 = 0 ), so ourassumptions do not rule out this counter-example a priori.
Revised: [Equation]

Original: This is a constrained problem with feasibleset x ∈ [ − 1 , 1] , the optimal solution is x ∗ = − 1 .
Revised: Since it is a constrained problem with feasible set x ∈ [ − 1 , 1] , the optimal solution is x ∗ = − 1 .

Original: Since they consider constrained problems, theirclaimed “divergence" actually means the iterates will stay in a huge region with the size of wholefeasible set. Here, we call it “non-convergence" to distinguish from our result of “diverge to infinity"in Proposition 3.3.
Revised: In(Reddi et al., 2018), they use notation C rather than n .

Original: They consider sampling f i in the cyclic order: f 0 , f 1 , f 2 .
Revised: They consider sampling f i in the cyclic order: f 0 , f 1 , f 2 .

Original: In (Reddi et al., 2018), Function (9) ispresented as an “online optimization problem with non-zero average regret".
Revised: In (Reddi et al., 2018), Function (7) ispresented as an online optimization problem.

Original: We choose to use theform of (9) since it is more consistent with our notation in Algorithm 1.
Revised: We choose to use the form of (7) since it is moreconsistent with our notation in Algorithm 1.

Original: We re-state their results asfollows.
Revised: We re-state their results as follows.

Original: Theorem D.1 (Theorem 2 in (Reddi et al., 2018)) .
Revised: Theorem C.1 (Theorem 2 in (Reddi et al., 2018)) .

Original: For any fixed ( β 1 , β 2 ) s.t.
Revised: For any fixed ( β 1 , β 2 ) s.t.

Original: β 1 < √ β 2 , there existsfunction (9) with large enough n , s.t.
Revised: β 1 < √ β 2 , there existsfunction (7) with large enough n , s.t.

Original: Adam will converge to highly sub-optimal solution x = 1 .
Revised: Adam will converge to highly sub-optimal solution x = 1 .

Original: We briefly re-state the non-convergent condition for this Theorem.
Revised: We briefly re-state the non-convergent condition for this Theorem.

Original: As stated in Equation (7), Appendix
Revised: As stated in Equation (7), Appendix

Original: B in (Reddi et al., 2018), for every fixed ( β 1 , β 2 ) , they need a “constant n that depends on β 1 and β 2 ".
Revised: B in (Reddi et al., 2018), for every fixed ( β 1 , β 2 ) , they need a “constant n that depends on β 1 and β 2 ".

Original: As such, they require different n to cause non-convergence on different ( β 1 , β 2 ) .
Revised: As such, they requires different n to cause non-convergence on different ( β 1 , β 2 ) .

Original: So the consideredfunction class is constantly changing.
Revised: So the consideredfunction class is constantly changing.

Original: For completeness, we further re-state Theorem 1 in (Reddi et al., 2018).
Revised: For completeness, we further re-state Theorem 1 in (Reddi et al., 2018).

Original: Theorem D.2 (Theorem 1 in (Reddi et al., 2018)) .
Revised: Theorem C.2 (Theorem 1 in (Reddi et al., 2018)) .

Original: For function (9) , when β 1 = 0 and β 2 = 1 / ( n 2 +1) ,
Revised: For function (7) , when β 1 = 0 and β 2 = 1 / ( n 2 +1) ,

Original: Adam will converge to highly sub-optimal solution x = 1 .
Revised: Adam will converge to highly sub-optimal solution x = 1 .

Original: This theorem considers choosing ( β 1 , β 2 ) after n .
Revised: This theorem considers choosing ( β 1 , β 2 ) after n .

Original: However, this result only show non-convergence ona single point ( β 1 , β 2 ) =
Revised: However, this result only show divergence on asingle point ( β 1 , β 2 ) =

Original: (0 , 1 / ( n 2 + 1)) .
Revised: (0 , 1 / ( n 2 + 1)) .

Original: This point lies somewhere on the left boundary of Figure3.
Revised: This point lies somewhere on the left boundary of Figure 3.

Original: It seems unclear how Adam’s behavior would change as we change the ( β 1 , β 2 ) to anywhere else.
Revised: It seems unclear how Adam’s behavior would change as we change the ( β 1 , β 2 ) to anywhere else.

Original: E Proof of Proposition 3.3
Revised: D Experimental Settings and More Experiments

Revised: D.1 Experimental Settings Here, we introduce our experimental settings that are omitted in the full script. • Experiments on function (2) .

Revised: We use Algorithm 1 with cyclic order f 0 , f 1 , f 2 and so on.

Revised: We report the optimality gap x − x ∗ after 50k iteration, or equivalently 50000 /n epochs.

Revised: Weuse ϵ

Revised: = 10 − 8 for numerical stability.

Revised: We use diminishing stepsize η k = 0 .

Revised: 1 / √ k , where k is

Revised: the index of epoch.

Revised: Unless otherwise stated, this setting applies to all the other experiments

Revised: on function (2).

Revised: In Figure 1 (a), we use n = 10 and initialization x = 1 .

Revised: We will report more results withdifferent n and different initialization in Appendix D.2.

Revised: • MNIST (Deng, 2012).

Revised: We use one-hidden-layer neural network with width =16.

Revised: We setbatchsize =1, weight decay =0, stepsize =0.0001 and train for 20 epochs.

Revised: We use ϵ = 10 − 8for numerical stability.

Revised: • CIFAR-10 (Krizhevsky et al., 2009).

Revised: We use ResNet-18 (He et al., 2016) as the architecture.

Revised: We choose batchsize =16, weight decay =5e-4 , initial stepsize=1e-3.

Revised: We use a stage-wise

Revised: constant learning rate scheduling with a multiplicative factor of 0.1 on epoch 30, 60 and 90.We use ϵ = 10 − 8 for numerical stability.

Revised: For MNIST and CIFAR-10, larger batchsize will bring similar pattern as that in Figure 1,but the phase transition will occur at some smaller β

Revised: • NLP.

Revised: The WikiText-103 dataset is a collection of over 100 million tokens extracted from theset of verified ‘Good’ and ‘Featured’ articles on Wikipedia.

Revised: The base model of Transformer

Revised: XL contains 16 self-attention layers.

Revised: In each self-attention layer, there are 10 heads.

Revised: Theencoding dimension of each head is set to be 41.

Revised: We set batchsize = 60, number of iteration= 200k, and initial stepsize = 0.00025.

Revised: We use cosine learning rate scheduler, which is apopular choice for training Transformer.

Revised: We use ϵ = 10 − 8 for numerical stability.

Revised: D.2 More experiments

Revised: Estimation on ρ in Theorem 3.1.

Revised: To ensure convergence, Theorem 3.1 requires β 2 ≥ γ 1 ( n ) =1 − O ((1 − β n 1 ) /n 2 ρ ) .

Revised: Now we estimate the constant ρ .

Revised: According to our definition in Remark G.8in Appendix G.1, ρ = ρ 1 ρ 2 ρ 3 , where ρ 1 , ρ 2 , ρ 3 as defined as follows.

Revised: [Equation]

Revised: ρ 2 ≥ | max i ∂ l f

Revised: [Equation]

Revised: x k

Revised: [Equation]

Revised: [Equation]

Revised: |1 n (cid:80)

Revised: [Equation]

Revised: f

Revised: [Equation]

Revised: 2 ; ρ 3 ≥

Revised: [Equation]

Revised: =

Revised: [Equation]

Revised: f

Revised: [Equation]

Revised: [Equation]

Revised: (cid:113) 1 n (cid:80) ni =

Revised: [Equation]

Revised: ∂ l f

Revised: [Equation]

Revised: x k

Revised: [Equation]

Revised: 0 ) |.

Revised: These definitions can also be seen in Appendix F.1.

Revised: In worst case, we have 0 ≤ ρ 3 ≤ √ nρ 1 ≤ n .

Revised: However, ρ is highly dependent on the problem instance f ( x ) and training process.

Revised: We now estimatehow ρ changes with Adam’s trajectory on MNIST and CIFAR-10.

Revised: We use β 1 = 0 .

Revised: 9 , β 2 = 0 .

Revised: On both datasets, we set batchsize to be 64, which brings n = 937 on MNIST and n = 781 on

Revised: CIFAR-10.

Revised: We collect ρ 1 , ρ 2 , ρ 3 along the training process and estimate their distribution density.

Revised: The results are shown in Figure 7 and 8.

Revised: [Figure]

Revised: [Figure]

Original: {+On both CIFAR-10 and MNIST, we observe that the maximal ρ 1 < 25 ≈ O ( √ n ) , ρ 2 < 
+}
Revised: {+On both CIFAR-10 and MNIST, we observe that the maximal ρ 1 < 25 ≈ O ( √ n ) , ρ 2 < 
+}

Revised: ≈ O ( n ) ,ρ 3 < 0 .

Revised: 1 ≈ O (1 / √ n ) .

Revised: Therefore, ρ = ρ 1 ρ 2 ρ 3 ≈ O ( n ) .

Revised: More experiments on function (2) .

Revised: Figure 1 (a) shows the optimality gap after 50k iterationswhen n = 10 and initialization x = 1 .

Revised: Here, we provide more relevant experiments.

Revised: First, wheninitialized at x = 1 , we run experiments with n = 5 , 10 , 15 , 20 .

Revised: The results are shown in Figure 9.

Revised: We observe that the blue region shrinks as n increases.

Revised: This matches our conclusion in Theorem 3.1:when n increases, the convergent threshold of β 2 increases, which means we need larger β 2 to ensureconvergence.

Revised: It also matches the conclusion in Theorem 3.3: when n increases, the divergence regionwill expand (more evidence can be seen in Appendix E).

Revised: When initializing at x = − 5 , we further demonstrate that the gradient norm of f ( x ) can dramaticallyincrease.

Revised: All the setting is the same as that in Figure 9 except for the change of initialization.

Revised: Starting at x = − 5 , the algorithm can touch the “quadratic" side of function (2) where the gradient isunbounded.

Revised: The results are shown in Figure 10.

Revised: We observe an similar pattern as that in Figure 9.

Revised: Asa result, the gradient norm of f ( x ) is large in the left bottom corner.

Revised: We further plot the change of gradient norm along the iterations.

Revised: We pick β 1 = 0 .

Revised: 1 and β 2 =0 , 0 . 1 , 0 . 3 , 0 . 5 , 0 .

Revised: 99 to see the phase transition when increasing β 2 .

Revised: The result is shown in

Revised: Figure 11.

Revised: When β 2 is small, the gradient norm of f ( x ) increases rapidly along the iteration.

Revised: Mostof them are even much larger than the upper bound of color bar in Figure 10.

Revised: As a result, there is aphase transition from diverge to converge when increasing β 2 from 0 to 1.

Revised: [Figure]

Revised: [Figure]

Revised: Adam converges to a exact stationary point under SGC where D 0 = 0 .

Revised: Since function (2)satisfies SGC, we further provide empirical evidence that Adam converges to a exact stationary pointunder SGC.

Revised: We use function (2) n = 20 and initialization x = − 5 .

Revised: We choose some large enough β 2to ensure the convergence.

Revised: As shown in Figure 12, We observe 0 gradient norm after Adam converges.

Revised: All the hyperparameter setting is the same as before.

Revised: Adam converges to a bounded region when D 0 > 0 .

Revised: Now we show that Adam may not convergeto an exact stationary point when D 0 > 0 .

Revised: Instead, it converge to a bounded region near the stationarypoint.

Revised: For this part, we re-state the example from (Shi et al., 2020).

Revised: Consider the following function.

Revised: [Equation]

Revised: Summing up f j ( x ) we get

Revised: [Equation]

Revised: This f ( x ) is lower bounded by − a 2 9 with optimal solution x ∗ = 0 .

Revised: In (Shi et al., 2020), they try

Revised: SGD, RMSProp and AMSGrad on this function, the result is shown in Figure 13.

Revised: They observe that

Revised: [Figure]

Revised: (a

Revised: [Equation]

Revised: [Figure]

Revised: [Figure]

Revised: [Equation]

Revised: [Figure]

Revised: n

Revised: [Figure]

Revised: e n t n o r m [0.1,0.99][0.1,0.9][0.1,0.7][0.1,0.5][0.1,0.3] [0.1,0.1][0.1,0] (

Revised: [Figure]

Revised: ) n =

Revised: [Figure]

Revised: RMSProp with diminishing stepsize may not converge to the optimal solution when D 0 > 0 , even ifthey set the initialization at x = x ∗ = 0 .

Revised: See more description in (Shi et al., 2020, Appendix A.4)

Revised: [Figure]

Revised: E Proof of Proposition 3.

Original: We restate our counter-example here.
Revised: We restate our counter-example here.

Original: Consider f ( x ) = (cid:80) n − 1 i =

Revised: Consider f ( x ) = (cid:80) n − 1 i =

Original: f i ( x ) for x ∈ R , we define f i ( x ) asfollows:
Revised: f i ( x ) for x ∈ R , we define f i ( x ) asfollows:

Original: [Equation]
Revised: [Equation]

Original: Summing up all the f i ( x ) , we can see thatf ( x ) =
Revised: Summing up all the f i ( x ) , we can see that

Original: (cid:26) x, x ≥ −12 ( x + 2) 2 − 32 , x < −is a convex smooth function with optimal solution x ∗ = − 2 and optimal value f ( x ∗ ) = − 3 / 2 .
Revised: [Equation] is a convex smooth function with optimal solution x ∗ = − 2 and optimal value f ( x ∗ ) = − 3 / 2 .

Original: However, we are going to show that, for any fixed n > 2 , there exists an orange region shown in
Revised: However, we are going to show that, for any fixed n > 2 , there exists an orange region shown in

Original: Figure 2, s.t., Adam with any( β 1 , β 2) combination in the yellow region diverge to x = ∞ ratherthan the optimal solution x = − 2 , causing the divergence.
Revised: Figure 2, s.t., Adam with any β 1 - β 2 combination in the yellow region converges to x = ∞ ratherthan the optimal solution x = − 2 , causing the divergence.

Original: Now we introduce the formal statement of
Revised: Now we introduce the formal statement of

Original: Proposition 3.3.
Revised: Proposition 3.3.

Original: Proposition E.1.
Revised: Proposition E.1.

Original: (Formal statement of Proposition 3.3.) Consider the convex function (10) for afixed n .
Revised: (Formal statement of Proposition 3.3.) Consider the convex function (8) for afixed n .

Original: Starting at the initialization x = 1 and initial stepsize η 1 , the iterates of Adam diverge toinfinity if the following holds:
Revised: Starting at the initialization x = 1 and initial stepsize η 1 , the iterates of Adam diverge toinfinity if the following holds:

Original: [Equation]
Revised: [Equation]

Original: The analysis is motivated by that in (Reddi et al., 2018, Theorem 1).
Revised: The analysis is motivated by that in (Reddi et al., 2018, Theorem 1).

Original: However, (Reddi et al., 2018,
Revised: However, (Reddi et al., 2018,

Original: Theorem 1) considers a simplified case with β 1 = 0 .
Revised: Theorem 1) considers a simplified case with β 1 = 0 .

Original: Here, we consider non-zero β 1 , especially forthose β 1 > √ β 2 .
Revised: Here, we consider non-zero β 1 , especially forthose β 1 > √ β 2 .

Original: To show the divergence, we aim to prove the following claim: (we denote x k,i asthe value of x at the k -th outer loop and i -th inner loop )
Revised: To show the divergence, we aim to prove the following claim: (we denote x k,i asthe value of x at the k -th outer loop and i -th inner loop )

Original: Claim: for any fixed n > 2 , there exists an orange region shown in Figure 2 s.t., Adam with anyβ 1 - β 2 combination in the orange region 
Revised: Claim: for any fixed n > 2 , there exists an orange region shown in Figure 2 s.t., Adam with anyβ 1 - β 2 combination in the orange region gives x k +1 , 0 > 1 as long as x k, 0 = 1 .

Original: Since the gradient stays constant when x > 1 , so x will go to infinity if the claim holds.
Revised: Since the gradient stays constant when x > 1 , so x will go to infinity if the claim holds.

Original: To prove thisclaim, we only need to analyze the trajectory of Adam within one particular outer loop, e.g., the k -thouter loop.
Revised: To prove thisclaim, we only need to analyze the trajectory of Adam within one particular outer loop, e.g., the k -thouter loop.

Original: We will show that x k +1 , 0 > 1 if this outer loop is initialized with x k, 0 = 1 .
Revised: We will show that x k +1 , 0 > 1 if this outer loop is initialized with x k, 0 = 1 .

Original: Similarly as(Reddi et al., 2018), we assume f i ( x ) are sampled in the order of f 0 ( x ) , f 1 ( x ) , · · · , f n − 1 ( x ) withinthe k -th outer loop.
Revised: Similarly as(Reddi et al., 2018), we assume f i ( x ) are sampled in the order of f 0 ( x ) , f 1 ( x ) , · · · , f n − 1 ( x ) withinthe k -th outer loop.

Original: Now let us prove the claim.
Revised: Now let us prove the claim.

Original: For function (10), the update rule of Adam is shown as follows.
Revised: For function (8), the update rule of Adam is shown as follows.

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: We decompose the total movement (cid:80) n − 1 i =

Revised: We decompose the total movement (cid:80) n − 1 i =

Original: δ k,i into three terms as follows.
Revised: δ k,i into three terms as follows.

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation] [Table] [Equation] We will show that for some β 1 and β 2 : ( I ) , ( II ) > 0 and ( III ) < 0 .
Revised: We will show that for some β 1 and β 2 : ( I ) , ( II ) > 0 and ( III ) < 0 .

Original: However, ( I ) and ( II )outweigh ( III ) , causing the divergence.
Revised: However, ( I ) and ( II )outweigh ( III ) , causing the divergence.

Original: First, we show that m k − 1 ,n − 1 < 0 when β 1 is small.
Revised: First, we show that m k − 1 ,n − 1 < 0 when β 1 is small.

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: when β 1 is small,we have (1 − β n − 1) > (1 − β 1 ) β n − 1n , which implies − m k − 1 ,n − 1 > 0 .
Revised: when β 1 is small,we have (1 − β n − 1) > (1 − β 1 ) β n − 1n , which implies − m k − 1 ,n − 1 > 0 .

Original: For thesechoices of β 1 , we have ( I ) > 0 .
Revised: For thesechoices of β 1 , we have ( I ) > 0 .

Original: Now we derive a lower bound for ( II ) .
Revised: Now we derive a lower bound for ( II ) .

Original: [Equation]
Revised: [Equation]

Original: The inequality is due to the fact that v k, 0 ≤ n 2 .
Revised: The inequality is due to the fact that v k, 0 ≤ n 2 .

Original: Since β j 2 n 2 is small when β 2 is small and j is closeto n , there exists some small β 2 such that β j 2 n 2 ≤ 0 .
Revised: Since β j 2 n 2 is small when β 2 is small and j is closeto n , there exists some small β 2 such that β j 2 n 2 ≤ 0 .

Original: 1 for at least one j < n .
Revised: 1 for at least one j < n .

Original: For these small enough β 2 , we keep the summand with j ≥ log β 2 (0 .
Revised: For these small enoughβ 2 , we keep the summand with j ≥ log β 2 (0 .

Original: 1 /n 2 ) and drop the rest.
Revised: 1 /n 2 ) and drop the rest.

Original: We have the following lowerbound for ( II ) .
Revised: We have the following lowerbound for ( II ) .

Original: [Equation]
Revised: [Equation]

Original: However, this lower bound only holds for the small β 2 .
Revised: However, these lower bound only holds for the small β 2 .

Original: With simple modification, we derive auniversal lower bound of ( II ) for any β 2 ∈ (0 , 1) .
Revised: With simple modification, we derive anuniversal lower bound of ( II ) for any β 2 ∈ (0 , 1) .

Original: [Equation]
Revised: [Equation]

Original: Now we derive a lower bound for ( III ) .
Revised: Now we derive a lower bound for ( III ) .

Original: [Equation]
Revised: [Equation]

Original: The remaining step is to show for small enough step size η 1 , the iterates will stay in the linear region,thus the above gradient holds for all iterates in the trajectory.
Revised: The remaining step is to show for small enough step size η 1 , the iterates will stay in the linear region,thus the above gradient holds for all iterates in the trajectory.

Original: As we initial x as x 0 = 1 , if for all m < n and k > 0 , (cid:80) mi =0{+=
+} δ k,i ≥ − 2 , we can conclude all iteratesinthe trajectory stayin the linear region.
Revised: As we initial x as x 0 = 1 , if for all m < n and k > 0 , (cid:80) mi {+=
+} δ i,k ≥ − 2 , we can conclude all iteratesin trajectory stay

Original: Because it holds that m k − 1 ,n − 1 ≤ 0 and v k − 1 ,n − 1 ≥ 0 , we havethe following result after droppingmost negative terms in the definition of δ k,i :
Revised: Because it holds that m k − 1 ,n − 1 ≤ 0 and v k − 1 ,n − 1 ≥ 0 , we have

Original: [Equation]
Revised: [Equation]

Original: Therefore, to make (cid:80) m i =

Revised: Therefore, to make (cid:80) mi =

Original: δ k,i ≥

Original: [Equation]

Original: mi =0 − η

Original: [Equation]

Original: 1 − β

Original: [Equation]

Original: β

Original: [Equation]

Original: 1 √

Original: [Equation] 1  − [Equation] ) β n 2 ≥ − 2 , we have
Revised: δ i,k ≥ (cid:80) mi =0 − η 1 (1 − β 1 ) β i 1 √ (1 − β 2 ) β n 2 ≥ − 2 , we have

Original: [Equation]
Revised: [Equation]

Original: To show the divergence, we want to show that there exists some β 1 and β 2 s.t.
Revised: To show the divergence, we want to show that there exists some β 1 and β 2 s.t.

Original: the both of thefollowing conditions hold:
Revised: the both of thefollowing conditions hold:

Original: [Equation]
Revised: [Equation]

Original: The proof of Theorem 3.3 is completed.
Revised: The proof of Theorem 3.3 is completed.

Original: With the help of Python, we visualize the region where( C1 ) and ( C2 ) hold.
Revised: With the help of Python, we visualize the region where( C1 ) and ( C2 ) hold.

Original: The results are shown in Figure 16.
Revised: The results are shown in Figure 14.

Original: We use orange color to indicate the regionwhere ( C1 ) holds.
Revised: We use orange color to indicate the regionwhere ( C1 ) holds.

Original: White color is used for the counter part.
Revised: White color is used for the counter part.

Original: As for ( C2 ) , we use the gray verticalline to indicate the line where (1 − β n − 1) =
Revised: As for ( C2 ) , we use the gray verticalline to indicate the line where (1 − β n − 1) =

Original: (1 − β 1 ) β n − 1n .
Revised: (1 − β 1 ) β n − 1n .

Original: Note that there are two solutions tothis equation: one solution is β 1 = 1 and the other solution lies in 0 < β 1 < 1 , this is why there aretwo vertical lines in the figure.
Revised: Note that there are two solutions tothis equation: one solution is β 1 = 1 and the other solution lies in 0 < β 1 < 1 , this is why there aretwo vertical lines in the figure.

Original: ( C2 ) holds on the left hand side of the left gray vertical line.0.
Revised: ( C2 ) holds on the left hand side of the left gray vertical line.

Original: 0.8 beta 0.

Original: b e t a(a) n =0.

Original: 0.8 beta 0.

Original: b e t a(b) n =0.

Original: 0.8 beta 0.

Original: b e t a(c) n =0.

Original: 0.8 beta 0.

Original: b e t a(d) n =

Original: The intersection of two regions will be the region where Adam diverges, which is actually the orangeregion in Figure 16.
Revised: The intersection of two regions will be the region where Adam diverges, which is actually the orangeregion in Figure 14.

Original: As n increases, ( C2 ) holds for a wide range of β 1 , so the grey vertical linesmove towards β 1 = 1 and finally get overlapped.
Revised: As n increases, ( C2 ) holds for a wide range of β 1 , so the grey vertical linesmove towards β 1 = 1 and finally get overlapped.

Original: In addition, the size of divergence region increaseswith n
Revised: In addition, the size of divergence region increaseswith n

Original: Relation with γ 1 ( n ) in Theorem 3.1 According to Theorem 3.1, γ 1 ( n ) is at least in the order of1 − O ( n − 2 ) .
Revised: Relation with γ 1 ( n ) in Theorem 3.1 According to Theorem 3.1, γ 1 ( n ) is at least in the order of1 − O ( n − 2 ) .

Original: Combining with Figure 16.
Revised: Combining with Figure 14.

Original: It is not hard to see that γ 1 ( n ) is always larger than theupper boundary of the orange region, so there is no contradiction.
Revised: It is not hard to see that γ 1 ( n ) is always larger than theupper boundary of the orange region, so there is no contradiction.

Original: Proof of Corollary 4.1.
Revised: Proof of Corollary 4.1.

Original: For any ( β 1 , β 2 ) ∈ [0 , 1) 2 , condition (C1), (C2) and (C3) can be satisfiedby some sufficiently large n .
Revised: For any ( β 1 , β 2 ) ∈ [0 , 1) 2 , condition (C1), (C2) and (C3) can be satisfiedby some sufficiently large n .

Original: Therefore, Adam will diverge and the proof is concluded.
Revised: Therefore, Adam will diverge and the proof is concluded.

Original: F Some More Notations and Useful Lemmas for Convergence Analysis
Revised: F Some More Notations and Useful Lemmas for Convergence AnalysisF.1 More notations

Original: F.1 More notations

Original: • We denote x k,i , m k,i , v k,i ∈
Revised: • We denote x k,i , m k,i , v k,i ∈

Original: R
Revised: R

Original: d as the value of x, m, v at the k -th outer loop and i -th inner loop.
Revised: d as the value of x, m, v at the k -th outer loop and i -th inner loop.

Original: Further, we denote x l,k,i , m l,k,i , v l,k,i ∈ R as the l -th component of x k,i , m k,i , v k,i .•
Revised: Further, we denote x l,k,i , m l,k,i , v l,k,i ∈ R as the l -th component of x k,i , m k,i , v k,i .

Original:  We denote η k as the stepsize at the k -th epoch (outer loop).
Revised: • We denote η k as the stepsize at the k -th epoch (outer loop).

Original: We will focus mainly on diminishingstep size, especially η k = η 1 √ nk , where n is the number of batches (inner loop).
Revised: We will focus mainly on diminishing step size, especially η k = η 1 √ nk , where n is the number of batches (inner loop).

Original: • We denote ∂ l f ( x ) =
Revised: • We denote ∂ l f ( x ) =

Original: ∂∂x l f ( x ) , ∂ l f j ( x ) =
Revised: ∂∂x l f ( x ) , ∂ l f j ( x ) =

Original: ∂∂x l f j ( x ) .
Revised: ∂∂x l f j ( x ) .

Original: Further, we will use τ k,i to index  the i -thrandomly chosen batch in the k -th epoch.
Revised: Further, we will use τ k,i to index [Table] the i -th randomly chosen batch in the k -th epoch.

Original: In this sense, we denote ∂ l f τ k,i ( x ) as  ∂∂x l f τ k,i ( x ) .
Revised: In this sense, we denote ∂ l f τ k,i ( x ) as [Table] ∂∂x l f τ k,i ( x ) .

Original: [Table]
Revised: [Table]

Original: • Given an epoch k , we denote α as the index of the coordinate with the greatest gradient:
Revised: • Given an epoch k , we denote α as the index of the coordinate with the greatest gradient:

Original: [Equation]
Revised: [Equation]

Revised: 0.8 beta 0.

Revised: b e t a(a) n =0.

Revised: 0.8 beta 0.

Revised: b e t a(b) n =0.

Revised: 0.8 beta 0.

Revised: b e t a(

Revised: [Figure]

Original:  • We define ρ 1 , ρ 2 , ρ 3 as constants satisfying the following condition for any l = 1 , · · · , d :
Revised: 0.8 beta 0. b e t a [Figure] • We define ρ 1 , ρ 2 , ρ 3 as constants satisfying the following condition for any l = 1 , · · · , d :

Original: [Table]
Revised: [Equation]

Original: ρ 1 , ρ 2 , ρ 3 are problem-dependent constants.
Revised: ρ 1 , ρ 2 , ρ 3 are problem-dependent constants.

Original: In worst case, we have 0 ≤ ρ 3 ≤ √ nρ 1 ≤ n .•
Revised: In worst case, we have 0 ≤ ρ 3 ≤ √ nρ 1 ≤ n .

Original:  We define the constant △ x := η 0 √ x L √ d √ 1 − β 2 1 − β 1 1 − β 1 √ β .
Revised: ρ 3 is larger when ∂ l f i ( x k, 0 ) are more aligned. • We define the constant △ x : [Equation]

Original: This constant will be used repeatedly, especiallyin Lemma F.2.•
Revised: This constant will be used repeatedly, especially in Lemma F.3.

Original:  We define Q k := △ 1 n √ n √ k 32 √
Revised: • We define Q k :

Original: [Equation]
Revised: [Equation]

Original: 1 − β 2 ) n β n 2 .

Original: This constant will be used repeatedly.•
Revised: This constant will be used repeatedly.

Original:  We define
Revised: • We define

Original: [Equation]
Revised: [Equation]

Original: This constant will be used repeatedly.
Revised: This constant will be used repeatedly.

Original: Note that δ 1 → 0 when β 2 → 1 .• E ( · ) means taking expectation over the whole trajectory.• We will use E k [ · ] as a shorthand for E [ · (cid:12)(cid:12) x k, 0 , x k − 1 ,n − 1 , · · · , x 1 , 0 ] , i.e. the conditional expectation given all the history up to x k, 0 .
Revised: Note that δ 1 → 0 when β 2 → 1 . • E ( · ) means taking expectation over the whole trajectory. • We will use E k [ · ] as a shorthand for E [ · (cid:12)(cid:12) x k, 0 , x k − 1 ,n − 1 , · · · , x 1 , 0 ] , i.e. the conditionalexpectation given all the history up to x k, 0 .

Original: Similarly, E k − 1 stands for the conditional expectation givenall the history up to x k − 1 , 0 .
Revised: Similarly, E k − 1 stands for the conditionalexpectation given all the history up to x k − 1 , 0 .

Original: We further introduce some ‘index’ notation of the stochastic gradients, these notations will be usefulin the proof.
Revised: We further introduce some ‘index’ notation of the stochastic gradients, these notations will be usefulin the proof.

Original: For all the stochastic gradient, we have
Revised: For all the stochastic gradient, we have

Original: • True index: every stochastic gradient has its own true index: f 1 ( · ) , f 2 ( · ) , · · · , f n − 1 ( · ) .
Revised: • True index: every stochastic gradient has its own true index: f 1 ( · ) , f 2 ( · ) , · · · , f n − 1 ( · ) .

Original: All these{ f i ( · ) } n − 1 i =0 are fixed once the optimization problem is formulated.
Revised: Allthese { f i ( · ) } n − 1 i =0 are fixed once the optimization problem is formulated.

Original: Note that

Original:  for a fixed x , we have the following relationship between f i and f :
Revised: Note that for a fixed x , we have the following relationship between f i and f :

Original: [Equation]
Revised: [Equation]

Original: • Random-shuffle index: At the k -th epoch, all the stochastic gradients are sampled in the order off τ k, 0 ( · ) , f τ k, 1 ( · ) , · · · , f τ k,n − 1 ( · ) .
Revised: • Random-shuffle index: At the k -th epoch, all the stochastic gradients are sampled in theorder of f τ k, 0 ( · ) , f τ k, 1 ( · ) , · · · , f τ k,n − 1 ( · ) .

Original: Since Algorithm 1 is sampling without replacement, there is an implicit bijective mapping between{ f i ( · ) } n − 1 i [-=
-] and { f τ k,i ( · ) } n − 1 i =0 .
Revised: Since Algorithm 1 is sampling without replacement, there is an implicit bijective mappingbetween { f i ( · ) } n − 1 i [-=
-]= and { f τ k,i ( · ) } n − 1 i =0 .

Original: Further, for a fixed x , we have the following useful property:n
Revised: Further, for a fixed x , we have the following usefulproperty:

Original: [Equation]
Revised: [Equation]

Original: • Index for m and v : As shown in Algorithm 1, we denote :
Revised: • Index for m and v : As shown in Algorithm 1, we denote :

Original: [Equation]
Revised: [Equation]

Revised: m l

Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Original: F.2 Some useful lemmas
Revised: F.2 Some useful lemmas

Original: We begin with proving some useful lemmas.
Revised: We begin with proving some useful lemmas.

Original: Lemma F.1.
Revised: Lemma F.1.

Original: For any β ̸ = 1 , we have
Revised: Let M and k be 2 integers with k > M , we have

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Revised: For the lower bound,

Revised: [Equation]

Revised: For the upper bound,

Revised: [Equation]

Revised: Lemma F.2.

Original:  Proof only involves basic calculation, we omit the proof here.
Revised: For any β ̸ = 1 , we have [Equation] Proof. Proof only involves basic calculation, we omit the proof here.

Original: Lemma F.2.
Revised: Lemma F.3.

Original: Under Assumption 2.1, if β 1 < √ β 2 , we have
Revised: Under Assumption 2.1, if β 1 < √ β 2 , we have

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We start with bounding | x l,k,i +1 − x l,k,i | .
Revised: We start with bounding | x l,k,i +1 − x l,k,i | .

Original: By the update formula of Adam, we have| x l
Revised: By the update formula of Adam, we have

Original: [Equation]
Revised: [Equation]

Original: Now, by Assumption 2.1, we have
Revised: Now, by the Lipschitz assumption of ∂ l f τ k,i ( x k,i ) (Assumption 2.1), we have(cid:12)(cid:12)

Original: [Equation]
Revised: [Equation]

Original: Proof is completed.
Revised: Proof is completed.

Original: Lemma F.3.
Revised: Lemma F.4.

Original: Under Assumption 2.1 and 2.2, for any integer i ∈ [0 , n − 1] , we have the following two results, 
Revised: Under Assumption 2.1 and 2.2, for any integer i ∈ [0 , n − 1] , we have the followingtwo results, [Equation]

Original: [Equation]
Revised: [Equation]

Original: where | ∂ α f ( x k, 0 ) | = max l ∈ [1 ,d ] | ∂ l f ( x k, 0 ) | .
Revised: where | ∂ α f ( x k, 0 ) | = max l ∈ [1 ,d ] | ∂ l f ( x k, 0 ) | .

Original: Proof.
Revised: Proof.

Revised: | ∂ l f τ k,j ( x k,j ) |

Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: where ( ∗ ) is because of the definition (see Appendix F.1): ρ 1 is a constant satisfying ρ 1 ≥
Revised: [Table] [Equation] [Table] [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: ni =
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: |
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original:  l f i ( x k, 0 ) | [Equation]
Revised: where ( ∗ ) is because of the definition (see Appendix F.1): ρ 1 is a constant satisfying ρ 1 ≥(cid:80) ni =1 | ∂ l f i ( x k, 0 ) | √ (cid:80) ni =1 | ∂ l f i ( x k, 0 ) | 2 ; ( ∗∗ ) is due to ∥∇ f ( x k, 0 )

Original: ( ∗∗ ) is due to ∥∇ f ( x k, 0 ) ∥ 22 ≤ d | ∂ α f ( x k, 0 ) | 2  .
Revised: ∥ 22 ≤ d | ∂ α f ( x k, 0 ) | 2 ; ( ∗∗ ) is because √ x + y ≤ 793 √ x + √ y , for any x, y ≥ 0 .

Original: The proof of (23) follows the same procedure as above.
Revised: The proof of (21) follows the same procedure as above.

Original: G Proof of Theorem 3.1G.1 A roadmap of the proof
Revised: G Proof of Theorem 3.1G.1 A roadmap of the proof

Original: Our proof is based on the Descent Lemma:
Revised: Our proof is based on the Descent Lemma:

Original: [Equation]
Revised: [Equation]

Original: The expectation E ( · ) is taken on the whole trajectory.
Revised: The expectation E ( · ) is taken on the whole trajectory.

Original: Summing both sides from the initializationk =
Revised: Summing both sides from the initializationk =

Original: t 0 to k = T , we have the following re-arranged inequality: (usually we set t 0 = 1 .)
Revised: t 0 to k = T , we have the following re-arranged inequality: (usually we set t 0 = 1 .)

Original: [Equation]
Revised: [Equation]

Original: To prove the convergence, we need an upper bound for L 2 E ∥ x k +1 , 0 − x k, 0 ∥ 22 , as well as a lower boundfor E ⟨∇ f
Revised: To prove the convergence, we need an upper bound for L 2 E ∥ x k +1 , 0 − x k, 0 ∥ 22 , as well as a lower boundfor E ⟨∇ f

Original: ( x k, 0 ) , x k, 0 − x k +1 , 0 ⟩ (and such a lower bound should be in the order of 1 √ k E ∥∇ f
Revised: ( x k, 0 ) , x k, 0 − x k +1 , 0 ⟩ (and such a lower bound should be in the order of 1 √ k E ∥∇ f

Original: ( x k, 0 ) ∥ ).
Revised: ( x k, 0 ) ∥ ).

Original: The upper bound for L 2 E ∥ x k +1 , 0 − x k, 

Revised: The upper bound for L 2 E ∥ x k +1 , 0 − x k,

Original: ∥ 22 is relatively easy to get: according to the update rule of
Revised: ∥ 22 is relatively easy to get: according to the update rule of

Original: Adam, this term is in the order of O ( η 2 k m l,k,i /v l,k,i )
Revised: Adam, this term is in the order of O ( η 2 k m l,k,i /v l,k,i )

Original: =
Revised: =

Original: O ( η 2 k ∥∇ f
Revised: O ( η 2 k ∥∇ f

Original: ( x k, 0 ) ∥ / ∥∇ f ( x k, 0 ) ∥ )
Revised: ( x k, 0 ) ∥ / ∥∇ f ( x k, 0 ) ∥ )

Original: =
Revised: =

Original: O ( 1 k ) .
Revised: O ( 1 k ) .

Original: Further, recall (cid:80) Tk = t 0 1 k ≤ log T +1 t 0 , so L 2 E ∥ x k +1 , 0 − x k, 0 ∥ 22 contributes to the log term in Theorem3.1.
Revised: Further, recall (cid:80) Tk = t 0 1 k ≤ log T +1 t 0 , so L 2 E ∥ x k +1 , 0 − x k, 0 ∥ 22 contributes to the log term in Theorem3.1.

Original: The proof is shown in Lemma F.2.
Revised: The proof is shown in Lemma F.3.

Original: However, the lower bound for E ⟨∇ f
Revised: However, the lower bound for E ⟨∇ f

Original: ( x k, 0 ) , x k 
Revised: ( x k, 0 ) , x k +1 , 0 − x k, 0 ⟩ = E (cid:104) (cid:80) dl

Original: [Equation]{+=1 (cid:80) n − 1 i =
+}
Revised: {+=1 (cid:80) n − 1 i =
+}

Original: requires sophisticated derivation, we explain as follows.
Revised: ∂ l f ( x k, 0 ) m l,k,i √ v l,k,i (cid:105)requires sophisticated derivation, we explain as follows.

Original: Before taking the expectation, we firstwork on every possible realization of (cid:80) dl
Revised: Before taking the expectation, we firstwork on every possible realization of (cid:80) dl

Original: =1 (cid:80) n − 1 i =

Revised: =1 (cid:80) n − 1 i =

Original: ∂ l f ( x k, 0 ) m l,k,i √ v l,k,i .
Revised: ∂ l f ( x k, 0 ) m l,k,i √ v l,k,i .

Original: We perform the followingdecomposition for every l ∈ [1 , d ] .
Revised: We perform the followingdecomposition for every l ∈ [1 , d ] .

Revised: n − 1 (cid:88) i =

Revised: ∂ l f ( x k, 0 ) m l,k,i √ v l,k,i = n − 1 (cid:88)

Original: [Equation]
Revised: i = ∂ l f ( x k, 0 ) √ v l,k,i ( ∂ l f i ( x k, 0 ) + m l,k,i − ∂ l f i ( x k, 0 )) = (cid:34) n − 1 (cid:88) i =

Revised: ∂ l f ( x k, 0 ) √ v l,k,i ∂ l f i ( x k, 0 ) (cid:35) (cid:124) (cid:123)(cid:122) (cid:125)

Revised: ( a ) + (cid:34) ∂

Revised: l f ( x k, 0 ) n − 1 (cid:88) i = 1 √ v l,k,i ( m l,k,i − ∂ l f i ( x k, 0 ))

Revised: (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( b ) .

Original: First of all, we introduce the following Lemma G.1 to further decompose ( a ) and ( b ) .
Revised: First of all, we introduce the following Lemma G.1 to further decompose ( a ) and ( b ) .

Original: The intuitionof this decomposition is as follows: by increasing β 2 , we can control the movement of the movingaverage factor v l,k,i , so the trajectory of Adam is easier to control.
Revised: The intuitionof this decomposition is as follows: by increasing β 2 , we can control the movement of the movingaverage factor v l,k,i , so the trajectory of Adam is easier to control.

Original: Lemma G.1.
Revised: Lemma G.1.

Original: Under Assumption 2.1, for those l satisfying max i | ∂ l f i ( x k, 0 )
Revised: Under Assumption 2.1, for those l satisfying max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k :=
Revised: | ≥ Q k :=

Original: △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 , we have the following lower bound for ( a ) in (25) :[Equation] [Equation] where δ 1 = (1 − β 2 )4 nρ 2 β n
Revised: △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 , we have the following lower bound for ( a ) in (23) :

Original: [Equation]
Revised: [Equation]

Original: ρ 2 , ρ 3 are constants satisfying ρ 2 ≥ | max i ∂ l f i ( x k, 0 ) | 2 1 n (cid:80) ni =1 | ∂ l f i ( x k, 0 ) | 2 ,ρ 3 ≥ |
Revised: (cid:12)(cid:12)(cid:12)(cid:12)

Original: [Equation]
Revised: [Equation]

Original: l f i

Original: [Equation]
Revised: [Equation]

Original: x k
Revised: where δ 1 = (1 − β 2 )

Original: [Equation]
Revised: [Equation]

Original: n
Revised: ρ 2 , ρ 3 are constants satisfying ρ 2 ≥

Original: [Equation]
Revised: [Equation]

Original: l f i
Revised: δ

Original: [Equation]
Revised: [Equation]

Original: k
Revised: 0 when β

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: | 2 .

Original: Note that δ 1 → 0 when β 2 → 1 .

Original: Under the same condition, we also have a lower bound for ( b ) in (25) :
Revised: Under the same condition, we also have a lower bound for ( b ) in (23) :

Original: [Equation]
Revised: [Equation]

Original: The proof can be seen in Appendix G.2.
Revised: The proof can be seen in Appendix G.2.

Original: Lemma G.1 decomposes both ( a ) and ( b ) : in the denominator,we approximate v l,k,i (which is changing with i ) by v l,k, 0 (which is fixed).
Revised: Lemma G.1 decomposes both ( a ) and ( b ) : in the denominator,we approximate v l,k,i (which is changing with i ) by v l,k, 0 (which is fixed).

Original: Accordingly, theapproximation error can be controlled by increasing β 2 (since δ 1 → 0 when β 2 → 1 ).
Revised: Accordingly, theapproximation error can be controlled by increasing β 2 (since δ 1 → 0 when β 2 → 1 ).

Original: However,
Revised: However,

Original: Lemma G.1 can only be applied to those l with max i | ∂ l f i ( x k, 0 )
Revised: Lemma G.1 can only be applied to those l with max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k , so we need to discuss thefollowing two cases to complete the whole picture.
Revised: | ≥ Q k , so we need to discuss thefollowing two cases to complete the whole picture.

Original: Case 1: unbounded gradient.
Revised: Case 1: unbounded gradient.

Original: Given x k, 0 , consider those l with max i | ∂ l f i ( x k, 0 )
Revised: Given x k, 0 , consider those l with max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k , we havethe following decomposition:
Revised: | ≥ Q k , we havethe following decomposition:

Original: [Equation]
Revised: [Table]

Original: ( a ) + (cid:34) ∂

Original: l f ( x k, 0 ) n − 1 (cid:88) i = 1 √ v l,k,i ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( b ) .{+− (cid:34) δ(cid:115) 
+}
Revised: (cid:123)(cid:122) (cid:125) ( a 1 ) {+− (cid:34) δ(cid:115) 
+}

Original: LemmaG.1 ≥ ∂ l f ( x k, 0 ) 2 √ v l,k,− δ 1 (cid:12)(cid:12)(cid:12)(cid:12) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:12)(cid:12)(cid:12)(cid:12) n − 1 (cid:88) i = | ∂ l f i ( x k, 0 ) | + ∂ l f ( x k, 0 ) √ v l,k,n − 1 (cid:88) i =( m l,k,i − ∂ l f i ( x k, 0 )) − δ 1 (cid:12)(cid:12)(cid:12)(cid:12) ∂
Revised: ρβ nn − 1 (cid:88) i =

Original:  l f  ( x k, 0 )√ v l,k, 0 (cid:12)(cid:12)(cid:12)(cid:12) n − 1 (cid:88) i = | m l,k,i − ∂ l fi ( x k, 0 ) | LemmaG [Table]
Revised: | ∂ l f i ( x k, 0 ) | (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( a 2 ) + (cid:34) ∂ l f ( x k, 0 ) √ v l,k,n − 1 (cid:88) i =

Original: (cid:20) ∂ l f  ( x [Table] [Table]
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: [Table] √ v l, [Table] ,(cid:21) (cid:124) (cid:123)(cid:122) (cid:125) ( a 1 ) − (cid:34) δ(cid:115) 2 ρβ nn{+
+}
Revised: (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( b 1 ) − (cid:34) δ(cid:115) {+
+}

Original: [Table] i = | ∂ l f i ( x k, 0 ) | (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( a 2 ) + (cid:34) [Equation]
Revised: ρβ nn − 1 (cid:88) i = | m l,k

Original: [Equation] (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( [Equation]
Revised: i − ∂ l f i ( x k, 0 ) | (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) ( b 2 ) , (26)

Original: Case 2: bounded gradient.
Revised: Case 2: bounded gradient.

Original: Given x k, 0 , consider those l with max i | ∂ l f i ( x k, 0 ) | ≤ Q k , the analysisdegenerates to the “bounded gradient” scenario, we have the following lower bound:
Revised: Given x k, 0 , consider those l with max i | ∂ l f i ( x k, 0 ) | ≤ Q k , the analysisdegenerates to the “bounded gradient” scenario, we have the following lower bound:

Original: [Equation]
Revised: [Equation]

Original: where F 1 :=
Revised: where F:= △ 1 n 2 √ n 32 √ [Equation] [Equation] − β

Revised: [Equation]

Revised: ) n

Original: [Equation]
Revised: [Equation]

Original: n
Revised: n

Original: [Equation]
Revised: [Equation]

Original:  n , ( ∗ ) is due to the fact that | ∂ l f ( x k, 0 ) | ≤n max i | ∂ l f i ( x k, 0 ) | ≤ nQ k .
Revised: − β 1 1 − β 1 √ β 2 n , ( ∗ ) is due to the fact that | ∂ l f ( x k, 0 ) | ≤n max i | ∂ l f i ( x k, 0 ) | ≤ nQ k .

Original: Combining Case 1 & 2 together, we have the following result (note that ( a 1 ) , ( a 2 ) , ( b 1 ) , ( b 2 ) varywith l .
Revised: Combining Case 1 & 2 together, we have the following result (note that ( a 1 ) , ( a 2 ) , ( b 1 ) , ( b 2 ) varywith l .

Original: A more precise notation will be ( a 1 ) l , etc.
Revised: A more precise notation will be ( a 1 ) l , etc.

Original: We drop the subscript for brevity).
Revised: We drop the subscript for brevity).

Revised: d

Original: [Equation]
Revised: [Equation]

Original: l large
Revised: l large

Original: [Equation]
Revised: [Equation]

Original: (cid:125) unboundedgradient [Equation]
Revised: (cid:125) unbounded gradient − (cid:88) l

Original: l small
Revised: small

Original: F [Equation] k
Revised: F 1 1 √ k

Original: [Equation] (cid:125) boundedgradient ≥ 
Revised: (cid:124) (cid:123)(cid:122) (cid:125) bounded gradient ≥ (cid:88) l large

Original: [Equation]
Revised: [Equation]

Original: l large

Original: [Equation]
Revised: [Equation]

Original: ) } (cid:124) (cid:123)(cid:122) (cid:125) unboundedgradient − d 
Revised: (cid:124) (cid:123)(cid:122) (cid:125) unbounded gradient − d (cid:88)

Original: [Equation] (cid:26)F
Revised: l = (cid:26)

Original: [Equation] √ k (cid:27) (cid:124) (cid:123)(cid:122) (cid:125) boundedgradient ≥   (cid:88)
Revised: F 1 1 √ k (cid:27) (cid:124) (cid:123)(cid:122) (cid:125) bounded gradient ≥   (cid:88)

Original: l large ( a 1 )   +   (cid:88)
Revised: l large ( a 1 )   +   (cid:88)

Original: l large ( b 1 )   −   (cid:88) l large { ( a 2 ) + ( b 2 ) }   − dF [-
 [Equation]-] where “ l large” stands for the gradient component in Case 1 and “ l small” indicates those in Case2 .
Revised: l large ( b 1 )   −   (cid:88) l large { ( a 2 ) + ( b 2 ) }   − dF [-
 [Equation]-]1 1 √ k , (27) where “ l large” stands for the gradient component in Case 1 and “ l small” indicates those in Case2 .

Original: We assume “ l large” is not an empty set, otherwise the analysis degenerates to the proof withbounded gradient assumption.
Revised: We assume “ l large” is not an empty set, otherwise the analysis degenerates to the proof withbounded gradient assumption.

Original: Now we take expectation on (29).
Revised: Now we take expectation on (27).

Original: The expectation is taken on all the possible trajectories up to thei -th iteration in k -th epoch.
Revised: The expectation is taken on all the possible trajectories up to thei -th iteration in k -th epoch.

Original: [Equation]
Revised: [Equation]

Original: In the following context, we will discuss how to bound all terms in (30), respectively.
Revised: In the following context, we will discuss how to bound all terms in (28), respectively.

Original: First andforemost, we derive a lower bound for E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .
Revised: First andforemost, we derive a lower bound for E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .

Original: Since E (cid:104)(cid:80) l large ( b 1 ) (cid:105) contains all thehistorical gradient information, it involves great effort to handle it.
Revised: Since E (cid:104)(cid:80) l large ( b 1 ) (cid:105) contains all thehistorical gradient information, it involves great effort to handle it.

Original: We will show that the lower boundof E (cid:104)(cid:80) l large ( b 1 ) (cid:105) will vanish when β 2 is large and k goes to infinity.
Revised: We will show that the lower boundof E (cid:104)(cid:80) l large ( b 1 ) (cid:105) will vanish when β 2 is large and k goes to infinity.

Original: Recall E (cid:104)(cid:80) l large ( b 1 ) (cid:105) =
Revised: Recall E (cid:104)(cid:80) [Equation]

Original: E (cid:104)(cid:80) l large ∂ l f ( [-x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =
-]
Revised: = [Equation] ( [-x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =
-]m [Equation]

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: − ∂ [Equation] f i ( x k [Equation]

Original: (cid:105) .
Revised: (cid:105) .

Original: When β 1 is large,m l,k,i contains heavy historical signals.
Revised: When β 1 is large,m l,k,i contains heavy historical signals.

Original: It seems unclear how large E (cid:104)(cid:80) l large ( b 1 ) (cid:105) would bewhen β 1 goes to 1.
Revised: It seems unclear how large E (cid:104)(cid:80) l large ( b 1 ) (cid:105) would bewhen β 1 goes to 1.

Original: Existing literatures (Zaheer et al., 2018; De et al., 2018; Shi et al., 2020)
Revised: Existing literatures (Zaheer et al., 2018; De et al., 2018; Shi et al., 2020)

Original: takea naive approach: they set β 1 ≈ 0 so that m k,i ≈ ∇ f τ k,i ( x k,i ) .
Revised: takea naive approach: they set β 1 ≈ 0 so that m k,i ≈ ∇ f τ k,i ( x k,i ) .

Original: Then we get δ ( β 1 ) ≈ 0 .
Revised: Then we get δ ( β 1 ) ≈ 0 .

Original: However,this naive method cannot be applied here since we are interested in practical cases where β 1 is largein [0 , 1) .
Revised: However,this naive method cannot be applied here since we are interested in practical cases where β 1 is largein [0 , 1) .

Original: We emphasize the following technical difficulties in bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) for anyβ 1 ∈ [0 , 1) :
Revised: We emphasize the following technical difficulties in bounding E (cid:104) (cid:80) l large ( b 1 ) (cid:105) for anyβ 1 ∈ [0 , 1) :

Original: • Issue (i)
Revised: Issue (i)

Original: In order to bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) , we need to first know how to bound its simpler 
Revised: In order to bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) , we need to first know how to bound its simpler version:

Original: version: E (cid:104)(cid:80) n − 1 i =

Revised: E (cid:104)(cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) .
Revised: (cid:105) .

Original: This term measures the difference between the current gradient and weighted historical gradients.
Revised: This term measures the difference between the currentgradient and weighted historical gradients.

Original: It seems unclear that how large this term could bewhen β 1 goes to 1.• Issue (ii) Even if we can bound E (cid:104)(cid:80) n − 1 i =

Revised: It seems unclear that how large this term couldbe when β 1 goes to 1. Issue (ii) Even if we can bound E (cid:104)(cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) , it is still unclear how to bound
Revised: (cid:105) , it is still unclear how to bound

Original: E (cid:104) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Revised: E (cid:104) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) , which is further multiplied by a random variable∂ l f ( x k, 0 ) √ v l,k, 0 .
Revised: (cid:105) , which is further multiplied by a random variable∂ l f ( x k, 0 ) [Table] l,k, 0 .

Original: • Issue (iii)
Revised: Issue (iii)

Original: Even if we can bound E (cid:104) ∂
Revised: Even if we can bound E (cid:104) ∂

Original: l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Revised: l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) , it is still different from E (cid:104)(cid:80) l large ( b 1 ) (cid:105) which contains additional operation “ (cid:80) l large " inside the expectation.
Revised: (cid:105) , it is still different from E (cid:104)(cid:80) l large ( b 1 ) (cid:105) which contains additional operation “ (cid:80) l large " inside the expectation.

Original: Note that the set “ l large" is a random variable which changes along different trajectories, so there is still non-negligible gap between bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) 
Revised: Note that the set “ l large" is a random variable which changes along different trajectories, so there is still non-negligible gap between bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) and

Original: [-andE (cid:104) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =
 ( m l,k,i − ∂ l f i ( x k, 0 )) (cid:105) .-]
Revised: [-andE (cid:104) ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =
 ( m l,k,i − ∂ l f i ( x k, 0 )) (cid:105) .-][Equation]

Original: To our best knowledge, there is no general approach to tackle the above issues.
Revised: To our best knowledge, there is no general approach to tackle the above issues.

Original: In the followingcontent, we will overcome difficulties (i), (ii) and (iii) in Step 1, 2 and 3 respectively.
Revised: In the followingcontent, we will overcome difficulties (i), (ii) and (iii) in Step 1, 2 and 3 respectively.

Original: In Step 1, wewill discuss how to bound E (cid:104)(cid:80) n − 1 i =

Revised: In Step 1, wewill discuss how to bound E (cid:104)(cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) , which is a simplified version of E [( b 1 )] .
Revised: (cid:105) , which is a simplified version of E [( b 1 )] .

Original: Bounding this term will shed light on bounding the whole term E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .
Revised: Bounding this term will shed light on bounding the whole term E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .

Original: Then in Step 2 and3, we will prove several technical lemmas to handle the effect of ∂ l f ( x k, 0 ) √ v l,k, 0 and “ (cid:80) l large ", by whichwe can tackle (ii) and (iii) .
Revised: Then in Step 2 and3, we will prove several technical lemmas to handle the effect of ∂ l f ( x k, 0 ) √ v l,k, 0 and “ (cid:80) l large ", by whichwe can tackle (ii) and (iii) .

Original: Combining all together we can bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) . Bounding 
Revised: Combining all together we can bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) . Bounding E (cid:104)(cid:80)

Original: [Equation]
Revised: [Equation]

Original: Step 1.
Revised: Step 1.

Original: We now introduce the key idea of bounding
Revised: We now introduce the key idea of bounding

Original: [-E (cid:104)(cid:80) n − 1 i =
-]
Revised: [-E (cid:104)(cid:80) n − 1 i =
-][Equation] As discussed in Appendix A, we analyze

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: [Figure] [Equation] f i ( x k, 0 ))

Original: (cid:105) .
Revised: (cid:105) for every epoch and find out some pattern across the successive epochs. In Appendix A, we distill our idea into a toy example with two balls in a box. Here,to better convey our message, we further describe a similar case with three balls. The “color-ball " model with three balls. Consider a box containing 3 balls labeled a 0 , a 1 , a 2 ∈ R(could be negative).

Revised: In each round (epoch), we randomly sample all the balls from the box withoutreplacement, then we put them all back.

Original: In Appendix A, we distill our idea into a toy example called "thecolor-ball model of the 1st kind" and thus we prove Lemma A.1.
Revised: We denote the 1st sampled label in the k -th epoch as a τ k, 0 ,and similarly for a τ k, 1 , a τ k, 2 . We define the following random variables: ( β ∈ [0 , 1] ) [Equation] [Equation]

Revised: Further

Original: This lemma is crucial for bounding
Revised: [Equation] we introduce the following deterministic constants. [Equation] For the simplicity of analysis

Revised: [Equation]

Revised: we denote

Revised: [Equation]

Revised: [Equation]

Original: {+2 i =
+}
Revised: {+2 i =
+}

Original: {+f i (cid:17) k − 1 + · · · + (cid:16)(cid:80) 2 i =
+}
Revised: {+f i (cid:17) k − 1 + · · · + (cid:16)(cid:80) 2 i =
+}

Original: {+f i (cid:17) 1 ,where (cid:16)(cid:80) 2 i =
+}
Revised: {+f i (cid:17) 1 ,where (cid:16)(cid:80) 2 i =
+}

Original: {+m i (cid:17) k and (cid:16)(cid:80) 
+}
Revised: {+m i (cid:17) k and (cid:16)(cid:80) 
+}

Original: {+i =
+}
Revised: {+i =
+}

Revised: f i (cid:17) k contain the summand in the k -th epoch, so on so for.

Original: {+Now, wetry to calculate E (cid:104)(cid:80) 2 i =
+}
Revised: {+Now, wetry to calculate E (cid:104)(cid:80) 2 i =
+}

Original: {+m i − (cid:80) 2 i =
+}
Revised: {+m i − (cid:80) 2 i =
+}

Revised: f i (cid:105) , where the expectation is taken on all the possible drawsuntil the end of k -th round (epoch).

Revised: Lemma G.2.

Revised: In the colorb-ball example above, we have

Revised: E

Revised: [Equation]

Revised: For general n = 1 , 2 , 3 , · · · , we have

Revised: E

Revised: [Equation]

Revised: Proof.

Original: {+Since E [ · ] = E [ E k [ · ]] , we first calculate E k (cid:104)(cid:80) 2 i =
+}
Revised: {+Since E [ · ] = E [ E k [ · ]] , we first calculate E k (cid:104)(cid:80) 2 i =
+}

Original: {+m i − (cid:80) 2 i =
+}
Revised: {+m i − (cid:80) 2 i =
+}

Revised: f i (cid:105) .

Revised: Since all the history before k -th epoch is fixed, we delegate the history part to later analysis.

Revised: Now, we focus on 890 (cid:16)(cid:80)

Revised: [Equation]

Original: E (cid:104)(cid:80) n − 1{+(cid:16)(cid:80) 
+} i [-=
-]
Revised: − {+(cid:16)(cid:80) 
+} i [-=
-][Equation] f

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: [Equation] (cid:17) k . As shown in Figure 15, there are (3)! =

Revised: 6 possible order in the k -thepoch.

Revised: To calculate E k ( · ) , we need to sum up all 6 possible shuffle order.

Revised: Figure 15 tells us: when all thehistory sampling orders are fixed, we have: (we use

Revised: = all 6 possible shuffle )(cid:88) s (cid:40)(cid:32) 2 (cid:88)

Revised: [Equation]

Revised: Since all the sampling order occur with the same probability 1 / 6 , we have:

Revised: [Equation]

Original: {+Now we move one step further to calculate E k − 1 E k (cid:104)(cid:80) 2 i =
+}
Revised: {+Now we move one step further to calculate E k − 1 E k (cid:104)(cid:80) 2 i =
+}

Original: {+m i − (cid:80) 2 i =
+}
Revised: {+m i − (cid:80) 2 i =
+}

Original:  (cid:105) .
Revised: f i (cid:105) .

Revised: Since all the historybefore ( k − 1) -th epoch is fixed, we do not discuss them for now.

Revised: In particular, we need to calculate

Revised: E k −(cid:40)

Revised: E k (cid:34)(cid:32) 2 (cid:88) i

Revised: = m i (cid:33) k − (cid:32) 2 (cid:88) i =

Revised: f i (cid:33) k

Revised: + (cid:32) 2 (cid:88) i

Revised: = m i (cid:33) k −− (cid:32) 2 (cid:88) i =

Revised: f i (cid:33) k −(cid:35)(cid:41) =

Revised: E k −(cid:40)

Revised: [Equation]

Revised: (cid:88) i

Original: We refer the readers to Appendix A for more explanation.
Revised: = a i ( [Equation] 1 − [Equation]

Revised: 2 ) + (cid:32) 2 (cid:88) i

Revised: = m i (cid:33) k −− (cid:32) 2 (cid:88) i =

Revised: f i (cid:33) k −(cid:41) .

Revised: We highlight the blue term since it is “the residue from the k -th epoch”, this term will be canceledout later.

Original: {+Similarly as Figure 15, we sum up all possible draws (cid:80) s (cid:16)(cid:80) 2 i =
+}
Revised: {+Similarly as Figure 15, we sum up all possible draws (cid:80) s (cid:16)(cid:80) 2 i =
+}

Revised: m

Revised: i (cid:17) k − 1 in Figure 16.

Revised: As shown in Figure 16, we have(cid:88) s   2 (cid:88)

Revised: i = a i ( − 2 6 β 1 − 4 6 β 2 ) + (cid:32) 2 (cid:88) i

Revised: = m i (cid:33) k −− (cid:32) 2 (cid:88) i =

Revised: f i (cid:33) k

Revised: −  = 2 (cid:88)

Revised: i = a i ( − 2 β 1 − 4 β 2 ) + 2 (cid:88)

Revised: i = a i (2 β 1 + 4 β 2 ) (cid:124) (cid:123)(cid:122) (cid:125)

Revised: canceloutthe k -thresidue + 2 (cid:88)

Revised: i = a i (6 β 3 + 4 β 4 + 2 β 5 ) − 2 (cid:88)

Revised: i = a i (6 β 3 + 6 β 4 + 6 β 5 ) (cid:124) (cid:123)(cid:122) (cid:125) createthe ( k − 1) -thresidue

Revised: = 2 (cid:88)

Revised: i = a i ( − 2 β 4 − 4 β 5 ) .

Revised: Equivalently,

Original: Lemma A.1 can provide insights in bounding E (cid:104)(cid:80) n − 1 i [-=
-]
Revised: E k −  E k   (cid:32) 2 (cid:88) i = m i (cid:33) k − (cid:32) 2 (cid:88) i [-=
-]= f i (cid:33) k

Revised: + (cid:32) 2 (cid:88) i

Revised: = m i (cid:33) k −− (cid:32) 2 (cid:88) i =

Revised: f i (cid:33) k −   

Revised: Figure i = a i ( − 2 6 β 4 − 4 6 β 5 )

Revised: From (29) we find out that: the residue from the k -th epoch will be cancelled out by part of the redterms in the ( k − 1) -th epoch; meanwhile, new residue term will be created.

Revised: Moreover, we observethe following pattern in the residue : 1) the order of β involves the largest & second largest order thatpossibly occur within the epoch.

Revised: 2) the coefficients are fixed to be − 2 6 and − 4 6 .

Revised: This is actually anuniversal pattern for every epoch from 1 to k .

Revised: We illustrate this pattern in the following Figure 17.

Original: {+Combining the above analysis and Figure 17, we observe the following pattern: when iterativelyplugging in E k − j , the residue term remaining in the ( k − j + 1) -th epoch will be erased (i.e. 908 (cid:80)i =
+}
Revised: {+Combining the above analysis and Figure 17, we observe the following pattern: when iterativelyplugging in E k − j , the residue term remaining in the ( k − j + 1) -th epoch will be erased (i.e. 908 (cid:80)i =
+}

Original: {+a i β ( j − 1) ( − 2 β 4 − 4 β 5 ) ), whereas a new residue term will be produced (i.e. (cid:80) 2 i =
+}
Revised: {+a i β ( j − 1) ( − 2 β 4 − 4 β 5 ) ), whereas a new residue term will be produced (i.e. (cid:80) 2 i =
+}

Revised: a i β j ( − 2 β 4 −

Revised: 4 β 5 ) ).

Revised: Since E [ · ] = E [ E 1 ( E k − 1 ( · · · E k ( · )))] , we derive the following result by repeatedly applyingthe steps in (29):

Revised: [Equation]

Revised: The above analysis also holds for general n = 1 , 2 , 3 , · · · , we just need to change the constants in thefollowing form:

Revised: [Equation]

Revised: Since the analysis for general “ n balls” is exactly the same as that of n = 3 , we omit the proof forbrevity.

Original: {+The color-ball toy example above can provide insights in bounding E (cid:104) (cid:80) n − 1 i =
+}
Revised: {+The color-ball toy example above can provide insights in bounding E (cid:104) (cid:80) n − 1 i =
+}

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) .
Revised: (cid:105) .

Original: However, there isstill certain gap between E (cid:104)(cid:80) n − 1 i =

Revised: However, there is still some gap between E (cid:104)(cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) and the quantities in Lemma A.1.
Revised: (cid:105) and the quantities in Lemma G.2.

Original: Weelaborate as follows:
Revised: We elaborate as follows:

Original: • To mimic the color-ball example, we need to expand (cid:80) n − 1 i =

Revised: • To mimic the color-ball example, we need to expand (cid:80) n − 1 i =

Original: ∂ l f i ( x k, 0 ) into an infinite sumsequence: (cid:80) n − 1 i =

Revised: ∂ l f i ( x k, 0 ) into an infinite sumsequence: (cid:80) n − 1 i =

Original: ∂
Revised: ∂

Original: l f i ( x k, 0 )(1 − β 1 )(1 + β 1 + · · · β ∞ 1 ) .
Revised: l f i ( x k, 0 )(1 − β 1 )(1 + β 1 + · · · β ∞ 1 ) .

Original: However, (cid:80) n − 1 i =

Revised: However, (cid:80) n − 1 i =

Original: m l,k,i is a finite sumsequence up to the order of β kn .
Revised: m l,k,i is a finitesum sequence up to the order of β kn .

Original: In contrast, both sequences in the color-ball example are “finitesum”.
Revised: In contrast, both sequences in the color-ball example are “finite sum”.

Original: As such, there is an error term caused by “finite sum v.s. infinite sum”.
Revised: As such, there is an error term caused by “finite sum v.s. infinite sum”.

Original: • When taking the expectation, the variable x in each possible trajectory is different.
Revised: • When taking the expectation, the variable x in each possible trajectory is different.

Original: In contrast, inthe color-ball example, { a i } 2 i =0 are fixed in all shuffling order (
Revised: In contrast, in the color-ball example, { a i } 2 i =0 are fixed in all shuffling order (so it is much

Original: so it is much easier to calculate the expectation by summing them up
Revised: easier to calculate the expectation by summing them up).

Original: To mimic the color-ball example, we repeatedly take conditional expectation at the beginning ofeach k -th epoch.
Revised: To mimic the color-ball example, we repeatedly take conditional expectation at the beginningof each k -th epoch.

Original: In this way, x k, 0 will be fix.
Revised: In this way, x k, 0 will be fix.

Original: Despite x k,i is still changing across the trajectory,we can transform x k i into x k, 0 by using Lipschitz property.
Revised: Despite x k,i is still changing across thetrajectory, we can transform x k,i into x k, 0 by using Lipschitz property.

Original: • In each trajectory of (cid:80) n − 1 i =

Revised: • In each trajectory of (cid:80) n − 1 i =

Original: m l,k,i , the variable x in the summand of m l,k,i varies with k andi ;
Revised: m l,k,i , the variable x in the summand of m l,k,i varies with k and

Original:  while the variable in ∂ l f i ( x k, 0 ) is fixed to be x k, 0 .
Revised: i ; while the variable in ∂ l f i ( x k, 0 ) is fixed to be x k, 0 .

Original: In contrast, in the color-ball example, { a i } 2 i = are the same across the epoch.
Revised: In contrast, in the color-ball example,{ a i } 2 i =0 are the same across the epoch.

Original: To mimic the final step in the color-ball example (Figure 5), at each step of conditional expectation,we need to simultaneously move the variables in m l,k,i and ∂ l f i ( x k, 0 ) (using Lipschitz property)so that they can match and cancel out with each other.
Revised: To mimic the final step in the color-ball example (Figure 17), at each step of conditionalexpectation, we need to simultaneously move the variables in m l,k,i and ∂ l f i ( x k, 0 ) (using Lipschitz property) so that they can match and cancel out with each other.

Original: This operation will introduce new errorterms and it is our duty to put them under control.
Revised: This operationwill introduce new error terms and it is our duty to put them under control.

Original: We omit the proof for bounding E (cid:104)(cid:80) n − 1 i =

Revised: We omit the proof for bounding E (cid:104)(cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) since this is not our actual goal.
Revised: (cid:105) since this is not our actual goal.

Original: Instead, we will directly use the above ideas to bound E [( b 1 )] .
Revised: Instead, we will directly use the above ideas to bound E [( b 1 )] .

Original: To do so, we need to further tacklethe issue (ii) and (iii) mentioned before.
Revised: To do so, we need to further tacklethe issue (ii) and (iii) mentioned before.

Original: We explain as follows.
Revised: We explain as follows.

Original: Bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) :
Revised: Bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) :

Original: Step 2.
Revised: Step 2.

Original: We now resolve issue (ii) , i.e., handle the effect of ∂ l f ( x k, 0 ) √ v l,k, 0 .
Revised: We now resolve (ii), i.e., handle the effect of ∂ l f ( x k, 0 ) √ v l,k, 0 .

Original: The key idea is as follows: when we calculate E [( b 1 )] , we sequentially take conditional expectation 952 E k ( · ) , E k − 1 ( · ) , etc.. When taking E k ( · ) , we will fix all the historical information up to k -th epoch, so ∂ l f ( x k, 0 ) √ v l,k, 0 can be regarded as a constant.
Revised: Thekey idea is as follows: when we calculate E [( b 1 )] , we sequentially take conditional expectation 942 E k ( · ) , E k − 1 ( · ) , etc.. When taking E k ( · ) , we will fix all the historical information up to k -th epoch,so ∂ l f ( x k, 0 ) √ v l,k, 0 can be regarded as a constant.

Original: In this sense, E k (cid:104)
Revised: In this sense, E k (cid:104)

Revised: [Table]

Revised: i − ∂

Original: [Equation]
Revised: [Equation]

Revised: f i

Original: [Equation]
Revised: [Equation]

Revised: x k, 0 ))

Original: can be calculated following the same idea as color-ball toy example.
Revised: (cid:105)can be calculated following the same idea as color-ball toy example.

Original: However, when taking E k − 1 ( · ) , ∂ l f ( x k, 0 ) √ v l,k, 0 will become a random variable which changes with differenttrajectories.
Revised: However, when taking E k − 1 ( · ) , ∂ l f ( x k, 0 ) √ v l,k, 0 will become a random variable which changes with differenttrajectories.

Original: In this case, the color-ball method cannot be applied.
Revised: In this case, the color-ball method cannot be applied.

Original: To fix this issue, we introducethe following lemma to change ∂ l f ( x k, 0 ) √ v l,k, 0 into ∂ l f ( x k − 1 , 0 ) √ v k − 1 , 0 , which can again be regarded as a fixedconstant when taking E k − 1 ( · ) .
Revised: To fix this issue, we introducethe following lemma to change ∂ l f ( x k, 0 ) √ v l,k, 0 into ∂ l f ( x k − 1 , 0 ) √ v k − 1 , 0 , which can again be regarded as a fixedconstant when taking E k − 1 ( · ) .

Original: Lemma G.2.
Revised: Lemma G.3.

Original: Suppose Assumption 2.1 holds and β 1 < √ β 2 .
Revised: Suppose Assumption 2.1 holds and β 1 < √ β 2 .

Original: For any integer j ∈ [0 , k ] , ifmax
Revised: For any integer j ∈ [0 , k ] , if

Original: [Equation]
Revised: [Equation]

Original: | ≥ Qk
Revised: Q

Original: [Equation]
Revised: [Equation]

Original: j (where
Revised: (where

Original: Q
Revised: Q

Original: [Equation]
Revised: [Equation]

Original: :=
Revised: β

Original: [Equation]
Revised: [Equation]

Original: 2 ), then we have the following result:
Revised: 2 ), then we have the following result:

Original: [Equation]
Revised: [Equation]

Original: (cid:16) 1 − (1 − β 2 )4 nρ 2 β n(cid:17)

Original: δ 1 , where δ 1 is defined in (16) .

Original: The proof of Lemma G.2 can be seen in Appendix G.3.
Revised: The proof of Lemma G.3 can be seen in Appendix G.3.

Original: To proceed, we combine Lemma G.2 and the“color-ball method of the 2nd kind", and thus we can prove Lemma A.2.

Original: This lemma is crucial forour current goal: bounding E ( b 1 ) .

Original: We refer the readers to Appendix A for more information.

Original: We emphasize that here are still the following gap between Lemma A.2 and our goal E ( b 1 ) .

Original: • We have the similar gap as discussed at the end of the Step 1 .

Original: • The condition in Lemma G.2 has requirement on the gradient norm, while this requirement is

Original: temporarily ignored in the color-ball method of the 2nd kind.

Original: • The result in Lemma G.2 has additional error terms other than O (1 / √ k ) .

Original: This is slightly different

Original: from the setting in Lemma A.2.

Original: It requires some technical lemmas to handle these gaps.

Original: We fill in these gaps in Lemma G.4.

Original: Thetechnical details can be seen Appendix G.4.

Original: Bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) :
Revised: Bounding E (cid:104)(cid:80) l large ( b 1 ) (cid:105) :

Original: Step 3.
Revised: Step 3.

Original: Now we shift gear to tackle (iii): handling the random variable“ l large".
Revised: Now we shift gear to tackle (iii): handling the random variable“ l large".

Original: We rewrite “ l large" into the indicator function as follows:
Revised: We rewrite “ l large" into the indicator function as follows:

Original: E  [Equation] l = I
Revised: E  (cid:88) l large ( b 1 )

Original: [Equation]
Revised: [Equation]

Original:  f[Equation] [Equation] v l,k,n − 1 (cid:88) i =
Revised: , where I k := I (cid:16) max i | ∂ l f i ( x k, 0 )

Original: ( m l,k,i − ∂ l f i ( x k, 0 )) (cid:35) , where I k :
Revised: | ≥ Q k :=

Original: = I (cid:16) max i | ∂ l f i ( x k, 0 ) | ≥ Q k := △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 (cid:17) is the indicator function ( I ( A ) =
Revised: △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 (cid:17) is the indicator function ( I ( A ) =

Original: 1when event A holds and I ( A )
Revised: 1when event A holds and I ( A )

Original: = 

Revised: =

Original: otherwise.)
Revised: otherwise.)

Original: Similarly as before, when taking E k ( · ) , “ I k " canbe regarded as a constant index.
Revised: Similarly as before, when taking E k ( · ) , “ I k " can beregarded as a constant index.

Original: Therefore, E k (cid:104)(cid:80) dl
Revised: Therefore, E k (cid:104)(cid:80) dl

Original: =

Revised: =

Original: I k ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Revised: I k ∂ l f ( x k, 0 ) √ v l,k, 0 (cid:80) n − 1 i =

Original: ( m l,k,i − ∂ l f i ( x k, 0 ))
Revised: ( m l,k,i − ∂ l f i ( x k, 0 ))

Original: (cid:105) canbe calculated following the same idea as color-ball model of the 1st kind.
Revised: (cid:105) can becalculated following the same idea as color-ball toy example.

Original: However, when taking980 E k − 1 ( · ) , I k will become a random variable which changes with different trajectories.
Revised: However, when taking E k − 1 ( · ) , I kwill become a random variable which changes with different trajectories.

Original: In this case,the color-ball method cannot be applied.
Revised: In this case, the color-ballmethod cannot be applied.

Original: Similarly as in Step 2 , we introduce the following lemma tochange I k into  I k − 1 (defined later), which can again be regarded as a fixed when taking E k − 1 ( · ) .
Revised: Similarly as before, we introduce the following lemma to change I k into 961 I k − 1 (defined later), which can again be regarded as a fixed when taking E k − 1 ( · ) .

Original: Lemma G.3.
Revised: Lemma G.4.

Original: Suppose Assumption 2.1 holds and β 1 < √ β 2 .
Revised: Suppose Assumption 2.1 holds and β 1 < √ β 2 .

Original: For 0 ≤ j ≤ k , we define
Revised: For 0 ≤ j ≤ k , we define

Original: I k − j := I (cid:16) max i | ∂ l f i ( x k − j, 0 ) | ≥ (cid:80) kp = k − j Q p (cid:17) , where Q k :=
Revised: I k − j := I (cid:16) max i | ∂ l f i ( x k − j, 0 ) | ≥ (cid:80) kp = k − j Q p (cid:17) , where Q k :=

Original: △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 , then we havethe following results.
Revised: △ 1 n √ n √ k 32 √ 2 (1 − β 2 ) n β n 2 , then we havethe following results.

Original: I k
Revised: I k

Original: = I   max i | ∂ l f i ( x k, 0 )
Revised: = I   max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k and max i | ∂ l f i ( x k − j, 0 )
Revised: | ≥ Q k and max i | ∂ l f i ( x k − j, 0 )

Original: | ≥ k (cid:88) p =
Revised: | ≥ k (cid:88) p =

Original: k − j
Revised: k − j

Original: Q p 
Revised: Q p 

Original: 
Revised: 

Original: + I   max i | ∂ l f i ( x k, 0 ) | ≥ Q k and max i | ∂ l f i ( x k − j, 0 ) |
Revised: + I   max i | ∂ l f i ( x k, 0 ) | ≥ Q k and max i | ∂ l f i ( x k − j, 0 ) |

Original: ≤ k (cid:88) p =
Revised: ≤ k (cid:88) p =

Original: k − j
Revised: k − j

Original: Q p 
Revised: Q p 

Original:  , (31)
Revised:  , (30)

Original: I   max i | ∂ l f i ( x k, 0 )
Revised: I   max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k and max i | ∂ l f i ( x k − j, 0 )
Revised: | ≥ Q k and max i | ∂ l f i ( x k − j, 0 )

Original: | ≥ k (cid:88) p =
Revised: | ≥ k (cid:88) p =

Original: k − j
Revised: k − j

Original: Q p 
Revised: Q p 

Original:  =
Revised:  =

Original: I   max i | ∂ l f i ( x k − j, 0 )
Revised: I   max i | ∂ l f i ( x k − j, 0 )

Original: | ≥ k (cid:88) p =
Revised: | ≥ k (cid:88) p =

Original: k − j
Revised: k − j

Original: Q p 
Revised: Q p 

Original:  =
Revised:  =

Original: I k − j .
Revised: I k − j .

Original: Proof.
Revised: Proof.

Original: Equation (31) is straightforward, we only prove (32) here.
Revised: Equation (30) is straightforward, we only prove (31) here.

Original: Under Assumption 2.1, we have| ∂ l f i ( x k, 0 ) − ∂ l f i ( x k − j, 0 ) | ≤ n △ n ( k − 1) + · · · + n △ n ( k − j ) ≤ Q k − 1 + · · · + Q k − j .
Revised: Under Assumption 2.1, we have| ∂ l f i ( x k, 0 ) − ∂ l f i ( x k − j, 0 ) | ≤ n △ n ( k − 1) + · · · + n △ n ( k − j ) ≤ Q k − 1 + · · · + Q k − j .

Original: To show the 988 second inequality, it requires comparing the value between Q k and n △ n ( k ) , which are both problemdependent constants.
Revised: To show the 967 second inequality, it requires comparing the value between Q k and n △ n ( k ) , which are both problemdependent constants.

Original: Here, the inequality holds when n △ n ( k ) ≤ Q k .
Revised: Here, the inequality holds when n △ n ( k ) ≤ Q k .

Original: If otherwise, we can alwaysdefine ˜ Q k :=
Revised: If otherwise, we can alwaysdefine ˜ Q k :=

Original: max { Q k , n △ n ( k ) } and the inequality still holds by changing all the Q k into ˜ Q k .
Revised: max { Q k , n △ n ( k ) } and the inequality still holds by changing all the Q k into ˜ Q k .

Original: Wetemporarily omit this step for now.
Revised: Wetemporarily omit this step for now.

Original: Since | ∂ l f i ( x k, 0 ) − ∂ l f i ( x [Equation] | ≤
Revised: Since | ∂ l f i ( x k, 0 ) − ∂ l f i ( x k − j, 0 ) | ≤

Original: Q k −+ · · · +
Revised: Q k −+ · · · +

Original: Q k − j , the event 993 (cid:110) max i | ∂ l f i ( x k − j, 0 ) | ≥ (cid:80) kp = k − j Q p (cid:111) implies the event { max i | ∂ l f i ( x k, 0 )
Revised: Q k − j , the event 972 (cid:110) max i | ∂ l f i ( x k − j, 0 ) | ≥ (cid:80) kp = k − j Q p (cid:111) implies the event { max i | ∂ l f i ( x k, 0 )

Original: | ≥ Q k } , so theproof is completed.
Revised: | ≥ Q k } , so theproof is completed.

Original: Now we are ready to bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .
Revised: Now we are ready to bound E (cid:104)(cid:80) l large ( b 1 ) (cid:105) .

Original: Combining Step 1, 2 and 3 together, we prove thefollowing Lemma G.4.
Revised: Combining Step 1, 2 and 3 together, we prove thefollowing Lemma G.5.

Original: Lemma G.4.
Revised: Lemma G.5.

Original: Under Assumption 2.1, consider β 1 < √ β 2 , when k is large such that: k ≥ 4 ;β ( k − 1)
Revised: Under Assumption 2.1, consider β 1 < √ β 2 , when k is large such that: k ≥ 4 ;β ( k − 1)

Original: n≤ β n 1 √ k − 1 , we have the following result:
Revised: n≤ β n 1 √ k − 1 , we have the following result:

Original: [Equation]
Revised: [Equation]

Original: where δ 1 = (1 − β 2 )4 nρ 2 β n+ (cid:18) 1 √ β n 2 −(cid:19) , which goes to 0 when β 2 goes to 1.
Revised: where δ 1 = (1 − β 2 ) [Equation] , which goes to 0 when β 2 goes to 1.

Original: Proof can be seen in Appendix G.4.
Revised: Proof can be seen in Appendix G.4.

Original: Now we bound the error terms: E (cid:104)(cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:105) .
Revised: Now we bound the error terms: E (cid:104)(cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:105) .

Original: Since they are multiplied by δ 1 , these two terms vanish when β 2 → 1 .
Revised: Since they are multiplied by δ 1 , these two terms vanish when β 2 → 1 .

Original: Bounding E (cid:104)(cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:105) .
Revised: Bounding E (cid:104) (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:105) .

Original: We bound these two terms in the following Lemma G.5.
Revised: We bound these two terms in the following Lemma G.6.

Original: Lemma G.5.
Revised: Lemma G.6.

Original: Given all the history up to x k, 0 , we denote α as the index of the coordinate withthe greatest gradient: α = arg max l =1 , 2 , ··· ,d | ∂ l f ( x k, 0 ) | .
Revised: Given all the history up to x k, 0 , we denote α as the index of the coordinate withthe greatest gradient: α = arg max l =1 , 2 , ··· ,d | ∂ l f ( x k, 0 ) | .

Original: Under Assumption 2.1 and 2.2, considerβ 1 < √ β 2 , when k is large such that: k ≥ 4 ; β ( k − 1) n
Revised: Under Assumption 2.1 and 2.2, considerβ 1 < √ β 2 , when k is large such that: k ≥ 4 ; β ( k − 1) n

Original: [Equation]
Revised: [Equation]

Original: k − 1 , we have the following results:
Revised: k − 1 , we have the following results:

Revised: [Equation]

Revised: E 

Original: [Equation]
Revised: [Equation]

Original: where δ 1 is defined in (16) , D 0 is defined in Assumption 2.2.
Revised: where δ 1 is defined in (14) , D 0 is defined in Assumption 2.2.

Original: Detailed proof of Lemma G.5 is shown in Appendix G.8.
Revised: Detailed proof of Lemma G.6 is shown in Appendix G.8.

Original: Bounding E (cid:110)(cid:80) l large ( a 1 ) + (cid:80) l large ( b 1 ) − (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:111) .
Revised: Bounding E (cid:110)(cid:80) l large ( a 1 ) + (cid:80) l large ( b 1 ) − (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:111) .

Original: With the help of Lemma G.4and G.5, we can bound E (cid:110)(cid:80) l large ( a 1 ) + (cid:80) l large ( b 1 ) − (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:111) .
Revised: With the help of Lemma G.5and G.6, we can bound E (cid:110)(cid:80) l large ( a 1 ) + (cid:80) l large ( b 1 ) − (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:111) .

Original: The intuition is asfollows:
Revised: The intuition is asfollows:

Original: • ( a 1 ) is in the order of | ∂ l f ( x k, 0 ) | , multiplied by some positive constant;• (cid:80) l large { ( a 2 )
Revised: • ( a 1 ) is in the order of | ∂ l f ( x k, 0 ) | , multiplied by some positive constant;

Original: + ( b 2 ) } vanishes when β 2 → 1 .• E (cid:110)(cid:80) l large ( b 1 ) (cid:111) vanishes when β 2 → 1 and k goes to infinity Therefore, E (cid:110)(cid:80) l [Equation]
Revised: • (cid:80) l large { ( a 2 ) + ( b 2 ) } vanishes when β 2 → 1 .

Original: [Equation] + [Equation] (cid:80) l large 
Revised: • E (cid:110)(cid:80) l large ( b 1 ) (cid:111) vanishes when β 2 → 1 and k goes to infinity.

Original: [Equation] + ( b [Equation] } (cid:111) is still in the order of | ∂ l f ( x k, 0 ) |(multiplied by a positive constant when β 2 is large).
Revised: Therefore, E (cid:110)(cid:80) l large ( a 1 ) + (cid:80) l large ( b 1 ) − (cid:80) l large { ( a 2 ) + ( b 2 ) } (cid:111) is still in the order of | ∂ l f ( x k, 0 ) |(multiplied by a positive constant when β 2 is large).

Original: More formal results are shown in Lemma G.6.
Revised: More formal results are shown in Lemma G.7.

Original: Lemma G.6.
Revised: Lemma G.7.

Original: Under Assumption 2.1 and 2.2, when the hyperparameters satisfy: i) β 1 < √ β 2 < 1 ,ii) β 2 is large enough such that A ( β 2 ) is small enough to satisfy (33) , where A ( β 2 ) is a non-negativeconstant that approaches 0 when β 2 approaches 1 .
Revised: Under Assumption 2.1 and 2.2, when the hyperparameters satisfy: i) β 1 < √ β 2 < 1 ,ii) β 2 is large enough such that A ( β 2 ) is small enough to satisfy (32) , where A ( β 2 ) is a non-negativeconstant that approaches 0 when β 2 approaches 1 .