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

We thank the reviewer for their positive review and careful reading of our paper, including the appendix. We address below all of the reviewer’s comments. Due to space limitations, we will include an additional "Official Comment" to answer all the reviewer's questions.

Practical application of power law scaling on $\\mathscr{D}\_i^2(0)$. First, it is important to note that this power law scaling is not on the initialization but rather on the distance from the initialization to the ground truth (signal) X^\\star. If one initializes at 0, which is reasonable to do in the least squares setting, then this is a power law scaling on the ground truth. In settings where the data has power law covariance, one can also see power law scaling in the ground truth. For more details on this, see, for example, reference $56$ mentioned in the paper, particularly their discussion of scaling laws in section 6.

Limitations of the paper. Added a “Limitations” paragraph. For the version going in the main paper, see “All Reviewers Rebuttal”.

Coupled ODEs in Eq (7). Thank you for pointing out the potential issue with Equation (7). You are correct that there was an oversight in the first term of Eq(7). We have revised the equation to include the missing term and corrected ODEs. Page 17 contains the correct ODEs for S and reflects the updated formulation. We appreciate your attention to detail and the opportunity to correct this error.

Scalings. Thank you for the questions. Please see the discussion for all reviewers, where we discuss other scalings beside $(1/d)$.

You also asked about whether a learning rate on the scale of $1/d$ is large enough for the algorithm to really move away from initialization. The key here is that we work in the regime where the number of iterations is proportional to dimension. Thus, even with learning rate of size $1/d$, one sees significant movement after order $d$ iterations. Indeed, for examples such as linear or logistic regression, this learning rate yields convergence of the risk in order $d$ iterations.

You are correct that, in certain cases (e.g. phase retrieval with a cold start), the algorithm can become stuck at initialization and require more than order $d$ iterations to converge. This is also captured in our result. The theorem is not a statement about the convergence of the algorithm to an optimal solution, but rather about the convergence of the SGD dynamics to a deterministic curve. Thus, our theorem can accurately predict the risk trajectory for order $d$ iterations, even when the risk does not converge to zero in that time. We’ll add that to capture the escape of phase retrieval from the high-dimensional saddle $O(d \log d)$ iterations are needed, and moreover the high-dimensional limit of the risk curve will contain non-concentrating and non-vanishing effects of SGD noise (this is implicit in Tan and Vershynin. Online Stochastic Gradient Descent with Arbitrary Initialization Solves Non-smooth, Non-convex Phase Retrieval. 2019).

Finally, you asked about our choice to bound the norm of $X$, independent of $d$. Because we are interested in attaining a high-dimensional limit of the dynamics, we need to scale things in such a way that the risk is dimension-independent. In our set-up the risk is always expressible as a (generally non-linear) function of the inner product $X^Ta$, so we need this inner product to be dimension-independent. Our choice of scaling achieves this (although one could also, for example, simultaneously adjust the scaling of $X$ and $a$ to achieve an equivalent result). There are, as you suggest, interesting problems in which the norm of $X$ may grow with $d$, but those are beyond the scope of this paper.

Typos. Thank you for pointing out these things. We corrected them.

Dear author(s), thank you for your general rebuttal and for the rebuttal concerning my specific comments. I understand all your points, and after the corrections/specifications you put, I will raise my score. Some quick points:

• strong anisotropy and verifying the approximate solution: you comments in italics are necessary in my opinion.
• Polyak comment: yes, please.
• scalings comment: see one of the questions missing (at the very end) in my long list to close this aspect.
• all other typos: thank you for fixing them.

last comment: This is a very deep paper, maybe a journal deserves it.

My score is higher. I hope the other reviewers will raise their very low grades. Some scores, in their very description, are too harsh with respect to the actual review.

My score reflects my opinion on the exact description provided by Neurips: moderate-to-high impact. I stand with the idea that this paper is definitely not a 3.

Good luck!!

Acknowledged, see full comment below. I believe you are missing the last three questions (one might be a typo therefore I ping you about this).

We thank the reviewer for their support of our paper. We answer some of the reviewer’s reservations below.

“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. “

We agree with the reviewer that the algorithms studied are idealized and/or not widely used in practice. However, traditional analysis of adaptive stochastic learning rate algorithms (e.g., idealized line search/Polyak, etc.) all give the same rates as the ones used in practice. The traditional theory does not distinguish between “good” practical algorithms and “bad” practical algorithms. This suggests we need a new type of analysis. Part of our goal was to provide such a new framework. While we only analyze “idealized algorithms”, we now understand problems for which these “idealized algorithms” do not perform well. This hopefully illuminates and begins closing the gap between practice and theory.

In future work, we want to derive the dynamics of more practical algorithms. We strongly believe this framework might explain theoretically why certain often used algorithms in practice perform better.

Focus on quadratic problems is a major limitation, especially since the analyzed methods are not typically used on quadratic problems.

We agree with the reviewer that “quadratic problems” are idealized, and we want in future work to extend the analysis of these algorithms to more complex losses. On the other hand, many of the algorithms were designed by approximating the “complex” loss with a quadratic. Moreover, it is surprising that we do not even know how these algorithms perform on these quadratics. In our opinion, we should understand this simple loss in detail first.

This said, our analysis extends beyond quadratics to include generalized linear models such as logistic regression and more. Moreover, we do provide some analysis for other losses beyond quadratics. For strongly convex problems (not necessarily quadratic) we find a bound on the limiting learning rate, see Section D1. Also, Figure 1 shows our predictions matching the iterates of SGD with AdaGrad Norm for binary logistic regression.

Response to Weaknesses

• Too much set-up in the main paper. Thank you for your suggestions. In the updated version of the paper, we plan to lighten the notation and move the least square results upfront; see our reply to all reviewers.
• Plot fonts too small. Thank you for pointing that out, in the next version of the paper we will update the figures and make the font larger. See the attached PDF with the revised figures.
• Paper ends abruptly. Added a “Conclusion” paragraph. See comment to “All Reviewers” for the specific paragraph.
• Discussion of implications and potential applications. Added a “Limitations” paragraph. The conclusion paragraph contains potential applications. See comment to “All Reviewers” for the specific paragraph.

The reviewer's comments required more space than 6000 characters to answer. We will include an additional "Official Comment" with the answers to the reviewer's questions.

Significance of results - Main contributions

This work introduces a framework that goes beyond traditional analysis of stochastic algorithms to utilize the high-dimensionality in the problem. It allows for a finer analysis of stochastic learning rate algorithms (e.g., AdaGrad, line search, Adam, RMSprop, etc).

We also emphasize that one of the main contributions of this work is an explicit deterministic expression for the dynamics of stochastic adaptive learning rate algorithms (Thm 2.1). That is, we can predict the behavior of the loss value and learning rate (both stochastic) at any iteration without ever running the stochastic algorithm (see Fig 1). We then analyze these dynamics to gain insights into the evolution of the stochastic learning rate at any iteration. These dynamics exactly agree when the problem is “high-dimensional” (number of parameters is large) and often numerically reproduce the dynamics on real data (see Fig. 6 on CIFAR-5m dataset).

Line search and Polyak only “idealized version”

It is true that the algorithms presented are idealized versions. However, we emphasize that a theoretical understanding of the behaviors of these idealized (stochastic) algorithms is not known $Note: the reference given is for a line search on a deterministic objective and less is known about the behavior of stochastic line search$ . Adding more practical versions, which are often more complex, only adds complications to the already unknown behaviors. As such, a natural starting point for developing a theory for the exact dynamics of stochastic learning rates is to understand the behavior of these idealized versions.

Line search methods do not work well is not new

We first politely disagree with the reviewer that line searches do not work in practice. For deterministic optimization problems, line search methods are widely used and are the default in many well known algorithms (e.g., see scipy documentation for BFGS (Armijo condition)). While we appreciate the reference that showed an example where line searches fail on a deterministic optimization problem, line searches do provably work for a large class of deterministic optimization, including the deterministic versions of the problems analyzed in this paper. The theory for deterministic line searches is well developed, but for stochastic optimization problems the theory is still in its infancy.

We do agree with the reviewer that on stochastic optimization problems, line search methods are not used in practice. In fact, one of the goals of this paper is to explain this! We show why they perform badly due to anisotropic data; in contrast, a standard minimax-optimal analysis says they are as good as SGD (see e.g. $50$ ). Moreover, we precisely quantify how different the convergence rates of a stochastic line search method will be from a tuned fixed learning rate SGD algorithm.

The theoretical framework established in this paper (Theorem 2.1) provides the tools (albeit more complex than the “textbook analysis”) to do a finer analysis of stochastic line search methods.

Complexity of the theory is difficult

The standard tools in textbooks for theoretically analyzing stochastic algorithms often give too crude estimates on their behavior in practice. As an example, the standard analysis would say that you need to use a decreasing learning rate with SGD, but nobody does this in practice. The point of this paper is to provide another tool that allows for a finer analysis of these algorithms – traditional minimax optimality isn’t capable of distinguishing these algorithms – so something more complicated is necessary.

Complexity of theory makes it hard to derive insights. Theory will not be extended by others.

We politely disagree with the reviewer. The related work paragraph (Page 4 Lines 106-110) describes past works which use a high-dimensional framework for analyzing dynamics of stochastic algorithms. These past works usually assume isotropic Gaussian data ($N(0,I)$) and do not consider stochastic learning rates for their algorithm, but are part of a growing body of literature based on this style of analysis.

We know that we are doing something quite different than the vast majority of optimization researchers, but we think there is a need for better ways to distinguish performance of stochastic algorithms. For example, the more traditional analysis of line search in Vaswani et al.
$50$ suggests SGD + line-search has no performance cost (their Theorems 1 and 4). We quantify the performance cost, as one incurs an extra factor of condition-number. This is a statement 99% of optimization-researchers should be able to understand, even if the methods are not to their taste, and one needs a more complicated setup to see it!

$50$ Vaswani, Mishkin, Laradji, Schmidt, Gidel, Lacoste-Julien. Painless Stochastic Gradient: Interpolation, Line-Search, and Convergence Rates. 2019.

Limitation paragraph. Please see the response to "All Reviewers" for the exact paragraph we will add to the main document.

We thank the reviewers for their comments. We will update the paper to emphasize the difference between our analysis and that of $14$ . We state below the main differences between this work and $14$ :

• First, the learning rate is by itself a stochastic process. In addition, since it carries all the history of the previous gradients, the iterates of SGD are no longer a Markovian process. We extend the proof of the previous paper to handle such situations by providing an additional concentration of the learning rate to its equivalent deterministic limit.

• Second, there are differences in the analysis of the Volterra equation on the least square problem. Having the learning rate as a nonlinear function of the loss makes the analysis much harder, and we had to introduce new techniques to study this new generalized form of the Convolutional Volterra equation (see Appendix D).

• Third, in this paper, we study in depth the effect of different covariance matrix structures. For the line search method, we analyze in detail the example of two distinct eigenvalues. For AdaGrad-Norm, we study different power law covariance scalings and zero eigenvalues, for which we had to utilize some nontrivial asymptotic methods.

The assumption of the data being Gaussian could be relaxed. Most of the concentration estimates we use to prove the deterministic equivalent are based on concentration inequalities, which are valid beyond Gaussian data. In addition, Figure 6 in the paper shows that we can predict multipass SGD dynamics on more realistic data (CIFAR-5m) which suggests that our theory extends beyond the Gaussian data and the streaming setting.

$14$ Elizabeth Collins-Woodfin, Courtney Paquette, Elliot Paquette, and Inbar Seroussi. Hitting the high-dimensional notes: An ODE for SGD learning dynamics on GLMs and multi-index models. arXiv preprint arXiv:2308.08977, 2023.

We address below the questions of the reviewer.

1. Isotropic covariance matrix $K$. We emphasize that our results hold for non-isotropic data covariance matrices $K$ (under mild assumption $||K||\_{op} \< C$, and the average eigenvalue of $K$ is bounded independent of $d$). We point to Reviewer ipsV who also noted the anisotropic assumption on $K$.

