Adaptive Methods through the Lens of SDEs: Theoretical Insights on the Role of Noise

We thank the reviewer for their very positive review.

We are pleased to hear that the reviewer found our work to provide an excellently sound and good contribution towards understanding adaptive optimizers and that it was excellently presented with a very clear interpretation of results.

Here is how we addressed the weaknesses and questions: we included them in the revised PDF.

W1. "Do the findings tell you anything about how the studied optimizers can be improved, tuned, etc.?": Our focus was to provide new insights into existing adaptive methods that are known to perform well in practice, even though the reasons for their effectiveness are not yet fully understood. These insights point out some important features of adaptive methods which would be important to preserve while developing new optimizers. While we do not currently have suggestions for developing other optimizers, this is certainly a direction we are actively considering.

W2. "Form of stochastic gradient appears unjustified. Is the noise Z being additive, and distributed as a Gaussian or heavy-tailed, accurate to practice?": While we acknowledge that this noise structure is not fully realistic, it is very common in the literature: See Lines 139-142 of the original submission for a selected collection of the references, with particular attention to Jastrzebski et al., who offer a justification to back this claim. We highlight that more complex and realistic noise structures were presented in the appendix of the original submission with analogous conclusions — see Section C.3 ALTERNATIVE NOISE ASSUMPTIONS.

Q1. "For Lemma 3.4 (line 199), since $Y_t \in R^d$, is the condition e.g. $|Y_t| > 3/2$ denoting the absolute value treated element-wise, like $|(Y_t)_i| > 3/2$, or does $|Y_t|$ denote the norm of $Y_t$?": It has to be treated element-wise and we clarified it in the updated version.

We thank the reviewer for their review.

We regret that our contribution was considered only fair and poorly presented, as this contrasts starkly with feedback from all three other reviewers who deemed our paper "well-written", "rich in valuable results for practitioners", and "excellently presented with clear interpretations of the findings". We however take this feedback seriously: We appreciate the points raised and outline below how we address these weaknesses and questions, which we incorporated into the revised PDF.

Respectfully, we do not believe the weaknesses highlighted by the reviewer to be substantial: They can be addressed through rewording or additional clarifications. None of the mentioned concerns compromise the validity, generality, or impact of our analysis. We have already updated the PDF incorporating the feedback of the reviewer (see discussion below) and are happy to discuss further with the reviewer if some concerns remain.

Here is how we addressed the weaknesses and questions: we included them in the revised PDF.

Weaknesses

W1. While this informal definition was inspired by Bernstein et al. (2018) (who first provided an extended theoretical analysis of SignSGD), introducing the SNR was intended to support the reader's intuition and was not meant as a formal concept. We have reworded the mentioned paragraphs accordingly.

W2. On the contrary, this claim is in fact supported: Eq. 4 shows that if the SNR is large (e.g., the gradient is much larger than its noise), the SDE of SignSGD becomes the classical ODE of SignGD. The left and center-left sides of Fig. 2 visually support this: Thank you for pointing this out --- We have reworded the mentioned paragraph accordingly.

W3. While we initially omitted certain details to avoid overloading the paper with technical information that could affect readability, we have added some level of detail to the caption.

W4 and W5. The relevance of SDEs in optimization was already discussed in a dedicated (large) paragraph in the Related Works while Appendix B already provided introductory material on SDEs. To address your point regarding the notation, we have expanded the dedicated paragraph.

Regarding the interpretation of our results, we find this comment puzzling. As highlighted by other Reviewers, our results were already complemented by a paragraph where we interpret them, even in dedicated “Remarks”, or “Conclusions” in gray boxes, specifically providing brief intuitive explanations. For instance, Reviewer QC2v stated that "Results are presented very clearly and interpretations are made clear". Below is a guide indicating where each result was interpreted. All mentioned lines refer to the original submission:

1. Lemma 3.4: This result is meant to break down and make Corollary 3.3 and Theorem 3.2 more digestible: Its interpretation was in lines 230-234 and visually supported in Fig. 1 and Fig. 2 (left).

2. Lemma 3.5 and Lemma 3.6: Their interpretation was in lines 245-250 together with the Remark in lines 254-261. Visually supported in Fig. 2 (center-left).

3. Lemma 3.7 and Lemma 3.8: Their interpretation was in lines 277-279. Visually supported in Fig. 2 (center-right).

