Dynamics of Stochastic Momentum Methods on Large-scale, Quadratic Models
Authors: Courtney Paquette
Elliot Paquette
Venue: NeurIPS 2021 Poster (2021)
Keywords: optimization
convex
random matrix theory
complexity
average-case

Abstract: We analyze a class of stochastic gradient algorithms with momentum on a high-dimensional random least squares problem. Our framework, inspired by random matrix theory, provides an exact (deterministic) characterization for the sequence of function values produced by these algorithms which is expressed only in terms of the eigenvalues of the Hessian. This leads to simple expressions for nearly-optimal hyperparameters, a description of the limiting neighborhood, and average-case complexity. 

As a consequence, we show that (small-batch) stochastic heavy-ball momentum with a fixed momentum parameter provides no actual performance improvement over SGD when step sizes are adjusted correctly. For contrast, in the non-strongly convex setting, it is possible to get a large improvement over SGD using momentum. By introducing hyperparameters that depend on the number of samples, we propose a new algorithm sDANA (stochastic dimension adjusted Nesterov acceleration) which obtains an asymptotically optimal average-case complexity while remaining linearly convergent in the strongly convex setting without adjusting parameters.

Dynamics of Stochastic Momentum Methods on Large-scale, Quadratic Models
Courtney Paquette
Elliot Paquette
Department of Mathematics and Statistics McGill University
Department of Mathematics and Statistics McGill University
Montreal, Quebec H2Y 2M5
Montreal, Quebec H2Y 2M5 elliot wwwootpppppeeeeeeegglll....
courtney paquettedmcgill..e
Abstract
We analyze a class of stochastic gradient algorithms with momentum on a high- dimensional random least squares problem. Our framework, inspired by random matrix theory, provides an exact (deterministic) characterization for the sequence of loss values produced by these algorithms which is expressed only in terms of the eigenvalues the Hessian. This leads to simple expressions for nearly-optimal hyperparameters, description of the limiting neighborhood, and average-case complexity.
As a consequence, we show that (small-batch) stochastic heavy-ball momentum with fixed momentum parameter provides actual performance improvement over SGD when step sizes are adjusted correctly. For contrast, in the non-strongly convex setting, is possible to get a large improvement over SGD using momentum. By introducing hyperparameters that depend on the number of samples, we propose new algorithm SDANA (stochastic dimension adjusted Nesterov acceleration) which obtains an asymptotically optimal average-case complexity while remaining linearly convergent in the strongly convex setting without adjusting parameters.
Methods that incorporate momentum and acceleration play an integral role in machine learning where they are often combined with stochastic gradients. Two of the most popular methods in this category are the heavy-ball method (HB) [Polyak, 1964] and Nesterov's accelerated method (NAG) [Nesterov, 2004]. These methods are known to achieve optimal convergence guarantees when employed with exact gradients (computed on the full training data set), but in practice, these momentum methods typically implemented with stochastic gradients. the influential work Sutskever et al. [2013), authors demonstrated empirical advantages of augmenting stochastic gradient descent (SGD) with momentum machinery and, as result, momentum methods are widely used for training deep neural networks. Yet despite the popularity of these stochastic momentum methods, the theoretical understanding of these algorithms remains rather limited.
In this paper, we study the dynamics of stochastic momentum methods (with batch size 1 and constant step size) rooted in heavy-ball momentum and Nesterov's accelerated gradient algorithms on a least squares problem. Our approach uses framework inspired by the phenomenology of random matrix theory (see Paquette al. [2021]). which gains explanatory power when the number of samples (n) and features (d) are large. A key contribution of this work is simple description of the exact dynamics for a class of stochastic momentum methods smooth, deterministic function (t) such that (*x) (k//n) (i) as - xo. This function de solves
high-dimensional limit; we construct
"Website courtneypaquette.. github.oo. Website elliotpaquette github.io.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Algorithm Generic stochastic momentum method. step sizes F1, I > 0 and momentum parameter (kk Initialize: zo Rd and Yo for 2 1, Select & E [n] uniformly and update
Given:
yk (1 A (k) 9k-1 + T1 Vi. (mx) and
#k41 3k. T2Vfi, Si, (xx)
yk
(0.2)
a Volterra integral equation:
(01)= F()) + L( K. (.))))))
(0.1)
Here Ftt) and (t) are explicit, see Theorem for a precise statement. This Volterra equation (0.1) gives accurate prediction of the behavior of stochastic methods, see Figure We then an- alyze these dynamics providing insight into step size and momentum parameter selections as well providing both upper and lower average-case complexity (i.e., the complexity of an algorithm ate.
100
function values
10-1
Volterra
averaged over all possible inputs) for the last iter-
10
64
n 256
1024
As we show in this work, both theoretically and empirically, (small batch size) SGD with heavy- ball momentum (SHB) for any fixed momentum parameter does not provide any acceleration over plain SGD on large-scale least square problems. conclude under an identification the param- eters that (argh)) = [ax²") for all k up to errors that vanish grows large (upper bounds of this nature have been observed before: see Kidambi al. [2018], Sebbouh et al. [2020], Zhang et al. (2019]). Thus while SHB may provide a speed-up effective increase in the learning rate, and this speed-up could be matched by appropriately adjusting the learning rate of SGD.
4096
10 10
10-
102
epochs
Figure 1: Concentration of stochastic heavy- ball (SHB) on a Gaussian random least squares problem (Sec. 2), d = n, an 80% confidence interval (shaded region) over 10 runs for each n, the parameters for SHB (Table 1) are (0, 7) (0.1, 0.08). The random least squares problem becomes non- non-random in the large limit and all runs of SHB converge deterministic function (t) (red) given by our Volterra equation (0.1).
over SGD, it is only due
The root SHB's ' failure to provide meaningful acceleration is that a fixed momentum parameter is not aggressive enough when is large. We propose modification of Nesterov (see Alg. and Table 1). The resulting algorithm, SDANA, matches the average-case complexity of SGD when the least-squares problem strongly convex and obtains average-case complexity of 1/k 3 in the convex setting.
algorithm that uses dimension-bain based
1 Motivation and related work
We consider the large finite-sum setting
min mERd - =
f(xx)
2 > [a.
A1 = b1l2) b1l²]
for data matrix A Ruxdd whose i-th row is denoted by a, E Rdxx and target vector b E RD (detailed in Section 2). We make the convention that the matrix A has max row norm equal to 1. Note we absorb some n--dependence into A and b by setting f1 = 1 (a.æ 6,)². We investigate a generic class of stochastic momentum algorithms (see Alg. and Table 1). Particularly, we introduce sub-class, denoted by SDA(71 (71, , 72, A), of Alg. which has parameters that are appropriately adjusted for large problems (large number of samples n and large model size d); refer to the dimension of

---

Table 1: Summary of the parameters for variety of stochastic momentum algorithms that fit within the framework of Alg. 1. denote the normalized trace by m = n
def >>>ー>> a The
default parameters chosen that its linear rate is no slower, by a factor of 4 than the fastest possible rate for am algorithm having optimized over all step size choices.
Alg.
Parameters
Methods
Default Parameters
Ti T2
A(k)
Stochastic gradient descent: SGD(?)
0
1
(Prop. E.4)
Stoch. gradient descent w/
momentum: SHB(), 0)
7
0
(see Fig. 2)
Stoch. dimension-adjusted heavy-ball: SDAHB (y,
7
0
0 -= 2 (Prop. E.3)
0
(This paper)
n
Stoch. dimension addjusted Nesterov's
accel. method: SDANA(1) (MAAAA(110 (/.72211 72. 8)
71
72
0
71 4m 72- m
1: +
0 (Cor. D.1)
(This paper)
the problem as n. The class SDA(11 72, A) is defined by setting in Alg.
===
'1 r2
72. and
(k)
def
1t (log
()))
(1.1)
where 71, 72 > 0 are step sizes and yp is smooth function that represents a momentum schedule.
Although we develop some theory for general (p, we are principally interested in the two cases:
(SDANA) A(k) = k+n + (t) = t) and (SDAHB) A(k) = + p(t) = "t (1.2)
To avoid confusion between different algo- rithms, we add superscripts indicating the algorithm (e.g., denote T2 7 the step size parameter for SGD). For all these algorithms, we are interested in:
10-1
last iterate function value
1.
An expression for the (deterministic) dy- namics of these algorithms when mulii- ple passes the data set are allowed.
10-4
0.25
0.50
0.75
This contrasts with the "streaming" or "online" setting where at each itera- tion one generates an independent never- before-used data point.
10-7
Volterra, SGD
Volterra, SDAHB sDAHB
102
100
10
2. A formula for choosing the hyperparam- eters and a discussion the dependence of these hyperparameters on number of features and samples.
theta
Figure 2: Equivalence of SGD and stochastic Heavy- Ball., For every 7-8d E {0.25, 0.50, 0.75}, we select
pair of parameters (7 sbb, gshb) so that 7 gas run SHB 3000 times with varying @shb on (2.1) with = 500,
00 We
Upper and lower bounds on the average- case complexity of the last iterute to a neighborhood; this neighborhood dis- appears entirely in the overparameter- ized regime, while in the underparam- eterized regime the limiting distance to optimality concentrates in the high- dimensional limit.
1000, and plot the value of the last iterate after 50 epochs. Small gshh matches SGD (orange, theory), illustrating their equivalence. With ghhb se 1, change is observed, giving a small improvement for some values of sdd Plotted against theory for SDAHB (red Volterra, see Thm. and App E), which the same algorithm SHB after change of parameters.
Why divide by n? A negative result. Throughout the literature, there are examples for which (small batch size) stochastic momentum methods such as SHB and stochastic Nesterov's accelerated method (SNAG) achieve performances equal to (or even worse) than small batch size SGD (sec
3

---

e.g., Kidambi et al. [2018], Sebbouh et al. [2020], Zhang et al. [2019] for heavy-ball and Assran and Rabbat [2020], Liu and Belkin [2020], Zhang et al. [2019] for Nesterov). We also observe this phenomenon (see (2.4) and App. E.4, Thm 5), and we illustrate this in Fig. 2. The stochastic heavy-ball method for any fixed step size and momentum parameters (Table 1) has the exact same dynamics vanilla SGD, which is to say, by setting the step size parameter in SGD to be > f the two algorithms have the same loss values provided the number of samples is sufficiently large, Le., f(xx = =0 x
As consequence of this, the average-cass complexity of SHB equals the last iterate complexity of SGD (This was observed [Sebbouh et al., 2020] with an upper bound, but our result shows an exact equivalence between last iterate SGD and SHB). Although App. E,4, Thm 5 (see also (2.4)) gives an unsatisfactory answer to stochastic heavy-ball with fixed and small batch-size, our analysis illuminates a path forward. Particularly, one must choose fiieenioon-depndeen parameters to achieve dynamics which differ from SGD.
Why divide by n? A positive result. Adapt- ing SHB for dimension, we arrive at stochas- tic dimension adjusted heavy ball (SDAHB).
500
SDANA SGD
While formally equivalent to SHB, we include
halting time
400
the dimension parameters to emphasize that any improvement in its performance for large n re-
300
quires Nonetheless, the speed-up for heavy ball is modest (see Fig. 2).
200
On the other hand, we show that dimen- sion adapted version of Nesterov acceleration, SDANA, has large improvement in the non- strongly convex case. Moreover, with a simple parameter choice (see the default parameters in Table 1), it will perform linearly in the strongly convex case, and competitively with learning- rate-tuned SGD (or SHB), while performing orders-of-magnitude faster (k-S as compared to SGD k ) for the non-strongly convex setting (see Fig. and Table 3). We believe this gives SDANA promise an algorithm outside of the least squares context, in situations in which loss landscapes can range between alternately curved and very flat, frequently observed in neural net- work settings (see Ghorbani et al. [2019], Li et al. [2018], Sagun et al. (2016)).
100
8.5
1.0
1.5
2.0
r
Figure 3: Convergence of SDANA. Halting time of SDANA vs SGD with default parameters on the Gaussian random least squares problem (2.1) with varying d and n = 1024. When the ratio r max ( Amin AA" Amin (m.n(AAAAA (AA")} + 0, SDANA requires significantly fewer iterations loss of 10 55 As the ratio r moves away from 1, the performance of SDANA matches SGD.
d/n + (in which case
reach
Related work. Recent works have established convergence guarantees for SHB in both strongly convex and non-strongly convex setting [Flammarion and Bach, 2015, Gadat al., 2016, Orvieto et al., 2019, Sebbouh et al., 2020, Yan et al., 2018]; the latter references having established almost sure convergence results. Specializing converge linearly (but not in I under an exactness assumption [Loizou and Richtárik, 2017] while under some additional assumptions on the noise of the stochastic gradients, [Can et al., 2019, Kidambi et al., 2018) show linear convergence to a neighborhood of the solution.
the setting of minimizing quadratics, the iterates of SHB
Convergence results for stochastic Nesterov's accelerated method (SNAG), under both strongly convex and non-strongly setting. have also been established. The works [Assran and Rabbat, 2020, Aybat et al., 2018, Can et al., 2019, Kulunchakov and Mairal, 2019] showed that SNAG converged at the optimal accelerated rate to a neighborhood of the optimum. Under stronger assumptions, convergence to the optimum at accelerated rate is guaranteed. Examples include the strong growth condition [Vaswani et al., 2019] and additive noise on the stochastic gradients [Laborde and Oberman, 2019].
The lack of general convergence guarantees showing acceleration for existing momentum schemes, such as beavy-ball and NAG, in the stochastic setting, has led many authors to design alternative

---

acceleration schemes -Allen-Zhu, 2017, Ghadimi and Lan, 2012, 2013a, Kidambi et al., 2018, Kulunchakov and Mairal, 2019, Liu and Belkin, 2020].
2 Random least squares problem
To formalize the analysis of a high-dimensionn typical least squares problem, we define the random least squares problem:
arg min
{ f(z)
1 Σ f((z) def 1
2 (a 1 ] with b def A + +
(2.1)
=ER*
The data matrix A random and we shall introduce assumptions on A as they are needed, but we suggest central example the Gaussian random least squares where each entry of A is sampled independently from a standard normal distribution with variance 2. We always make the assumption that each row a; E Rdx1 is centered and is normalized so that max; {EE [|]| =.
As for the target b AT 77, assume it comes from generative model corrupted by noise, where T is signal and r7 is noise.
Assumption (Initialization, signal, and noise). The initial vector ay E Rd chosen so that xy is independent of the matrix A. The noise is centered and has i.i.d. entries, independent of A. The signal and noise are normalized so that
1 - x R and
[1l =
Note that deterministic 2p ã satisfies this assumption. The vectors 2g and r7 arise as a result of preserving constant signal-to- -ignal-tonnoise ratio in the generative model. Such generative models with this scaling have been used in numerous works [Gerbelot et al., 2020, Hastie et al., 2019, Mei and Montanari, 2019].
For the data matrix A we introduce the Hessian matrix H A" A and its symmetrization H AAT. Let Å₁ 2 ... 2 An be the eigenvalues of the matrix H. Up to appending zeros, this is the same ordered sequence of eigenvalues as those of the Hessian. Define the empirical spectral measure (ESM) of H, eh by the formula
f (()))Hr(d)))
1 Σ
9(x))
for any continuous function g IR + R.
(2.2)
This gives the interpretation for the empirical spectral measure as the distribution an eigenvalue of H chosen uniformly random.
Diffusion approximation. Our analysis will use a diffusion approximation to analyze the SDA(1, 72, A) class of stochastic momentum methods (see (1.1)) the random least squares setup (2.1). We call the approximation homogenized SGD:
71 ptt) / d(s) where d2, def VJ(X)) dt+ ( ((( (2.3)
dx, def
7/ d2,
and with initial conditions given by X, = Zy- The process (B, 2 0) is a d-dimensional standard Brownian motion. Here time is scaled in such way that t= 1 represents pass over the dataset or n calls to the stochastic oracle. Similar SDEs have appeared frequently in the theory around SGD, see e.g., Li et al. [2017, 2019], Mandt et al. [2016].
The advantage of the homogenized SGD diffusion is that we are able to give an explicit representation of the expected loss values on least squares problem, even at finite n.
Theorem (Volterra dynamics at finite n). Let EH [-] be the conditional expectation where is held fixed. There non-negative functions Ft)) and (t)(t) for s,t 2 O depending the spectrum of so that for all 2
E [x(, = F(1)+ + S ( for all t t 0.
(2.4)
5

---

Table 2: Summary of the convolution kernel for the Volterra equations (2.4) for all algorithms considered. The convolution kernel (below) for SDANA an approximation the true kernel. The forcing terms (L) for SGD and SDAHB are similar to the kernel whereas the forcing term for SDANA is defined only by solving 3rd-order ODE.
(t)
Methods
Kernel, (t)
SGD(?)
22126-2xA(6- 2yA((-6)
SDAHB(7;
2n e (1- -cos(l (o(((())/W)
4X7 f
(This paper)
(1 SDANA (7, 72, 0) X - -Arr((t-s)
(1- -oscos((t S) + 0)
tan
(1)
(w-271) 1271
(-- 231 )- 2
(This paper)
471 7/
The forcing function and kernel K. are given by
1 " F(t) = (RA, + and (C. = Σ (t) i=1
The functions G(A) and (AA) are solutions of an initial value problem with a 3-rd order ODE which depend on the hyperparaaeetess (71, 72, A) (Note, there is a 1-to-1 relationship with sp, see (1.1)).
We refer to the supplemental materials for the explicit third-order ODE (see Theorem 4 for full details). The expression (2.4) is Volterra integral equation, which can be analyzed explicitly, and has a relatively simple theory, especially in the case that the kernel is of convolution type (i.c. (t) I(t s) for some function Σ); see Table for kernels. We also note that in the case of SGD(r) (s is unused), the functions G and K become particularly simple
(t) = -2- and (t) > 2TM(-s)
Comparing homogenized SGD to the SDA class. When A is random matrix, we can compare the diffusion (2.3) SDA (1.1) when n and d are large. The argument based on the results of Paquette al. [2021], and we do only in the case of SGD:
Theorem (Concentration of SGD). Suppose that A is left-orthogonally invariant random matrix,
meaning that for any orthogonal matrix E Ruxnn OA A. Suppose further that the noise vector is independent of A and that it satisfies
law
El || = O(n--1p2 for any c,p > 0.
Fix 7 < 2n( H)-1, the convergence threshold of SGD(T). There is an absolute constant ε >> 0 and constant c(T, xx that with p min (d, n).
Pr( sup |E - f (()))) > e(T, e(T,At)) (t)) 5 0S1ST
We expect that this theorem can be generalized, include the entire SDA class. We also expect that the orthogonal invariance assumption can be relaxed somewhat (for example to include classes of non non-Gause Gaussian isotropic features matrices), but not entirely: the left singular vectors need have some degree of isotropy for the result to hold. The numerical results show very good general agreement with theory and demonstrate the validity of the approximation: see Figures 1, 2, and as well as Figure 6 on real data. Nonetheless, it is of great theoretical interest to establish the theorem in greater generality. We show a heuristic derivation App. B.
Main results
In this section, and in light of the Thm. 1, we outline how to use this Volterra equation (2.4) to produce average-case analysis, nearly optimal hyperparameters, and exact expressions for the neighborhood and convergence thresholds. For additional details, see Supplementary Materials.

---

Table 3: Asymptotic average-case convergence guarantees for (x R dim (ker (H))
(last iterate) under default parameters (see Table 1) for the isotropic features model. The norm of the kernel is controlled by two values: the normalized trace of the matrix mass of the spectral measure (empirical or limiting) at which we denote by p = dim(ker( ) Average-case complexity is strietly better than the worst-case complexity, in some cases by factor Y vs. 7²- As in Paquette et al. [2021], the worst-case rates in non-strongly convex setting have dimension dependent constants due the distance to the optimum | æoll Rs d which appears in the bounds. SDANA obtains an accelerated average-case rate in the non-strongly convex case over SGD while matching the average-case rate of SGD in strongly convex regime. These rates are achieved without changing hyperparameters in SDANA. For worst-case rates, see [Bottou et al., 2018, Theorem 4.6] [Ghadimi and Lan, 2013b, Theorem 2.1]: X can be replaced by the max-f2-row-norm
2n(1- [21]
- ||a,1| / /n and the
Non-strongly
Kernel, |||||
Strongly convex
convex
SnD(n)
Worst
exp( >xp(->1>* - + 3 (A+)22)
(R+ R d)
3 m (Eq. (C.4)
exp( -)
Rt Rt-3/2 2 + Rt 1/2
SGD(1)
Avg
(Lem. C.4)
(Eq. (C.6))
20 m (Eq. (F.11))
20-0
+07)
Rt-3/2
Rt 3/2 + Rt 1/2
SDAHB (7,
(Prop. E.3)
(Eq. (E.13)
m1 (1 p) + 2 m
exp(- (1- 1 321 2 +471
RI + Rt-1 (Prop. D.3)
SDANA( 02.0)
2y2
(Eq. (D.25))
(Cor. D.1)
3.1
Convolution Volterra convergence analysis: convergence threshold and neighborhoud
For all algorithms considered (SDANA, SDAHB, SHB, SGD), the Volterra equation in Theorem can be expressed in a simpler form, that is, convolution--ype type Volterra equation
= F(I)+ + /. I( ((x.)) for all 1.2 0.
=
(3.1)
The forcing function F and the convolution kernel Z are non-negative functions that depend on the spectrum H and SDA parameters (see Table 2 for the kernels of various algorithms). In the case of SDANA, the kernel is in fact not a convolution Volterra equation, but it can be approximated by one so that it matches the non con-convolttion equation
oc.
First, the forcing function F will in all cases be bounded, and in fact it will converge as t + oo to a deterministic value,
Ru (10)) Rdii(ker(HH))
F(I) (-+00 2
(3.2)
2n
Here Hh is the empirical spectral measure (2.2) which exists for even non-random matrices. It
follows that the solution of (3.1) remains bounded if the norm ||]|| = fo I(t) di is less than 1.
Theorem 3 (Convergence threshold and limiting loss). If the norm || the algorithm is
convergent in that
R dioororr( H))
E (((,)) bn
(limiting loss).
2n(1
[[]I]
r-1.00
This theorem gives a convergence threshold for all algo- rithms in Table based only the norm of the kernel of the Volterra equation, which easily computable (see Table 3).
r=122
density
11.50
11.75
--2,50
3.2 Average case analysis
magnitude
Limiting spectral measures. Average-case complexity looks at the typical behavior of an algorithm when some
Figure 4: The Marchenko- -Pastur law(r).
Varying r d/n.
7

---

SDANA (strangly convex, r=0.7)
SDANA (non-strongly convex, r=1)
Vottera
10
function values
64
function values
10-2
Voleee
1001
252
1724
10
0046
10-4
10-2
256
10
1028
4006
10a
10 epochs
104
10a
10
10-
10a
epochs
spachs
100
Figure 5: Concentration of SDANA. 80% confidence interval on 10 runs with default parameters on Gaussian random least squares problem (2.1). d/n= r, with noise = 0.01 and signal R = 1. The convolution-type Volterra equation (red, (3.1)) predicts the performance of SDANA and it reflects the oscillatory trajectories typically seen in momentum methods due to overshooting. Because the convolution Volterra only an approximation to the kernel, there always an initial mismatch between actual runs of SDANA and the Volterra solution. As + co, the convolution-tye Volterra equation better approximates SDANA. For more details on numerical simulations see App.
of inputs are chosen at random. To formulate an average case analysis that representative of what is seen in a large scale optimization problem, we will take limit of the empirical spectral measure as n and are taken to infinity. So, we suppose that the following holds:
Assumption (Spectral limit). Let be an n x d matrix drawn from a family of random matrices such that the number of features, d, tends infinity proportionally the size of the data set, n, so that + +rE E (0, oo); and suppose these random matrices satisfy the following.
The eigenvalue distribution of H = AAT converges to a deterministic limit μ with compact support. Formally, the empirical spectral measure (ESM) converges weakly to μ,, in that for all bounded continuous g : + R
9(a)
Pr R-000
s(A)u(dA).
(3.3)
==1
2. The largest eigenvalue X of H converges to the largest element X in the support of i.e. H
X
Pr
X
This assumption is typical in random matrix theory. An important example is the featurrss model, which is a random n x matrix whose every entry is sampled from a common, mean variance z distribution with fourth moment Odd-²), such as a Gaussian N 4). In this case, the ESM p of H AAT converges to the Marchenko-Pastie law (see Figure 3.1):
(A) def do(A) max r, r, +
TV (A )-+ - A)
2Td (1 })² and A (1 Ty2
(3.4)
where
More generally, the convergence of the spectral measure of matrices drawn from a consistent ensemble well studied in random matrix theory, and for many random matrix ensembles the limiting spectral measure is known. In the machine learning literature, it has been shown that the spectrum of the Hessians neural networks share characteristics with the limiting spectral distributions found in classical random matrix theory [Behrooz et al., 2019, Dauphin et al., 2014, Granziol et al., 2020, Liao et al., 2020, Martin and Mahoney, 2018, Papyan, 2018, Pennington and B., 2017, Sagun et al., 2016].
Complexity analysis. The forcing function and the convolution kernel both converge under As- sumption 2, and the result is that
lim Ex(s(x,) = (t) where = + 1 00 It - s))b(s) ds for all n=
t 2 0.
8

---

The forcing function and interaction kernel are given as integrals against the limit measure (such as (3.4) the case of isotropic features) and
(1) def / (RA (()(() and (1) del Jo (O)(a)))).
The kernel norm still determines the convergence properties of the Volterra equation. In particular, to have convergence of the algorithm, we need that ||Zμ||| < and jud. like in the finite-n case (see Lemma C.I)
(0}{
(t) (00) def
2(1 IZ,
To discuss average-case rates, we use the function (t). We consider separately the regimes when the problem is strongly convex and not. Having taken the limit, say the problem is strongly convex if the intersection of the support of μ with (0, oc) is closed. Intuitively, this says there a gap between O and the next smallest eigenvalue À of the hessian (strictly speaking it allows vanishing fraction of the eigenvalues to approach 0).
Spectrum of mNIST
In the non-strongly convex case the average-cayy complexity relatively simple to compute. The rate of convergence of (t) (0)) - is only determined by the rate of of convergence of the forcing function (t)
1.5
Fy (co). This decays like - where 3 in turn is controlled by the exponent a at which u((0,e)) el) x e as s + (see Lemma C.2). Particular if there are more small eigenvalues, the rate is slowed. In Table 3, the rates are reported for typical behavior for random matrix distributions with "hard-edge." such Marchenko- Pastur with aspect ratio r =
JinII
1. This is the
10
15
Empirical Volterra
109
SGD/sDAHB
SDANA
In the strongly convex case, X > O, the kemee I, plays a larger role, that it may slow down the convergence rate. In particular, if it exists, we define the Malthusian exponent X as the solution of
1055 runctionn
10
20
10
10 (I,,
25
75
100
epochs
Figure 6: SDANA & SGD vs Theory on MNIST. MNIST (60000 28 28 images) [LeCun et al., 2010] is reshaped into 10 matrices of dimen- sion 1000 4704, representing 1000 samples of groups of 6 digits (preconditioned to have cen- tered rows of norm-l). First digit of each 6 chosen to be the target b. Algorithms were run 10 times with default parameters (without tuning) to solve (2.1). 80%c-confidence interval is dis- played. Volterra (SDANA) is generated with eigen- values from the first MNIST data matrix (top pane, X = 0.041). Volterra predicts the convergent be- havior of SDANA in this non-idealized setting. SDANA outperforms equivalent SGD/SDAHB. See also Appendix F.
The rate of convergence of to (at exponen- tial scale) will then be the slower of (t)t) and A** (see Lemma C.3). In the case of SGD on Marchenko-Pastur, the exact value of X* is worked out in exact form in Paquette et al. [2021]. By bounding these Malthusian expo- nents, we produce the rate guarantees in Table (see Cor. D.I and Prop. E.3 for the bounds in the Appendix). The default parameters are chosen so that its linear rate is no slower, by factor of than the fastest possible rate for algorithm having optimized over all step size choices. This is achieved by lower bounding the Malthusian exponent at the default param- eters and upper bounding the optimal rate by minimizing (t) (oc) over all convergent parameters.
Conclusions from the analysis of homogenized SGD. of SDAHB with parameter choices no @odahb and
The SHB algorithm a special instance ny ny sdalb, By evaluating the kernels

---

for SDAHB in the large n limit, it is easily seen that the homogenized SGD equations for X x satisfy
sbb and
[m([(X F [[(;
0,
(3.5)
11-900
see (Thm. 5). On the other hand SDAHB with default parameters is always strictly faster (for sufficiently large 8) than tuned SGD, but its linear rate is never more than factor of 2 faster than SGD (Prop. E.5). It also does not substantially improve over SGD in non-strongly convex case. In contrast, the dimension adjusted Nesterov acceleration (SDANA) greatly (and provably, using homogenized SGD) improves over SGD in non-strongly convex case (see Prop. D.3), while remaining linear (and nearly as fast as SGD) the convex case (Cor. D.I). Furthermore, the predictions of homogenized SGD are born out even on real data (Fig. 6), which is a non-idealized setting that does verify the assumptions we imposed for the theoretical analysis.
Future directions. We would like to explore the applicability of homogenized SGD to other datasets and other convex losses as well as generalizing the theoretical setting under which homoge- nized SGD applies (see the discussion below Thm. 2). Moreover, we would like to test and extend SDANA to non-convex problems and extend homogenized SGD to non-convex settings.
Funding Transparency Statement. C. Paquette's research was supported by CIFAR AI Chair, MILA. Research by Paquette was supported by a Discovery Grant from the Natural Science and Engincering Council (NSERC). Additional revenues related to this work: C. Paquette has part-time employment at Google Research, Brain Team, Montreal, QC.
References
Z. Allen-Zhu. Katyusha: The first direct acceleration of stochastic gradient methods. The Journal of Machine Learning Research, 18(1):8194-82) 8244, 2017.
Asmussen. Applied probability and queues, volume 51 of Applications of Mathematics (New York). Springer-Verlag. New York, second edition, 2003, Stochastic Modelling and Applied Probability.
M. Assran and M. Rabbat. On the Convergence of Nesterov's Accelerated Gradient Method in Stochastic Settings. In Proceedings of the 37th International Conference on Machine Learning (ICML)), 2020.
N. Aybat, A. Fallah, M. Gurbuzbalaban, and A. Ozdaglar. Robust Accelerated Gradient Methods for Smooth Strongly Convex Functions. arXiv preprints urXiv: 10579, 2018.
G. Behrooz, S. Krishnan, and Y. Xiao. An Investigation into Neural Net Optimization via Hessian Eigenvalue Density. Proceedings of the 36th International Conference on Machine Learning (ICML), 2019.
L. Bottou, F.E. Curtis, and Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223-311, 2018
B. Can, M. Gurbuzbalaban, and Zhu. Accelerated Linear Convergence of Stochastic Momentum Methods in Wasserstein Distances. arXiv preprints
2019.
E. Coddington and N. Levinson. Theory of ordinary differential equations. McGraw-Hil Book Company, Inc., New York-Toorotto-London,, 1955.
Y.N. Dauphin, R. Pascanu, C. Gulcehre, K. Cho, S. Ganguli, and Y. Bengio. Identifying and attacking the saddle point problem in high-dimensional non-convex optimization. In Advances in Neural Information Processing Systems (NeurIPS), 2014.
N. Flammarion and F. Bach. From Averaging to Acceleration, There is Only Step-size. In Proceedings The 28th Conference on Learning Theory (COLTT), volume 40 of Proceedings of Machine Learning Research, pages 658-695 PMLR, 03-06 Jul 2015.
S. Gadat, Panloup, and S. Saadane. Stochastic Heavy Ball. arXiv preprints arXiv: 160iv:1609.00 042904428. 2016.
10

---

C. Gerbelot, A. Abbara, and F. Krzakala. Asymptotic Errors for High Sigh--Dmeesional Convex Penalized Linear Regression beyond Gaussian Matrices. In Proceedings of Thirty Third Conference on Learning Theory (COLT), volume 125 of Proceedings of Machine Learning Research, pages 1682-1713. PMLR, 2020.
Ghadimi and Lan. Optimal stochastic approximation algorithms for strongly convex stochastic composite optimization I: A generic algorithmic framework. SIAM J. Optim., 2264):1469-1492, 2012. ISSN 1052-6234, doi: 10. 11371 10848864.
Ghadimi and Lan. Optimal stochastic approximation algorithms for strongly convex stochastic composite optimization, II: Shrinking procedures and optimal algorithms. SIAM J. Optim., 23(4): 2061-2089, 3a. doi: 10. 10.1137/110844 10848876.
Ghadimi and G. Lan. Stochastic first- and zeroth-order methods for nonconvex stochastic program- ming. SIAM J. Optim., 23(4):2341-2366 2013b.
B. Ghorbani, S. Krishnan, and Y. Xiao. An investigation into neural net optimization via hessian eigenvalue density. International Conference on Machine Learning (ICMI.), pages 2232-2241.
PMLR, 2019.
D. Granziol, S. Zohren, and S. Roberts. Learning Rates as Function of Batch Size: A Random Matrix Theory Approach to Neural Network Training. arXiv preprint arXiv:2006.00002 09092, 2020.
G. Gripenberg. On the resolvents of nonconvolution Volterra kernels. Funkcial. Ekvac., 23(1):83-955 1980. ISSN 0532-8721. URL http: http://wwwwwthhh.t///////////hht math kobe-u. ac. attp:aacccpp////eeppp//p//// httppp////////////e//eee////
xml.
G. Gripenberg, S.O. Londen, and O. Staffans. Volterra Integral and Functional Equations. Encyclo- pedia of Mathematics and its Applications. Cambridge University Press, 1990.
Hastie, A. Montanari, S. Rosset, and R.J. Tibshirani. Surprises high-dimensional ridgeless least squares interpolation. arXiv preprint arXiv:
2019.
R. Kidambi, P. Netrapalli, P. Jain, and S. Kakade. On the insufficiency of existing momentum schemes for stochastic optimization. In 2018 Information Theory and Applications Workshop (ITA), pages 1-9, 2018. doi: 10. 200 8 8503173.
A. Kulunchakov and J. Mairal. A generic acceleration framework for stochastic composite optimiza- tion. In Advances in Neural Information Processing Systems (NeurIPS), volume 32, 2019.
M. Laborde and A. Oberman. Nesterov's 's method with decreasing learning rate leads to accelerated stochastic gradient descent. arXiv preprints arXiv: 1908.0766,
Y. LeCun, C. Cortes, and C. Burges. "mnist" handwritten digit database, 2010. URL http: //yannleccnn lecun. wtt.ccoooc/ceeeeeeeee mxxdb/mnist...
C. Li, H. Farkhoor, Liu, and J. Yosinski. Measuring the intrinsic dimension of objective landscapes. 1804.8044.88 arXiv:/804.000 2018
arXiv preprint arXiv:
Q Li, C. Tai, and W. E. Stochastic Modified Equations and Adaptive Stochastic Gradient Algorithms. In Proceedings of the 34th International Conference on Machine Learning (ICLR), volume 70, pages 2101-2110 10, 2017.
Q. Li, C. Tai, and W. E. Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms Mathematical Foundations. Journal of Machine Learning Research, 20(40):1-47, 2019.
Z. Liao, R. Couillet, and M. Mahoney, A Random Matrix Analysis of Random Fourier Features: Beyond the Gaussian Kernel, a Precise Phase Transition, and the Corresponding Double Descent. arXiv preprint arXiv:2006.05013. 2020.
C. Liu and M. Belkin. Accelerating SGD with momentum for pee-parramttrrieed learning. In Proceedings of the 37th International Conference on Machine Learning (ICML), 2020.
N. Loizou and Richtárik. Momentum and Stochastic Momentum for Stochastic Gradient, Newton, Proximal Point and Subspace Descent Methods. arXiv preprint 1rXiv:1712.09677, 2017.

---

S. Mandt, M. Hoffman, and D. Blei. A variational analysis of stochastic gradient algorithms. In International conference on machine learning (ICML), 2016.
C.H. Martin and M. W. Mahoney. Implicit self-regularizatiinn in deep neural networks: Evidence from random matrix theory and implications for learning. arXiv preprint arXiv: 2018.
S. Mei and A. Montanari. The generalization emror of random features regression: Precise asymptoties and double descent curve. arXiv preprint 1r0.v:190..05355. 2019.
Y. Nesterov. Introductory lectures on convex optimization. Springer, 2004.
A. Orvieto, J. Kohler, and A. Lucchi. The Role of Memory in Stochastic Optimization. In Conference on Uncertainty in Artificial Intelligence (UAI 2019). 2019.
V. Papyan. The full spectrum deepnet hessians at scale: Dynamics with SGD Training and Sample Size. arXiv preprint urXiv:1811 Ttf:::8111.000000 1.07062, 2018.
C. Paquette, K. Lee, F. Pedregosa, and E. Paquette. SGD in the Large: Average-case Analysis, Asymptotics, and Stepsize Criticality. arXiv preprint arXiv:2 2021.
J. Pennington and Yasaman B. Geometry of Neural Network Loss Surfaces via Random Matrix Theory. In Proceedings of the 34th International Conference Machine Learning (ICML), volume 70 of Proceedings of Machine Learning Research, pages 2798-2806, 2017.
B.T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computa- tional Mathematics and Mathematical Physics, 04, 1964.
L. Sagun, L. Bottou, and Y. LeCun. Eigenvalues of the Hessian in Deep Learning: Singularity and Beyond. arXiv preprint arXiv: 16Xiv:1611.07476 2016.
O. Sebbouh, R. Gower, and A. Defazio. Almost sure convergence rates for Stochastic Gradient Descent and Stochastic Heavy Ball. arXiv preprint arXiv:2006.07867, 2020.
I. Sutskever, J. Martens, G. Dahl, and G. Hinton. On the importance initialization and momentum in deep learning. Proceedings of the 30th International Conference on Machine Learning (ICML), volume 28, pages 1139-1147 117, 2013.
S. Vaswani, F. Bach, and M. Schmidt. Fast and Faster Convergence of SGD for Ever-Parameterizer Models and an Accelerated Perceptron. In Proceedings of the Twenty-Secone International Conference Artificial Intelligence and Statistics (ICML), volume 89 of Proceedings of Machine Learning Research, pages 1195-1204, PMLR, 2019.
Y. Yan, T Yang, Z. Li, Q. Lin, and Yi Yang. A Unified Analysis of Stochastic Momentum Methods for Deep Learning. In Proceedings of the Twenty-Seenthh Intelligence, IJCAI-18, pages 2955 1961. International Joint Conferences on Artificial Intelligence Organization, 7 2018. doi: 10. 20124496/jaaa22118444
Seventh International Joint Conference Artificial
G. Zhang, L. Li, Z. Nado, J. Martens, S. Sachdeva, G. Dahl, C. Shallue, and R. Grosse. Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model, arXiv preprint 1rXiv:1907.044 2019.
12


Reviews:
[{"review_id": "XbhAdNdwEis", "reviewer": "Reviewer_7r4P", "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:\n\n(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.\n\n(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.\n\n(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.\n\n(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.", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "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.\n\nEven 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:\n\n(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.\n\n(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.\n\n(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?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "W9KZjxqCBv1", "reviewer": "Reviewer_meKg", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is in general well-written. \n\nThe 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.\n\nUnder 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.\n\nThe 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).\n\nReference mentioned above:\n\n- 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.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "None.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "IdS94g0pJ99", "reviewer": "Reviewer_ZKu7", "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", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "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.\n\nI 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\n\nThen 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.\n\nAlso, 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.\n\nAlso, 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?\n\nI may recommend acceptance if the authors sufficiently address my questions.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I think the limitation can be better discussed. See above.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2 hours."}, "extra_text": {}}, {"review_id": "2pIyjTlYCh4", "reviewer": "Reviewer_BCD3", "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. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "To my knowledge, the results appear to be novel and I tend towards acceptance. \n\n[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). \n\n[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.\n\nBach and Moulines (2013): Non-strongly-convex smooth stochastic approximation with convergence rate O(1/n)\n\nRakhlin et al (2012): Making gradient descent optimal for strongly convex stochastic optimization\n\nGe et al (2019):  The Step Decay Schedule: A Near Optimal, Geometrically Decaying Learning Rate Schedule for Least Squares", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

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)