Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper presents results on stochastic gradient methods on high dimensional least squares problems and presents a characterization of the loss of the final iterate run with a class of stochastic dimension adjusted (SDA) based momentum/Nesterov acceleration methods. The methods guarantee convergence to a neighborhood of the optimal solution where the size of the neighborhood shrinks to zero as the limit of the number of samples increases to infinity. 
main_review: To my knowledge, the results appear to be novel and I tend towards acceptance. 

[1] Can assumptions on scale of initialization, independence of noise and noise variance (assumption 1) be relaxed? I believe most standard analysis of SGD in these settings work without relying on these assumptions (see the work of Bach and Moulines (2013) for an example). 

[2] Can the authors obtain results similar to what they have presented, but with comparison to L(w^*) instead of the neighborhood of the solution? In principle, one can obtain results that converge to the solution using ideas like iterate averaging [Bach and Moulines 2013]/suffix averaging[Rakhlin et al. 2012]? Alternatively, if we wish to obtain guarantees on the final iterate, there needs to be some decay on the stepsizes [Ge et al. 2019]? I am interested as to what technical challenges exist in a path towards this result.

Bach and Moulines (2013): Non-strongly-convex smooth stochastic approximation with convergence rate O(1/n)

Rakhlin et al (2012): Making gradient descent optimal for strongly convex stochastic optimization

Ge et al (2019):  The Step Decay Schedule: A Near Optimal, Geometrically Decaying Learning Rate Schedule for Least Squares
limitations_and_societal_impact: Yes.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This paper analyzes the behavior of stochastic gradient algorithms on a high dimensional least-squares problem. Under the proposed unified framework, it concludes that stochastic momentum methods with a fixed momentum parameter does not improve the convergence of SGD. The authors then propose a new algorithm with momentum parameters depending on sample size, which achieves an optimal complexity asymptotically.
main_review: The paper is in general well-written. 

The authors propose an unified framework for stochastic gradient descent algorithms with momentum and table 1 gives a clear summary of different methods under this framework. Under the assumption of random least-squares model, they consider a diffusion approximation method (*homogenized SGD*) to anaylze the behavior of stochastic momentum algorithms. Theorem 1 describes the dynamics of  homogenized SGD and Theorem 2 links this approximation to SGD. However, as also mentioned by the authors, the high probability bound in Theorem 2 only includes the case of SGD and it would be better to consider the general cases and gives a solid theoretical support of their proposed sDANA and sDAHB algorithms.

Under the  stochastic momentum framework, the authors are able to express the result in Theorem 1 in a simple convolution–type Volterra equation and analyze the limiting behavior of these algorithms and the new algorithm sDANA obtains an accelerated average-case rate in the non-strongly convex case.

The theoretical proof in this paper requires delicate technique and I appreciate the effort authors made to solve these ODEs explicitly. My only concern is that most techniques and theoretical results are direct extension of Paquette et al. (2021). The link between SGD and the homogenized SGD is important but incomplete to the general cases, and Theorem 2 is also a direct citation from Paquette et al. (2021).

Reference mentioned above:

- C. Paquette, K. Lee, F. Pedregosa, and E. Paquette. SGD in the Large: Average-case Analysis, Asymptotics, and Stepsize Criticality. *arXiv preprint arXiv:2102.04396*, 2021.
limitations_and_societal_impact: None.
review_title: null

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: summary: The authors study exact asymptotics of the dynamics of momentum-accelerated SGD for a least-squares regression model with a "mean-field" regression design matrix A, in the scaling limit where n,d -> infinity proportionally. The paper has the following contributions:

(1) The authors introduce a d-dimensional diffusion approximation, eq. (7), which they conjecture to describe the asymptotic behavior of a class of momentum-accelerated multi-pass SGD algorithms, under certain assumptions for the design. The conjecture is proven rigorously for (non-momentum-accelerated) SGD and orthogonally invariant designs A in Theorem 2, using a small adaptation of the argument by Paquette et al '11. A heuristic derivation of this approximation for the more general class of momentum-accelerated algorithms of interest is given in Appendix B.1, and the conjecture is supported strongly by numerical results on both Gaussian designs and regression designs derived from MNIST images.

(2) The authors prove that the expected loss value at any epoch t that is predicted by this diffusion equation admits a form given by a Volterra integral equation, Theorem 1, which depends on the eigenvalues of the Hessian A'A. Assuming the correctness of the approximation in (1), the convergence rates of various momentum acceleration schemes may then be analyzed and compared by analyzing the long-time behavior of this Volterra integral equation.

(3) The authors carry out this analysis of the Volterra integral equation for the heavy-ball and Nesterov acceleration schemes, and establish either exponential rates of convergence or polynomial rates of convergence to the limit loss, depending on whether d/n -> 1. (This limit loss is not zero in the studied regime, and is instead characterized by Theorem 3.) The results of this analysis are summarized in Table 3. This analysis is more straightforward for the heavy-ball method, where the kernel defining the Volterra equation is of convolution form. For the Nesterov-accelerated scheme, the authors show that the long-time dynamics may be approximated by a convolutional kernel (Appendix C.3) and hence analyzed via this approximation.