4. Lemma 3.9: Its interpretation was in lines 280-283. Visually supported in Fig. 2 (right).

5. We summarized Section 3 with a Conclusion to collect the findings of the section.

6. Lemma 3.10 and Corollary 3.11 had a dedicated Conclusion where we interpreted them in lines 317-322.

7. Lemma 3.13: We introduced it in lines 362-370 and interpreted it in lines 406-409. Visually supported in Fig. 5.

8. Lemma 3.14: We interpreted it in the Conclusion in lines 447-453. Visually supported in Fig. 5 (right).

We now added a dedicated paragraph to Theorem 3.12, which we recognized was not properly interpreted, but visually well-supported in Fig. 3 and 4.

W6.

1. Let us recall that AdamW is defined as Adam + decoupled weight decay: Since we study such an optimizer, we need to study how decoupled weight decay affects the dynamics of AdamW as decoupled weight decay plays a fundamental role in modern LLMs [1,2].

2. On the contrary, this claim is, in fact, supported: Eq. 13 shows that the bound on the expected loss of AdamW is itself bounded with respect to the noise level $\sigma$. This was clearly stated in lines 415, 442, 449, and 479, and we have further clarified this point in the updated version.

3. As mentioned, this was provided in Eq. 13. In contrast, Eq. 14 showed that Adam on the $L^2$-regularized loss function does not enjoy such a stabilization: See Lines 412 to 420 of the original submission.

[1] D'Angelo, Francesco, et al. "Why Do We Need Weight Decay in Modern Deep Learning?." NeurIPS 2024.

[2] Xiao, Lechao. "Rethinking Conventional Wisdom in Machine Learning: From Generalization to Scaling.", 2024

We thank the reviewer for their very positive review.

We are pleased to hear that the reviewer found our work 1) to provide a "well-written, mathematically sound, [and] good contribution" towards understanding adaptive optimizers; 2) to show "a rich set of results" that are "insightful for practitioners"; and 3) to support the theoretical results with experiments on a variety of DNNs.

Here is how we addressed the weaknesses and questions: we included them in the revised PDF.

Weaknesses

1. "The paper's reliance on a simplified model and unrealistic noise design limits its applicability.": We would like to clarify that contrary to the reviewer's statement, none of our SDEs derivation is restricted to quadratic functions or Gaussian noise: The theory applies to general smooth functions and general noise structures (see Assumption C.3 and lines 134-142 of the original submission, respectively). Regarding the applicability of our theory, we highlight, for instance, that although our novel scaling rule for AdamW was derived for strongly convex functions, we validated its effectiveness on a Pythia-like 160M LLM trained on $2.5B$ tokens from the SlimPajama dataset (see Appendix F.8 and Figure 14 of the original submission): This demonstrates the broad applicability of our theory beyond the theoretical setting.

2. "The toy model is too simple and far from practical.": We highlight that the only results specific for quadratic functions are Lemma 3.7 and Lemma 3.14, where we derive stationary distributions, which are calculated on quadratic functions in line with the existing literature (see lines 142 to 145 of the original submission for references, especially [5,7,8,9,10] whose joint citation count is 3000+). All other results are valid for general strongly convex functions. For the sake of simplicity and to avoid crowding the main paper, we already included their generalization to general non-convex smooth functions in the Appendix (see Lemma C.17 and Lemma C.50 of the original submission).

3. "The noise design is artificial and unrealistic. It doesn't capture practical batch noise, which depends on the loss, parameters, and changes over time. The current noise function uses a time-independent, identical distribution.": While we agree that this is an artificial design, this assumption is very common and accepted in the literature, especially that of SDEs for optimization (see lines 139-142 of the original submission for reference, especially [1,2,3,4,5,6,7] whose joint citation count is 1700+). We employ this design in the main paper only for didactical reasons, but more complex and realistic noise structures are presented in the appendix with analogous conclusions — see Section C.3 ALTERNATIVE NOISE ASSUMPTIONS of the original submission: It is quite rare to tackle these many noise structures, thus we are more complete than other papers. Therefore, we kindly disagree with the reviewer: Our setup is not restrictive and is more general or, at least, in line with the literature.

4. "Experiments use full batch with synthetic (Gaussian) noise instead of mini-batches, which is consistent with the paper's problem setup but unrealistic.": As we discussed above, these assumptions are widely accepted in the literature (see lines 139-142 of the original submission for reference, especially [1,2,3,4,5,6,7] whose total citation count is 1700+) and already bring novel insights into the dynamics of these optimizers.