2. Writing of the manuscript. Please see some specific changes in the “All Reviewers” section of the response. We appreciate the reviewer’s suggestions.

3. New developments from previous work $14$ . See discussion to all reviewers.

4. Beyond non-quadratic losses. In Appendix D, we show that under a restricted secant inequality, similar conclusions to the quadratic case can be derived for AdaGrad-Norm. Now the results we get are essentially the same as one gets for quadratics. We weren’t so keen to generalize the whole paper to something like losses satisfying an RSI if, at the end of the day, we just get the same answers as for quadratics. Identifying interesting and important non-quadratics with different phenomenology for adaptive algorithms is an interesting direction of future research (e.g. nonconvex problems like retrieval or problems exhibiting implicit bias). Using Theorem 2.1, the ODEs are already in place, and so there is a framework for finer analysis of adaptive algorithms which was missing from the literature. We believe these methods can eventually be used to provide better or more informed adaptive stepsize methods.

5. High-dimensional vs low-dimensional problem extensions. Thanks for these points – they are very good. Please see the discussion in the rebuttal to all reviewers. We’ll remark that the \\text{Tr}(K^2) \= o(d) condition you observed is implied by \\text{Tr}(K) \= o(d), so long as we keep bounded operator norm of $K$, and so it falls into the regime where SGD and the ODEs both need to be scaled differently. We can add a discussion of this to the paper.

   Part of the reason we wrote this paper is that we needed to demonstrate that it was actually possible to analyze these ODEs and derive meaningful conclusions about optimization algorithms. Especially, we needed to analyze these ODES in the presence of anisotropy, which can heavily degrade the value of an adaptive stepsize algorithm (e.g. line searches work great in the isotropic case – and are essentially optimal – but they are bad in the presence of anisotropy. In contrast, adagrad norm does much better, although still fails in some power law-type setups). So half of this paper is analyzing various algorithms and half is formulating a precise theorem which says where the ODE comes from.

   We hope this paper could be a starting point for us (or, gladly, other researchers!) to refine the method further, for example:

• Finding the optimal covariance conditions under which these ODEs hold and making a fully non-asymptotic comparison.

• Understanding the extra correction terms to the ODEs in non-high-dimensional setups, such as the classical source/capacity conditions. (The rank-1 problem you mentioned, in contrast, is not well-aligned with the goals of this type of analysis, which we think could be summarized as: how to give meaningful quantification of adaptive algorithms on anisotropic problems).

• Digging deep into other adaptive stepsize strategies to see if they perform better in anisotropic setups.

• Showing a version of Theorem 2.1 for non-Gaussian data (high-dimensionality also gives a path to showing universality of the ODEs).

  It could be, for example, that Adagrad norm on classical capacity source setups could be analyzed in the way we did and have interesting conclusions.

6. SGD power-law spectrum rate. We thank the reviewer for the references. Under our set-up, constant stepsize SGD would be $\\mathcal{R}(t) \\asymp t^{\\beta \+\\delta \-2}$, and there is no transition at $\\beta+\\delta \=1$. We will add this to our paper. The references $1$ , $2$ , $3$ suggested by the reviewer do consider similar assumptions on the spectrum and eigenvalues of the covariance matrix, but the papers referenced are not in the high-dimensional regime (diverging $\\text{Tr}(K)$). All $1$ , $2$ , $3$ contain results which formally correspond to $\\beta \> 1$, so a quantitative comparison is not really possible ( $3$ discusses our regime here, but does not seem to have rates in this regime). Qualitatively, they all demonstrate similar results as seen here (formally similar to how the loss undergoes a phase transition as the ‘source’ parameter varies in source/capacity setups). In the next version of the paper, we will add clarification regarding the scaling with constant stepsize together with comparison and citations of the relevant works.

7. Learning rate phase transition. This is a great question and worthy of future research. While we did not keep track of the exact finite dimensional $d$ effects, one could in principle use our analysis to theoretically quantify when the asymptotics for the learning rate and the risk take over.

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).