The High Line: Exact Risk and Learning Rate Curves of Stochastic Adaptive Learning Rate Algorithms

Summary

The paper is a continuation of a broad line of work exploring deterministic limits of stochastic gradient descent. In particular, the author(s) follow one specific branch of the many, which derives a solution to a coupled integro-differential equation by passing on the complex space. The novelty lies mainly in the fact that they derive, for least squares, deterministic limits of dynamics that have adaptive step-sizes, for non identity covariance of the observed dataset. Industrially, this is important as the idealized scenario of fixed step-sizes is __very much__ idealized, and realistic datasets have structure. Thanks to the general expressions (under 6 assumptions that are claimed to be not super-restrictive), they can proceed to study the various differences between standard schedules. The limiting learning rate can be derived, as well as the scaling of the risk along time. This is done for Line Search, Polyak step-size, AdaGrad Norm. Interestingly, their bounds suggest that for covariances with largely separated eigenvalues (say, power law spectrum), there should be a change in the phenomenology (i.e. a phase transition). They verify this for a power law covariance, which exhibits three phases depending on the parameters of the problem (notably the starting distance from $X^{\star}$). All is accompanied by experiments corroborating the predictions.

Strengths

The motivation in lines 17-31 is very well stated. As a reader, I would be inclined to continue parsing the paper. - Proposition 4.3 states that with a null initialization $X_0 = 0$ and subject to a condition on the covariance (which roughly says that $\lambda_{\min}$ can be small but that there is a very small number of vanishing eigenvalues), if the ground truth signal is spread out, then the learning rate is never zero. For symmetric activations (say, phase retrieval), and constant learning rate, I think you cannot really start at zero because you have a symmetry, you need at least a tiny bit of overlap. From this result, it appears that you can start from the most useless guess and still get away with it, provably. Interesting. - The higher arching aim of the work is very important, as knowing that the behavior of SGD is predictable is pivotal for understanding how our models work. - The techniques are of independent interest and are a continuation of a line of work that is systematically, step by step, developing a global answer to these questions. - The result that Polyak's step-size converges to the inverse of the average eigenvalue is yet another validation of a method originated from convex optimization holding in more generality. This is reassuring.

Weaknesses

Weaknesses are questions and questions are weakenesses in some sense. Once the author(s) engage, I will reorder things that get answered and things that remain as "concerns". - Proposition 4.4 supposes that $\mathscr{D}_i^2(0)$ could potentially have a different scaling at each site $i$. I understand that the assumption includes an uninformative initialization $\delta = 0$. However, when is it really that one can think of a setting in which you can have an initialization that scales like $\lambda_i^{-\delta}$ for non trivial $\delta$? The result is very aesthetic, but I did not think through of its realistic value. Anyways, it is very nice so I am __not__ criticizing the math. Maybe I am just interested in understanding if you have in mind something. - In the Checklist, you state that you have been "__very careful__ \[...\] to include __all__ the assumptions needed". I agree with this statement. However, as a reader, I do not agree with this being the limitation of your work, as in some sense you claim over the whole text that your assumptions cover standard adaptive step-sizes (hence you hiddenly claim that this theory is all we need). I would be very happy to see a paragraph where you clearly state your limitations. Where does your theory fail? Where does the technique __not__ extend naturally? If you want to sell me your paper, I do not believe its value is in the assumptions. On a first aspect, the ideas and the techniques are of independent interest. Secondly, it would be sad if the bottleneck were the assumptions. - Your equation (7) is found across some other works. One of the earliest appearances is (Yoshida and Okada, 2019). However, from a first check, I would guess that you are missing a term on the first object. To be fair, I quickly tried to re-derive it and I would guess that you need a further $H_{2, t}\mathscr{V}_{12, i}(t)$ in the first expression. Am I wrong? If so, please correct me. It looks like the equations are missing some symmetry of the dynamics. - I am concerned about the scaling. My understanding of the literature is that scalings are very important. Your assumptions require that $\lVert K\rVert$ is bounded independently of the dimension and that $\lVert X_0\rVert$ is bounded, independently of the dimension. Moreover, your step-size is $\frac{1}{d}\mathfrak{g}_k$. Therefore, at the beginning you move by very little, I think by $\propto \frac{1}{d}$. What is done in some other references is very different. Let me take an example, and correct me if I am wrong. 1. You can derive heuristically deterministic dynamics for step-size $\mathfrak{g}_k$ (not normalized), starting from a standard Gaussian standard gaussian vector $X_0$, with data-points being i.i.d. Gaussians $a\sim \mathcal{N}\left(0,\frac{1}{d}I_d\right)$. Now, the step-size is not divided by dimension, the $X_0\sim \mathcal{N}(0, I_d)$ is not bounded in norm by something independent of $d$, and the $a$ matches your assumptions. In this case however, the signal is far stronger, as we are removing two normalizations you indeed have. Below my questions regarding this. 2. Can you clarify why it makes sense to choose your scaling? Apart from it allowing the proofs. 3. In particular, how is it not possible to allow for norms dependent on $d$, and in parallel, how is it that the signal is enough to be not stuck at the initialization? 4. Does this implicitly mean that you are exploring a slow regime? Say, a Gradient Flow? What happens if in your experiments you just change the scalings to something larger in magnitude? #### Typos __NOTE__: I am including for completeness typos in "weaknesses". Please do not count them as such. - (line 560) In the reminder of $\Omega$, you are missing a $\max$. - (equation 20) there is no $\Gamma(t)$ in the expression, right? - (lines 1140-1141) the way of quoting "noise" and "variance" is not the correct one, TeX-wise.

