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

Authors: Enea Monzio Compagnoni, Tianlin Liu, Rustem Islamov, Frank Norbert Proske, Antonio Orvieto, Aurelien Lucchi
Primary area: optimization
Keywords: Stochastic Differential Equations, Stochastic Optimization, Adaptive Methods

Abstract:
Despite the vast empirical evidence supporting the efficacy of adaptive optimization methods in deep learning, their theoretical understanding is far from complete. This work introduces novel SDEs for commonly used adaptive optimizers: SignSGD, RMSprop(W), and Adam(W). These SDEs offer a quantitatively accurate description of these optimizers and help illuminate an intricate relationship between adaptivity, gradient noise, and curvature. Our novel analysis of SignSGD highlights a noteworthy and precise contrast to SGD in terms of convergence speed, stationary distribution, and robustness to heavy-tail noise. We extend this analysis to AdamW and RMSpropW, for which we observe that the role of noise is much more complex. Crucially, we support our theoretical analysis with experimental evidence by verifying our insights: this includes numerically integrating our SDEs using Euler-Maruyama discretization on various neural network architectures such as MLPs, CNNs, ResNets, and Transformers. Our SDEs accurately track the behavior of the respective optimizers, especially when compared to previous SDEs derived for Adam and RMSprop. We believe our approach can provide valuable insights into best training practices and novel scaling rules.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1730315619687

thread_type:
official_review

<thread_object>
creator: reviewer_qc2v

<content>
summary:
This work derives a novel SDE (of weak approximation order 1, under Gaussian gradient noise Z) for the SignSGD optimizer, and using this SDE, uncovers three phases of dynamics with distinct behaviours. They show theoretically that the effect of noise on SignSGD differs from SGD, and that SignSGD is more resilient to high levels of noise. Further, they show that when Z is heavy-tailed noise, the associated SDEs for SignSGD take a very similar form, further highlighting SignSGD's resilience to high levels of noise. They then derive SDEs (of weak approximation order 1) for AdamW and RMSpropW under much less restrictive assumptions than previous work, and find that decoupled weight decay has an important role in stabilization of dynamics at high noise levels. They also provide empirical validations of their theoretical results.

strengths:
Results are presented very clearly and interpretations are made clear

Use numerical integrators to verify that their derived SDEs match the associated discrete-time optimizers across a diverse range of relevant setups

Technical assumptions appear much less restrictive than other works

weaknesses:
Mostly theoretical, with practical implications of findings underexplored; do the findings tell you anything about how the studied optimizers can be improved, tuned, etc.?

Form of stochastic gradient appears unjustified. Is the noise Z being additive, and distributed as a Gaussian or heavy-tailed, accurate to practice?

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730606594370

thread_type:
official_review

<thread_object>
creator: reviewer_2dqw

<content>
summary:
The work is centered on introducing novel SDEs for adaptive optimizers. Under weak assumptions, they derive SDE for SignSGD and compare it against vanilla SGD to show three different phases in the dynamics of SignSGD. They also show an inverse relationship between noise and convergence rate for SGD which is in contrast to the quadratic effect on loss & variance observed for SGD. They also analyze decoupled weight decay and compare it against SignSGD.

strengths:
The paper has introduced A rigorous mathematical framework of SDEs for SignSGD and its comparison against commonly used adaptive optimizers like ADAMW. They also characterize a relationship between noise and convergence rates for SignSGD and SGD. Also, the framework is used to understand decoupled weight decay efficacy and infer the effect of noise on such approaches. Their analysis is followed by relevant experiments.

weaknesses:
Major weakness

1. Line 245-246. The definitions of 'signal' and 'noise' need to be explicitly specified in this context. The current characterization of 'signal' in these lines does not make sense. $Y_t$ line 199 defined signal-to-noise ratio but I don't see what is being referred to in lines 245-246 and what part is being referred to as "large". The entire paragraph here should be rephrased to better connect it to the surrounding content as it currently seems disconnected, or poorly worded.

2. Line 245: "SNR is large, meaning SignSGD behaves like SignGD": This claim is unsupported. You should provide the reasoning or evidence that supports this claim, such as referencing relevant equations, prior results, or additional explanation. 

