Exact risk curves of signSGD in High-Dimensions: quantifying preconditioning and noise-compression effects

Dear Reviewer:

Thank you for your comments and for taking the time to evaluate our manuscript.

Note that the stochastic gradient algorithm itself does not care whether we use a minibatch gradient or not.

Your point is taken regarding the similarities between heavy ball momentum and conjugate gradient. We avoided discussing stochastic gradient methods at this line, because there is considerable nuance in showing acceleration for stochastic methods. We would be happy if you have a reference in mind. To our knowledge, outside of minibatch SGD (where indeed batch is relevant), there is not a stochastic gradient + momentum setup with provable speed-up.

The extra conjecture adds nothing essentially to the paper. Adam is simply a learning rate setting for the stochastic gradient algorithm… A careful study of Tsypkin's classic 1971 work on Adaptation and learning in automatic systems will make this much clearer.

Thank you for the Tsypkin reference! It's interesting to see some modern algorithms described so early in the literature.

The point of including the conjecture was to show that our methodology can be used to take Adam - which has a stochastic learning rate - and compute a deterministic equivalent - which becomes an arbitrarily good approximation in high dimensions - that can be explicitly written in terms of the risk and the data covariance. We hope to formally prove the novel conjecture in future work, which would allow us to quantify the difference between SignSGD, Adam and SGD

[1] Malladi et al., On the SDEs and Scaling Rules for Adaptive Gradient Algorithms: NeurIPS 2022.

[2] Kingma et al., ADAM: A Method for Stochastic Optimization: ICLR 2015

Dear Reviewer QUHk,

First, thank you for the exceptionally thorough review that you have performed of our paper. Regardless of the outcome of the review process, your feedback has already helped us improve the paper.

Much of your review discusses the Weak-Approximation framework. We agree that a comparison between our work and the Weak-Approximation framework will improve our paper. The Weak-Approximation approach is quantitatively and qualitatively different from our approach, and we would like to explain the details here. We will make a more detailed, itemized response to your points later, but we wanted to begin discussion of this substantial point of your critique first.

The Weak--approximation SDE is the following:

(Malladi et al. -- Adam SDE, Def 4.4). Let $c_1 \triangleq\left(1-\beta_1\right) / \eta^2, c_2 \triangleq\left(1-\beta_2\right) / \eta^2$. Let $\gamma_1(t) \triangleq 1-\exp \left(-c_1 t\right)$ and $\gamma_2(t) \triangleq 1-\exp \left(-c_2 t\right)$. Define the state of the SDE as $\boldsymbol{X}_t=\left(\boldsymbol{\theta}_t, \boldsymbol{m}_t, \boldsymbol{u}_t\right)$ and the dynamics as

where, $\boldsymbol{B}_t$ is $d$-dimensional Brownian motion.

To recover SignSGD from ADAM, we would take $\beta_2=0$ and $\beta_1=0$. The noisy gradient model which most closely reproduces our setup has $\boldsymbol{\Sigma}(\boldsymbol{\theta}_t) = \boldsymbol{K}{2\mathcal{P}(\theta_t)}$ with $2\mathcal{P}(\theta_t)$ as in (Eqn 2) of our paper. (Note Malladi et al also has RMS-Prop, with a slightly simpler SDE. But even recovering the RMS-Prop SDE from the ADAM SDE requires taking a nontrivial limit $c_1 \to \infty$. We will add an appendix to the paper in which we take these limits; even after doing so, however, the resulting SDE is substantially different from ours).

In comparison, our SDE is given by (Eqn 7): $\mathrm{d} \Theta_t=-\eta_t \frac{\varphi\left(\mathcal{R}\left(\Theta_t\right)\right)}{\sqrt{2\mathcal{R}\left(\Theta_t\right)}}\overline{\mathbf{K}}\left(\Theta_t-\theta_*\right)\mathrm{d} t+\eta_t \sqrt{\frac{\mathbf{K}_\sigma}{\pi d}} \mathrm{d}\mathbf{B}_t$

The SDEs are different in many ways. In the comment below, we discuss in detail some of these details.

3. Conceptual Accuracy and Veracity concerns

A number of comments are regarding phrasing and accuracy of some statements. We address these below indexed by the bullet point where these comments were made:

