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: The authors study the dynamics of accelerated gradient methods applied to two prototypical, non-convex problems: recovering the ground state of a mixed p-spin model, and recovering a planted signal from a spiked matrix-tensor model. The two models are defined according to the following loss function  $$ \mathcal{L}(x) = - \eta_p(\Delta_p) \cdot \sum_{i_1, \ldots, i_p}^N T_{i_1, \ldots, i_p} \cdot x_{i_1} \ldots x_{i_p} - \eta_2(\Delta_2) \cdot \sum_{i,j}^N Y_{i,j} \cdot x_{i} x_{j} $$  where $T \in (\mathbb{R}^N)^{\otimes p}$ is a tensor, $Y \in \mathbb{R}^{N \times N}$ is a matrix and $\eta_p(\Delta_p)$ and $\eta_2(\Delta_2)$ are normalization constants dependent on $\Delta_p$ and $\Delta_2$ respectively. When $T$ and $Y$ possess entries that are i.i.d Gaussians with mean 0 and variance $\Delta_p$ and $\Delta_2$ respectively, we have the mixed $p$-spin model. When $T$ and $Y$ are given by a Gaussian perturbation of vector $x^*$ on the $\sqrt{N}$ radius $N$-dimensional sphere, we have the spiked matrix-tensor model. Finding the ground state for mixed $p$-spin, and recovering the planted signal from the spiked tensor model then correspond to finding $x$ which minimizes loss $\mathcal{L}$. Note in the latter case that $x$ which minimizes loss will also have a large overlap with $x^*$ hence why one considers this problem as recovering a planted signal.  This work uses the cavity method to derive mean dynamics for iterates produced by accelerated gradient methods. They construct a sequence of equations that track correlation between iterates maintained by Nesterov’s acceleration and Polyak’s Heavy Ball Momentum. These equations are parameterized by different time steps of the algorithm; for example, the equation for $C_x[t, t’]$ corresponds to the overlap of $x$ at time $t$ with that at time $t’$.  The dynamic mean field equations can be used to heuristically determine computational thresholds for statistical signal recovery in the following sense. For an instance of finding a planted signal from a spiked matrix-tensor model given by samples of $T$ and $Y$, consider the instance’s signal-to-noise ratio $\frac{\Delta_p}{\Delta_2}$. For a fixed value of $\Delta_p$ and as one increases $\Delta_2$, it becomes increasingly difficult to recover the planted spike $x^*$. At a certain value of $\Delta_2$, say $\delta^*$, an iterative method solving the instance may exhibit a phase transition: for $\Delta_2 < \delta^*$ the method will find $x$ that has large overlap with $x^*$, while for $\Delta_2 > \delta^*$, the method only converges to $x$ that has vanishingly small (w.r.t. $N$) overlap with $x^*$.   The authors use the mean dynamic equations to compute statistical recovery thresholds for Nesterov acceleration and Heavy Ball momentum. The cavity method calculations are heuristic, and so the authors support their estimates with experiments that suggest the true threshold for these two methods are very close to their estimates. The results also suggest that the threshold for accelerated methods align closely with the threshold achieved by vanilla gradient descent. As a separate results, the authors demonstrate an equivalence between the dynamics of Polyak’s Heavy Ball Momentum and the dynamics of a massive particle under the effect of a certain potential field.