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

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 . 1 Motivation and related work . 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 ] .

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.