#8) We have changed the phrasing to “an important characterizing feature of SignSGD is its effect on the covariance of the signed stochastic gradients”

#13) While the landscape might remain constant (the problem itself has not changed) the landscape has effectively changed from the point of view of the optimizer. Given that we are writing about the optimizer it seems fair to make this claim.

#18) Good point, we will change how we refer to interpolation. Here we truly just mean noiseless.

#23) Thank you for pointing this out. We mistakenly referenced the wrong paper. This has been fixed.

#24) Yes this scheduler is bounded so long as the noise has non-zero variance (as then $\mathcal{P} \geq \mathcal{v}$ where $\mathbb{E}[\epsilon^2] = \mathcal{v}$). We have assumed this. In Remark 1, we point out that if we consider the no-noise scenario then we must assume the risk is bounded away from $0$ and again the scheduler is bounded)

#25) In the case of infinite variance, SGD does not even converge. We have added a citation.

#26) Our main theorems (1 and 2) show that the risk curves of signSGD and signHSGD concentrate in fact become identical as dimension goes to infinity. Thus, it is consistent to refer to training signSGD when we discuss the dynamics of signHSGD since they are referring to identical processes in the high-dimensional limit.

#27) We added a forward reference to the missing equation.

#28) We say why we would do this in the text. We also do not follow what rescaling by risk does with regards to variance. The sign function properly controls for gradients with infinite variance.

#29) This sentence was intended to provide a reason why one may want to keep the schedule from signSGD, as it adjusts for ‘vanishing gradients’. We have clarified the text.

#30) These are excellent questions: we will look into an expanded discussion and adding simulations.

#31) The risk $\mathcal{R}$ is defined independent of the noise. So if noise variance goes to infinity $\mathcal{R}$ will not.

#32)

• Kurtosis is a broader concept than the Pearson definition in terms of the fourth moment, which is of course useless in the case of distributions with infinite 4th moment. See Balanda and Mcgillivray. "Kurtosis: a critical review"
• Eq (19) is a dominated convergence exercise.

#33) We believe it is clearer as written.

#34) This scheduler depends on $R_t$ the deterministic equivalent of $\mathcal{R}(\boldsymbol{\Theta}_t)$ which is independent of the noise. This is indeed the optimal schedule, which confirms the discussion earlier. Note the difference between $P$-risk and $R$-risk.

#35) See our other responses regarding Malladi et al.

#36) We do not like the idea of combining them, as they contain different information. (The suggestion to provide experimental validation is well-taken, and we will look into it).

#37) In our setup, we have $\mathcal{R}(\boldsymbol{\theta_k}) = (\boldsymbol{\theta_k} - \boldsymbol{\theta_k}^*)^T \mathbf{K} (\boldsymbol{\theta_k} - \boldsymbol{\theta_k}^*)$. The Hessian of this function is $\mathbf{K}$, the covariance.

From Appendix point #1) we do not actually drop this term. Equation 44 is a simplification of Equation 36 to be index $i$ independent, we justify this in Lemma 1. This section is to build intuition for the reader to see where the argument is going. From Appendix point #6) as we show the dynamics of signSGD and signHSGD converge. It’s valid to refer to either converging to the other.

Thank you again for your exceptionally detailed review of our manuscript. Your suggestions have improved the structure and clarity of our manuscript. We have collected your points into a few key critiques and address each below. We’ll also reference specific bullet point numbers of your review when relevant.

1. Inclusion of Related Works, Literature Review, and Limitations (Detailed feedback points 1-5)

We appreciate the reviewer’s recommendation to expand the discussion on related work and provide a clearer comparison with the Weak-Approximation (WA) framework [2] and Malladi et al.'s SDE derivations [1]. In response, we have added Appendix F, A new section comparing our approach with the frameworks discussed by [1] and [2], highlighting key differences in focus, methodology, and results.

We summarize some of key points here:

