Dynamics of Stochastic Momentum Methods on Large-scale, Quadratic Models
Year: 2021

Abstract:
We analyze a class of stochastic gradient algorithms with momentum on a highdimensional 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 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. 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 are typically implemented with stochastic gradients. In the influential work Sutskever et al. [2013], the authors demonstrated empirical advantages of augmenting stochastic gradient descent (SGD) with the momentum machinery and, as a 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 a framework inspired by the phenomenology of random matrix theory (see Paquette et al. [2021]), which gains explanatory power when the number of samples (n) and features (d) are large. A key contribution of this work is a simple description of the exact dynamics for a class of stochastic momentum methods in the high-dimensional limit; we construct a smooth, deterministic function (t) such that f(xk) ! (k/n) as n ! 1. This function solves ⇤Website courtneypaquette.github.io . †Website elliotpaquette.github.io . 35th Conference on Neural Information Processing Systems (NeurIPS 2021). Algorithm 1 Generic stochastic momentum method. Given: step sizes 1, 2 > 0 and momentum parameter (k) > 0 Initialize: x0 2 R and y0 = 0 for k 1, Select ik 2 [n] uniformly and update yk = (1 (k))yk 1 + 1rfik(xk) and xk+1 = xk 2rfik(xk) yk (0.2) a Volterra integral equation:

text: We analyze a class of stochastic gradient algorithms with momentum on a highdimensional 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 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.
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 are typically implemented with stochastic gradients. In the influential work Sutskever et al. [2013], the authors demonstrated empirical advantages of augmenting stochastic gradient descent (SGD) with the momentum machinery and, as a 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 a framework inspired by the phenomenology of random matrix theory (see Paquette et al. [2021]), which gains explanatory power when the number of samples (n) and features (d) are large. A key contribution of this work is a simple description of the exact dynamics for a class of stochastic momentum methods in the high-dimensional limit; we construct a smooth, deterministic function (t) such that f(xk) ! (k/n) as n ! 1. This function solves
⇤Website courtneypaquette.github.io . †Website elliotpaquette.github.io .
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Algorithm 1 Generic stochastic momentum method. Given: step sizes 1, 2 > 0 and momentum parameter (k) > 0 Initialize: x0 2 Rd and y0 = 0 for k 1, Select ik 2 [n] uniformly and update
yk = (1 (k))yk 1 + 1rfik(xk) and xk+1 = xk 2rfik(xk) yk (0.2)
a Volterra integral equation:
(t) = F (t) +
Z t
0 Ks(t) (s) ds. (0.1)
Here F (t) and Ks(t) are explicit, see Theorem 1 for a precise statement. This Volterra equation (0.1) gives an accurate prediction of the behavior of stochastic methods, see Figure 1. We then analyze these dynamics providing insight into step size and momentum parameter selections as well as providing both upper and lower average-case complexity (i.e., the complexity of an algorithm averaged over all possible inputs) for the last iterate.
As we show in this work, both theoretically and empirically, (small batch size) SGD with heavyball momentum (SHB) for any fixed momentum parameter does not provide any acceleration over plain SGD on large-scale least square problems. We conclude under an identification of the parameters that f(xshb
k ) = f(xsgd k ) for all k up to errors
that vanish as n grows large (upper bounds of this nature have been observed before: see Kidambi et al. [2018], Sebbouh et al. [2020], Zhang et al. [2019]). Thus while SHB may provide a speed-up over SGD, it is only due to an effective increase in the learning rate, and this speed-up could be
matched by appropriately adjusting the learning rate of SGD.
The root of SHB’s failure to provide meaningful acceleration is that a fixed momentum parameter is not aggressive enough when n is large. We propose a new algorithm that uses a dimension-based modification of Nesterov (see Alg. 1 and Table 1). The resulting algorithm, SDANA, matches the average-case complexity of SGD when the least-squares problem is strongly convex and obtains an average-case complexity of 1/k3 in the convex setting.