Questions

Please note that some of the questions below are very ingenuous. - (lines 49-57) This is more of an impression. Why is $\lambda_{\min}(K) > C > 0$ for you a situation of strong anisotropy? The power-law case is anisotropic in that the minimum eigenvalue is not bounded below by a constant as $d\to\infty$. Put another way, I do not understand how the __minumum__ eigenvalue being bounded below from zero would describe the whole behavior of the covariance profile of the data-points $a$. I understand that this might be related to the fact that the optimal strategy uses the trace so you say something in the line of "if the minimum eigenvalue is very well distant from the average then there is anisotropy". At the very least, it appeared to me as a reader that you somehow __define__ strong anisotropy as this eigenvalue condition, while in fact it is a consequence of the algorithmic behavior that you can tune the eigenvalues in this "anisotropic" way and observe that "greed can be arbitrarily bad". - The original introduction of the Polyak step-size is that of a step-size that uses the optimal value of the function, and the value of the current iterate. In your case, what you call Polyak step-size has a connection with the original method, which is related to your assumptions, as you find that eventually it gets to a $\gamma_{\infty}$ that takes care of the geometry of the function. From a pedagogical point of view, it would be nice to comment further on this. - Theorem 2.1 states that the concentration is for any $\epsilon \in \left(0,\frac{1}{2}\right)$, line 48 says "for any inverse power of d", and Theorem B.1, Corollary B.1 cited below Theorem 3.1 in line 195 give a result for existance of an $\epsilon$ such that the approximate solution of the integro-differential equation is verified. Can you shed light on these differences? Since it is your main Theorem, it would be nice to have an explicit discussion of what goes on where. More in detail, why do you get from existance in the appendix to a whole interval of $\epsilon$ (that by the way gives a rather slow decay wrt to say $\epsilon > 1$), to claim in text "for any inverse power"? I have to be honest here and I did not spend time re-doing all the proofs. - Related to the point above, in equation (27), your supremum is over $T\wedge \hat{\tau}_M$. In line 531 you state that you will not specify which stopping time it is when it is "clear from context", but for (27), given that you want to prove (26), my natural question is if stability is for $\hat{\tau}_M(S)$ or $\hat{\tau}_M(\mathscr{S})$, as you only mention one. At the very least, it is not clear from context here which stopping time it is in (27). - Lemma C.1 requires as condition that $S(t, z)$ is an approximate solution over the $\xi$-mesh. How do you guarantee that this condition holds in your case? If ever, it is not evident from text that this is the case. If it is a direct consequence of arguments in reference 14, I would appreciate it being made explicit. - Can you clarify the sentence of lines 239-240? Ignore the "sufficiently anisotropic" aspect which I already discussed above. The idea that smaller learning rates lead to under-performance is not entirely true. In the simplest setting if you over-shoot with your learning rate you will never converge. I believe here this is tightly linked to your Table 1, but there line search and Polyak have the same convergence rate for the risk. Where is the difference? - In Figure 1, how many runs do you make? It would be nice to know. - The same question above but for any other plot where you show error bars. - (line 1125) You choose as scaling for the initialization $X_0\sim\mathcal{N}\left(0,\frac{I_d}{\sqrt{d}}\right)$ and the ground truth $X^{\star} = \frac{1}{\sqrt{d}}\mathbf{1}_d$. The norm of the latter is $1$, but the former has norm concentrating at $\sqrt{\sqrt{d}}$, and this violates your assumptions (not bounded by a constant independent of $d$). Is it a typo? Or is it that the experiment does not match the assumptions? If I look at the scaling of Figure 3, line 1099, you indeed sample $X_0\sim \mathcal{N}\left(0,\frac{I_d}{d}\right)$. - Proposition 4.4 is nice. What happens when $\beta +\delta > 2$? My soundness score is due to the many questions. I hope I will raise it after the discussion. The overall score follows the same principle.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