We believe that although our paper also derives an SDE approximator for a discrete optimizer, the main results are fundamentally different. For one, we believe the weak-approximation framework they propose is unable to directly recover a corresponding signSGD SDE from their ADAM SDE formulation. For instance, In order to construct their ADAM SDE, they estimate the normalizing constants $(1-\beta_1^k)$ in continuous time by $1-exp(-c_1 k \eta^2)$ where $c_1 = (1-\beta_1)\eta^2$. This estimator fails when $\beta_1 =0$, meaning that recovering a signSGD SDE is not simply a matter of selecting the appropriate hyperparameters, as is the case when recovering the signSGD optimizer from the Adam optimizer. We discuss this further in Appendix F of our updated manuscript.

We apply a different method of analysis to study optimization algorithms as compared to [1] and [2]. As a result, we arrive at different approximation bounds. Malladi shows a comparison between the SDE and the optimizer as the learning rate $\eta$ goes to $0$. In contrast, we show that the exact dynamics of our SDE and signSGD closely track each other in the high-dimensional limit. This is what allows us to study the behavior of the actual optimizer (signSGD) by studying its SDE. Our goal is not only to derive SDEs for signSGD but to also gain insight into how adaptive methods like signSGD and eventually Adam, behave in the limit of large problem sizes. The aspect of dimensionality is not addressed in [1] or [2] thus requires a different set of tools.

It is true that the weak approximation framework allows for a wider selection of loss functions and noise models. However, our approach on just quadratic statistics and linear regression, although restrictive, allows for more in-depth analysis of the optimizer. For example, the exact effects of the noise can be seen in $\phi$ as well as the transformation of the covariance $\boldsymbol{K}$ induced by the sign operator. Our focus on linear regression and quadratic statistics is certainly also a limitation, and believe an extension to GLMs is a natural direction of future research. However, even in the simple setting of linear regression, extracting these dynamics was a non-trivial task.

references

[1] - Li et al. (2019) "Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations".

[2] - Malladi et al. (2022) "On the SDEs and Scaling Rules for Adaptive Gradient Algorithms".

Please let us know if we can clarify anything further. Your detailed review was quite helpful, and we'd love to know if we were able to address some of your criticisms (particularly those involving comparisons to the weak-approximation framework).

31-This discussion is all in the case of finite variance, where there is an interesting tradeoff and where there is a point to comparing to SGD (the whole of section 4 is about comparison).

Saying more about SignSGD in the infinite variance case is a reasonable goal for future research.

33-Thanks for pointing out the typo about $R \to 0$. Indeed the comment was made with unimodal distributions in mind, your example is a clear counterexample to the statement about kurtosis. In fact, we think (19) is already relatively clear in terms of what the tradeoffs are regarding when it is large or small, so we removed the reference to kurtosis.

38-Sorry we did not reply directly to this point. We assume you refer to the role of conditioning in choosing which optimizer to use: after selecting learning rates correctly, signSGD gains from the diagonal preconditioning when the averaged condition number of $\overline{\mathbf{K}}$ decreases from the averaged condition number of $\mathbf{K}$. We added an explicit sentence to this effect in the text.

39-As above, sorry we did not reply directly to this point. We have discussed in the paper the difficulty of the analysis of $\mathbf{K}_\sigma$ for general covariances $\mathbf{K}$. Aside from the trivial case where $\mathbf{K}$ is diagonal, we do not have a compelling non-diagonal example which would serve as a starting point for further analysis.

Minor Comments:

We appreciate the feedback regarding the formatting and have indeed made many changes as a result. We have already addressed most of the minor comments in the revision. With regards to some of the formatting and language suggestions that we did not implement, we respectfully disagree with the suggestions and have chosen to keep the current version.

Appendix:

Thank you for the feedback! Regarding your previous comments:

1-We have reformatted this section of the appendix to be more readable. Our intent was to provide motivation for the upcoming lemmas. We have also added an overview prior to the appendix to summarize the main content of each section.

2-We note that we do not specify a bound on the risk, hence this may be taken to be as large as one likes. Since the risk cannot blow up to infinity in finite time, if one chooses an appropriate upper bound, with some probability, the stopping time would not be observed within a given time horizon. We show in Lemma 8 for signHSGD (also see Line 1385 in updated) that one may choose a bound so that with overwhelming probability, the bound is never exceeded within the given time horizon. The concentration result of Theorem 1 between signHSGD and signSGD allows us to extend this bound to signSGD. This effectively removes the bounded risk assumption. This can be found in the original manuscript and has been clarified in red in the updated.