heading: 1 Motivation and related work
text: We consider the large finite-sum setting
min x2Rd
⇢ f(x) = 1
n
nX
i=1
fi(x) = 1
2
nX
i=1
(aix bi)2 = 1
2 kAx bk2 ,
for data matrix A 2 Rn⇥d whose i-th row is denoted by ai 2 Rd⇥1 and target vector b 2 Rn (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 1
n fi(x) = 1 2 (aix bi) 2. We investigate a generic class of stochastic momentum algorithms (see Alg. 1 and Table 1). Particularly, we introduce a sub-class, denoted by SDA( 1, 2, ), of Alg. 1 which has parameters that are appropriately adjusted for large problems (large number of samples n and large model size d); we refer to the dimension of
the problem as n. The class SDA( 1, 2, ) is defined by setting in Alg. 1
1 = 1
n , 2 = 2, and (k)
def = 1 n (log')0( k n ), (1.1)
where 1, 2 > 0 are step sizes and ' is a smooth function that represents a momentum schedule. Although we develop some theory for general ', we are principally interested in the two cases:
(SDANA) (k) = ✓ k+n $ '(t) = (1 + t) ✓ and (SDAHB) (k) = ✓ n $ '(t) = e✓t. (1.2)
To avoid confusion between different algorithms, we add superscripts indicating the algorithm (e.g., we denote 2 = sgd, the step size parameter for SGD). For all these algorithms, we are interested in:
1. An expression for the (deterministic) dynamics of these algorithms when multiple passes on the data set are allowed. This contrasts with the "streaming" or "online" setting where at each iteration one generates an independent neverbefore-used data point.
2. A formula for choosing the hyperparameters and a discussion of the dependence of these hyperparameters on number of features and samples.
3. Upper and lower bounds on the averagecase complexity of the last iterate to a neighborhood; this neighborhood disappears entirely in the overparameterized regime, while in the underparameterized regime the limiting distance to optimality concentrates in the highdimensional limit.
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 (see
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 as vanilla SGD, which is to say, by setting the step size parameter in SGD to be sgd = shb
✓shb ,
the two algorithms have the same loss values provided the number of samples is sufficiently large, i.e., f(xsgd
k ) = f(xshb k ).
As a consequence of this, the average-case complexity of SHB equals the last iterate complexity of SGD (This was observed in [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 dimension-dependent parameters to achieve dynamics which differ from SGD.
Why divide by n? A positive result. Adapting SHB for dimension, we arrive at stochastic dimension adjusted heavy ball (SDAHB). While formally equivalent to SHB, we include the dimension parameters to emphasize that any improvement in its performance for large n requires it. Nonetheless, the speed-up for heavy ball is modest (see Fig. 2).
On the other hand, we show that a dimension adapted version of Nesterov acceleration, SDANA, has a large improvement in the nonstrongly 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 learningrate-tuned SGD (or SHB), while performing orders-of-magnitude faster (k 3 as compared to SGD k 1) for the non-strongly convex setting (see Fig. 3 and Table 3). We believe this gives SDANA promise as 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]).
Related work. Recent works have established convergence guarantees for SHB in both strongly convex and non-strongly convex setting [Flammarion and Bach, 2015, Gadat et al., 2016, Orvieto et al., 2019, Sebbouh et al., 2020, Yan et al., 2018]; the latter references having established almost sure convergence results. Specializing to the setting of minimizing quadratics, the iterates of SHB converge linearly (but not in L2) 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.
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 an 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 heavy-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].

heading: 2 Random least squares problem
text: To formalize the analysis of a high–dimensional, typical least squares problem, we define the random least squares problem:
argmin x2Rd
n f(x) = 1
n
nX
i=1
fi(x) def =
1
2
nX
i=1
(aix bi)2 o , with b def= Aex+ ⌘. (2.1)
The data matrix A is random and we shall introduce assumptions on A as they are needed, but we suggest as a central example the Gaussian random least squares where each entry of A is sampled independently from a standard normal distribution with variance 1
d . We always make the assumption
that each row ai 2 Rd⇥1 is centered and is normalized so that maxi {E [kaik2]} = 1. As for the target b = Aex + ⌘, we assume it comes from a generative model corrupted by noise, where ex is signal and ⌘ is noise. Assumption 1 (Initialization, signal, and noise). The initial vector x0 2 Rd is chosen so that x0 ex 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
E kx0 exk22 = R dn and E [k⌘k 2 2] = eR.
Note that deterministic x0 ex satisfies this assumption. The vectors x0 ex and ⌘ arise as a result of preserving a constant signal-to-noise 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 fH = ATA and its symmetrization H = AAT . Let 1 . . . n 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 , µH by the formula
Z g( )µH(d ) def = 1
n
nX
i=1
g( i) for any continuous function g : R ! R. (2.2)
This gives the interpretation for the empirical spectral measure as the distribution of an eigenvalue of H chosen uniformly at random.
Diffusion approximation. Our analysis will use a diffusion approximation to analyze the SDA( 1, 2, ) class of stochastic momentum methods (see (1.1)) on the random least squares setup (2.1). We call the approximation homogenized SGD:
dXt def = 2 dZt
1
'(t)
Z t
0 '(s) dZt, where dZt
def = rf(Xt) dt+ q 2 n f(Xt)r2(f) dBt, (2.3)
and with initial conditions given by X0 = x0. The process (Bt : t 0) is a d–dimensional standard Brownian motion. Here time is scaled in such a way that t = 1 represents one 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 a least squares problem, even at finite n. Theorem 1 (Volterra dynamics at finite n). Let EH [·] be the conditional expectation where H is held fixed. There are non-negative functions F (t) and Ks(t) for s, t 0 depending on the spectrum of H so that for all t 0
EH [f(Xt)] = F (t) + Z t
0 Ks(t)EH [f(Xs)] ds, for all t 0. (2.4)
The forcing function F and kernel K are given by
F (t) = 1
n
nX
i=1
(R i + eR)G( i)(t) and Ks(t) = 1
n
nX
i=1
K ( i) s (t).
The functions G( ) and K( ) are solutions of an initial value problem with a 3-rd order ODE which depend on the hyperparameters ( 1, 2, ) (Note, there is a 1-to-1 relationship with ', see (1.1)).
We refer to the supplemental materials for the explicit third–order ODE (see Theorem 4 for full details). The expression in (2.4) is a 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.e. Ks(t) = I(t s) for some function I); see Table 2 for kernels. We also note that in the case of SGD( ) (' is unused), the functions G and K become particularly simple
G ( )(t) = e 2 t and K( ) s (t) = 2 2e 2 (t s).
Comparing homogenized SGD to the SDA class. When A is a random matrix, we can compare the diffusion (2.3) to SDA (1.1) when n and d are large. The argument is based on the results of Paquette et al. [2021], and we do it only in the case of SGD: Theorem 2 (Concentration of SGD). Suppose that A is a left-orthogonally invariant random matrix, meaning that for any orthogonal matrix O 2 Rn⇥n,OA law= A. Suppose further that the noise vector ⌘ is independent of A and that it satisfies
E [k⌘kp1] = O(n✏ p/2) for any ✏, p > 0. Fix < 2n(trH) 1, the convergence threshold of SGD( ). There is an absolute constant " > 0 and a constant c(T, +
H ) so that with p = min{d, n},
Pr( sup 0tT
|EH [f(Xt)] f(x[nt])| > c(T, +H)p " | + H )  p ".
We expect that this theorem can be generalized, to include the entire SDA class. We also expect that the orthogonal invariance assumption can be relaxed somewhat (for example to include classes of non–Gaussian isotropic features matrices), but not entirely: the left singular vectors need to 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 5 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 in App. B.

heading: 3 Main results
text: 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.

heading: 3.1 Convolution Volterra convergence analysis: convergence threshold and neighborhood
text: For all algorithms considered (SDANA, SDAHB, SHB, SGD), the Volterra equation in Theorem 1 can be expressed in a simpler form, that is, as a convolution–type Volterra equation
EH [f(Xt)] = F (t) + Z t
0 I(t s)EH [f(Xs)] ds for all t 0. (3.1)
The forcing function F and the convolution kernel I are non-negative functions that depend on the spectrum of 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–convolution equation as t ! 1. First, the forcing function F will in all cases be bounded, and in fact it will converge as t ! 1 to a deterministic value,
F (t) ! t!1 eRµH({0}) 2 = eR dim(ker(H)) 2n . (3.2)
Here µH 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 kIk = R1 0 I(t) dt is less than 1. Theorem 3 (Convergence threshold and limiting loss). If the norm kIk < 1, the algorithm is convergent in that
EH [f(Xt)] ! t!1 eR dim(ker(H)) 2n(1 kIk) (limiting loss).
Figure 4: The Marchenko-Pastur law(r). Varying r = d/n.
This theorem gives a convergence threshold for all algorithms in Table 1 based only on the norm of the kernel of the Volterra equation, which is easily computable (see Table 3).

heading: 3.2 Average case analysis
text: Limiting spectral measures. Average-case complexity looks at the typical behavior of an algorithm when some
of its inputs are chosen at random. To formulate an average case analysis that is representative of what is seen in a large scale optimization problem, we will take a limit of the empirical spectral measure as n and d are taken to infinity. So, we suppose that the following holds:
Assumption 2 (Spectral limit). Let A be an n⇥ d matrix drawn from a family of random matrices such that the number of features, d, tends to infinity proportionally to the size of the data set, n, so that d
n ! r 2 (0,1); and suppose these random matrices satisfy the following.
1. 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 ! R
1
n
nX
i=1
g( i) Pr !
n!1
Z 1
0 g( )µ(d ). (3.3)
2. The largest eigenvalue + H of H converges to the largest element + in the support of µ, i.e. + H Pr ! n!1 + .
This assumption is typical in random matrix theory. An important example is the isotropic features model, which is a random n⇥ d matrix A whose every entry is sampled from a common, mean 0, variance 1
d distribution with fourth moment O(d 2), such as a Gaussian N(0, 1 d ). In this case, the
ESM µH of H = AAT converges to the Marchenko-Pastur law (see Figure 3.1):
dµMP( ) def = 0( )max{1 r, 0}+ r
p ( )( + )
2⇡ 1[ , +] ,
where def= (1 q
1 r )2 and + def= (1 + q 1 r )2 .
(3.4)
More generally, the convergence of the spectral measure of matrices drawn from a consistent ensemble is 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 of 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 Assumption 2, and the result is that
lim n!1
EH [f(Xt)] = (t) where (t) = Fµ(t) + Z 1
0 Iµ(t s) (s) ds for all t 0.
The forcing function and interaction kernel are given as integrals against the limit measure (such as (3.4) in the case of isotropic features) and
Fµ(t) def =
Z 1
0 (R + eR)G( )(t)µ(d ) and Iµ(t) def =
Z 1
0 K
( )(t)µ(d ).
The kernel norm still determines the convergence properties of the Volterra equation. In particular, to have convergence of the algorithm, we need that kIµk < 1 and just like in the finite-n case (see Lemma C.1)
(t) ! t!1
(1) def= eRµ({0})
2(1 kIµk) .
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, we say the problem is strongly convex if the intersection of the support of µ with (0,1) is closed. Intuitively, this says there is a gap between 0 and the next smallest eigenvalue of the hessian (strictly speaking it allows a vanishing fraction of the eigenvalues to approach 0).
In the non-strongly convex case the average-case complexity is relatively simple to compute. The rate of convergence of (t) (1) ! 0 is only determined by the rate of of convergence of the forcing function Fµ(t) Fµ(1). This decays like t where in turn is controlled by the exponent ↵ at which µ((0, "]) ⇣ ✏↵ as " ! 0 (see Lemma C.2). Particular if there are more small eigenvalues, the rate is slowed. In Table 3, the rates are reported for ↵ = 12 . This is the typical behavior for random matrix distributions with a “hard-edge,” such as Marchenko–Pastur with aspect ratio r = 1.
In the strongly convex case, > 0, the kernel Iµ plays a larger role, in that it may slow down the convergence rate. In particular, if it exists, we define the Malthusian exponent ⇤ as the solution of
Z 1
0 e ⇤ tIµ(t) = 1.
The rate of convergence of to 0 (at exponential scale) will then be the slower of Fµ(t) and e ⇤t (see Lemma C.3). In the case of SGD on Marchenko–Pastur, the exact value of ⇤ is worked out in exact form in Paquette et al. [2021]. By bounding these Malthusian exponents, we produce the rate guarantees in Table 3 (see Cor. D.1 and Prop. E.3 for the bounds in the Appendix). The default parameters are chosen so that its linear rate is no slower, by a factor of 4 than the fastest possible rate for an algorithm having optimized over all step size choices. This is achieved by lower bounding the Malthusian exponent at the default parameters and upper bounding the optimal rate by
minimizing Fµ(t) Fµ(1) over all convergent parameters.
Conclusions from the analysis of homogenized SGD. The SHB algorithm is a special instance of SDAHB with parameter choices n✓shb = ✓sdahb and n shb = sdahb. By evaluating the kernels
for SDAHB in the large n limit, it is easily seen that the homogenized SGD equations for Xshb t and Xsgd
t satisfy
|EH [f(Xshbt )] EH [f(X sgd t
)]| ! n!1 0, (3.5)
see (Thm. 5). On the other hand SDAHB with default parameters is always strictly faster (for sufficiently large ✓) than tuned SGD, but its linear rate is never more than a 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) in the convex case (Cor. D.1). Furthermore, the predictions of homogenized SGD are born out even on real data (Fig. 6), which is a non-idealized setting that does not 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 homogenized 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 E. Paquette was supported by a Discovery Grant from the Natural Science and Engineering Council (NSERC). Additional revenues related to this work: C. Paquette has part-time employment at Google Research, Brain Team, Montreal, QC.

Review summaries:
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.

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.