None.

Summary

Analyzes the behavior of AdaGrad-Norm, Line-search and Polyak step size using ODEs.

Strengths

This paper provides an interesting characterization of the behavior of several optimization methods using ODE tools, which are not as heavily used in optimization theory as they should be. This work seems technically sound, although I didn't review the (long) Appendix. I am in favor of acceptance, although with reservations. The methods studied are interesting to theoreticians, but they are rarely if ever used in practice, which limits the practical appeal of this work. Especially the idealized special cases for the line search and Polyak step size. The focus on quadratic problems is also a major limitation, especially since the analyzed methods are not typically used on quadratic problems. There is definitely potential for followup work to build on this work.

Weaknesses

- Paper spends a lot of time setting up a general framework, which it then applies to specific settings. The amount of notation introduced seems over the top for a conference paper. For instance, the results for least-squares can be presented without establishing the full general framework, which could be left to the Appendix. Section 1.2 and 2 are just .... a lot to get through. - Paper ends very abruptly. - Some plot fonts are too small to read when printed. - Further discussion of the implications and potential applications would be good.

Questions

See above.

Rating

7

Confidence

2

Soundness

4

Presentation

2

Contribution

2

Limitations

See above

Summary

### Update after rebuttal After reading the feedback carefully, I updated my score for the paper. ### Original review In this work, the authors study SGD and its adaptive variants in the setting of noisy generalized linear models. Assuming the covariance matrix of the data is bounded by a dimension-independent constant and a few more similar assumptions, the authors establish that the expected risk of SGD and the expectation of adaptive learning rates both converge to deterministic curves described by an ODE. I found the paper rather hard to follow and I couldn't understand the significance of the obtained results. The theory requires a lot of simplifications: the prediction model is linear; the learning rates that are studied are "idealized", so they are not the same as those used in practice; Adagrad-norm isn't really a method popular in practice; the ODE solution is not trivial in the case of least squares, as equation (11) is implicit, so it requires further assumptions to make a statement about the solution. I did not understand what insight we can draw from all of this.

Strengths

1. The theory reveals that line search leads to slower convergence that Polyak stepsize. 2. In some simplified cases, there exist exact expressions that describe the dynamic of the studied methods. 3. The theory is also supported by numerical simulations.

Weaknesses

1. The theory for line search and Polyak stepsize is only for the idealized version. 2. The insight that line search methods do not work particularly well is not exactly new, see for instance J Bolte, E Pauwels "Curiosities and counterexamples in smooth convex optimization" 3. The complexity of the theory and the discussion is such that it is very difficult to draw any insight, and it seems the theory wouldn't be extended by others because of how convoluted it is. ### Minor issues Abbreviation "SGD" is used before (line 23) it is introduced (line 30) "We fully expect" -> "We expect" "worst case convergence guarantees" -> "worst-case convergence guarantees" "strongly-convex" -> "strongly convex"

Questions

1. What's the intuition for why $h$ exists and why would it be well-behaved? I can see a discussion in Appendix B, but it doesn't provide any high-level intuition. 2. Why do you say in Appendix B that $R(X)$ is "an expectation of a Gaussian vector"? 3. The authors wrote "We shall remove this condition by approximation in what follows", where is it done? 4. The authors say they need "mild assumptions on the learning rate", but why are they mild?