3-See 2) in detailed.

4-We believe it is more suitable to move the technical lemmas to a later section of the appendix, as their proofs are largely self-contained and do not directly relate to the core ideas of the main theorem. Given the technical nature of our main result, we consider it more important to present the key part of the appendix—the main proof—at its forefront.

5-This is defined in equation 51 in the original manuscript and equation 56 in updated.

6-This is in fact interchangeable. Provided that both signSGD and signHSGD have the same starting point, they are defined independently of each other, i.e. one does not need to run signSGD to derive signHSGD.

7-Thank you for the catch! In the original manuscript, this was defined in Line 1519 in appendix A.3 after having been moved to a later point in the appendix. It is now defined near the beginning of appendix A.2.

8-Good catch! This has been fixed.

9-In the original manuscript, $\nu_k$ was defined at the start of the appendix A.1 in Line 661 and $\mathbf{v}_k$ was defined in Line 742.

11-In our setting (and indeed, any linear regression setting with well-behaved covariance) the learning rate must be scaled as 1/d; we invite you to conduct numerical explorations in our code (for example in anisotropic_signHSGD.ipynb, change the line ‘lr=0.5’ to ‘lr = 0.5 * d’ and vary the ‘d’).

In neural network settings, the equivalent scaling dimension is typically the width of fully connected blocks --- not the number of parameters. For standard parameterizations, optimal (and often, largest stable) learning rates usually scale as 1/width. This is a well known phenomenon; see e.g. https://arxiv.org/pdf/2203.03466 figure 1, left panel. In your NanoGPT example, the embedding layers have dimension 768, implying a learning rate on the scale of 1e-3 - comparable to the supplied 6e-4.

The difference between the linear regression setting and the neural network setting reflects the fact that the important quantity is the effective dimension. In our problem and notation the effective dimension is $tr(K)/\|\|K\|\|_{op}$. See https://arxiv.org/pdf/2406.11733 for a discussion (where it is called "intrinsic dimension"). In Appendix C.3 of that paper, this quantity is computed for ResNet18 on CIFAR10 and is found to be ~ 1e3.

13-We’ve updated the wording. When you apply the sign-function to the gradient, you are no longer using the original least-squares geometry. So sure, the original landscape ‘exists’ but it is also not the relevant geometry for understanding the algorithm. One point of the analysis we do is to allow a comparison between signSGD, and a langevin type diffusion (signHSGD) whose gradient term is interpretable. So in that sense the relevant optimization landscape has $(\overline{\mathbf{K}} + \overline{\mathbf{K}}^T)$ as a Hessian and gradients are flatted by the $\varphi$ function.

15-One component of the noise is inherent in the setup; our stochastic gradients are always noisy, as they come from a statistical setup. You can get an ODE by sending stepsize to 0, but it need not be the same one as when you have access to the exact gradient.

If you do a batch version of our setup, taking the batch average pre-sign function, it can converge to [3], because the stochastic gradients become exact preactivation for sufficiently large batch.

16-We have made changes to figures 1, 2 and 3). Do they improve the visibility? Besides changing colors, we decreased the number of runs, to allow the confidence intervals of signSGD and signHSGD to fluctuate more (and hence not be on top of each other).

17-This change was made in the last manuscript. The longer time horizon plots are in the appendix. Here we changed the axes to correspond to signSGD iterations.

19-Indeed, we could optimize for $p$ if the constant $c(2p,K)$ were known. But we don’t actually know the precise values of $c(2p,K)$. See Lemma 12 and Corollary 2.

20-Theorem 2 is the Theorem? The SGD theorem is in Collins-Woodfin et al, which is cited. We’re not quite sure to which statements you are referring. We agree that all mathematical statements we are making should be stated formally (barring trivialities); so we have cleaned up the appendix, including putting more statements inside of formal lemma claims (see below)

21-Ok

22-Noted

23-See our response to 15 above. We chose this specific SGD-as-SDE formulation for comparison because it most closely matches our chosen setting for signSGD

