Adam with model exponential moving average is effective for nonconvex optimization

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.

Paper

References (62)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer m9Ez4/10 · confidence 2/52024-07-08

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.

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. [1] Qinzi Zhang and Ashok Cutkosky. Random scaling and momentum for non-smooth non-convex optimization. In International Conference on Machine Learning. PMLR, 2024.

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. 2. 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? 3. Could the author explain the meaning of choosing such a comparator $u_t$ in Lemma 7. 4. 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})$. 5. Could the authors explain the details of how they get Corollary 12 from the results of the definition $T$ in Theorem 11. 6. 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? 7. 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.

Rating

4

Confidence

2

Soundness

3

Presentation

2

Contribution

2

Limitations

No limitations.

Reviewer m9Ez2024-08-13

I thank the authors for their response. I still have some resolved questions. Since the author claimed they had dropped the notation of $w_t$, which plays a significant role in their original proof, I can not check the correctness of their results in their new version. Despite this significant problem, I also feel confused about their explanations for the following points: - Since they drop the variable $w_t$, why they still need $\alpha_t$ for their further proof? At least in their original version, $\alpha_t$ was introduced in the definition of $w_t$, and now, I even don't know where is this $\alpha_t$ first introduced. - The authors might misunderstand my concerns about line 466. Actually, I want to know why $E[\nabla F(w_t)-g_t|z_t]=0$. - The authors's explanation of implications and motivations does not convince me, several definitions or settings, seems too far-fetched and specific. While the authors consider clipping a minor modification, I have never used it or seen it used by others in practice. Similar questions also exist for the definition of comparator $u_t$ and learning rate $\eta$. As I stated in my review, it seems only given such specific forms, the authors get an updating rules share a similar structure with Adam. However, I can not see any direct connections between these frameworks and Adam. Due to these questions, I have to keep my original rating score. However, I want to stress that I'm not familiar with online optimization. Specifically, my comments about these intermediate variables might be unfair and it's welcome to point it out with more convincing evidence. Moreover, I keep my confidence at 2 and if other reviewers feel more confident in their judgment, I would defer to them.

Authorsrebuttal2024-08-13

Thank you

- To clarify the new scheme, the update is now $\mathbf{x_t}= \mathbf{x_{t-1}} + \alpha_t \mathbf{z_t}$, where $\alpha_t$ is sampled from an exponential distribution (which is a very light-tailed distribution with mean 1 and variance 1). So $\alpha_t$ is still used. Via an elementary argument involving fundamental theorem of calculus and the pdf $p(\alpha) = \exp(-\alpha)$, we obtain $\mathbb{E}[F(\mathbf{x_t}) - F(\mathbf{x_{t-1}})] = \mathbb{E}[\langle \mathbf{g_t}, \mathbf{x_t} - \mathbf{x_{t-1}}\rangle]$. This identity says that the function gap is *exactly equal* to the linearization of the function gap. In a sense, the randomization makes the first-order Taylor approximation perfectly correct. In the current submission, we sample $\mathbf{g_t}$ at $\mathbf{w_t}$ to ensure this identity. However, changing $\alpha_t$ to be exponentially distributed and multiplying the update $\mathbf{z_t}$ by $\alpha_t$ means that this is no longer necessary. In short, although $\mathbf{w_t}$ plays a significant role, the entirety of that role is contained in establishing the identity $\mathbb{E}[F(\mathbf{x_t}) - F(\mathbf{x_{t-1}})] = \mathbb{E}[\langle \mathbf{g_t}, \mathbf{x_t} - \mathbf{x_{t-1}}\rangle]$. This other option keeps the same identity and so works just as well. That said, if the reviewers prefer to keep the results the same as in the submitted version for which all of the analysis is completely available, we are of course happy to do so. - Regarding why $\mathbb{E}[ F(\mathbf{w_t}) - \mathbf{g_t}|\mathbf{z_t}]=0$: this is is because $\mathbf{g_t}$ is a standard stochastic gradient oracle evaluated at $\mathbf{w_t}$ and as such $\mathbb{E}[\mathbf{g_t} | \mathbf{w_t}]=F(\mathbf{w_t})$. As a concrete illustration, suppose that $F=\mathbb{E}[f(x,r)]$, where $x$ is the model parameter, $r$ is a randomly sampled data point, and $f$ is the loss of the model on the data point. Then $\mathbf{g_t}$ can be taken to be $\nabla f(\mathbf{w_t}, r_t)$ for some i.i.d. sample $r$. This formalism exactly captures the standard approach used in training. In this case, since $r_t$ is independent of $\mathbf{w_t}$ and $\mathbf{z_t}$, we have $\mathbb{E}[ F(\mathbf{w_t}) - \mathbf{g_t} | \mathbf{z_t}]=0$. - Regarding the clipping: You're right that this is not too common in current practice. However, we feel this is more an analytical artifact than a real change in the algorithm because the clipping should not be active in most iterations. The reason for this is that intuitively as the algorithm convergence we expect the gradients to be dominated by noise so that they look roughly like mean-zero random variables. In this case, the clipping will not occur. On the other hand, if the clipping *does* occur frequently, then an alternative analysis based upon https://jmlr.org/papers/v18/17-079.html would actually show that we converge incredibly fast. This in turn would put us in the regime in which gradients are dominated by noise, and so clipping would stop occurring.