Rating

5

Confidence

3

Soundness

3

Presentation

1

Contribution

3

Limitations

The authors could do a better job when discussing the limitations of their work. There is no explicit limitations section and in the ckecklist, the authors state that they already did everything needed by clearly outlining the required assumptions. However, I think it'd benefit the paper if the authors made an explicit statement on how some aspects of the theory are too restrictive.

Summary

This paper studies Stochastic Gradient Descent (SGD) training with adaptive step sizes. The setting is a generalization of single-index models. There is a ground truth vector X^*, and the model must find a vector X such that a loss L(X) = E_{a,epsilon}[f(<X,a>, <X^*,a>, epsilon)] is minimized. Here epsilon is a random additive error, and a is distributed according to N(0,K) for a potentially anisotropic covariance K. This paper shows that (1) in the high-dimensional regime, the training dynamics converge to a deterministic limit given by an ODE. It studies these dynamics in two significant cases of adaptive step-size algorithms: (2) exact-line-search (i.e., greedily decreasing the risk optimally at each step), and (3) AdaGrad-Norm. For exact-line search, it is shown that this can be very suboptimal if the data covariance is highly anisotropic. For AdaGrad-Norm, it is shown that if the data covariance has lower-bounded eigenvalues the optimal step size (within a constant factor) is reached. However, for harder problems with worse conditioning of the data covariance, AdaGrad-Norm can be overly pessimistic.

Strengths

This paper provides a fine-grained analysis of practical optimization algorithms, and also proves a general theorem on convergence to deterministic dynamics that I believe could be used to analyze dynamics for different adaptive-learning-rate algorithms in the future. There is also a strong match between theory and experiments. I believe it is an important step in the direction of understanding these practical optimization algorithms.

Weaknesses

In terms of techniques, Theorem 2.1 on convergence to an ODE is based in large part on modifying a previous analysis of [Collins-Woodfin ’23] to the case of adaptive step sizes, so the techniques are not completely novel. The data is also assumed to be Gaussian, although anisotropy is allowed.

Questions

I am satisfied with the presentation in the paper, and do not have questions.

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

Yes

Summary

The paper focuses on analytical analysis of dynamics of a class of optimization algorithms with adaptive learning rate applied to linear models with Gaussian data. The class of algorithms includes AdaGrad-Norm, RMSprop-Norm, Polyak stepsize and line search, but doesn't include, for example, Adam, classical RMSprop and AdaGrad. Also, authors consider $\frac{1}{d}$ scaling of learning rate with dimension $d$ of the problem. First, the authors establish convergence of the optimization trajectory to a deterministic ODE in the limit $d\to\infty$. Then, they focus on analysing the resulting dynamics for quadratic loss functions and specific algorithms - Plolyak stepsize. Obtaining learning rate. - Exact line search. Obtaining learning rate for two cases of data covariance $K$: isotropic and with two distinct eigenvalues - AdaGrad-Norm. $O(t^{-\frac{1}{2}})$ learning rate asymptotic for noisy observations. Also, learning rate and risk asymptotic for isotropic and power-law data covariance $K$. The theoretical results are validated with numerical experiments on synthetic data.

Strengths

- The main result (convergence to deterministic continuous-time ODE) is valid for general loss functions. The precise characterization of the optimization trajectory, as opposed to upper/lower bounds, is typically harder to obtain for non-quadratic loss functions - For AdaGrad Norm algorithm, the authors characterized the dynamics of risk and learning rate for several distinct types of covariance matrix spectrum. However, both these points are achieved only partially (see *weaknesses* for details). Completing them would significantly improve the paper.

Weaknesses

