Adam with model exponential moving average is effective for nonconvex optimization

Thank you for your detailed comments! We answer your questions one by one below.

1. 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.

   • 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].

   • 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.

2. 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.

3. 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.

4. You are correct. Thanks for catching the typo.

5. 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. SInce $1-\beta = (\frac{\epsilon}{10\alpha})^2$ and $\alpha=G+\sigma$, it follows that $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)$ as desired.

6. 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 $\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,$ 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. In general, we will detail all the arguments regarding expectations in our final version.

7. 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. Although 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. We 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.

We would be glad to address any remaining concern!

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.

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.

Regarding 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.

Thank you for your constructive comments! Let us address your questions one by one.

• (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:

  • 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.
  • 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]$.

• (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.

• (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.

• (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).