[1] Li, Qianxiao, et al. "Stochastic modified equations and adaptive stochastic gradient algorithms." ICML, 2017.

[2] Mertikopoulos, Panayotis, et al. "On the convergence of gradient-like flows with noisy gradient input." SIAM Journal on Optimization 2018.

[3] Raginsky, Maxim, et al. "Continuous-time stochastic mirror descent on a network: Variance reduction, consensus, convergence." CDC IEEE, 2012.

[4] Zhu, Zhanxing, et al. "The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects." ICML 2019.

[5] Mandt, Stephan, et al. "A variational analysis of stochastic gradient algorithms." ICML 2016.

[6] Ahn, Sungjin, et al.. "Bayesian posterior sampling via stochastic gradient Fisher scoring." ICML, 2012.

[7] Jastrzebski, Stanislaw, et al. "Three factors influencing minima in SGD." ICANN, 2018.

[8] Ge, Rong, et al. "Escaping from saddle points—online stochastic gradient for tensor decomposition." CLT, 2015.

[9] Levy, Kfir Y. "The power of normalization: Faster evasion of saddle points." (2016).

[10] Jin, Chi, et al. "How to escape saddle points efficiently." ICML, 2017.

[11] Xie, Zeke, et al. "Positive-negative momentum: Manipulating stochastic gradient noise to improve generalization." ICML, 2021.

Continuation of Rebuttal on Weaknesses

5. "The paper is overly dense, with little intuition provided for the theorems and proofs. It also relies heavily on in-line equations, hindering readability.": We highlight that we already complemented our results with interpretations in the original submission (see reply to Reviewer 2DqW for a precise guide indicating where each result is interpreted). To address your point, we added the proof to three key results: Other results follow similar strategies. Additionally, we reduced the use of in-line equations and used the full $10$ pages to reduce the density of the paper.

Questions

Q1. While this is the notation of the seminal paper by Li et al. (2017) and has become standard in the literature, we have redefined this in our paper as well.

Q2. See the answer to Weakness 4 above. We add here that the reason why we inject noise with this structure is that we want to empirically verify our SDEs on a variety of neural architectures and datasets: To our knowledge, we are the first to take on such a challenge to this extent --- We welcome future work to improve our initial effort in this direction.

Q3. Please find our insights for this setup below and the details in Appendix C.5 of the updated version together with other additional alternative (realistic) noise structures.

Q4. We added a sketch of the proof of some key results to guide the reader's intuition.

Linear Regression with Gaussian design

In this section, we adopt the notation of Paquette et al.. As per Eq. 5 of Paquette et al., the loss function can be rewritten as

$

f(\theta)= \frac{1}{2} \langle D(W \theta-b),(W \theta-b)\rangle, \quad \text { where } D=\operatorname{diag}\left(j^{-2 \alpha}\right) \in \mathbb{R}^{v \times v}, \text{ and} 1 \leq j \leq v.

$

Without loss of generality, we define $\phi:= W \theta-b$, which implies that

$

f(\theta)= \frac{\phi^{\top} D \phi}{2}, \quad \text { where } D=\operatorname{diag}\left(j^{-2 \alpha}\right) \in \mathbb{R}^{v \times v}, \text{ and } 1 \leq j \leq v.

$

The stochastic gradient is unbiased and its covariance is

$B \Sigma(\phi) = (\phi^{\top} D \phi) D + D \phi \phi^{\top} D = 2 f(\phi) D + \nabla f(\phi) \nabla f(\phi)^{\top},$ where $B$ is the batch size.