The main drawback of the manuscript is its writing. The paper mostly focuses on the introduction and setting, leaving too little space for actual results and their discussion. Out of 9 pages: - Five pages are reserved for introduction and setup. - Two pages are occupied by Section **2**, which mainly introduces notations and demonstrates them on the quadratic problem example. Indeed, it contains Theorem **2.1**, but that one is quite short - Only the last two pages contain the main results of the paper. This leaves very little space for discussing them. Some of the discussions are present in the (quite lengthy) *main contributions* part, but, for readability purposes, it would be nicer to have this discussion together with or after the results. - Any kind of concluding section is absent Regarding the figures, only figure **1** is referenced in the text while figures **2**,**3**, and **4** are not referenced. This adds an extra effort to interpret them and put into right context. **New developments from previous work**. As mentioned by the authors, the current paper follows the technical framework developed in [14], extending it to adaptive algorithms. Then, it would be good if the manuscript mentioned more clearly which results/notions were already obtained in [14] and which are new additions of the current paper. - In particular, Theorem **2.1** of the current paper seems to be very close to Theorem 1.1 of [14] which can also handle non-constant learning rates. If two theorems are indeed close, it would be interesting to discuss new aspects introduced by adaptive algorithms. - Assumptions **1-6**, and the notations of sections **1.1**, **1.2**, **2** look very close to that of [14]. If that is the case, it would be interesting to discuss more specifically their differences. Also in that case, sections **1.1**, **1.2**, **2** could be compressed by partially moving them in appendix, leaving more space for the results and their discussion. **Results for non-quadratic losses**. Although the framework, as introduced in section **2** is designed for general non-quadratic problems, the specific results are given only for quadratic problems in the main paper. Some results for non-quadratic problems are given in appendix sections **D.1**, **E**, **F.1**, **G.1** (mainly based on results of [14]). In such form of these results, it is hard to extract interesting conclusions, or determine the contributions compared to the previous literature.

Questions

