Adam with model exponential moving average is effective for nonconvex optimization
Authors: Kwangjun Ahn
Ashok Cutkosky
Venue: NeurIPS 2024 poster (2024)
Keywords: Adam
exponential moving average
EMA
nonconvex optimization.
Primary area: optimization

Abstract: In this work, we offer a theoretical analysis of two modern optimization techniques for training large and complex models: (i) adaptive optimization algorithms, such as Adam, and (ii) the model exponential moving average (EMA). Specifically, we demonstrate that a clipped version of Adam with model EMA achieves the optimal convergence rates in various nonconvex optimization settings, both smooth and nonsmooth. Moreover, when the scale varies significantly across different coordinates, we demonstrate that the coordinate-wise adaptivity of Adam is provably advantageous. Notably, unlike previous analyses of Adam, our analysis crucially relies on its core elements---momentum and discounting factors---as well as model EMA, motivating their wide applications in practice.

Adam with model exponential moving average is effective for nonconvex optimization
Kwangjun Ahn Microsoft Research Cambridge, MA vranggjumaimmmlcooooff..cccc
Ashok Cutkosky Boston University Boston, MA ashok@cutkosky.com com
Abstract
In this work, offer theoretical analysis of two modern optimization techniques for training large and complex models: (i) adaptive optimization algorithms, such as Adam, and (ii) the model exponential moving average (EMA). Specifically, we demonstrate that clipped version of Adam with model EMA achieves the optimal convergence rates various nonconvex optimization settings, both smooth and nonsmooth. Moreover, when the scale varies significantly across different coor- dinates, we demonstrate that the coordinate-wei wise adaptivity of Adam provably advantageous. Notably, unlike previous analyses of Adam, our analysis crucially relies on its core elements -ooeentm and discounting factors as well as model EMA, motivating their wide applications in practice.
Introduction
In neural network training, the training loss F Rd R is often optimized using an iterative optimization algorithm which starts with the initial iterate xo and then updates during each iteration t = 1,2,... follows:
X; = = Ze
where denotes the increment chosen by the algorithm during the 1-th iteration. One of the most popular optimization algorithms is Adam [Kingma and Ba. 2014]. Adam has gained significant attention due to its effectiveness in training Transformer-based language models [Zhang et al., 2020a, Kunstner et al., 2023, Jiang et al., 2023, Pan and Li, 2023, Ahn et al., 2024a, Kunstner et al., 2024, Zhang et al., 2024a].
The model exponential moving average (EMA) [Polyak and Juditsky, 1992, Ruppert, 1988] is an optimization technique that has gained popularity in conjunction with Adam for various recent applications. EMA maintains an exponential moving average of the model iterates, Xt, which contributes to the stabilization of these iterates. There has been resurgence of interest in this technique due to its effectiveness training high-quality generative models [Yaz et al., 2018, Karras et al., 2019, Song et al., 2021b, Dhariwal and Nichol, 2021, Nichol and Dhariwal, 2021, Song et al., 2021a. la, Balaji al., 2022, Karras et al., 2022, Rombach et al., 2022, Kang al., 2023, Karras et al., 2023]. Moreover, recent work by Block et al. [2024] demonstrates the effectiveness of EMA for both language modeling and imitation learning applications.
In this work, we theoretically study the effectiveness of these two modern optimization techniques. Our main results can informally summarized as follows.
Theorem (Informal). CLIPPED-ADAN with the EMA its iterates achieves the optimal conver- gence rate for nonconvex optimization both for smooth and nonsmooth settings (Section 5). The coordinate-wise adaptivity of Adam is particularly effective when the scale varies across different coordinates (Section 6).
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Our main results are based on the online-to- coahe-to-noneee no--nononvex conversion framework of Cutkosky et al. [2023], which chooses the increment z, based on an online learner of choice. In particular, our approach is quite different than the previous analyses of Adam (see Section 1.1 below). Notably, our analysis relies on the key components of Adam (momentum and adaptive learning rate) as well as EMA of the iterates, offering new, theoretical insight into their success, See Section 7 for a more detailed discussion.
At a high level, our analysis combines the main insights from the two recent works: Zhang and Cutkosky [2024] and Ahn et al. [2024b]. We first carefully modify the corpuutted-tonooncooe conversion framework (Lemma 7) of Zhang and Cutkosky [2024] which converts an online learner that achieves good discounted regret (Definition 6) into a good noncovex optimizer. We then combine it with the main insight of Ahn et al. [2024b] that an effective discounted online learner can be designed based on scale-free Dollow-the-Regulaaie
Leader (FTRL) [Orabona and Pál, 2018]. In particular, the way we arrive at Adam is similar to Ahn et al. [2024b]: choosing a discounted version of FTRL in the discounted-tor sescounted-to-noonoovyee tonooconvee conversion leads Adam.
1.1 Related work
Even though Adam is widely used in deep learning, our theoretical understanding of its inner workings, especially the importance of its core components-moneee incomplete, as pointed out by Ahn et al. [2024b]. Most theoretical work Adam and its variations focus on characterizing the convergence rate for convex or smooth nonconvex functions, where methods like SGD already achieve the minimax optimal convergence rate. [Reddi et al., 2018, Zhou et al., 2019, Chen et al., 2019, Zou et al., 2019, Alacaoglu et al., 2020, Guo et al., 2021, Défossez et al., 2022, Zhang et al.. 2022, Li et al., 2023, Wang al., 2023]. In fact, even the most recent results [Li et al., 2023, Wang et al., 2023] are not reflective of practice, in the sense that Adam's 's convergence rate worsens with momentum [Wang et al., 2023, $6) or is no better than that of SGD [Li et al., 2023, $7]. A notable exception is Crawshaw et al. [2022], which demonstrates the advantages of momentum under the generalized smoothness conditions of Zhang et al. [2020a]. However, the algorithm they analyze is signSGD, which differs significantly from the original Adam. In contrast, as we will explore in the subsequent sections, momentum and discounting factors are important in our analysis. See Section 7 for more details.
and discounting factors-remains
We also highlight that our analysis relies on model EMA, technique widely used in practice as mentioned above (also see a recent work by Block et al. [2024]). It is worth noting that EMA (or model averaging in general) has shown to have generalization benefits in practice [Tarvainen and Valpola, 2017, Izmailov et al., 2018]. In this paper, we study EMA from an optimization perspective, and show that the of EMA leads to optimal guarantees for nonconvex optimization. Interestingly. EMA naturally derives from the discounted-to- -o-llnnne conversion (see Algorithm 1), which, we believe, provides new theoretical insights into this practical method.
The use of EMA also represents a significant departure from most non-convex optimization analyses. While EMA is classical technique in the convex setting, theoretical analyses in the non-convex setting typically randomly select an iterate as the "final output" of the optimizer, rather than using EMA. This random selection is intuitively extremely impractical (indeed, average it actually wastes half of the computation), and is never performed in real implementations.
Our analysis follows a line of work studying convergence guarantees for non-smooth non-convex optimization. Our particular convergence criterion similar to finding the Goldstein stationary points [Goldstein, 1977] that were first studied in the context of modern machine learning by [Zhang et al., 2020b), and has seen much subsequent interest [Tian and So, 2022, Jordan et al., 2023, Davis et al., 2020]. Other notions of convergence are also massonable-comme envelope, or imposing a weak convexity condition [Davis et al., 2018, 2022a].
alternatives involve the Moreau
2 Setting for nonconvex and nonsmooth optimization
Throughout this paper, unless specified otherwise, ||-I| denotes the L2 norm. Following [Cutkosky et al., 2023], we consider optimizing loss function F that satisfies the following conditions, accessing information about F through a stochastic gradient oracle STOGRAD: : Rdd - Rdd for the set of randomness Z.

---