Under this assumption, we have that for $Y_t := \frac{\sqrt{B}(\Sigma(\phi_t))^{-\frac{1}{2}}\nabla f(\phi_t)}{\sqrt{2}}$ and $\mathcal{S}(\phi_t)=\mathbb{E}[(Sign(\nabla f_{\gamma}(\phi_t))(Sign(\nabla f_{\gamma}(\phi_t))^{\top}]$ the SDE of SignSGD becomes:

$

d \phi_t = - Erf \left( Y_t \right) dt + \sqrt{\eta} \sqrt{\mathcal{S}(\phi_t) - Erf \left(Y_t \right) Erf \left(Y_t \right)^{\top}} d W_t.

$

As a consequence, Lemma 3.5 becomes:

Let $f$ be as above, $f_t:=f(\phi_t)$, and $\mathcal{L}_{\tau} := Tr(D)$. Let $\mu$ be the minimum eigenvalue of $D$, and $L$ be its maximum one. Then, during

1. Phase 1, the loss will reach $0$ before $t_* = 2 \sqrt{\frac{S_0}{\mu}}$ because $f_t \leq \frac{1}{4} \left( \sqrt{\mu}t - 2 \sqrt{f_0}\right)^2$;
2. Phase 2 with $\beta := \frac{\eta}{2} \mathcal{L}_{\tau}$ and $\alpha:= \frac{ m \mu \sqrt{B}}{\sqrt{2 L}}$,

$

\mathbb{E}[S_t] \leq \frac{\beta^2 \left( \mathcal{W}\left( \frac{(\beta + \sqrt{S_0} \alpha)}{\beta} \exp\left(-\frac{\alpha^2 t - 2 \sqrt{S_0} \alpha}{2 \beta} - 1 \right) \right) + 1 \right)^2}{\alpha^2} \overset{t \rightarrow \infty}{\rightarrow} \frac{\beta^2}{\alpha^2}; $

3. Phase 3 with $\beta := \frac{\eta}{2} \mathcal{L}_{\tau}$ and $\alpha:= \sqrt{\frac{2}{\pi}} \frac{\mu \sqrt{B}}{\sqrt{L}}$,

$

\mathbb{E}[S_t] \leq \frac{\beta^2 \left( \mathcal{W}\left( \frac{(\beta + \sqrt{S_0} \alpha)}{\beta} \exp\left(-\frac{\alpha^2 t - 2 \sqrt{S_0} \alpha}{2 \beta} - 1 \right) \right) + 1 \right)^2}{\alpha^2} \overset{t \rightarrow \infty}{\rightarrow} \frac{\beta^2}{\alpha^2},

$

where $\mathcal{W}$ is the Lambert $\mathcal{W}$ function.

We thank the reviewer very much for their dedication to the review process and for taking the time to carefully study our manuscript.

We are pleased to hear that the reviewer found our work to provide a "solid foundation for understanding the dynamics of these optimizers", providing "intriguing insights", with "direct implications for better training practices in deep learning". Importantly, the reviewer noted that the presentation is good, with a "clear background" that "enhances the readability and understanding of the main results", which are "supported by extensive experimental evidence across various neural network architectures".

Here is how we addressed the weaknesses in the updated PDF:

1. "What does $\mathbb{P}$ represent in Theorem 3.2?": As per Appendix B, all our analyses take place on a probability space $\left(\Omega, \mathcal{F}, \mathbb{P}\right)$: Therefore, $\mathbb{P}$ represents the probability measure of the probability space --- We added this to the paragraph dedicated to notation. In that equation, we wanted to measure the probability that the stochastic gradient is (element-wise) smaller than $0$.

2. "How is the error function defined in Corollary 3.3?": The error function (also called the Gauss error function) is a function erf: $\mathbb{C} \rightarrow \mathbb{C}$ defined as:

   $$\operatorname{erf} z=\frac{2}{\sqrt{\pi}} \int_0^z e^{-t^2} \mathrm{~d} t.$$

   Despite it being a known object in probability, we included its definition in the manuscript for clarity.

3. "Does Lemma 3.7 explicitly establish the stationary distribution of SignSGD?": Yes. Point 1 of Lemma 3.7 shows that the expected value of the iterates converges to $0$. Point 2 of the same lemma shows that the covariance matrix of the iterates converges to $\frac{\eta}{2} \left( \sqrt{\frac{2}{\pi}} I_d + \frac{\eta}{\pi} H \Sigma^{-\frac{1}{2}}\right)^{-1} H^{-1} \Sigma^{\frac{1}{2}}$. Therefore, we have that the stationary distribution of SignSGD is:

$$\left(\mathbb{E} [X_\infty], Cov[X_\infty]\right)=\left(0, \frac{\eta}{2} \left( \sqrt{\frac{2}{\pi}} I_d + \frac{\eta}{\pi} H \Sigma^{-\frac{1}{2}}\right)^{-1} H^{-1} \Sigma^{\frac{1}{2}} \right).$$

We clarified this in the updated version.