Title: SDEs for Adaptive Methods: The Role of Noise
Venue: NeurIPS 2024
Accepted: no
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. In this work, we introduce novel SDEs for commonly used adaptive optimizers: SignSGD, RMSprop(W), and Adam(W). Our SDEs offer a quantitatively accurate description of these optimizers and help bring to light 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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work derives SDEs for adaptive gradient methods and study the role of gradient noise. The analysis starts from theoretically driving the SDE for SignSGD and highlight its significant difference from SGD. The work further generalize the SDE analysis for AdamW and RMSpropW, two popular adaptive optimizers with decoupled weight decay and reveal key properties of weight decay. Finally, the work integrates the derived SDEs with Euler-Maruyama to confirm that the SDEs faithfully track their respective optimizers with various modern neural networks.\n\nSTRENGTHS:\n-The theoretical results are novel. To my knowledge, this is a first SDE analysis for SignSGD with quantitatively accurate descriptions.\n\n-The theoretical analysis reports some novel properties in terms of gradient noise and convergence. These properties are interesting.\n\n-The proofs seem complete and reasonable.\n\n-A useful theory should be quantitatively verifiable. This work definitely make it. The experiments that SDEs fit the empirical results with various optimizers and models are informative and impressive.\n\nWEAKNESSES:\n-It seems that the reported theoretical results and insights cannot directly lead to some theory-inspired and improved methods. This raise a question on the significance of this work.\n\n-While this work did literature review, some important references are still missing, such as [1] on analyzing Adam using SDEs. As weight decay plays a key role in the results, it may be helpful to review recent papers analyzing novel or overlooked properties of weight decay.\n\n\nReference:\n\n[1] Xie, Z., Wang, X., Zhang, H., Sato, I., & Sugiyama, M. (2022, June). Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum. In International conference on machine learning (pp. 24430-24459). PMLR.\n\nQUESTIONS:\n- Please see the weaknesses.\n\n- Could you please explain more how L2 regularization and decoupled weight decay behaves differently in your results?\n\nLIMITATIONS:\nThis work discussed the limitations in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-23"}

comments: SUMMARY:
The authors derive SDE for signSGD and Adam(W). The experiments show that the algorithm will converge toward the limit of the theorem indicates.

STRENGTHS:
The authors propose "accurate" SDEs for algorithms Sign-SGD and Adam(W).

WEAKNESSES:
1. In Remark after Lemma 3.6, the authors claim that Sign-SGD is (almost) linear in $\sigma_{max}$. However, with $\Delta$ either in Phase 2 or Phase 3, there should be $\sigma_{max}^2$ in the final bound.

2. All the stationarity holds when Hessian is the same from $X_0$ to $X_t$ and convergence holds for strongly convex. However, the hessian changes a lot during network training.

QUESTIONS:
1. Notation of $W_t$ is not defined.  What is $W_t$?

2. How can we extend Lemma 3.13 to convex setting (or even nonconvex case)?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

comments: SUMMARY:
This paper derives SDEs for SignSGD, RMSprop, and Adam.
The analysis offers insights into the convergence speed, stationary distribution, and robustness to heavy-tail noise of adaptive methods.

STRENGTHS:
- The derived SDE for SignSGD exhibits three different phases of the dynamics.

- The analysis reveals the difference between SignSGD and SGD in terms of the asymptotic expected loss, the robustness of noise variance, etc.

- The analysis of AdamW provides insights into the different roles of noise, curvature, and weight decay.

WEAKNESSES:
Refer to Questions and Limitations.

QUESTIONS:
- What learning rate (lr) do the experiments in Figure 4 use? Within what range of lr does this SDE align well with the original algorithm (experimentally)?

- Could the authors intuitively explain why the asymptotic expected loss of SignSGD is proportional to $\sigma_{\max}$ instead of $\sigma_{\max}^2$?

- How can the derived SDE explain the loss spike phenomenon of SignSGD/AdamW?

- Many works about SGD noise ([1][2][3]) admit the noise structure $\mathbb{E}(g_i-g)(g_i-g)^{\top}\sim \mathcal{L} H$. What conclusions (such as those related to the training phases) can be derived from the SDE if we change the noise assumption in Corollary 3.3 to $\mathbb{E}(g_i-g)(g_i-g)^{\top}\sim\mathcal{L} H$?


[1] Ziyin et al. Strength of minibatch noise in SGD.

[2] Wojtowytsch. Stochastic gradient descent with noise of machine learning type. part II: Continuous time analysis.

[3] Wu et al. The alignment property of SGD noise and how it helps select flat minima: A stability analysis.

LIMITATIONS:
The SDE for AdamW is limited to quadratic functions.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-23