Assumption 2. Let F: :R + R be differentiable function with the following properties:
Let A = F(xo) - infx F(x).
For any two points x and y, F(y) - F(x)= = So (VF(x + t(y x)).y x) dt.
Lipschitzness. F is G-Lipshitz, for any point x, VF(x)] G.
Stochastic gradient variance. For any point x, the stochastic gradient +- STOGRAD(X, r) for randomness E 2 satisfies E(g] = Fx and E lig - VF(x)] oo.
The Lipschitz continuity condition is standard assumption in nonconvex nonsmooth settings. However, as we will discuss Section 6, one of the key insights from our results is that Adam enables us to adapt to the Lipschitz constants coordinate-wis wise without requiring prior knowledge of these constants. Note that we almost certainly need some form of structural assumption on the "difficulty" of the loss function; thus, relaxing the Lipschitz assumption would likely come at the cost of another assumption, such as smoothness.
The second condition, called well-behavedess in [Cutkosky et al., 2023, Definition 1J, is mild regularity condition. For any locally Lipschitz function F. applying an arbitrarily small perturbation to the function is sufficient to ensure this condition [Cutkosky et al., 2023, Proposition 2].
For the notion of optimality, follow Zhang and Cutkosky [2024] and consider the following notion of stationarity for nonconvex and nonsmooth functions. This notion is a slight relaxation of the notion of Goldstein stationarity point [Goldstein, 1977], which was further studied by recent works [Zhang et al., 2020a, Davis et al., 2022b, Tian et al., 2022, Jordan et al., 2023].
Definition 3 ((A, stttttionaayy point). Suppose F Rd R is differentiable. We say is (X, e) -ttatioinyy point of if (VF)))) s, where
[[F(()))
inf )EP(R"). Fy--py]=x
{IE(VF())]I+ +X A E x xI²}
To further motivate this definition, we remark that -sstatitti points retain the desirable prop- erties of Goldstein stationary points. Specifically, the following result [Zhang and Cutkosky, 2024, Lemma 2.3] demonstrates that, akin to Goldstein stationary points, (X, e)sstationary points can be reduced to first-order stationary points with appropriate choices of À when the objective function is smooth or second-order smooth.
Lemma 4. If F L-smooth, then an (L'e-1, "e)-",,)-tattonan point x of F satisfies F < 2s. Moreover, if F is H seccond-odder-smooth then an (H/2,)) point x of F satisfies [[)]( S 28.
Moreover, as shown by [Zhang and Cutkosky, 2024, Lemma 2.4), (A, -Statssta points can also be reduced to Goldstein stationary points when F is Lipschitz.
Lemma Suppose F is G-Lipschitz. For any X, 20
(X, -statioaa point is (6,e")-
Goldstein stationary point, where s' = x -e.
Now we design algorithms that find (A, -statsnta points efficiently.
3 Discounted--o-noocorree conversion: online learning of increments
Our main results are built on the malin-to-nnoccovve conversion framework of Cutkosky et al. [2023]. At its core, this framework involves selecting the increment z, using online learner, as discussed by Ahn et al. [2024b]. Specifically, we follow a variant developed by Zhang and Cutkosky [2024], which carefully incorporates the discounting factor in the conversion process. Note that make slight modifications to the version proposed by Zhang and Cutkosky [2024] as follows. Here Exp(1) denotes the exponential random variable with mean 1.
Given Algorithm 1, it turns out we need to design an online learner that minimizes the discounted regret, formally defined below. It is worth noting that discounted regret has been recently studied with the goal of better adapting online learners to dynamic environments [Ahn et al., 2024b, Zhang et al., 2024b, Jacobsen and Cutkosky, 2024].

---

Algorithm Discounted--o-nnncoovxx conversion (choosing increments via online learning)
Input:
Initial point E N, online learning algorithm A, and discounting factor , E (0, 1)
for t= 1,2...:7 :....T... do
Receive 2, from A // choose the increment using an online learner
Update X; f- x + a, where at Exp())(1) i.dd.
Compute gr - STOGRAD(Xx, Tr) with freshly sampled randomness rt
Send (! (z) = (5- gt. z) to A
// Maintain exponential moving average (for output only): : (Equivalently, x, I- E.=1 8 xx,)
Update X, X- X T X
end for
Definition 6 (Discounted regret). For comparator u, the B-discounted regret is defined as
Regret, (u)
|0] (u) == 3 Σ (6.9 (z.)
(.9 uu)
8
go
u)
The discounted regret of an online learner A can be used to upper bound the norm of averaged gradients, as shown in the following result.
Lemma (Discounted conversion), Suppose that F satisfies Assumption 2. Then for the comparator sequence chosen as U, == D Σ: F((x)
Algorithm gives
Ei VPiix.
A
2G
E
E
E
S
+
3yT + 1oVI-3
1-
DT
(1-
1
Regreth 10) (ur) (1 (1-33 (11-3) B) Σ Regret [o] (m))
+
8 E
DT
where is distributed over {x,)_=1 as Ply: = x.) = gt-s 1-8 T-Br for 8 = 1, ...
The proof combines the techniques of [Cutkosky al., 2023, Theorem 7] and [Zhang and Cutkosky, 2024, Theorem 3.3]. See Appendix A for details.
We briefly explain how Lemma can be used find a (X, e)-stationary point (Definition 3). Recall that (X, e )-stationnarit essentially requires producing point x Ely] such that both |[[VF(y) and E Elyy- x are small.
Given this context, Lemma states that long as the discounted regret of the online learner A is small, we can ensure that the EMA iterates X, = [8[y:] serve good candidates for (A, e -)-stationarity. since the term E |[E, Fy:))) can be kept small. The remaining task is bound the variance term, E E - x which will be addressed later in Lemma 10.
Moreover, the comparator roughly models the update direction that an oracle algorithm with perfect knowledge of the loss would select. In the proof of Lemma 7, we demonstrate that moving along the u, direction effectively decreases the loss value, which forms the basis for establishing convergence guarantee.
Thanks to the discounted- coaetted--t--onnccoo nonconvex conversion, the task now reduces to designing online learner that achieves low discounted regret.
Scale-free Follow-the-Reguaaiiee--Leee (FTRL)
In this section, we introduce an algorithmic component, called the wwwp:wedd--e---ggaaaii-eee (FTRL), a powerful online learning technique with various applications [Gordon, 1999, Kalai and Vempala, 2005, Shalev-Shwartz and Singer, 2006, Abernethy al., 2008, Nesterov, 2009, Hazan and Kale, 2010].

---

For the setting, consider the online linear optimization (OLO) setting, where during each round t= I....... T. and online learner chooses 2₁, and then the linear loss 4(-) (v.-) is revealed by the environment. Here the goal of the online learner is to minimize the regret defined as E, (v,, u),
where u1 is the comparator in hindsight. For this setting, FTRL presented in Algorithm
Algorithm 2 Nollow-the-Reeguariizee
(FTRL)
Require: Regularizers {vr IR the domain C Rd
1: for = =,,,.... do
2: Update 2, f- argmin ZED p/ (2) + + (x(x))
3: Receive the next loss (x(+) (Ver.)
4:
end for
The key insight of Ahn et al. [2024b] and Zhang et al. [2024b] is that in order to design an online learner for discounted regret, is important that the online learner scale-free described below. In particular, following Ahn et al. [2024b], we consider gradient adaptive scale-free FTRL algorithm called scale-free FTRL [Orabona and Pál, 2018].
We will focus on the case where D = Pp the d-dimensional Iz-ball of radius D > 0. Scale-free FTRL is given by Algorithm 2 with the following choie:
1
D
t (-)
11-11 and
"
V Eiillval'
Then using the clipping operator clipp( (x) x min( (1/1x 1), write down the update rule more explicitly as follows:
1
Σ ==
Ze f- argmin xcBp
2
+
(v., z) cllpp D
(SCALE-FREE FTRL)
7/
6=1 |v
Here, if the denominator is zero, i.e., VI = Vt-1 O, then we set + O. Note that this algorithm scale-free in the sense that when the loss sequence scaled by scalar c > 0, the updates remain the same.
Let us now present the regret bound of SCALE- FTRL.
Lemma 8 (Gradient-adaptive regret bound). For any > O, loss sequence VA.T and comparator E Rd s.f. ||u|| < D, SCALE-FREE FTRL guarantees the following regret bound:
T
T
(v.. Σ
u) S 4D
Σ
|
We note that Lemma 8 follows (with slightly worse constant) from [Orabona and Pál, 2018, Theorem 1]. and the version we invoke is here due to [Ahn et al., 2024b, Theorem A.
Recall from Lemma 7 that online learner for the discounted--o--noncosse conversion (Algorithm 1) needs to have low discounted regret. To achieve this, following Ahn al. [2024b] and Zhang et al. [2024b], we simply substitute V. + B into SCALE- FTRL, resulting in the update
Z clipp D Σ
g-gg. g.
(3-FTRL)
2-
3
g+
Here again, if the denominator is zero. i.e., g1 following result characterizes the discounted regret guarantee of 3-FTRL.
g:-1 O. then we set 0. Then, the
Theorem 9 (Discounted regret bound). Let B E (0, 1}. For any T > 0, loss sequence gi.r and comparator E gd s.t. |u|| < D, 3-FTRL guarantees the following static regret bound
Rf (u) S 4D
( |g
We next use this result to design algorithm for nonconvex optimization.
5

---

Discounted-FTRL leads to adaptive nonconvex optimization
In this section, warm-up. let see the implications of choosing A B-FTRL in Algorithm 1. First, let us obtain a bound on the expected discounted regret. By Theorem 9 together with Jensen's inequality, have the following regret bound for any 1, .......:2.
T
T Σ
E Regret,
( (u.)
& AD E
pa(-
g 2 <DD
E
Since "|gl:1 S G2 + oz and
V1-3 it follows that
V1-82
| Regret (u.)
4D +
4D(G+
E
(1)
V1-
VI
5.1 From gradient-adaptive regret to nonconvex optimization
In order to obtain nonconvex optimization guarantees in terms of the (x,,)-stationane (Definition 3), we need to handle the variance term. Following [Zhang and Cutkosky, 2024, Lemma 3.2], the variance term can be bounded as follows.
Lemma 10 (Variance bound). Using the notations of Lemma for any L 1.2,..., T, 3-FTRL satisfies
D
E k
Ely. - x11 -x,1 12
(1 -1-
Proof. From [Zhang and Cutkosky, 2024, Lemma 3.2]. it follows that E EΣ lly, - x, I 12
(1-9)a E EΣ |1 Now since ||z,|| D for all t 1.2..... T, after dividing each side by T. we get the desired inequality.
Plugging the regret bound (1) into Lemma 7 and combining it with Lemma 10, we arrive at the following optimization guarantee in terms of the (A, -)ttttiionn See Section B. for proof.
Theorem 11. Suppose that F satisfies Assumption 2 and consider any X > For > 0, choose A = 8- 8-FTRL in Algorithu 1 with the following parameters:
-11) D
B1= and T
1
4/
12C
β = 8-1-
10C
421/2
1-3
c3/2
Then we have En- (F [
(1+ ( In other words, randomly chosen model EMA X, (>, (1 G+a )e)- stationary point, in expectation.
5.2
Optimality and gradient adaptivity
Here, we discuss several notable aspects the guarantee provided in Theorem 11.
5.2.1
Optimality
As shown in [Zhang and Cutkosky, 2024, Corollary 5.1]. the lower bound on the iteration com- plexity for finding (A, (Attt-stationary point is Q((G + -x12227777/ provided that À < G4 Theorem 11 implies that setting C = G + a achieves this optimal iteration complexity.
A
Corollary 12. In Theorem 11, choosing C G + leads to the following iteration complexity for finding a (X, stttttioonyy point:
(G (G+o)AA
(G (G+
O max
g//2
In particular, treating G, and A as constants, as long X 2 E, this leads to the optimal complexity of of 1

---

In light of Lemma 4, the above optimal complexity can be converted into the optimal complexities for smooth settings.
Corollary 13 (Smooth settings). Corollary 12 implies the following optimal iteration complexity for smooth settings. Choosing X = O) implies the optimal complexity of (( ) for smooth loss functions [Arjevani et al., 2023]. Similarly, with X = O(1), it achieves the optimal iteration complexitiy of of (0(== for second-order smooth loss functions [Arjevani et al., 2020J.
of
We next discuss the benefits of using the gradient-adaptive regret bound (Theorem 9) by considering the case where we do not have knowledge of G.o. o.
5.2.2 Gradient adaptivity
A remarkable consequence of Theorem 11 is that, due to the gradient-adaptin regret bound of Theorem 9, the final convergence guarantee has better dependence on G. a in the case when we do not have knowledge of them. For concreteness, in the following discussion, we treat G, 0, A as constants, and focus on the regime X ž e.
First, our Theorem 11 with C' = 1 (since we do not know o, leads to the following iteration complexity for finding a (A, -sttatnttt point:
o (IG ((++o)?AANV2=-11)
The price we pay for not knowing G, relative the lower bound is a multiplicative factor of (G + oo)+0)3/2 o/3/2, To see the benefit of this adaptive regret approach, let us consider the guarantees given by Zhang and Cutkosky [2024]. Their approach is based on choosing online gradient descent for A in Algorithm I, when the learning rate is not properly tuned with the knowledge of G and a, it would
lead to the following (suboptimal) discounted regret bound:
[8] (u:) So DIG (DIG+ o0)
E R
VI-3 (
Then, the resulting iteration complexity becomes 0(2A1/(( ((-a)))) )-72), which is equal to
O + ((G=o))'A\)) ->12/>//2--777 This is larger than the complexity due to adaptive approach by
a multiplicative factor of (G + 7/7/
Next, we build on the results from this section and consider a better approach to design an adaptive nonconvex optimizer.
6 Coordinate-wise adaptivity via (clipped-)Adan)
In this section, we consider the setting where the Lipschitzness constants vary across different coordinates, which is empirically observed to be reflective of practical neural network training (see, e.g. [Crawshaw et al., 2022, Zhuang et al., 2022]). Formally, we consider the following setting. Assumption 14. Under the same setting as Assumption 2, we replace the last two conditions with the following coordinate-wise version:
For each coordinate = 1.2.....4 d, there is a Lipschitzness constant G; > and a variance constant o; >> 0 such that Vx, |0, [x)) < G; and the stochastic gradient g + STOGRAD(X: satisfies E[g[i]] Fx,) and EE g 0, Fx)) s of (Here, O,F denotes the partial deriva- live of F w.r.f. the i-th coordinate.)
Let G == (G1,..., Ga) and o (o1...., Oa). Then, the above condition implies the last two conditions in Assumption with G = |GG||, and o 1
As we mentioned before, the previous approaches [Cutkosky et al., 2023, Zhang and Cutkosky, 2024] choose the online learner to be online gradient descent, and hence choosing the learning rate requires the knowledge of G, o; for all i. However, for neural network training, d is equal to the number parameters the network, so tuning them individually computationally infeasible. We instead consider running 8-FTRI. coordinate-wist in Algorithm 1, which will automatically adapt to each coordinate. We begin with an important observation that such an approach fact leads to popular optimizer widely used in practice.
7

---

6.1 Coordinate-wise discounted FTRL corresponds to (clipped-)AAamm
For notational simplicity, a coordinate among ...,.... d, and let us denote the iterate by dis the stochastic gradient by gt, and the update by 24- Then the resulting optimizer becomes:
E. $i-"0 go 2++1 = -lipp D
(CLIPPED-ADAM)
E' $2 92
where B1 = 8 and B2 = _ Here, again if the denominator is zero, i.e., if 91 .....-.......... 9r 0. then we set the update be zero, i.e., 2(+1 O. Note that this almost exactly the Adam optimizer [Kingma and Ba, 2014], except that now we add clipping to control the variance of the iterates relative to their EMA. Notably, CLIPPED-ADAM retains one of the most important properties of Adam: it is mndee-ivvrrint. The scale invariance causes the optimizer to make updates of the same magnitude on each coordinate even when the scale differs across different coordinates.
In practice, we expect that the clipping operation will effectively be a no-op. This is because, when the algorithm is converging (even if the convergence somewhat slow), the gradients are likely to behave as approximately mean-zero random variables (due to factors such as stochastic noise, unstable training trajectories, etc.). such cases, standard concentration inequalities imply that (Bt- and hence, the clipping has no effect.
E 3t go S
We also remark that CLIPPED-ADI ADAM does not consider the "bias correction" terms in the original updates of Adam [Kingma and Ba, 2014]. However, note that the bias correction terms are coordinate- independent, and they can be merged into the scalar D.
6.2 Nonconvex optimization guarantees of CLIPPED-ADAM
We next discuss the theoretical guarantees CLIPPED- for nonconvex and nonsmooth opti- mizaton. Inspired by [Duchi et al., 2010, McMahan and Streeter, 2010], where the coordinate-we wise online learners lead to regret bounds with respect to the Lı norms of stochastic gradients, we consider the following variant of Definition in the same vein [Cutkosky et al., 2023, Section 4].
Definition 15 ((A, -staiooorr point). Suppose F: 12 + is differentiable. We say x is -itttioona point of F V & E, where
|[V
inf DEP(R"), Ey-ply]=x
(I +X Ely x1
Using the fact 11:111 S Vd I1-1122- one can connect the two notions of (A, e)-staioonay points. Lemma 16. A
point is (X, point.
In order to obtain the guarantee in terms of Lı-norm, consider the coordinate-wise version of discounted-to-onling conversion, in the same vein as [Cutkosky et al., 2023, Appendix G]. See Section A.1 for details.
Lemma 17 (Li-variant of Lemma 7). Suppose that F satisfies Assumption 14. Consider the 8- for i = 1,2,.... d.
comparator sequence chosen as defined as u/l0] - E. -0,[(x.) Then, Algorithm 1 gives
A
2|
E
E VF()))
<
|اl|| v1 /1-3
DT
(1- (1-B)1 B)T
T Σ
+
B R
Regret (8] (ur)) (ur))
+ +11 (1-13) (1-3))
F
Regret
M (m:)
DT
t=1
where Yt is distributed over {x Ply: Xs) = g
1-3 1- for 1,2, ...
8

---

Next, let us consider the (expected) regret bound. Fix a coordinate = 1....,d d. Then, by the one-dimensional version Theorem together with Jensen's inequality, have the following regret bound for any / = 1, ... T
Regret ! (w.( S 4D
T Σ
gac
E |
4D +o))
-1-
Hence, taking the sum over all coordinates
l,....d, d. we obtain
Regret [8] (w.)
4D G +all
E
(2)
VI-3
Combining these together, one get the following guarantee in terms of the L1 norm. See Section B.2 for proof.
Theorem 18. Suppose that F satisfies Assumption 14 and consider any X >> O. For C > O, choose the coordinate-wise optimizer CAIPPEDDADAN Algorithm with the following parameters:
(1- B) and T
1/2
12C
B B=
10C
4d2/2
1- max
e/2
E
Then we have Ei~11y E
)s. In other words, a randomly chosen model EMA X, a (A, (1 || )E) LL--stitionar point, in expectation.
A < (1+ | (
6.3 Benefits of coordinate-wise adaptivity CLIPPED-ADAM
In this section, discuss the benefits of coordinate- wise adaptivity by examining the guarantee from Theorem 18 and compare with that of Theorem 11. We begin with the (X, (x,e))La-stationaay point guarantee due to Theorem 18. We consider the scenario where 3 is carefully tuned by making the optimal choice of C.
Corollary 19. Theorem 18, choosing C ||G + oll1 leads to the following iteration complexity for finding (A, E)-E) -taatioary point:
|G + ol? /2
IG
max
77/2
In order to better appreciate the benefits of coordinate-wise adaptivity, let us compare the above iteration complexity with that of Theorem 11.
For concreteness, we treat G = |G]|| and a = ||اه| as constants throughout, and more importantly, we assume that the coordinates are heterogeneous in the sense that
|| + oll; ||G + +
(3)
The assumption (3) roughly says that a few coordinates of G + take much larger values than the rest; if all the coordinates of G have similar magnitudes, then ||G oll = VdllG a ll In the case Z e, Corollary 19 implies that the iteration complexity is
O[]G + ol A2/22
Next, let consider the counterpart that does not adapt to each coordinate separately. In this case, we apply Lemma 16, which tells us that to find a (A,=)-Lı stationary point it suffices to find a (2/Vd, eeddde//V))saaial -tttstatio point. Then, from Corollary 12, the iteration complexity is O(IG + o A( -7/)), i.e.,
O[|G + o 2/22
(
Hence, when (3) holds, (4) can be lower than (5) by a multiplicative factor of d, showing the benefits of coodiintte-wiss adaptivity.

---

7 Discussion
Our analyses of Adam based on the discounted-tor -o-olee conversion is quite different than the previous ones. As discussed in Section 1.1, the previous analyses often result in guarantees that are not quite reflective of practice-e,, the rates get better without momentum and the rates are no better than that of non-adaptive methods. In contrast, our analyses and results highlight the role of practical components as highlighted below.
Momentum, In order obtain a low discounted regret, any sensible online learner should integrate the past history of stochastic gradients gi:r to make the decision 24++1- Such online learners under the sescounted-to-onlin next increment 4++1- In particular, non-momentum methods would correspond to aggressive online learners that only use the last gradient g: to make the decision Z This perspective provides new insights into understanding the role of momentum, echoed by Ahn et al. [2024b].
conversion lead momentum methods that integrate gi:t to obtain the
Adaptive learning rates. The adaptive learning rate due to 3-FTRL leads to gradient-adaptive regret bound (Theorem 9), which important to obtain a better Lipshitzness dependence (Sec- tion 5.2) as well as the coordinate- wise adaptivity for nigh-dimension settings (Section 6.3). Our analysis offers theoretical benefits of adaptive learning rate from discounted regret perspective.
conversion (Algorithm naturally leads to guarantees in terms of the model EMA, X,. At high level (see Appendix A precise details), this is because for dynamic environment, it is important discount the losses such that online learners adapt to changing environments. The appearance of model EMA in the discounted-to toneedd-tooonccovee conversion provides new perspective on its role.
Model EMA. Lastly, the discountedd-t--nononyy
Our analyses and results have several limitations and raise several interesting questions. Firstly, CLIPPED-AD ADAM does not precisely match the original Adam algorithm, warranting further investiga- tion into the original Adam update. Specifically, our analysis suggests choosing 31 = 3 and 32 5 which does not align with the commonly used practical choices. Understanding the exact roles of these practical choices for 31 and B₂ would be valuable.
In Section 5.2, we observed that our iteration complexity for finding a (A, --ttatiiiia point is O + 0)7/221/26-7/2, when G and a are unknown. Investigating whether this complexity is optimal presents another intriguing direction for future research.
Lastly, from practical standpoint, developing more advanced online learner for discounted regret and designing algorithm that surpasses Adam in practicality would have significant practical implications.
Funding Acknowledgments
AC is supported by NSF grant number CCF-2211718.
References
Jacob D. Abernethy, Elad Hazan, and Alexander Rakhlin. Competing in the dark: An efficient algorithm for bandit linear optimization. In Conference on Learning Theory (COLT), pages 263-274. Omnipress, 2008
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention (maybe) all you need (to understand transformer optimization). In The Twelfth International Conference Learning Representations, 2024a. URL https::/openreet net./forummddd-0u1541111rr
Kwangjun Ahn, Zhiyu Zhang, Yunbum Kook, and Yan Dai. Understanding Adam optimizer via online learning updates: Adam is FTRL in disguise. In International Conference on Machine Learning. PMLR, 2024b.
Ahmet Alacaoglu, Yura Malitsky, Panayotis Mertikopoulos, and Volkan Cevher. A new regret analysis for Adam-type algorithms. In International conference on machine learning (ICML), pages 202-210. PMLR, 2020.
10

---

Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Ayush Sekhari, and Karthik Sridharan. Second-order information non-convex stochastic optimization: Power and limitations. In Conference on Learning Theory, pages 242-299, 2020.
Yossi Arjevani, Yair Carmon, John Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming, 199(1): 165-214, 2023.
Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Qinsheng Zhang, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, et al. eDiff-I: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01322 2022.
Adam Block, Dylan Foster, Akshay Krishnamurthy, Max Simchowitz, and Cyril Zhang. Butterfly effects of SGD noise: Error amplification in behavior cloning and autoregression. In The Twelfth International Conference on Learning Representations, 2024. URL attps://openrevieee net/ forumiid=CgPs0419900
Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of A class of Adam-type algorithms for non-convex optimization. In International Conference on Learning Representations (ICLR), 2019.
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in neural information processing systems, 35:9955-9968, 2022.
Ashok Cutkosky, Harsh Mchta, and Francesco Orabona. Optimal stochastic non-smooth non-convex optimization through online-to- toon---t--on-ccnwee conversion. In International Conference on Machine Learning (ICML), volume 202, pages 6643-6670. PMLR, 2023.
Damek Davis, Dmitriy Drusvyatskiy, Kellie J MacPhee, and Courtney Paquette. Subgradient methods for sharp weakly convex functions. Journal of Optimization Theory and Applications, 179:962-982, 2018.
Damek Davis, Dmitriy Drusvyatskiy, Sham Kakade, and Jason D Lee. Stochastic subgradient method converges tame functions. Foundations of computational mathematics, 20(1): 20(1):119-154 194, 2020.
Damek Davis, Mateo Díaz, and Dmitriy Drusvyatskiy. Escaping strict saddle points of the moreau envelope in nonsmooth optimization. SIAM Journal on Optimizaaiin,, 32(3):1958-1988 2022a.
Damek Davis, Dmitriy Drusvyatskiy, Yin Tat Lee, Swati Padmanabhan, and Guanghao Ye. A gradient sampling method with complexity guarantees for lipschitz functions high and low dimensions. In Advances in Neural Information Processing Systems, volume 35, pages 6692-6703. Curran Associates, Inc., 2022b.
Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of Adam and AdaGrad. Transactions on Machine Learning Research (TMLR), 2022.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
J. Duchi, F. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. In Conference on Learning Theory (COL.T), pages 257-269, 2010.
AA Goldstein. Optimization of lipschitz continuous functions. Mathematical Programmingg 13: 14-22, 1977.
Geoffrey J. Gordon. Regret bounds for prediction problems. In Conference Learning Theory (COLT), pages 29-40. ACM, 1999.
Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianbao Yang. A novel convergence analysis for algorithms of the Adam family. arXiv preprint arX1v:2112.03459, 2021.
Flad Hazan and Satyen Kale. Extracting certainty from uncertainty: Regret bounded by variation in costs. Machine Learning, 10:165-188, 2010.

---

Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pages 885. Association For Uncertainty in Artificial Intelligence (AUAI), 2018.
Andrew Jacobsen and Ashok Cutkosky. Online linear regression in dynamic environments via discounting, In International Conference on Machine Learning. PMLR, 2024.
Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li. How does adaptive optimization impact local neural network geometry? In Advances in Neural Information Processing Systems (NeurIPS), 2023. URL
htttps://openrevieee www//teeuortteeee dig.diddggIIIIaaaa
Michael Jordan, Guy Kornowski, Tianyi Lin, Ohad Shamir, and Manolis Zampetakis. Deterministic nonsmooth nonconvex optimization. In The Thirty Sixth Annual Conference on Learning Theory, pages 4570-4597. PMLR, 2023.
Adam Kalai and Santosh Vempala. Efficient algorithms for online decision problems. Journal of Computer and System Sciences (JCSS), 71(3):291-37,, 2005.
Minguk Kang, JunJJunnaa Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung Park. Scaling up gans for text-to-image synthesis. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 10124-10134, 2023.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4401-4410, 2019.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in Neural Information Processing Systems, 35:26565-26577, 2022.
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. arXiv preprint arXiv:2312.02666, 2023.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2014.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between SGD and Adam on Transformers, but sign descent might be. In International Conference on Learning Representations (ICLR), 2023.
Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Hcavy-tailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint arXiv:2402. 0eti:2422.9444 19449, 2024.
Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie. Convergence of Adam under relaxed as- sumptions. In Advances in Neural Information Processing Systems (NeurIPS), 2023, URL https://openreview
H. Brendan McMahan and Matthew Streeter. Adaptive bound optimization for online convex optimization. In Proceedings of the 23rd Annual Conference on Learning Theory (COLT), pages 256, 2010.
Yurii Nesterov. Primal-dual subgradient methods for convex problems. Mathematical programming 120(1):221-259, 2009,
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models.
In International conference on machine learning, pages 81 8162-8171. PMLR, 2021.
Francesco Orabona and Dávid Pál. Scale-free online learning. Theoretical Computer Science, 716: 50 69, 2018.
Yan Pan and Yuanzhi Li. Toward understanding why Adam converges faster than SGD for Trans- formers. arXiv preprint arXiv:2306.00l 00204, 2023.
12

---

Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM journal on control and optimization, 30(4):838 855, 1992.
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference Learning Representations, 2018.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pages 10684-10695, 2022.
David Ruppert. Efficient estimations from slowly convergent robbins-monn monro process. Technical report, Cornell University Operations Research and Industrial Engineering, 1988.
Shai Shalev-Shwartz and Yoram Singer. Online learning meets optimization in the dual. In Conference on Learning Theory (COLTT), volume 4005, pages 423-437 Springer, 2006.
Jiaming Song, Chenlin Meng, and Stefano Ennon. Denoising diffusion implicit models. In Inter- national Conference Learning Representations, 2021a. URL https:///openrvvion .ett
foruniid-StliiarCCi
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b. URL https: /ttp::/://////ee/////// net/ met/forum?id PXTIG12RRES
Antti Tarvainen and Harri Valpola. Mean teachers are better role models: Weight -iegtt-averagd consistency targets improve semi-supervised deep learning results. Advances in neural information processing systems, 30, 2017.
Lai Tian and Anthony Man-Cho So. No dimension free deterministic algorithm computes approxi- mate stationarities of lipschitzians. arXiv preprint arXiv:2210.06907, 2022.
Lai Tian, Kaiwen Zhou, and Anthony Man-Cho So. On the finite-time complexity and practical com- putation of approximate stationarity concepts of lipschitz functions. In International Conference on Machine Learning, pages 21360-21379. PMLR, 2022.
Bohan Wang, Jingwen Fu, Huishuai Zhang, Nanning Zheng, and Wei Chen, Closing the gap between the upper bound and lower bound of Adam's 's iteration complexity. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Yasin Yaz, Chuan-Sheng Sheng Foo, Stefan Winkler, Kim-Hui Yap, Georgios Piliouras, Vijay Chandrasekhar, et al. The unusual effectiveness of averaging in GAN training. In International Conference Learning Representations, 2018
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient elipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Representations (ICLR), 2020a.
Jingzhao Zhang. Hongzhou Lin, Stefanie Jegelka, Suvrit Sra, and Ali Jadbabaie. Complexity of finding stationary points of nonconvex nonsmooth functions. In International Conference on Machine Learning, 2020b.
Qinzi Zhang and Ashok Cutkosky. Random scaling and momentum for non-smooth non-convex optimization. In International Conference on Machine Learning. PMLR, 2024.
Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. Adam can converge without any modification on update rules. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pages 28386-28399, 2022.
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why trans- formers need adam: hessian perspective. arXiv preprint arXiv:2402. 16788, 2024a.
Zhiyu Zhang, David Bombara, and Heng Yang. Discounted adaptive online prediction. In Interna- tional Conference on Machine Learning. PMLR, 2024b.
13

---

Zhiming Zhou, Qingru Zhang, Guansong Lu, Hongwei Wang, Weinan Zhang, and Yong Yu. Adashift: Decorrelation and convergence of adaptive learning rate methods. In International Conference on Learning Representations (ICLR), 2019.
Zhenxun Zhuang, Mingrui Liu, Ashok Cutkosky, and Francesco Orabona. Understanding AdamW through proximal methods and freen-freenes. Transactions on Machine Learning Research, 2022.
Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for conver- gences of Adam and RMSProp. In The IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pages 11127-11135, 2019.
14

---

Appendix
Proof of discounted--oonnonoovex conversion (Lemma 7)
15
A.1 Proof of the coordinate-wise version (Lemma 17)
17
Proof of main theorems
18
B.1 Proof of Theorem 11
18
B.2 Proof of Theorem 18
19
A Proof of discounted-tt-nnonoovvx conversion (Lemma 7)
Note first that via change of summation, we get
T
gn (Fx,) F(x,
gn (F)((())) F(x--1))
T
= Σ (1 8 (F(())))((x))) F(x--1))
1=1
F(XT) - F(xo) Σ
(T-(+((((x)) F(x:-1))
1=1
Rearranging the above together with the fact F(xo) - F(XT) S F(xo)
-infx F(x) =: A, we get
T
- A E
- (1) - -((x-))))((
+E
(T-(t((((x)) F(x--1))
B
We also recall the following fact about exponential random variable due [Zhang and Cutkosky, 2024, Lemma 3.1].
Lemma 20. Let a ~ Exp(A) for some À > 0, then
Fi az) F(x)) E az), (z) z)/Ax
By Lemma 20 with À 1, together with X, = x + a follows that F[ - F F))
(g:, 2,)]
This identity indicates that the function gap is exactly equal to the linearization of the function gap. In this sense, the randomization renders the first-order Taylor approximation perfectly accurate.
Now, with this result, we will address each term separately.
Analysis of A
Note that for each t < n, since X, Xt-1 a.Z, for a Exp(1), Lemma 20 yields
[x(((x)) F [x) = E. (V F Xxx kx)
E (VF (x)) un) + E. (VF (x), -u,) g: 2z
E (VF un)) (VF x
u + E (g:, 2i U
= E (VF u.)
E V (x
gi. u
E (g. Zt un
15

---

where the last line follows from the fact ([(()((x)) - g., z,)l = 0. More specifically, note that the randomness in the stochastic gradient oracle independent of the randomness due to a,. Since (())((x) g+) = 0, it follows that
-
(xx) ge zt)l = ((
gt), z =0 0,
where the inner expectation is with respect to the randomness in the stochastic gradient oracle and the outer is with respect to all other quantities.
Now let us handle each term.
Note that using the definition of yn. we have
8 Fxx,,, -D En 8r (xx)
E
3 (1 3) (VF()))) un) = -1
12 1== 3 - [F(x)]]
1=1
1=1
n
1 3
Σ
1-8 f 8t (x((
- E
9t (n-1x((x)), -D
1- Br
Br
Σ
gn 3V
I- -
- - - E
E (y
S DD E VF
+ DG3
Therefore, summing over = 1
T, we obtain:
DG
E
BT- - B) (VF(x)) un) S -DE Σ
E VVy)))
+ 1-3
2 For the second term, using Cauchy-Schwarte inequality, we have
8 (VF(x)) - ge: unn) S gm (VF(x) - -80)
E E
Using the bounded variance assumption on the stochastic gradient oracle, we have
32(n-t) [F(()) grl1 o² 32
3m (VF(x) gi)
Therefore, summing over I,..., T. and using the fact that following bound on the second term:
T-ge < T-B
we get the
T B)
oD
E
3" (1 3) (VF(x))
u,, -
Σ (1
SoDT SODTv/-2
-- B.
V1
B2
m=1
n=1
3 Lastly, for the third term, we have
T
T
E
ga 3) (g. -un) = B3) E
E 3nn--gg,2 un)
n=
t=
1=1
10 Regret (ux)
B)
Analysis of B
Note that for each f, since X; +- X + 0z for Exp())) Lemma 20 yields
E(F(x:)= - F(xx-1)) FF E (gr, 2x)
= (g., z, ut + E (gr. 1T) $ < (g., 2, -ur) + D(G
16

---

Thus,
E
8 (F(x))
F(x)-1))
3 13t BB.-"B..21 gi, -ur) + 3--
D(G+
3 ERegrett
(9) (ur)] +
1 1-3
Combining and B
A
Combining the above analyses and rearranging, it follows that
DE
(
DG SA+
1-3 + +oDT /1 -
Regret [0] (u
y+
D(G+
B EReggret 19] (ur)] 1 1-3
Dividing both sides by DT, we get the desired result.
A.1 Proof of the coordinate-wise version (Lemma 17)
The proof closely follows that of Lemma 7. In particular, with F(xo) iinf F(x), we have
- A <
g" (1 (F((x) F(x--1))
+E
Σ
TT (F
F(x(-1))
LI= t=1
1=1
B
We begin with the term
Using the same decomposition as before, we have
B[F(x,) - F(x, E (g., 2 u + E (g:, ur) = E (gr, - ur)
Σ
Eg Mullur
1=1
E (g:, 24 -ur) + D Σ
G (.)
Thus.
T
T
d
3 - ++1 (F(x)) F(x--1)
= BB
gi, 2: - -ur) + TT Σ D(G
1=1
+ D D d レ (G, +
S E Regret (uT))
Moving onto the we again use the same decomposition:
El F(x--)) 1)] (VF(())) u. + (VF(x)) ge, - u,,) + + (gi, U,))
As before, let us handle each term one by one separately.
17

---

Note that using the definition of yn, for each coordinate i = 1,..., d, we have
8 vF((x)) ) (-D)
- 1= 3 0, (x,))
E0
8 (1
(1
1=1 B Ba- 0 F(x)))
n 1 B
8" (VF(F(()))) )
E 1-
1-9 Rr (VF(x)
(1 - (1- EE
Σ
1 -gp
De 1-3 Bm- BF (F((x)
- - (*) EE E Fy DDE E 0,F((n) + DG,9"
Therefore, summing over = 1...., and then I,...
T, we obtain:
T
T
D D G
E
5th (1 B) (VF(x), un) - DE
Σ
E VF(y))
n=l
2
For each coordinate = 1,...,d.. d, we have
E
(0,F(xx) - g.(((- u. S
E
Σ 3" (0,F(x)) - g
E u M
Using the coordinate-wise bounded variance assumption on the stochastic gradient oracle,
af
E 8 (0,
- grlb)) E 32(n 0,F((x,) -gllil? S 1- 32
Therefore, summing over n = 1..... T, and using the fact that following bound on the second term:
1
S T-8 we get the
1-
T
T
D DE d
E
gm gn-t(1 B) (VF(x)) - gi, unn) S
(1
3)
2i=1 0i
V1 B2
t=1
R=1
d
< DT
Σ 9i i=1
V1 - 3.
3
We use the same manipulation as before:
n 9 gn 3) (g.. Ze u,)) = (1 B)E
Regret [3] (u:).
E
Combining the above, we get the desired result in Lemma 17.
B Proof of main theorems
B.1 Proof of Theorem 11
By Definition 3, since Bly.] = Xg, holds that
|V
V
+ +AE ly. -x11²)
WTT
Vi
yt
We begin with the second term (the variance term). By Lemma 10, we have
AD2
A E Elyy - x11 < 12 (1 3)2
18

---

Hence, by choosing D (1-8)=4/2 follows that A [eTrT] By, lly. - x,1 Se. 4X172
Next consider the first term (the norm of the averaged gradients). Plugging the regret bound (1) into Lemma 7, we get
E b-
E
EVF(Y)
A
+
4(G o0) tov 1- 3+
+4(G+
Y1
DT
(1-3)1 3)T
TVI-B
4AA1/2
6G + 50
+ (4G + 50) V 1 B
122 (1 -1- -YT
where the last line follows since VI-3 3 and D (1-3)e1/2
12-72 Choosing 8 = 1 Tnc )2, the
last term is bounded by C+a 20 e. Moreover, choosing T (1 (1-3)-3 max (44/2/2--3/2, 12Ca
the first and second terms are bounded by and (+a
2C" e, respectively. This concludes the proof.
B.2 Proof of Theorem 18
By Definition 15, since E[y:] = X,, it holds that
E [~22]
< E b-[1]
(VF()))
+ + lyy x, 12 Y1
We begin with the second term (the variance term). This time, given that now each coordinate of update z, is bounded by D, i.e., |2[0l| S D, applying the variance bound due Lemma 10 coordinate- wise implies:
A E E - x1112 - 12 AdD2 t~ 7
(1 3)2
Hence, by choosing D (1-8)=*/2
4/7/2x1/2: follows that E Ey, lly: xel1 Se.
Next consider the first term (the Lı-norm of the averaged gradients). Plugging the regret bound (2)
into Lemma 17, and doing similar manipulations the proof of Theorem 11, we get
VF(y))
<
A
6G + 50 1l
+ ||4G + 50 v1-B
2~ T Dc
DT
3yT
1/2
||6G + 50
+
+ ||4G + 50 11 v1 - B,
B/
(1 B)T
where the last line follows since D 13-/2 /217- Choosing 8 = 1 T0C: )² the last term is bounded by |G++oll Moreover, choosing T (1 B)- 11 max { 12 -3/2 12Cs-1), the first
2C
and second terms bounded by and |G 20
E, respectively. This concludes the proof.
19

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction accurately reflect the paper's contributions regarding the analysis of Adam with EMA for nonconvex and nonsmooth optimization.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine to include aspirational goals motivation long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The paper discuss the limitations of the work in Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
The assumptions and the proofs provided both in the main text and the appendix.
20

---

Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [NA]
Justification: This paper is a theory paper and does not have any experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-soures models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification:
This paper is a theory paper and does not have any experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
21

---

Please see the NeurIPS code and data submission guidelines Sttee/CCoeSuubiiscee ssiopolicyy for more details.
While we encourage the release code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https /nwps..ippaa/pp/// cc/public/ /eidees/CodeSubmt wwwieeSSummiiiiiiiioi
for more details.
The authors should provide instructions data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reprodueible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [NA]
Justification: This paper is a theory paper and does not have any experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: This paper is a theory paper and does not have any experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1 I-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of eтrors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
22

---

8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the computer resources (type of compute workers, memory, time of execution) needed reproduce the experi- ments?
Answer: [NA]
Justification: This paper is theory paper and does not have any experiments.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't make it into the paper).
Code Of Ethics
Does the research conducted in the paper conform, every respect, with the NeurIPS Code of Ethics htttps:/nneut ips.
Question:
Answer: [Yes]
Justification:
The research adheres to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., s if there a special consideration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper is theory paper, and we do not forsee any direct societal implications arising from the research.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to par- ticular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, improvement in the quality of generative models could be used generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
legitimate to point out that
The authors should consider possible harms that could arise when the technology being used intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
23

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper a theory paper and does not have any experiments.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The paper accurately credits the existing results by citing the papers that inspired the techniques used in this work.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets some datasets. Their licensing guide can help determine the license of a dataset.
curated licenses for
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's
creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification:
The paper does not introduce any new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
24

---

The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable), You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then much detail possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should paid least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Sub- jects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review.
25


Reviews:
[{"review_id": "7ZEjuqqGjH", "reviewer": "Reviewer_BY48", "summary": "This paper proposes a variant of Adam involving per-iteration clipping of the updates (which look like Adam's updates) and EMA-based weight averaging. More specifically, two versions of this variant are proposed -- one is a global (w.r.t. all the coordinates) version and one is a per-coordinate version. The two main technical tools used in developing the proposed algorithm are (i) the online-to-nonconvex framework of references Cutkosky et al. [2023] and Zhang and Cutkosky [2024] wherein the update directions are chosen with an online learner, and (ii) choosing the online learner to be scale-free Follow-the-Regularized-Leader (FTRL) based on reference Ahn et al. [2024b]. The proposed algorithm attains the optimal convergence rate for both smooth and non-smooth nonconvex settings (although w.r.t. a different notion of stationary points for the non-smooth case). In addition, the benefit of the per-coordinate version is discussed.", "questions": "**1.** The purpose of Lemma 10 is not immediately clear to me. Is it to replace $\\mathbf{y}_t$ in Lemma 7 with $\\overline{\\mathbf{w}}_t$?\n\n**2.** It seems to me that clipping is always activated in $\\beta$-FTRL and CLIPPED-ADAM(?) This is because with the choice of $\\beta_1 = \\beta$ and $\\beta_2 = \\beta^2$ in CLIPPED-ADAM for instance, $\\sum_{s=1}^t \\beta_1^{t-s} g_s \\geq \\sqrt{\\sum_{s=1}^t \\beta_2^{t-s} g_s^2}$. \n\n**3.** Is the Lipschitz assumption necessary?", "limitations": "Discussed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**1.** The proposed algorithm is shown to attain the optimal convergence rate for both smooth and non-smooth non-convex problems. \n\n**2.** I like the discussion on the benefit of coordinate-wise updates.\n\n**3.** The results of this paper also shed some light on why EMA-based weight averaging can be useful in practice (although it doesn't seem strictly important in the context of the proposed algorithm; see Weaknesses).", "weaknesses": "**1.** It appears to me that this paper's main convergence results (Theorems 11 and 19) can hold even with $\\mathbf{y}_t$ as defined in Lemma 7. Since Theorems 11 and 19 are *not* w.r.t. the last *averaged* iterate (which is what is used in practice), it seems that EMA-based averaging is not strictly necessary in the context of the proposed algorithms, i.e., the results hold even with the unaveraged iterates appropriately sampled (as described in Lemma 7).\n\n**2.** I understand that this is a theoretical paper but because a new algorithm has been proposed involving clipping and EMA-based averaging, it would have been nice to show some empirical results at least in some simple settings. This would have made the proposed algorithm more appealing and convincing to practitioners.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments! Let us address your questions one by one.\n \n- (Weakness 1  / Question 1) We indeed need to choose $\\mathbf{\\bar{w}}_t = \\mathbb{E}[\\mathbf{y}_t]$ as the output, and this is crucial to achieve a  $(\\lambda,\\epsilon)$-stationarity, our main notion of convergence. This notion of convergence basically requires to output a point $\\mathbf{w} = \\mathbb{E}[\\mathbf{y}]$ such that both $\\mathbb{E}[\\nabla F(\\mathbf{y})]$ and $\\mathbb{E}\\|\\mathbf{y}-\\mathbf{w}\\|^2$ are small. In other words, we need to find a point $\\mathbf{w}$ for which **one can find a set of points around $\\mathbf{w}$ whose averaged gradients is small**. In light of this, the role of the two lemmas are as follows:  \n\n   - **Lemma 7** gives the set of points $\\mathbf{y}_t$ that guarantees the smallness of $\\mathbb{E}[\\nabla F(\\mathbf{y}_t)]$, and hence we need to output its expectation $\\mathbf{\\bar{w}}_t = \\mathbb{E}[\\mathbf{y}_t]$ as the output.\n  - **Lemma 10** ensures that $\\mathbb{E}\\|\\mathbf{y}_t-\\mathbf{\\bar{w}}_t\\|^2$ is small; in other words, it ensures that $\\mathbf{y}_t$ is closely clustered around its average $\\mathbf{\\bar{w}}_t= \\mathbb{E}[\\mathbf{y}_t]$. \n\n- (Weakness 2) As discussed in the introduction section, Adam with EMA is very successful for large complex neural network models. Our main goal was to understand why such a method is very effective. We consider the clipping as a minor modification, and given that, our main message is justifying the success of Adam with EMA. Hence, we do not think running further experiments is necessary, because proposing a new algorithm is not the main scope of this work.   \n\n- (Question 2) Actually, the inequality likely goes the other way so that clipping is a no-op. In practice, we expect the $g_s$ sequences to be random (due to stochastic noise, unstable training trajectories etc), for which the standard concentration inequalities yield $\\sum_{s=1}^t \\beta^{t-s}g_s \\lesssim \\sqrt{\\sum_{s=1}^t (\\beta^{t-s}g_s)^2}$. Hence, in practical settings, we expect the clipping to be unnecessary.\n\n- (Question 3) The Lipschitzness condition is indeed a standard assumption in the nonconvex nonsmooth settings. However, we note that one of the main messages of our main results is that Adam lets us adapt to the Lipschitzness constants (**coordinate-wise**) without the knowledge of them. Note that we almost certainly would require *some* kind of structural assumption on the \"difficulty\" of the loss function, so if we were to dispense with Lipschitzness it would likely be at the cost of some other assumption (such as smoothness).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dDQyqTEQkJ", "reviewer": "Reviewer_DXTx", "summary": "The authors show that with clipping and model exponential average, (Clipping) Adam can perform better than SGD.", "questions": "If change Lemma 7 with a different formulation will the conclusion change?", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Give the analysis of (Clipping) Adam with model exponential average.", "weaknesses": "1.  The paper is hard to follow.\n\n    (i) In Algorithm 1, it would be better to add \"output: $\\bar{w}_T$\"\n    \n    (ii) it is unclear what $v_s$ is in Definition 6. \n\n2. All of the following results are based on Lemma 7, I am not sure whether Lemma 7 is suitable for SGD.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments. Please see the other responses for some clarifications to the mathematics that we will include in the final paper. We hope this will make the paper easier to follow.\n\nRegarding Lemma 7: it is indeed an important part of our analysis, and as we show it does indeed allow us to analyze Adam - in fact it arguably allows us to derive Adam from first principles. A similar argument using a less advanced online learning algorithm could be made to analyze SGD with momentum instead.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KBNJ5aIAvM", "reviewer": "Reviewer_m9Ez", "summary": "In this paper, the authors analyze the online learning framework 'discounted-to-nonconvex conversion', and propose a bound for the expectation of the exponential average of the output gradient. Moreover, they find that if the online learner follows a specific loss structure(FTRL), then the output of this online learner would follow a clip-Adam rules.", "questions": "1. In line 109, the description of the Algorithm 1, could the author explain the reason they introduce the notations $\\mathbf{x_t}$ and $\\mathbf{w_t}$, since from my perspective, $\\mathbf{x_t}$ represent the iterations of coefficients of the loss function. Then the stochastic gradient $\\mathbf{g_t}$ should also be computed by $\\mathbf{x_t}$, instead of $\\mathbf{w_t}$. I also wonder about the implication of $\\alpha_t$ in the definition of $\\mathbf{w_t}$, since it only benefits the further proof. \n\n2. I might be wrong, but there might be typos in the formula of Definition 6. According to the definition of $\\ell_t^{[\\beta]}(z)$ in Algorithm 1,  the RHS of formula might be $\\beta^t\\sum_{s=1}^t \\langle \\beta^{-s} \\mathbf{v_s},  \\mathbf{z_s} -\\mathbf{\\mu}\\rangle$. Meanwhile, I also wonder whether $\\mathbf{v_s}$ is actually \\mathbf{g_s} defined in Algorithm 1 since I never find the definition of $\\mathbf{v_s}$. Moreover, could the author explain the meanings and implications of such a definition?\n\n3. Could the author explain the meaning of choosing such a comparator $u_t$ in Lemma 7.\n\n4. I might be wrong, but as I checked the Lagrangian function of constrained optimization in line 137, I guess the RHS should be $-\\text{clip}_D(\\eta \\sum \\mathbf{v_s})$. \n\n5. Could the authors explain the details of how they get Corollary 12 from the results of the definition $T$ in Theorem 11. \n\n6. Could the authors explain why line 466 holds since I'm not sure about the independence among $\\mathbf{z_t}$ and $\\mathbf{g_t}$. Moreover, could the author explain about the randomness in their formulas and which variables they take expectation w.r.t during the proof?\n\n7. Could the author further discuss the implications of their main result? Since from my perspective, it looks like they use the exponential discounted loss and choose FTRL learner in Discounted-to-nonconvex conversion, then they could get an updating rules share a similar structure with Adam. However, this result is far-fetched, and I have no idea what it implies.", "limitations": "No limitations.", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The result of this paper is solid.", "weaknesses": "Since this paper is built upon the 'discounted-to-nonconvex conversion' framework proposed in [1], focusing on the perspective of online learning, many definitions, concepts, and details require further elaborations and interpretations, especially for readers who are unfamiliar with this topic. I will list all my confusions and questions in the 'Questions' part and update my rating according to the response from the authors.\n\n[1] Qinzi Zhang and Ashok Cutkosky. Random scaling and momentum for non-smooth non-convex optimization. In International Conference on Machine Learning. PMLR, 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed comments! We answer your questions one by one below.\n\n1. You are correct that $\\mathbf{g_t}$ is computed at $\\mathbf{x_t}$ in practice. **In our new version, we actually fix this issue**, and in our new main results, $\\mathbf{g_t}$ is computed at $\\mathbf{x_t}$, and there is no need to introduce $\\mathbf{w_t}$ anymore. This is based on the \"exponential scaling\" technique due to [Zhang and Cutkosky, 2024], and we will update this in our final version. \n\n   - The role of $\\alpha_t$ is for the identity $F(\\mathbf{x_{t}}) - F(\\mathbf{x_{t-1}}) \\leq \\mathbb{E}[\\langle \\mathbf{g_t},~\\mathbf{x_t}-\\mathbf{x_{t-1}} \\rangle]$ to hold, which is crucial for the discounted-to-nonconvex conversion to work. When $F$ is convex, this holds without the random scaling $\\alpha_t$, but it turns out for the nonconvex setting, we need this random scaling. In our final version, where we fix the afornmentioned issue regarding $\\mathbf{w}_t$, we choose $\\alpha_t$ to be an exponential random variable, inspired by [Zhang and Cutkosky, 2024].\n \n   - Regarding $\\alpha$ being beneficial for the further proofs, sorry it's an overload of notations; the $\\alpha$'s that appears in the later part of the analysis should've been a separate quantity. We changed $\\alpha$'s in the later part to $C$ in our final version.\n\n2. Thank you for catching the typos. It should be ineed $\\beta^t\\sum_{s=1}^t \\langle \\beta^{-s} \\mathbf{g_s},  \\mathbf{z_s} -\\mathbf{u}\\rangle$. We corrected it in our final version.\n\n3. The comparator $\\mathbf{u}_t$ roughly models the \"update direction an oracle algorithm with perfect knowledge of the loss would have chosen.\" In the proof of Lemma 7, we show that along such a direction, the loss value decreases, and that is how we could come up with the convergence guarantee.\n\n\n4. You are correct. Thanks for catching the typo.\n \n5. Let's start with $$T = \\frac{1}{1-\\beta} \\max\\{\\frac{5\\Delta \\lambda^{1/2}}{\\epsilon^{3/2}}, \\frac{12\\alpha}{\\epsilon}\\}$$ in Theorem 11. \nSInce $1-\\beta = (\\frac{\\epsilon}{10\\alpha})^2$ and $\\alpha=G+\\sigma$, it follows that\n $$T = \\frac{100\\alpha^2}{\\epsilon^2} \\max\\{\\frac{5\\Delta \\lambda^{1/2}}{\\epsilon^{3/2}}, \\frac{12\\alpha}{\\epsilon}\\} = O\\left(   \\max\\{\\frac{(G+\\sigma)^2\\Delta \\lambda^{1/2}}{\\epsilon^{7/2}}, \\frac{(G+\\sigma)^3}{\\epsilon^3}\\} \\right)$$\n as desired.\n \n6. Thank you for your question. We will detail the argument in our final version. In our theorem statements, the expectations are taken over both randomness in the gradient oracle as well as the internal randomness of the algorithm from choosing $\\alpha$. Note that the randomness in the stochastic gradient oracle is independent of the randomness due to $\\alpha_t$. Since $\\mathbb{E}[\\nabla F(\\mathbf{w_t})- \\mathbf{g_t}] =0$, it follows that\n$$\\mathbb{E}[\\langle \\nabla F(\\mathbf{w_t})- \\mathbf{g_t}, \\mathbf{z_t} \\rangle]   = \\mathbb{E}[ \\langle \\mathbb{E}[ \\nabla F(\\mathbf{w_t})- \\mathbf{g_t}], \\mathbf{z_t} \\rangle] = 0,$$\nwhere the inner expectation is with respect to the randomness in the stochastic gradient oracle and the outer is with respect to all other quantities.\nIn general, we will detail all the arguments regarding expectations in our final version.\n\n7. As discussed in the introduction section, Adam with EMA is very succesful for large complex neural network models. Our main goal was to understand why such a method is very effective. \nAlthough we do have a minor modification to the original Adam (such as clipping), our analysis shows that Adam with EMA achieves an optimal convergence rate, and it is particularly benefical when the Lipshitzness constants are unknown and heterogenous across different coordinates. \nWe believe our results to be a significant progress towards understanding the success of Adam (with EMA) since before our work, it is unclear from the theoretical standpoint as to why Adam is much more effective than SGD.\n\nWe would be glad to address any remaining concern!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper examines the benefit of model exponential moving average when the model is trained with Adam. In particular, the paper uses a discount-to-nonconvex conversion framework with a scale-free FTRL learner to obtain non convex optimization guarantees. This resultant algorithm resembles a clipped variant of Adam with model EMA, thus providing optimization guarantees for it. While there are differences with the version of Adam used in practice, the results are nevertheless insightful. The reviews for the paper were mostly positive except for one reviewer with low-confidence providing a lower score. While some of their concerns were valid, I believe the rebuttal mostly addresses them. I think this paper will be valuable addition to the conference. I recommend acceptance.
Decision: Accept (poster)