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