3. Figure 2 referred to in line 249: Hard to understand. Please provide details about the loss function/landscape, dimensionality, and any other relevant parameters used to generate Figure 2.  Lacks description of the setting used to plot figures.

4. In general, (not commenting on technical soundness), the paper is very hard to follow and poorly structured. 
- Need to provide brief intuitive explanations before each lemma, highlighting its significance and key takeaways in order to understand what the authors want to convey.
- In section 3, authors should be more rigorous about introducing notation. There are unnecessary references to related work mentions which should be pushed somewhere else.  A comprehensive introduction of notation in the upcoming lemmas or some useful background should be introduced. Since the main argument is around novel SDEs, more introductory material on SDEs and their relevance to optimization algorithms before presenting the main results.

5. Section 3.1.1 - authors should ensure all new notation is properly introduced before it is used in lemmas, either by adding definitions as terms first appear or by including a notation section at the beginning of 3.1.1.

6. Line 412: Decoupled weight decay. 
- Clarify why decoupled weight decay is considered key in this context
- Provide a more explicit link between their analysis and the claim about stabilization at high noise levels. Currently, the authors' analysis to support this is not evident.
- Include specific equations or results that demonstrate the stabilization effect of decoupled weight decay

7. The authors should include key experimental results in the main paper, selecting the most important figures or tables that support their main claims, while potentially moving more detailed results to an appendix if space is an issue.

Not sure how to comment on the rigor of mathematical soundness of SDEs analysis but overall it was very hard to follow what the authors wanted to show/convey.

questions:
I am not quite sure what is the key message (takeaway) of SDE analysis for decoupled weight decay analysis. They are known to be more effective than their vanilla counterparts. Is there a key message here that helps improve our understanding of adaptive optimizers? Also, is the analysis valid only for convex quadratics in this case?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730950293323

thread_type:
official_review

<thread_object>
creator: reviewer_ahk4

<content>
summary:
This paper investigates the learning dynamics of adaptive optimizers in machine learning.  The authors analyze a simplified model, $f(x) = x^THx$, with noisy gradients defined as $g_\gamma = \nabla f(x) + \gamma$, where $\gamma$ is sampled from a distribution like a Gaussian. Focusing on SGD, sign-SGD,  and AdamW, they demonstrate that the learning dynamics can be effectively approximated by a stochastic differential equation (SDE) of the form $X_t = \text{function}(X_t)$. This allows them to analyze the first and second moments of $X_t$ for these optimizers as $t$ approaches infinity.  The authors support their theoretical findings with experiments on practical models like ResNets and small Transformers, trained on MNIST, CIFAR, and Shakespeare datasets.

While generally well-written, mathematically sound and provide many insights, the paper's reliance on a simplified model and unrealistic noise design limits its applicability.


--- updates --- 

Thank you for the clarifications. While I believe this work warrants a score of 7, I understand that's not an available option.  Considering that other reviewers have mostly given scores of 6, I am raising my score to 8 to better reflect the paper's merit.

strengths:
- generally well-written 

- mathematically seems sound, though i didn't check the proof 

- papper contains rich set of results that may be insightful for practicioners. 

- theoretical results are supported by real networks, though on

weaknesses:
* The toy model is too simple and far from practical. I would expect at least a linear regression model with a random (Gaussian) design.
* 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.
*  Experiments use full batch with synthetic (Gaussian) noise instead of mini-batches, which is consistent with the paper's problem setup but unrealistic.
* The paper is overly dense, with little intuition provided for the theorems and proofs. It also relies heavily on in-line equations, hindering readability.

questions:
* The notation $\nabla f_\gamma$ is unclear and could be misinterpreted as $\nabla (f_\gamma)$. 
* All algorithms in the paper use full batch gradients with added noise, which is misleading and doesn't accurately reflect real-world implementations of algorithms like Adam.
* It would be beneficial to generalize the results to mini-batch linear regression with random Gaussian design, similar to the approach in https://arxiv.org/abs/2405.15074.
* The paper should provide more intuition behind the proof of Theorem 3.2 to help readers understand the key ideas without having to delve into the appendix.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731082045951

thread_type:
official_review

<thread_object>
creator: reviewer_2nwx

<content>
summary:
This paper presents a theoretical analysis of adaptive optimization methods—specifically SignSGD, AdamW, and RMSpropW—through the lens of stochastic differential equations (SDEs), unveiling the intricate interplay between adaptivity, gradient noise, and curvature. Key contributions include:

- **SignSGD Dynamics:** Identification of three distinct phases in SignSGD’s dynamics, with noise inversely affecting both convergence rates of the loss and the iterates.
- **Enhanced SDE Models:** Derivation of new and improved SDEs for AdamW and RMSpropW, leading to a novel batch size scaling rule and an examination of the stationary distribution and stationary loss value in convex quadratic settings.

The derivation of new SDEs for SignSGD, AdamW, and RMSpropW provides a solid foundation for understanding the dynamics of these optimizers.

strengths:
Strengths includes:
- **Intriging Insights:** Identifying three distinct phases in SignSGD dynamics and the inverse relationship between noise and convergence rates offer valuable perspectives on optimizer behavior.
- **Practical Implications:** Introducing a novel batch size scaling rule and examining stationary distributions have direct implications for better training practices in deep learning.
- **Clear background:** The preliminaries and background information provided in the appendix enhance the readability and understanding of the main results.
- **Validation:** Theoretical findings are supported by extensive experimental evidence across various neural network architectures, including MLPs, CNNs, ResNets, and Transformers.

weaknesses:
- **Unclear Terminology and Notation:** Certain terms and symbols used in the paper are not clearly defined, which may lead to confusion. What does $\mathbb{P}$ represent in Theorem 3.2? How is the error function defined in Corollary 3.3? Does Lemma 3.7 explicitly establish the stationary distribution of SignSGD?

questions:
See weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732006264656

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Answer

comment:
Dear Reviewers and AC,

We sincerely appreciate your thorough reviews, insightful comments, and interesting questions regarding our paper: Your feedback has greatly contributed to the finalization of our work.


We are pleased that three out of four reviewers found our paper to be "**well-presented**" and to provide a "**solid (mathematical) foundation**" for understanding the dynamics of SignSGD, Adam(W), and RMSprop(W). These reviewers appreciated our theoretical contributions which, quoting their words, "**offer valuable perspectives on optimizer behavior**" and "**have direct implications for better training practices in deep learning**", noting that our "**intriguing**" findings are "**well-supported by extensive experimental evidence across various architectures**" and datasets. Reviewer QC2v described our **presentation as excellent**, highlighting the clarity of our results and interpretations (e.g. stating "**Results are presented very clearly and interpretations are made clear**").

We of course take the criticisms very seriously and we devoted time to addressing them below. Specifically, the weaknesses and questions mainly concern relatively **minor issues** around notation, the need for some rewording, the request for some additional technical details, and the intuition behind some proofs. All of these are properly taken care of in the updated version.

Finally, all concerns regarding assumptions or simplistic setups have been properly addressed: We referred reviewers to very well-cited papers from the literature that use our same (if not more restrictive) assumptions. Importantly, we added the study of a linear regression model with a random Gaussian design (suggested by Reviewer ahk4) to showcase the flexibility and generality of our approach.

To highlight the (relevant) changes made to the original submission, we marked them using **blue** in the revised version.
Once again, we are thankful to the reviewers for their constructive feedback. We look forward to the upcoming author-reviewer discussion period.

Thank you for your attention.

Best regards,

The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732415334599

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Kind Reminder

comment:
Dear Reviewers,  

As the discussion period approaches its conclusion, we kindly encourage your engagement in providing feedback or addressing any points from our rebuttal.

We greatly appreciate your time and look forward to continuing this discussion.  

Best regards,  

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces novel stochastic differential equations (SDEs) to model adaptive optimization methods like SignSGD, AdamW, and RMSpropW, providing insights into their dynamics and the role of noise, adaptivity, and curvature in convergence. The authors support their theoretical findings with experiments on various neural network architectures, demonstrating the accuracy of their SDE models compared to previous approaches. A weakness is the reliance on simplified or unrealistic models and noise assumptions in some analyses, which may limit the direct applicability to real-world deep learning scenarios; however, the strength lies in the introduction of a rigorous SDE framework for analyzing adaptive optimizers, supported by empirical validation, offering a deeper understanding of their behavior and potentially informing improved training practices. Given these considerations, I recommend accepting this paper, as it provides valuable insights into adaptive optimization methods.