Reviewer DXTx5/10 · confidence 3/52024-07-12

Summary

The authors show that with clipping and model exponential average, (Clipping) Adam can perform better than SGD.

Strengths

Give the analysis of (Clipping) Adam with model exponential average.

Weaknesses

1. The paper is hard to follow. (i) In Algorithm 1, it would be better to add "output: $\bar{w}_T$" (ii) it is unclear what $v_s$ is in Definition 6. 2. All of the following results are based on Lemma 7, I am not sure whether Lemma 7 is suitable for SGD.

Questions

If change Lemma 7 with a different formulation will the conclusion change?

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

N/A

Reviewer BY487/10 · confidence 3/52024-07-13

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.

Strengths

**1.** The proposed algorithm is shown to attain the optimal convergence rate for both smooth and non-smooth non-convex problems. **2.** I like the discussion on the benefit of coordinate-wise updates. **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). **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.

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$? **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}$. **3.** Is the Lipschitz assumption necessary?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Discussed.

Reviewer DXTx2024-08-12

Thank you for the response I have raised my score.

Reviewer BY482024-08-12

Thanks for the rebuttal; I'll keep my score. Regarding the direction of the inequality, I believe my direction is correct (with exact inequality and not inequality up to constants) if all the $g_s$'s are non-negative (this can be seen by squaring both sides).

Authorsrebuttal2024-08-13

Thank you

You're right of course that if all the $\mathbf{g_t}$s have the same sign, then the clipping will occur. Our observation is that this is extremely unlikely to happen. The reason is the following: imagine that the algorithm is converging (even if it is converging somewhat slowly). In such a case, we would expect the gradients to be roughly mean-zero random variables (e.g. https://arxiv.org/pdf/2111.06328 or https://arxiv.org/pdf/1508.00882), in which case the clipping should be expected to not occur. In fact, the case in which clipping does occur frequently results in "too good to be true" results that should provide further intuitive evidence of it being unlikely. For example, https://jmlr.org/papers/v18/17-079.html shows that if the prediction of FTRL is clipped on every round then in fact we would suffer only *logarithmic* regret rather than the standar $\sqrt{T}$ regret bound, and http://proceedings.mlr.press/v97/cutkosky19b.html shows that in the case that $\sum_{t=1}^T \mathbf{g_t}\ge \Omega(\sqrt{T})$ (even if clipping doesn't occur every single round), we can obtain preconditioning "for free" without any extra computation cost.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC