Preconditioned Gradient Descent for Over-Parameterized Nonconvex Matrix Factorization

1 Introduction . Numerous problems in machine learning can be reduced to the matrix factorization problem of recovering a low-rank positive semidefinite matrix M⋆ ⪰ 0 , given a small number of potentially noisy observations [ 1–7 ] . In every case , the most common approach is to formulate an n×n candidate matrix M = XXT in factored form , and to minimize a nonconvex empirical loss f ( X ) over its n× r low-rank factor X . But in most real applications of nonconvex matrix factorization , the rank of the ground truth r⋆ = rank ( M⋆ ) is unknown . It is reasonable to choose the rank r of the model XXT conservatively , setting it to be potentially larger than r⋆ , given that the ground truth can be exactly recovered so long as r ≥ r⋆ . In practice , this will often lead to an over-parameterized regime , in which r > r⋆ , and we have specified more degrees of freedom in our model XXT than exists in the underlying ground truth M⋆ . Zhuo et al . [ 8 ] recently pointed out that nonconvex matrix factorization becomes substantially less efficient in the over-parameterized regime . For the prototypical instance of matrix factorization known as matrix sensing ( see Section 3 below for details ) it is well-known that , if r = r⋆ , then ( classic ) gradient descent or GD Xk+1 = Xk − α∇f ( Xk ) ( GD ) 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . converges at a linear rate , to an ϵ-accurate iterate in O ( κ log ( 1/ϵ ) ) iterations , where κ = λ1 ( M ⋆ ) /λr∗ ( M ⋆ ) is the condition number of the ground truth [ 9 , 10 ] . But in the case that r > r⋆ , Zhuo et al . [ 8 ] proved that gradient descent slows down to a sublinear convergence rate , now requiring poly ( 1/ϵ ) iterations to yield a comparable ϵ-accurate solution . This is a dramatic , exponential slow-down : whereas 10 digits of accuracy can be expected in a just few hundred iterations when r = r⋆ , tens of thousands of iterations might produce just 1-2 accurate digits once r > r⋆ . The slow-down occurs even if r is just off by one , as in r = r⋆ + 1 . It is helpful to understand this pheonomenon by viewing over-parameterization as a special , extreme case of ill-conditioning , where the condition number of the ground truth , κ , is taken to infinity . In this limit , the classic linear rate O ( κ log ( 1/ϵ ) ) breaks down , and in reality , the convergence rate deterioriates to sublinear . In this paper , we present an inexpensive preconditioner for gradient descent . The resulting algorithm , which we call PrecGD , corrects for both ill-conditioning and over-parameterization at the same time , without viewing them as distinct concepts . We prove , for the matrix sensing variant of nonconvex matrix factorization , that the preconditioner restores the convergence rate of gradient descent back to linear , even in the over-parameterized case , while also making it agnostic to possible ill-conditioning in the ground truth . Moreover , PrecGD maintains a similar per-iteration cost to regular gradient descent , is stable under noise , and converges linearly to an information theoretically optimal error bound . We also perform numerical experiments on other variants of nonconvex matrix factorization , with different choices of the empirical loss function f . In particular , we consider different ℓp norms with 1 ≤ p < 2 , in order to gauge the effectiveness of PrecGD for increasingly nonsmooth loss functions . Our numerical experiments find that , if regular gradient descent is capable of converging quickly when the rank is known r = r⋆ , then PrecGD restores this rapid converging behavior when r > r⋆ . PrecGD is able to overcome ill-conditioning in the ground truth , and converge reliably without exhibiting sporadic behavior . 2 Proposed Algorithm : Preconditioned Gradient Descent . Our preconditioner is inspired by a recent work of Tong et al . [ 11 ] on matrix sensing with an illconditioned ground truth M⋆ . Over-parameterization can be viewed as the limit of this regime , in which λr ( M⋆ ) , the r-th largest eigenvalue of M⋆ , is allowed to approach all the way to zero . For finite but potentially very small values of λr ( M⋆ ) > 0 , Tong et al . [ 11 ] suggests the following iterations , which they named scaled gradient descent or ScaledGD : Xk+1 = Xk − α∇f ( Xk ) ( XTk Xk ) −1 . ( ScaledGD ) They prove that the scaling allows the iteration to make a large , constant amount of progress at every iteration , independent of the value of λr ( M⋆ ) > 0 . However , applying this same scheme to the over-parameterized case with λr ( M⋆ ) = 0 results in an inconsistent , sporadic behavior . The issues encountered by both regular GD and ScaledGD with over-parameterization r > r⋆ can be explained by the fact that our iterate Xk must necessarily become singular as our rank-r model XkXTk converges towards the rank-r ⋆ ground truth M⋆ . For GD , this singularity causes the per-iteration progress itself to decay , so that more and more iterations are required for each fixed amount of progress . ScaledGD corrects for this decay in per-iteration progress by suitably rescaling the search direction . However , the rescaling itself requires inverting a near-singular matrix , which causes algorithm to take on sporadic values . A classical remedy to issues posed by singular matrices is ℓ2 regularization , in which the singular matrix is made “ less singular ” by adding a small identity perturbation . Applying this idea to ScaledGD yields the following iterations Xk+1 = Xk − α∇f ( Xk ) ( XTk Xk + ηkIr ) −1 , ( PrecGD ) where ηk ≥ 0 is the damping parameter specific to the k-th iteration . There are several interpretations to this scheme , but the most helpful is to view η as a parameter that allows us to interpolate between ScaledGD ( with η = 0 ) and regular GD ( in the limit η → ∞ ) . In this paper , we prove for matrix sensing that , if the k-th damping parameter ηk is chosen within a constant factor of the error Clb∥XkXTk −M⋆∥F ≤ ηk ≤ Cub∥XkXTk −M⋆∥F , where Clb , Cub > 0 are abs . const . ( 1 ) then the resulting iterations are guaranteed to converge linearly , at a rate that is independent of both over-parameterization and ill-conditioning in the ground truth M⋆ . With noisy measurements , setting ηk to satisfy ( 1 ) will allow the iterations to converge to an error bound that is well-known to be minimax optimal up to logarithmic factors [ 14 ] . We refer to the resulting iterations ( with a properly chosen ηk ) as preconditioned gradient descent , or PrecGD for short . For matrix sensing with noiseless measurements , an optimal ηk that satisfies the condition ( 1 ) is obtained for free by setting ηk = √ f ( Xk ) . In the case of noisy measurements , we show that a good choice of ηk is available based on an approximation of the noise variance . 3 Background and Related Work . Notations . We use ∥ · ∥F to denote the Frobenius norm of a matrix and ⟨· , ·⟩ is the corresponding inner product . We use ≳ to denote an inequality that hides a constant factor . The big-O notation Õ hides logarithimic factors . The gradient of the objective is denoted by ∇f ( X ) ∈ Rn×r . The eigenvalues are assumed to be in decreasing order : λ1 ≥ λ2 ≥ · · · ≥ λr . The symmetric , linear variant of matrix factorization known as matrix sensing aims to recover a positive semidefinite , rank-r⋆ ground truth matrix M⋆ , from a small number m of possibly noisy measurements y = A ( M⋆ ) + ϵ , where A ( M⋆ ) = [ ⟨A1 , M⋆⟩ , ⟨A2 , M⋆⟩ , . . . , ⟨Am , M⋆⟩ ] T , in which A is a linear measurement operator , and the length-m vector ϵ models the unknown measurement noise . A distinguishing feature of matrix sensing is that A is assumed to satisfy the restricted isometry property [ 14 , 15 ] . Throughout this paper , we will always assume that A satisfies RIP with parameters ( 2r , δ ) . Definition 1 ( RIP ) . The linear operator A satisfies RIP with parameters ( 2r , δ ) if there exists constants 0 ≤ δ < 1 and m > 0 such that , for every rank-2r matrix M , we have ( 1− δ ) ∥M∥2F ≤ 1 m ∥A ( M ) ∥2 ≤ ( 1 + δ ) ∥M∥2F . A common approach for matrix sensing is to use a simple algorithm like gradient descent to minimize the nonconvex loss function : f ( X ) = 1 m ∥∥y −A ( XXT ) ∥∥ = 1 m ∥∥A ( M⋆ −XXT ) + ϵ∥∥2 . ( 2 ) Recent work has provided a theoretical explanation for the empirical success of this nonconvex approach . Two lines of work have emerged . Local Guarantees . One line of work studies gradient descent initialized inside a neighborhood of the ground truth where X0XT0 ≈ M⋆ already holds [ 10 , 16–19 ] . Such an initial point can be found using spectral initialization , see also [ 18 , 20–23 ] . With exact rank r = r⋆ , previous authors showed that gradient descent converges at a linear rate [ 9 , 10 ] . In the over-parameterized regime , however , local restricted convexity no longer holds , so the linear convergence rate is lost . Zhuo et al . [ 8 ] showed that while spectral initialization continues to work under over-parameterization , gradient descent now slows down to a sublinear rate , but it still converges to a statistical error bound of Õ ( σ2nr⋆/m ) , where σ denotes the noise variance . This is known to be minimax optimal up to logarithmic factors [ 14 ] . In this paper , we prove that PrecGD with a damping parameter ηk satisfying ( 1 ) also converges to an Õ ( σ2nr⋆/m ) statistical error bound . Global Guarantees . A separate line of work [ 13 , 24–31 ] established global properties of the landscapes of the nonconvex objective f in ( 2 ) and its variants and showed that local search methods can converge globally . With exact rank r = r⋆ , Bhojanapalli et al . [ 24 ] proved that f has no spurious local minima , and that all saddles points have a strictly negative descent direction ( strict saddle property [ 32 ] , see also [ 28 , 33 ] ) . In the over-parameterized regime , however , we are no longer guaranteed to recover the ground truth in polynomial time . Other related work . Here we mention some other techniques can be use to solve matrix sensing in the over-parameterized regime . Classically , matrix factorization was solved via its convex SDP relaxation [ 14 , 15 , 34–36 ] . The resulting O ( n3 ) to O ( n6 ) time complexity [ 37 ] limits this technique to smaller problems , but these guarantees hold without prior knowledge on the true rank r⋆ . Firstorder methods , such as ADMM [ 38–40 ] and soft-thresholding [ 41 ] , can be used to solve these convex problems with a per-iteration complexity comparable to nonconvex gradient descent , but they likewise suffer from a sublinear convergence rate . Local recovery via spectral initialization was originally proposed for alternating minimization and other projection techniques [ 21 , 23 , 34 , 42–45 ] . These also continue to work , though a drawback here is a higher per-iteration cost when compared to simple gradient methods . Finally , we mention a recent result of Li et al . [ 46 ] , which showed in the over-parameterized regime that gradient descent with early termination enjoys an algorithmic regularization effect .

Summary: The authors consider low-rank matrix recovery problems, where the $n\times n$ matrix variable $M$ is factorized as $M=XX^T$, for some $X$ of size $n\times r$, with $r\ll n$. It is known from previous work that, although they are typically non-convex, these problems can oftentimes be solved through simple gradient descent over $X$, applied to a suitable loss function.  However, gradient descent tends to converge very slowly when either the ground true matrix is ill-conditioned, or the problem is overparameterized (that is, $r$ is strictly larger than the rank of the ground true matrix). In Reference [11], Tong, Ma and Chi have proposed a gradient preconditioner which solves the slow convergence issue when it is due to ill-conditioning. However, the method does not work when the problem is overparameterized.  In this article, the authors propose a refined preconditioner, which accommodates both ill-conditioning and overparameterization. In the particular setting of RIP matrix sensing with least-squares cost, they show that, with this preconditioner, gradient descent converges linearly in a neighborhood of the solution, up to the unavoidable statistical error.