(4) A few practical insights from this analysis are a prescription of the scaling of the step size and momentum parameters in n that are needed to obtain a speed-up over non-accelerated SGD, and the dependence of this speed-up on the choice of acceleration scheme (heavy-ball vs. Nesterov) and on the ratio d/n.
main_review: I think this paper addresses an important question---the correct scaling-limit approximation of momentum-accelerated SGD---albeit in the context of a simple linear model. It would have been ideal if the paper made rigorous the arguments in Appendix B.1 to provide a more complete version of Theorem 2, and in the absence of this result, perhaps the authors can give a sense in Appendix B.1 on how large the technical challenges might be to establish this theorem for the full sDA class.

Even in the absence of this result, I think the proposal of the form of the Volterra dynamics and the analyses of these dynamics that the authors provide are sufficient for publication in NeurIPS. I provide a few comments on exposition below:

(1) I find Figure 2 hard to interpret---is it saying that for large n * theta, sDAHB is doing worse than SGD because it has a larger function value in the last iterate? I find the wording of the caption unclear about what the figure is showing.

(2) Precise statements of some of the main high-level take-aways of the paper---for example, the convergence rates summarized in Table 3 and the equivalence of sHB and SGD in Eq. (14)---cannot be found in the main text, and are instead delegated to the appendix. Actually, even in the appendix I cannot find where many of the claims of Table are established. (For example the analyses in the "non-strongly convex" setting in Appendices C and D have some dependence on the scaling exponent alpha, but where is alpha in Table 3?) I would like to ask the authors to formalize the discussion on page 9 and more precisely state the results that are referred to in Table 3 (i.e., what are the bounds for psi(t) under what assumptions), preferably in the main text, but otherwise in the appendix.

(3) My understanding is that the authors refer to the setting d/n -> 1 as "non-strongly-convex" and all other settings of d/n larger or smaller than 1 as "strongly convex". I don't know if this may cause confusion or is the usual nomenclature---perhaps it's more common to describe any setting d <= n as strongly convex and d > n as non-strongly-convex, which is simply based on whether the quadratic loss f(x) is strongly convex in x?
limitations_and_societal_impact: N/A
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper studies the momentum SGD method (and some variants) in the high-dimensional setting. The paper, based on its analysis, proposes a new task-dependent algorithm for optimization
main_review: Understanding momentum in SGD is an important theoretical topic in deep learning, and studying this in the setting of high-dimensional regression is fairly novel.

I say "fairly novel" because I feel that the authors did not discuss the related works sufficiently to help me judge the novelty of the present work. For example, is this work the first and only work that deals with SGD in a high-dimensional setting? To my knowledge, the answer is no. The discrete-time SGD for high-dimensional regression has been analyzed in Ziyin2021, Section 4.4: arxiv.org/abs/2102.05375

Then the important question would be how does the result in this paper compare with the existing previous work. My first feeling is that the present work (in theoretical analysis) is more restrictive than Ziyin2021. For example, Ziyin2021 studies discrete-time SGD directly, while this work relies on continuous-time. Then why is this work novel? I am not suggesting that this work is not novel, but that the authors need to clearly discuss the relationship between the previous work to make its contribution stand out.

Also, one major weakness in the theory of the present work is that it relies on the assumption that the noise \eta is i.i.d. and independent of the data matrix A. However, this is not the case for SGD noise. Ziyin2021 has a very detailed study about noise here. See equation 13 of Ziyin2021; the noise eta in high-dimensional regression setting has a non-trivial noise structure, which is dependent on A and not i.i.d. I can accept more limiting assumptions than previous works, but the authors need to discuss this limitation clearly in the paper and convince the readers why the result may still be important despite this limitation. However, this assumption is given in the present paper in passing in assumption 1 without sufficient explanation. Also, I consider this a significant flaw because if the primary goal of the present work is to understand momentum in the setting of minibatch noise, then it must make sure it assumes a reasonable noise.

Also, I note an interesting difference in the result between the present paper and Ziyin2021. The present paper suggests that the step size of SGD should scale with 1/n, but Ziyin2021 Theorem 5 suggests that the step size should scale with 1/d. What causes this difference? How do I understand it?

I may recommend acceptance if the authors sufficiently address my questions.
limitations_and_societal_impact: I think the limitation can be better discussed. See above.
review_title: null

Meta-review: This paper presents several interesting results concerning stochastic gradient methods on high-dimensional least-squares problems. The characterization of the role played by the momentum coefficient in terms of convergence is sharp. The new task-dependent algorithm matches the optimal complexity.
Decision: Accept (Poster)