24-As you write, this scheduler has problems as the risk goes to 0 (this is part of the discussion in Section 4.1, and is one of its undesirable features). For HSGD, it will actually cause the risk to not go to 0, even in the presence of 0 noise (as the risk gets small, the LR will go way above the descent threshold, causing the risk to increase). The observed dynamics of the risk will stabilize at a positive level. (In the isotropic case, using the ODEs in (25), one can easily verify this).

If one runs normalized SGD and performs an exponential moving average of the gradients, like in Cutosky and Mehta ‘20, with bandwidth B, then in a limit of $B \to \infty$ and $d \to \infty$ with $B \ll d$ one should recover this scheduled HSGD as a high-dimensional equivalent.

26-In Figure 1) we do confirm that the dynamics of signSGD, signHSGD, and signODE match as predicted in Theorems 1) and 2). What more is missing?

28-In the $P$-risk you have already taken an expectation over the noise, so either the variance is finite (like we are assuming in this section) or it is infinite and this makes no sense. If you are attempting to infer the $P$-risk, of course your point is good, but that's not what is being discussed here.

Dear reviewer, we thank you for taking the time to thoroughly look over our response.

We have completed experimental validation of Theorems 3 and 4. For validation of Theorem 3 please see Figure 6 of the newly revised manuscript. We have also updated Figure 1 to include the limit risk value of Theorem 3 when the noise is Gaussian as well as the CIFAR10 dataset. Theorem 4 is validated in Figures 7 and 8. The code has been updated to include these new experiments.

We now respond to your overall comments:

1. In our high-dimensional framework, we are dealing with a sequence of problems, where the complexity increases as the dimension grows. If we were to fix the dimension, in some sense, our work amounts to performing a kind of weak approximation. So there are strong parallels.

However, there are aspects of our framework which are also not present in a fixed-dimension setup. In particular, the risk curves concentrate around ODEs, which encode the effect of the randomness. This is not present in weak approximation, and it is a key result of this paper. We also hope this illustrates that there is some theory that we have to build from the ground-up, as this is nowhere in [1] or [2].

So to summarize: we are not dismissing Weak Approximation. But it’s also not a tool which is built for what we’re doing, which is to understand sequences of problems of growing dimension. Nor do the key results of our theory appear in weak approximation (risk curve concentration). They’re just different; they have different conclusions, different assumptions, and even different proof methods.

2. i) We don’t claim to check our assumptions on datasets. We have checked the predictions – that the ODEs match the behavior of the optimizer – on real datasets (See Figure 1). Our Assumptions, especially 3 and 5, are designed with random problems in mind – it’s not even clear what they would mean for a single dataset.
   ii) The suggestion of forming random projections is smart. We would be optimistic.
   iii) The main goal of this paper is to derive the high-dimensional limit of SignSGD. This is not a weak approximation, and these assumptions are important and reasonable from the point of view of matrices in high dimensions.

Your point about checking assumptions is reasonable, and assumptions should be challenged. There are lots of directions to take this work. One would be to find non-asymptotic assumptions (e.g. in terms of ‘intrinsic dimension’ – see the discussion below), which could be applied to fixed problems. We think your request makes sense for a paper attempting to make those non-asymptotic comparisons, instead of trying to check these assumptions (which are not engineered for that task) on general empirical problems.

Detailed feedback:

1,2-These are of course good things to have in the paper, but there are space considerations. Discussion in the appendix will be ignored by everyone except a small fraction of the audience. We would need to sacrifice the content of the paper to do this, and these do not raise to the level of concern as to sacrifice that. So in our judgment, it's simply inappropriate to add these.

The most significant limitations of this paper are built into this setup (e.g. linear regression), and we are certainly not hiding these. How does the theorem really change outside of linear regression? We don’t have anything smart to say, and it’s certainly a valuable direction for future research. Again, we don’t think this type of discussion adds enough to the paper to warrant cutting something else.

Also, the discussion we’re having is preserved on openreview. Many researchers benefit from openreview by going to the discussion after making an effort at reading the paper, and so there’s value in already having them here.

4-This relates back to 1) in main point