**High dimensional vs. low dimensional problems**. Some covariance matrices $K$ satisfy assumption **1** but have $\frac{\mathrm{Tr} K^2}{d}\to0$ and therefore produce trivial ODE dynamics coinciding with full-batch Gradient Flow (as can be seen from eq. 11), for example - rank 1 with non-trivial eigenvalue $\lambda=1$ (extreme case). - with power-law spectrum $\lambda_i=i^{-\alpha}, \alpha>1$ (e.g. classical *capacity* and *source* conditions). Indeed, for these two examples, the dynamics converge to full-batch Gradient Flow due to $\propto \frac{1}{d}$ scaling of learning rate in eq. 3, while learning rate independent of $d$ would be more natural. - Could the techniques developed in this paper handle such problems? - If no, is it possible to formulate a criterion on covariance matrix $K$ to distinguish non-trivial ODE dynamics from trivial ones (i.e. coinciding with full-batch gradient flow)? For general non-quadratic problems. **Results for power-law spectrum** The last section provides variety of results in case of power-law distribution of eigenvalues (with exponent $\beta$), and target (with exponent $\delta$). In the non-vanishing learning rate phase $\beta+\delta<1$, how the obtained rate $\mathcal{R}(t)\sim t^{\beta+\delta-2}$ compares with other results for SGD under power-law spectrum (e.g. [1](https://arxiv.org/abs/2006.08212),[2](https://arxiv.org/abs/2102.03183),[3](https://arxiv.org/abs/2206.11124))? **Learning rate phase transition**. For power-law spectrum, learning rate asymptotic exhibits transition at $\beta+\delta=1$, as demonstrated on fig. **4**. Interestingly, on this figure we can see that in both phases learning rate reduces similarly up until it reaches very small value ($\gamma\sim 10^{-4}$) at quite late optimization times ($t\sim10^{3}$). Can this dynamic transition (w.r.t. time $t$) be explained theoretically, for example, with some quantities characterizing transition scale with $d$?

Rating

5

Confidence

3

Soundness

3

Presentation

1

Contribution

2

Limitations

The paper surely doesn't have any negative societal impact.

Proposed conclusion/limitations paragraphs

**Conclusion paragraph (to be added):** *This work studies stochastic adaptive optimization algorithms when data size and parameter size are large, allowing for nonconvex and nonlinear risk functions, as well as data with general covariance structure. The theory shows a concentration of the risk, the learning rate and other key functions to a deterministic limit, which is described by a set of ODEs. The theory is then used to derive the asymptotic behavior of the AdaGrad-Norm and idealized exact line search on strongly convex and least square problems, revealing the influence of the covariance matrix structure on the optimization. A potential extension of this work would be to study other adaptive algorithms such as D-adaptation, DOG, and RMSprop which are covered by the theory. Studying the asymptotic behavior of the risk and the learning rate may improve our understanding of the performance and scalability of these algorithms on more realistic data. Another important application of the theory would be to analyze the ODEs presented here on nonconvex problems.* **Limitations (to be added):** *The current form of the theory is limited to Gaussian data, though many parts of the proof can be extended easily beyond Gaussian data. The main ODE comparison theorem is also only tuned for analyzing problem setups where the trace of the covariance is on the order of the ambient dimension; when the trace of the covariance is much smaller than ambient dimension, other stepsize scalings of SGD are needed. In addition, the analysis is limited to the streaming stochastic adaptive methods. We conjecture that a similar deterministic equivalent holds also for multi-pass algorithms at least for convex problems. This has already been shown in the least square problem for SGD with a fixed deterministic learning rate \[42,44\]. Lastly, numerical simulations on real datasets (e.g., CIFAR-5m) suggests that the predicted risk derived by our theory matches the empirical risk of multipass SGD beyond Gaussian data (see for example Figure 6).*

(Response to questions)

We thank the reviewer for finding typos, which have now been fixed. We answer below the reviewer's additional questions. **Intuition and existence of $h$ and why is $R(X)$ is “an expectation of a Gaussian vector”** The intuition for $h$ comes from the fact that $a^TX$ and $a^TX^{\\star}$ are correlated Gaussians with covariance given by $\[X|X^{\\star}\]^T K \[X|X^{\\star}\]$. Here $[X|X^{\\star}]$ means one concatenates $X$ and $X^{\\star}$. Since $a \\sim N(0,K)$, for any fixed vector $X$, then $y \= a^TX \\sim N(0, X^TKX)$. Both $y \= a^T X$ and $y^{\\star} \= a^TX^{\\star}$ are Gaussians, but they are correlated normal since they both depend on the same $a$. A simple computation shows that the correlation between $y$ and $y^{\\star}$ is given by $\[X |X^{\\star}\]^T K \[X | X^{\\star}\]$ where $\[X|X^{\\star}\]$ is the concatenated vector. This is to say that $(y, y^{\\star})$ is Gaussian and $(y, y^{\\star}) \\sim N(0, \[X|X^{\\star}\]^T K \[X|X^{\\star}\]$. (See Section 1.1). Suppose $\epsilon =0$ as otherwise this will just add another integral. From Eq (1), using the density of multivariate Gaussian $(N(0, \[X|X^{\\star}\]^T K \[X|X^{\\star}\])$, $$ R(X) \= \\int f(y,y^{\\star}, \\epsilon) e^{(y, y^{\\star})^T G (y, y^{\\star})} \\, dy dy^{\\star} $$ where G is the 2 x 2 matrix, $\[X|X^{\\star}\]^T K \[X|X^{\\star}\]$. Integrating out $y$ and $y^{\\star}$, $\\int f(y,y^{\\star}, \\epsilon) e^{(y, y^{\\star})^T G (y, y^{\\star})} \\, dy dy^{\\star}$ is a function of $\[X|X^{\\star}\]^T K \[X|X^{\\star}\]$. This function is $h$. Moreover, this also explains why $R(X)$ is an “expectation of a Gaussian vector” as $(y, y^{\\star}) \= (a^TX, a^T X^{\\star})$ is distributed as a multivariate Gaussian with covariance $\[X|X^{\\star}\]^T K \[X|X^{\\star}\]$ and you are taking the expectation over $(y, y^{\\star}, \epsilon)$. **“We shall remove this condition by approximation in what follows”** The reviewer is absolutely right: we did not end up removing the condition, although there is a straightforward way to do so, and we will correct this. In short, one creates an $f\_\\epsilon$ which is an approximation to $f$ formed by convolving with an isotropic Gaussian of variance $\\epsilon$. This is $C^2$ and has bounded second derivatives (as $f$ was smooth). We then will take limits as $\\epsilon \\to 0$. We note that Lemma B.1 was intended as discussion, and it is not used in the main theorem. Theorem B.1 generalizes to the more standard L-smooth f, at the cost of additional complexities in the formulation (discussed below the Lemma), and is intended to show why we didn’t simply use the more standard L-smooth f assumptions. **“Mild assumptions on the learning rate” Why mild?** The learning rate assumptions are mild since they encompass many adaptive learning rates used in practice; they do exclude learning rates which are non-concentrating. Among non-concentrating stepsizes, one can further consider those which bias the gradient and those which do not bias the gradient. Those which bias the gradient require a totally different theory (e.g. scaling the gradient to be norm-1 or things like gradient clipping). Those which do not bias the gradient would have a similar theory as those which we consider, but with an additional variance term in the ODEs (e.g. RMS-prop with a very short window of averaging, with respect to dimension). So we chose this case to keep the theory *as simple as possible* (which we know is aligned with the reviewer’s desires\!)

(Response to questions)

**Answers to the reviewers’ questions** * **Strong anisotropy.** When we say “strong anisotropy” we are not referring to the lower bound on the eigenvalues. Rather, we refer to context where, for fixed $\\lambda\_{\\min}$, we have $\\lambda\_{\\min}/(\\text{Tr}(K^2)/d)\\ll (\\text{Tr}(K)/d)^{-1}$ or in other words, $\\text{Tr}(K)/d\\ll\\text{Tr}(K^2)/d$. This is in contrast to the isotropic case where $\\text{Tr}(K)/d=\\text{Tr}(K^2)/d$. *We can make this more clear by writing $\\text{Tr}(K)/d\\ll\\text{Tr}(K^2)/d$ in parentheses after the phrase “strong anisotropy.”* * **Polyak stepsize.** Indeed, we are using the derivation of Polyak stepsize (see, e.g., \[1\]) as motivation to construct what should be the “idealized” stochastic version and not the $( R(X\_k)-R(X^\*) ) / ||\\nabla R(X\_k)||^2$. We will add commentary about the derivation of the Polyak step size in the paper and how it deviates from the classical Polyak stepsize. We appreciate the reviewer pointing this out. \[1\] Elad Hazan and Sham Kakade. “Revisiting the Polyak Step Size”. * **Concentration as an inverse power of d.** The concentration for the distance from SGD statistics to their deterministic equivalent is indeed $d^{-\\varepsilon}$ for $0\<\\varepsilon\<1/2$, so the best bound we obtain for this concentration is slightly worse than $1/\\sqrt{d}$. The comment in line 48 is that the concentration occurs with **probability** better than any inverse power of d. In other words, the probability that our statistic will be less than $d^{-\\varepsilon}$ is at least $1-d^{-C}$ for arbitrarily large C. In the theorem, we describe this more succinctly in the phrase “with overwhelming probability.” However, the definition of overwhelming probability is not introduced until line 74, so we do not use it in the comment on line 48\. * **Stopping times.** This is indeed unclear, and thank you for pointing it out. It is meant to be the minimum of the two stopping times, as in the preceding equation and in the Proposition that is cited from \[14\]. *We can clarify this in the updated version of the paper.* * **Verifying $S(t,z)$ is an approximate solution.** Indeed, this is a very important condition to verify, and doing so is the primary purpose of Section C.2. The way in which this is utilized with regard to the mesh is briefly described in Section C.1.3 where we prove Proposition C.1 (pending the details that are in Section C.2). *Based on your question, we can add a sentence at the location that you referenced summarizing the flow of the proof and where these details can be found.*

Rebuttal update

Dear authors, thank you for your response, including the general rebuttal. Based on the rebuttal, I am willing to increase the score to 5. Let me clarify the reasoning behind it, so that the authors and AC have a possibility to calibrate the score. The positive things are - proposed changes in writing are good - the discussion of relation with [14] adds clarity and transparency - the comments regarding different scaling of learning rate with dimension $d$, as well as additional discussion of power-law rates, help to better identify the positioning and relation of the current paper with other approaches. What prevents me from raising the score further (e.g. 7) is that the paper still explores just a little bit of the behaviors of adaptive algorithms. The theorem 2.1 is a solid foundation for studying adaptive algorithms. However, it really pays off with its subsequent application to different adaptive optimization scenarios. The applications in the current manuscript (sections 3 and 4) still feel as kind of short and not properly explored. The same goes to the application to non-quadratic problems - authors comment that they are roughly the same as quadratic problems, which does not sound satisfying.