Analytical Study of Momentum-Based Acceleration Methods in Paradigmatic High-Dimensional Non-Convex Problems

1 Introduction . In many computer science applications one of the critical steps is the minimization of a cost function . Apart from very few exceptions , the simplest way to approach the problem is by running local algorithms that move down in the cost landscape and hopefully approach a minimum at a small cost . The simplest algorithm of this kind is gradient descent , that has been used since the XIX century to address optimization problems [ 1 ] . Later on , faster and more stable algorithms have been developed : second order methods [ 2 , 3 , 4 , 5 , 6 , 7 ] where information from the Hessian is used to adapt the descent to the local geometry of the cost landscape , and first order methods based on momentum [ 8 , 9 , 10 , 11 , 12 ] that introduce inertia in the algorithm and provably speed up convergence in a variety of convex problems . In the era of deep-learning and large datasets , the research has pushed towards memory efficient algorithms , in particular stochastic gradient descent that trades off computational and statistical efficiency [ 13 , 14 ] , and momentum-based methods are very used in practice [ 15 ] . Which algorithm is the best in practice seems not to have a simple answer and there are instances where a class of algorithms outperforms the other and vice-versa [ 16 ] . Most of the theoretical literature on momentum-based methods concerns convex problems [ 17 , 18 , 19 , 20 , 21 ] and , despite these methods have been successfully applied to a variety of problems , only recently high dimensional non-convex settings have been considered [ 22 , 23 , 24 ] . Furthermore , with few exceptions [ 25 ] , the majority of these studies focus on worst-case analysis while empirically one could also be interested in the behaviour of such algorithms on typical instances of the optimization problem , formulated in terms of a generative model extracted from a probability distribution . The main contribution of this paper is the analytical description of the average evolution of momentumbased methods in two simple non-convex , high-dimensional , optimization problems . First we consider the mixed p-spin model [ 26 , 27 ] , a paradigmatic random high-dimensional optimization problem . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Furthermore we consider its spiked version , the spiked matrix-tensor [ 28 , 29 ] which is a prototype high-dimensional non-convex inference problem in which one wants to recover a signal hidden in the landscape . The second main result of the paper is the characterization of the algorithmic threshold for accelerated-methods in the inference setting and the finding that this seems to coincide with the threshold for gradient descent . The definition of the model and the algorithms used are reported in section 2 . In section 3 and 4 we use dynamical mean field theory [ 30 , 31 , 32 ] to derive a set of equations that describes the average behaviour of these algorithms starting from random initialization in the high dimensional limit and in a fully non-convex setting . We apply our equations to the spiked matrix-tensor model [ 29 , 33 , 34 ] , which displays a similar phenomenology as the one described in [ 24 , 35 ] for the phase retrieval problem : all algorithms have two dynamical regimes . First , they navigate in the non-convex landscape and , second , if the signal to noise ratio is strong enough , the dynamics eventually enters in the basin of attraction of the signal and rapidly reaches the bottom of the cost function . We use the derived state evolution of the algorithms to determine their algorithmic threshold for signal recovery . Finally , in Sec . 5 we show that in the analysed models , momentum-based methods only have an advantage in terms of speed but they do not outperform vanilla gradient descent in terms of the algorithmic recovery threshold . 2 Model definition . We consider two paradigmatic non-convex models : the mixed p-spin model [ 32 , 36 ] , and the spiked matrix-tensor model [ 28 , 29 ] . Given a tensor T ∈ ( RN ) ⊗p and a matrix Y ∈ RN×N , the goal is to find a common low-rank representation x that minimizes the loss L = − 1 ∆p √ ( p− 1 ) ! Np−1 N∑ i1 , ... , ip=1 Ti1 , ... , ipxi1 . . . xip − 1 ∆2 1√ N N∑ i , j=1 Yijxixj , ( 1 ) with x in the N -dimensional sphere of radius √ N . The two problems differ by the definition of the variables T and Y . Call ξ ( p ) and ξ ( 2 ) order p tensor and a matrix having i.i.d . Gaussian elements , with zero mean and variances ∆p and ∆2 respectively . In the mixed p-spin model , tensor and matrix are completely random T = ξ ( p ) and Y = ξ ( 2 ) . While in the spiked matrix-tensor model there is a low-rank representation given by x∗ ∈ SN−1 ( √ N ) embedded in the problem as follows : Ti1 ... ip = √ ( p− 1 ) ! Np−1 x∗i1 . . . x ∗ ip + ξ ( p ) i1 ... ip , Yij = x∗i x ∗ j√ N + ξ ( 2 ) ij . ( 2 ) These problems have been studied both in physics , and computer science . In the physics literature , research has focused on the relationship of gradient descent and Langevin dynamics and the corresponding topology of the complex landscape [ 32 , 37 , 38 , 36 , 39 , 27 , 40 ] . The state evolution of the gradient descent dynamics for the mixed spiked matrix-tensor model has been studied only more recently [ 33 , 34 ] . All these works considered simple gradient descent dynamics and its noisy ( Langevin ) dressing . In this work we focus on accelerated methods and provide an analytical characterization of the average performance of these algorithms for the models introduced above . In order to simplify the analysis we relax the hard constraint on the norm of the vector x and consider x ∈ RN while adding a penalty term to L to enforce a soft constraint µ4N ( ∑ i x 2 i −N ) 2 , so that the total cost function is H = L+ µ4N ( ∑ i x 2 i −N ) 2 . Using the techniques described in detail in the next section we write the state evolution for the following algorithms : • Nesterov acceleration [ 9 ] starting from y [ 0 ] = x [ 0 ] ∈ SN−1 ( √ N ) x [ t+ 1 ] = y [ t ] − α∇H ( y [ t ] ) , ( 3 ) y [ t+ 1 ] = x [ t+ 1 ] + t t+ 3 ( x [ t+ 1 ] − x [ t ] ) . ( 4 ) given α the learning rate of the algorithm . • Polyak ’ s or heavy ball momentum ( HB ) [ 8 ] starting from y [ 0 ] = 0 and x [ 0 ] ∈ SN−1 ( √ N ) , given the parameters α , β y [ t+ 1 ] = βy [ t ] +∇H ( x [ t ] ) , ( 5 ) x [ t+ 1 ] = x [ t ] − αy [ t+ 1 ] ; ( 6 ) • gradient descent ( GD ) starting from x [ 0 ] ∈ SN−1 ( √ N ) x [ t+ 1 ] = x [ t ] − α∇H ( x [ t ] ) . ( 7 ) This case has been considered in [ 27 , 33 ] with the constraint ∑ i x 2 i = N . The generalization to the present case in which constraint is soft is a straightforward small extension of these previous works . We will not compare the performance of these accelerated gradient methods to algorithms of different nature ( such as for example message passing ones ) in the same settings . Our goal will be the derivation of a set of dynamical equations describing the average evolution of such algorithms in the high dimensional limit N →∞ . 3 Dynamical mean field theory . We use dynamical mean field theory ( DMFT ) techniques to derive a set of equation describing the evolution of the algorithms in the high-dimensional limit . The method has its origin in statistical physics and can be applied to the study of Langevin dynamics of disordered systems [ 30 , 31 , 43 ] . More recently it was proved to be rigorous in the case of the mixed p-spin model [ 44 , 45 ] . The application to the inference version of the optimization problem is in [ 29 , 33 ] . The same techniques have also been applied to study the stochastic gradient descent dynamics in single layer networks [ 46 ] and in the analysis of recurrent neural networks [ 47 , 48 , 49 ] . The derivation presented in the rest of the section is heuristic and , as such , it is not fully rigorous . Making our results rigorous would be an extension of the works [ 44 , 45 ] where path-integral methods are used to prove a large deviation principle for the infinite-dimensional limit . Our non-rigorous results are checked against extensive numerical simulations . The idea behind DMFT is that , if the input dimension N is sufficiently large , one can obtain a description of the dynamics in terms of the typical evolution of a representative entry of the vector x ( and vector y when it applies ) . The representative element evolves according to a non-Markovian stochastic process whose memory term and noise source encode , in a self-consistent way , the interaction with all the other components of vector x ( and y ) . The memory terms as well as the statistical properties of the noise are described by dynamical order parameters which , in the present model , are given by the dynamical two-time correlation and response functions . In this first step of the analysis we obtain an effective dynamics for a representative entry xi ( and yi ) . The next step consists in using such equations to compute self-consistently the properties of the corresponding stochastic processes , namely the memory kernel and the statistical correlation of the noise . In Fig . 1 we anticipate the results by comparing numerical simulations with the integration of the DMFT equations for the different algorithms : on the left we observe the evolution of the loss , on the right we observe the evolution of the radius of the vector x , defined as the L2 norm of the vector ||x||2 . We find a good agreement between the DMFT state evolution and the numerical simulations . We compare Nesterov acceleration with the heavy ball momentum in the mixed p-spin model Fig . 1 , and in the spiked model Fig . 3 . Nesterov acceleration allows for a fast convergence to the asymptotic energy without need of parameter tuning . In Fig . 2 we compare the numerical simulations for the HB algorithm and the DMFT description of the corresponding massive momentum version for several control parameters . DMFT equations In the following we describe the resulting DMFT equations for the correlation and response functions . The details of their derivation for the case of the Nesterov acceleration are provided in the following section , while we leave the other cases to the supplementary material ( SM ) . The dynamical order parameters appearing in the DMFT equations are one-time or two-time correlations , e.g . Cxy [ t , t′ ] = ∑ i xi [ t ] yi [ t ′ ] /N , and response to instantaneous perturbation of the dynamics , e.g . Rx [ t , t′ ] = ( ∑ i δxi [ t ] /δHi [ t ′ ] ) /N by a local fieldH [ t′ ] ∈ RN where the symbol δ denotes the functional derivative . In this section we show only the equations for the mixed p-spin model and we discuss the difference and the derivation of the equations for the spiked tensor in the SM . From the order parameters we can evaluate useful quantities that describe the evolution of the algorithms . In particular in Figs . 1,2,3 we show the loss , the radius , and the overlap with the solution in the spiked case ( Fig . 3 ) : • Average loss L [ t ] = − α ∆pCx [ t , t ] p 2 t∑ t′′=0 Rx [ t , t ′ ] Cx [ t , t ′ ] p−1 − α ∆2Cx [ t , t ] t∑ t′′=0 Rx [ t , t ′ ] Cx [ t , t ′ ] ; ( 8 ) • Radius √ Cx [ t , t ] ; • Define mx [ t ] = 1N ∑ i xi [ t ] x ∗ i an additional order parameter for the spiked matrix-tensor model ( more details are given in the SM ) , the overlap with ground truth is x [ t ] · x∗ ||x|| = mx [ t ] √ Cx [ t , t ] Nesterov acceleration . It has been shown that this algorithm has a quadratic convergence rate to the minimum in convex optimization problems under Lipschitz loss functions [ 9 , 50 ] , thus it outperforms standard gradient descent whose convergence is linear in the number of iterations . The analysis of the algorithm is described by the flow of the following dynamical correlation functions Cx [ t , t ′ ] = 1 N ∑ i xi [ t ] xi [ t ′ ] , ( 9 ) Cy [ t , t ′ ] = 1 N ∑ i yi [ t ] yi [ t ′ ] , ( 10 ) Cxy [ t , t ′ ] = 1 N ∑ i xi [ t ] yi [ t ′ ] , ( 11 ) Rx [ t , t ′ ] = 1 N ∑ i δxi [ t ] δHi [ t′ ] , ( 12 ) Ry [ t , t ′ ] = 1 N ∑ i δyi [ t ] δHi [ t′ ] . ( 13 ) The dynamical equations are obtained following the procedure detailed in section 4 . Call Q ( x ) = x2/ ( 2∆2 ) + x p/ ( p∆p ) , Cx [ t+ 1 , t ′ ] = Cxy [ t , t ′ ] − αµ ( Cy [ t , t ] − 1 ) Cy [ t , t′ ] + α2 t′∑ t′′=0 Rx [ t ′ , t′′ ] Q′ ( Cy [ t , t ′′ ] ) + + α2 t∑ t′′=0 Ry [ t , t ′′ ] Q′′ ( Cy [ t , t ′′ ] ) Cxy [ t ′ , t′′ ] ; ( 14 ) Cxy [ t+ 1 , t ′ ] = Cy [ t , t ′ ] − αµ ( Cy [ t , t ] − 1 ) Cxy [ t , t′ ] + α2 t′∑ t′′=0 Ry [ t ′ , t′′ ] Q′ ( Cy [ t , t ′′ ] ) + + α2 t∑ t′′=0 Ry [ t , t ′′ ] Q′′ ( Cy [ t , t ′′ ] ) Cy [ t ′ , t′′ ] ; ( 15 ) Cxy [ t ′ , t+ 1 ] = 2t+ 3 t+ 3 Cx [ t+ 1 , t ′ ] − t t+ 3 Cx [ t , t ′ ] ; ( 16 ) Cy [ t ′ , t+ 1 ] = 2t+ 3 t+ 3 Cxy [ t+ 1 , t ′ ] − t t+ 3 Cxy [ t , t ′ ] ; ( 17 ) Rx [ t+ 1 , t ′ ] = Ry [ t , t ′ ] + δt , t′ − αµ ( Cy [ t , t ] − 1 ) Ry [ t , t′ ] + α2 t∑ t′′=t′ Ry [ t , t ′′ ] Ry [ t ′′ , t′ ] Q′′ ( Cy [ t , t ′′ ] ) ; ( 18 ) Ry [ t ′ , t+ 1 ] = 2t+ 3 t+ 3 Rx [ t+ 1 , t ′ ] − t t+ 3 Rx [ t , t ′ ] . ( 19 ) The initial conditions are : Cx [ 0 , 0 ] = 1 , Cy [ 0 , 0 ] = 1 , Cxy [ 0 , 0 ] = 1 , Rx [ t+ 1 , t ] = 1 , Ry [ t+ 1 , t ] = 2t+3 t+3 . The equations show a discretized version of the typical structure of DMFT equations . We can observe : terms immediately ascribable to the dynamical equations ( 3,4 ) and summations whose interpretation is less trivial without looking into the derivation . They represent memory kernels that take into account linear response theory for small perturbations to the dynamics ( e.g . the last term of Eq . ( 14 ) ) and a noise whose statistical properties encode the effect of all the degrees of freedom on a representative one ( e.g . the second last term of Eq . ( 14 ) ) . Heavy ball momentum . The DMFT equations are obtained analogously to previous ones , Cy [ t+ 1 , t ′ ] = βCy [ t , t ′ ] + µ ( Cx [ t , t ] − 1 ) Cxy [ t , t′ ] + α t′∑ t′′=0 Ry [ t ′ , t′′ ] Q′ ( Cx [ t , t ′′ ] ) + α t∑ t′′=0 Rx [ t , t ′′ ] Q′′ ( Cx [ t , t ′′ ] ) Cxy [ t ′′ , t′ ] ; ( 20 ) Cxy [ t ′ , t+ 1 ] = βCxy [ t ′ , t ] + µ ( Cx [ t , t ] − 1 ) Cx [ t , t′ ] + α t′∑ t′′=0 Rx [ t ′ , t′′ ] Q′ ( Cx [ t , t ′′ ] ) + α t∑ t′′=0 Rx [ t , t ′′ ] Q′′ ( Cx [ t , t ′′ ] ) Cx [ t ′ , t′′ ] ; ( 21 ) Cxy [ t+ 1 , t ′ ] = Cxy [ t , t ′ ] − αCy [ t+ 1 , t′ ] ; ( 22 ) Cx [ t+ 1 , t ′ ] = Cx [ t , t ′ ] − αCxy [ t′ , t+ 1 ] ; ( 23 ) Ry [ t+ 1 , t ′ ] = βRy [ t , t ′ ] + 1 α δt , t′ + µ ( Cx [ t , t ] − 1 ) Rx [ t , t′ ] + α t∑ t′′=0 Rx [ t , t ′′ ] Rx [ t ′′ , t′ ] Q′′ ( Cx [ t , t ′′ ] ) ; ( 24 ) Rx [ t+ 1 , t ′ ] = Rx [ t , t ′ ] − αRy [ t+ 1 , t′ ] . ( 25 ) with initial conditions : Cx [ 0 , 0 ] = 1 , Cy [ 0 , 0 ] = 0 , Cxy [ 0 , 0 ] = 0 , Ry [ t+ 1 , t ] = 1/α , Rx [ t+ 1 , t ] = −1 . Fig . 2 shows the consistency of theory and simulations . Mappings between discrete update equation and continuous flow for both heavy ball momentum and Nesterov acceleration have been proposed in the literature . In the SM we considered the work [ 51 ] that maps HB to second order ODEs in some regimes of α and β . This mapping establishes the equivalence of the algorithm to the physics problem of a massive particle moving under the action of a potential . This problem has been studied in [ 52 ] but the result is limited to the fully under-damped regime where there is no first order derivative term , corresponding therefore to a dynamics that is fully inertial and which never stops due to energy conservation . In the SM we obtain the dynamical equations for arbitrary damping regimes , and we recover the equivalence established in [ 51 ] comparing the results from the two DMFTs formulations . Gradient descent . A simple way to obtain the gradient descent DMFT is by taking the limitm→ 0 in the DMFT of the massive momentum description of HB . We get Cx [ t+ 1 , t ′ ] = Cx [ t , t ′ ] − αµ ( Cx [ t , t ] − 1 ) Cx [ t , t′ ] + α2 t′∑ t′′=0 Rx [ t ′ , t′′ ] Q′ ( Cx [ t , t ′′ ] ) + α2 t∑ t′′=0 Rx [ t , t ′′ ] Q′′ ( Cx [ t , t ′′ ] ) Cx [ t ′ , t′′ ] ; ( 26 ) Rx [ t+ 1 , t ′ ] = Rx [ t , t ′ ] + δt , t′ + α 2 t∑ t′′=0 Rx [ t , t ′′ ] Rx [ t ′′ , t′ ] Q′′ ( Cx [ t , t ′′ ] ) − αµ ( Cx [ t , t ] − 1 ) Rx [ t , t′ ] . ( 27 ) with initial conditions : Cx [ 0 , 0 ] = 1 , and Rx [ t+ 1 , t ] = 1 . Apart from the µ-dependent term , these equations are a particular case of the ones that appear in [ 36 , 37 ] and we point to these previous references for details .

Summary: This work studies two types of high-dimensional non-convex optimization problems: (1) the mixed p-spin model; (2) the spiked matrix tensor model. Authors studied a high-dimensional, mean-field type limit of momentum-accelerated optimization dynamics solving these two types of non-convex problems. Mean-field type dynamical equations for Polyak's heavy-ball and Nesterov's method are derived.