7-The diffusion term for signHSGD is trivially Lipschitz (for SDE's the relevant Lipschitz-ness is in the spatial variable), and this diffusion term is constant in space. Theorem 5.2.5 of Karatzas Shreve applies directly. It applies as well to the HSGD, although we are not proving anything about it.

Linear regression can have multiple minimizers in overparameterized settings, you are correct. Under assumption 2, the coefficient is Lipschitz regardless of the existence of multiple minimizers. Without assumption 2, you have to stop the SDE when you reach a small risk (as we do in Theorem 6).

8-This was changed in the last manuscript, text currently reads: "As we will see, an important characterizing feature of signSGD is its effect on the covariance of the signed stochastic gradients." Please elaborate.

... (continued)

References

[1] -Li et al. (2019) "Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations".

[2] - Malladi et al. (2022) "On the SDEs and Scaling Rules for Adaptive Gradient Algorithms".

38. I see: Unfortunately, the information you mention is typically unavailable since $K$ and similar matrices are unknown in practice. Are there maybe practical insights that could be drawn? For example, can you recommend when a practitioner should choose signSGD over SGD based on observable behaviors in deep learning?

39. To clarify, there is no formal statement or experiment backing this claim?

Regarding the appendix: It has improved significantly. While there is room for further refinement, the major rework is promising.

Regarding experiments: Nice job verifying Theorem 3. However, Theorem 4 presents a bound, and your verification focuses solely on $R_{\infty}$, which is essentially verifying Theorem 3 again. Shouldn't you have plotted the curve described by Eq. 29 as a function of $t$?

Unfortunately, many other insights are not verified. For example, the effects of the schedulers are not verified --- Is SGD with the scheduler in Eq. 14 actually behaving like signSGD to some extent? The list goes on and it would be easier to pinpoint them all if your insights were formally conjectured in formal statements.

I am not asking for more experiments: I am simply acknowledging that there are some formal results with some degree of verification, some others without, and quite some informal insights without any experiments to support them.

To Conclude

I acknowledge the major changes and triple revision made to the paper: These include images, additional literature review, rewording of unclear paragraphs, and experiments. On top of this, great progress has been made on the appendix which is now significantly different (and better) than the original submission.

However, the following concerns persist:

1. Assumptions: These appear non-standard in the literature, with insufficient justification for their validity.
2. Formality vs. Intuition: Many insights remain qualitative or intuitive, leaving them open to counterexamples. For a theoretical paper, formal rigor should take precedence.
3. Experimental Validation: Some results and insights lack experimental support. While I understand it may now be too late to address this, it is a notable limitation.

Post Scriptum

I want to highlight that I enjoyed this back-and-forth. I like what you are doing in this paper and believe it can shine. I would have not invested so many hours in this review process if I did not value this paper and your dedication.

If my feedback has come across as overly critical or adversarial at times, please know that my intention has always been to be the kind of reviewer I would have appreciated having: someone who thoroughly reads and evaluates the work, thoughtfully interacts with it, and strives to help the paper evolve into the best possible version for both the authors and the broader community.

I wish you every success moving forward, and I will now discuss this paper with the other reviewers and the Area Chair before finalizing my recommendation. Thank you again for your dedication and hard work!

Dear Reviewer, Thanks for your comments; we address your concerns with the approximations here.

Discuss any theoretical or empirical evidence the authors have for how well the continuous approximation matches discrete SignSGD over long time horizons.

The constant $C(\overline{\mathbf{K}},\epsilon)$ in Theorem 1 controls the timescale of good approximation. The key point here is that $C(\overline{\mathbf{K}},\epsilon)$ is independent of dimension. The rescaled time $T$ in the theorem is $k/d$, where $k$ is the number of iterates; therefore the approximation is good for $O(d \log(d))$ iterates. This is a key benefit of our theoretical analysis; by using the high dimensional limit, we have a bound that does NOT require a small learning rate (such as in the provided reference [1], see also [2] and [3]). We provide a more detailed comparison with approaches [2] and [3], and how our approach avoids certain limitations of those works, in Appendix F of our revision. We believe this is highly related to the issues raised in your reference [1].

Figure 1 in our original submission shows the correspondence between our theory and the true stochastic optimizer over a time period of 5000 iterates; the original figure was labeled in rescaled time (iterates/$d$). We have corrected the figure, and have added a plot in the appendix which extends this another factor of 10 with good agreement (Figure 5), following your suggestion.

Clarify if there are specific regimes or conditions where the authors expect the approximation to break down for long-time behavior.

As stated by the theorem and validated by our experiments, we expect that for high dimensions the approximations are good for timescales long compared to the overall relaxation time of the dynamics. One case that might require a very high dimension for good approximation would be power-law features and covariance structure; this is studied in [4] for the SGD case.

Consider adding a discussion of the limitations of the continuous approximation approach, particularly for long-time dynamics, to their paper.

We have added some more discussion about other continuous approximations. Regarding very-long-time distributional behavior, such as convergence to the invariant distribution -- we can provide some information, such as Theorem 3, which describes the limit risk for constant stepsize. It's an interesting question if the invariant distribution of signHSGD meaningfully relates to the invariant distribution of signSGD, beyond this, but we don't have much to say.

[1] Kushner, H. J. (1982). A cautionary note on the use of singular perturbation methods for “small noise” models. Stochastics: An International Journal of Probability and Stochastic Processes, 6(2):117–120.

[2] Li et al., Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations: Journal of Machine Learning Research, 2019

[3] Malladi et al., On the SDEs and Scaling Rules for Adaptive Gradient Algorithms: NeurIPS 2022.

[4] Paquette et al. 4+ 3 Phases of Compute-Optimal Neural Scaling Laws. arXiv preprint arXiv:2405.15074 (2024)

Thank you very much for the rebuttal. You can find the detailed reply below the thread of our interactions.

1. I took a closer look at the paper by Malladi, and indeed, deriving the SDE of signSGD from that of Adam and/or RMSprop is not feasible. One would need to derive the SDE for signSGD "from scratch" to enable a clear comparison, which would be an interesting challenge to undertake (if it is even possible to derive an SDE for signSGD under WA, which is not immediately obvious).

2. This is a cumbersome claim. In my opinion, the fact that WA does not require any constraints on the dimensionality inherently makes this a non-issue. I would like to highlight that [1,2] perform experiments to validate their SDEs on PreResNet32, VGG19, and ResNet-50, which have up to 138M parameters, using learning rates that are not infinitesimal but rather standard, such as $10^{-3}$. Therefore, I would advise caution when making these comments and suggest refraining from dismissing a theoretical framework that has been successfully employed and experimentally validated in various setups for years. This is particularly important considering that in your experiments, $d$ ranges from $d=50$ to $d=750$, which can easily be classified as low-dimensional relative to the DL models mentioned above.

3. I agree that the comparison is difficult to perform, and demonstrating this adds value to your paper. It does not detract from WA but rather shows that deriving the SDE for signSGD directly from Adam is not the appropriate approach.

Finally, I look forward to seeing your new experiments to substantiate your claims, as these are still missing.

I would like to add an additional point that I initially overlooked: I apologize for this oversight, but it is an important issue. The experimental details are quite poorly described. It would be best to specify all the hyperparameters, design choices, and related details. While one could read your code to find this information, it is not standard practice, as your code may become unavailable in the future.

[1] Li, Zhiyuan, Sadhika Malladi, and Sanjeev Arora. "On the validity of modeling SGD with stochastic differential equations (SDEs)." Advances in Neural Information Processing Systems 34 (2021): 12712-12725.
[2] Malladi, Sadhika, et al. "On the SDEs and scaling rules for adaptive gradient algorithms." Advances in Neural Information Processing Systems 35 (2022): 7697-7711.

Dear reviewers, AC,

We have uploaded a second revision.

There are many minor updates, which are marked in red, which we will not summarize here.

Besides these, we made the following more substantial changes

• We added an Appendix D with additional supporting experiments (3 sets, which are designed to add numerical verification to Theorems 1/2, Theorem 3 and Theorem 4 respectively. We have also updated the anonymous github linked in the paper with the code to generate these figures.
• We made a pass through Appendix A and B to improve the presentation of the proofs, putting all statements inside of formal lemmas, and adding additional details where we thought they were lacking.
• We added details to Appendix H on the how the experiments were conducted.
• We updated the figures in the main text, with the intention of making them clearer to colorblind readers and improving the content of Figure 1.