Preconditioned Gradient Descent for Over-Parameterized Nonconvex Matrix Factorization
Year: 2022

Abstract:
In practical instances of nonconvex matrix factorization, the rank of the true solution r is often unknown, so the rank r of the model can be overspecified as r > r. This over-parameterized regime of matrix factorization significantly slows down the convergence of local search algorithms, from a linear rate with r = r to a sublinear rate when r > r. We propose an inexpensive preconditioner for the matrix sensing variant of nonconvex matrix factorization that 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. Classical gradient descent in a neighborhood of the solution slows down due to the need for the model matrix factor to become singular. Our key result is that this singularity can be corrected by l2 regularization with a specific range of values for the damping parameter. In fact, a good damping parameter can be inexpensively estimated from the current iterate. The resulting algorithm, which we call preconditioned gradient descent or PrecGD, is stable under noise, and converges linearly to an information theoretically optimal error bound. Our numerical experiments find that PrecGD works equally well in restoring the linear convergence of other variants of nonconvex matrix factorization in the over-parameterized regime.

heading: 1 Introduction
text: 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.

heading: 2 Proposed Algorithm: Preconditioned Gradient Descent
text: 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.

heading: 3 Background and Related Work
text: 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.

heading: 4 Sublinear Convergence of Gradient Descent
text: In order to understand how to improve gradient descent in the over-parameterized regime, we must first understand why existing methods fail. For an algorithm that moves in a search direction D with step-size α, it is a standard technique to measure the corresponding decrement in f with a Taylor-like expansion
f(X − αD) ≤ f(X)− α ⟨∇f(X), D⟩︸ ︷︷ ︸ linear progress +α2 (L/2)∥D∥2F︸ ︷︷ ︸ inverse step-size
(3)
in which L is the usual gradient Lipschitz constant (see e.g. Nocedal and Wright [47, Chapter 3]). A good search direction D is one that maximizes the linear progress ⟨∇f(X), D⟩ while also keeping the inverse step-size (L/2)∥D∥2F sufficiently small in order to allow a reasonably large step to be taken. As we will show in this section, the main issue with gradient descent in the over-parameterized regime is the first term, namely, that the linear progress goes down to zero as the algorithm makes progress towards the solution.
Classical gradient descent uses the search direction D = ∇f(X). Here, a common technique is to bound the linear progress at each iteration by a condition known as gradient dominance (or the Polyak-Łojasiewicz or PL inequality), which is written as
⟨∇f(X), D⟩ = ∥∇f(X)∥2F ≥ µ(f(X)− f⋆) where µ > 0 and f⋆ = min X f(X). (4)
Substituting the inequality (4) into the Taylor-like expansion (3) leads to
f(X − αD) ≤ f(X)− α∥∇f(X)∥2F + α2(L/2)∥∇f(X)∥2F f(X − αD)− f⋆ ≤ [1− µα(1− αL/2)] · (f(X)− f⋆). (5)
Here, we can always pick a small enough step-size α to guarantee linear convergence:
Q = 1− µα+ µα2L/2 < 1 =⇒ f(Xk)− f⋆ ≤ Qk[f(X0)− f⋆]. (6)
In particular, picking the optimal step-size α = 1/L minimizes the convergence quotient Q = 1 − 1/(2κ), where κ = L/µ is the usual condition number. This shows that, with an optimal step-size, gradient descent needs at most O(κ log(1/ϵ)) iterations to find an ϵ-suboptimal X .
Matrix sensing with exact rank r = r⋆ is easily shown to satisfy gradient dominance (4) by manipulating existing results on (restricted) local strong convexity. In the over-parameterized case r > r⋆, however, local strong convexity is lost, and gradient dominance can fail to hold. Indeed, consider the following instance of matrix sensing, with true rank r⋆ = 1, search rank r = 2, and A set to the identity
f(X) = ∥XXT − zzT ∥2F where X = [ 1 0 0 ξ ] and z = [ 1 0 ] . (7)
We can verify that ∥∇f(X)∥2 = 4ξ2[f(X) − f⋆], and this suggests that f satisfies gradient dominance (4) with a constant of µ ≤ 2ξ2. But ξ is itself a variable that goes to zero as the candidate XXT approaches to ground truth zzT . For every fixed µ > 0 in the gradient dominance condition (4), we can find a counterexample X in (7) with ξ < √ µ/2. Therefore, we must conclude that gradient dominance fails to hold, because the inequality in (4) can only hold for µ = 0.
In fact, this same example also shows why classical gradient descent slows down to a sublinear rate. Applying gradient descent Xk+1 = Xk − α∇f(Xk) with fixed step-size α to (7) yields a sequence of iterates of the same form
X0 = [ 1 0 0 ξ0 ] , Xk+1 = [ 1 0 0 ξk+1 ] = [ 1 0 0 ξk − αξ3k ] ,
from which we can verify that f(Xk+1) = (1−αξ2k)4 ·f(Xk). As each k-th XkXTk approaches zzT , the element ξk converges towards zero, and the convergence quotient Q = (1− αξ2k)4 approaches 1. We see a process of diminishing returns: every improvement to f worsens the quotient Q, thereby reducing the progress achievable in the subsequent step. This is precisely the notion that characterizes sublinear convergence.

heading: 5 Linear Convergence for the Noiseless Case
text: To understand how it is possible make gradient descent converge linearly in the over-parameterized regime, we begin by considering gradient method under a change of metric. Let P be a real symmetric, positive definite nr × nr matrix. We define a corresponding P -inner product, P -norm, and dual P -norm on Rn×r as follows
⟨X,Y ⟩P def = vec(X)TPvec(Y ), ∥X∥P
def = √ ⟨X,X⟩P , ∥X∥P∗ def = √ vec(X)TP−1vec(X),
where vec : Rn×r → Rnr is the usual column-stacking operation. Consider descending in the direction D satisfying vec(D) = P−1vec(∇f(X)); the resulting decrement in f can be quantified by a P -norm analog of the Taylor-like expansion (3)
f(X − αD) ≤ f(X)− α⟨∇f(X), D⟩+ α2(LP /2)∥D∥2P (8) = f(X)− α(1− α(LP /2))∥∇f(X)∥2P∗ (9)
where LP is a P -norm gradient Lipschitz constant. If we can demonstrate gradient dominance under the dual P -norm,
∥∇f(X)∥2P∗ ≥ µP (f(X)− f⋆) where µP > 0 and f⋆ = min f(X), (10)
then we have the desired linear convergence
f(X − αD)− f⋆ ≤ [1− µPα(1− αLP /2)] · (f(X)− f⋆) (11) = [1− 1/(2κP )] · (f(X)− f⋆) with α = 1/LP , (12)
in which the condition number κP = LP /µP should be upper-bounded. To make the most progress per iteration, we want to pick a metric P to make the condition number κP as small as possible.
The best choice of P for the fastest convergence is simply the Hessian ∇2f(X) itself, but this simply recovers Newton’s method, which would force us to invert a large nr × nr matrix to compute the search direction D at every iteration. Instead, we look for a preconditioner P that is cheap to apply while still assuring a relatively small condition number κP . The following choice is particularly interesting (the Kronecker product ⊗ is defined to satisfy vec(AXBT ) = (B ⊗A)vec(X))
P = (XTX + ηIr)⊗ In = XTX ⊗ In + ηInr,
because the resulting D = ∇f(X)(XTX + ηI)−1 allow us to interpolate between regular GD and the ScaledGD of Tong et al. [11]. Indeed, we recover regular GD in the limit η → ∞, but as we saw in Section 4, gradient dominance (10) fails to hold, so the condition number κP = LP /µP grows unbounded as µP → 0. Instead, setting η = 0 recovers ScaledGD. The key insight of Tong et al. [11] is that under this choice of P, gradient dominance (10) is guaranteed to hold, with a large value of µP that is independent of the current iterate and the ground truth. But as we will now show, this change of metric can magnify the Lipschitz constant LP by a factor of λ−1min(X
TX), so the condition number κP = LP /µP becomes unbounded in the over-parameterized regime.
Lemma 2 (Lipschitz-like inequality). Let ∥D∥P = ∥D(XTX + ηIr)1/2∥F . Then we have
f(X +D) ≤ f(X) + ⟨∇f(X), D⟩+ 1 2 LP (X,D)∥D∥2P
where
LP (X,D) = 2(1 + δ)
[ 4 +
2∥XXT −M⋆∥F + 4∥D∥P λmin(XTX) + η +
( ∥D∥P
λmin(XTX) + η )2] Lemma 3 (Bounded gradient). For the search direction D = ∇f(X)(XTX + ηI)−1, we have ∥D∥2P = ∥∇f(X)∥2P∗ ≤ 16(1 + δ)f(X).
The proofs of Lemma 2 and Lemma 3 follows from straightforward linear algebra, and can be found in the Appendix. Substituting Lemma 3 into Lemma 2, we see for ScaledGD (with η = 0) that the Lipschitz-like constant is bounded as follows
LP (X,D) ≲ ( ∥XXT −M⋆∥F /λmin(XTX) )2 . (13)
In the exact rank case r = r⋆, the distance of X from singularity can be lower-bounded, within a “good” neighborhood of the ground truth, since λmin(XTX) = λr(XTX) and
∥XXT −M⋆∥F ≤ ρλr(M⋆), ρ < 1 =⇒ λr(XTX) ≥ (1− ρ)λr(M⋆) > 0. (14) Within this “good” neighborhood, substituting (14) into (13) yields a Lipschitz constant LP that depends only on the radius ρ. The resulting iterations converge rapidly, independent of any illconditioning in the model XXT nor in the ground-truth M⋆. In turn, ScaledGD can be initialized within the good neighborhood using spectral initialization (see Proposition 6 below).
In the over-parameterized case r > r⋆, however, the iterate X must become singular in order for XXT to converge to M⋆, and the radius of the “good” neighorhood reduces to zero. The ScaledGD direction guarantees a large linear progress no matter how singular X may be, but the method may not be able to take a substantial step in this direction if X becomes singular too quickly. To illustrate: the algorithm would fail entirely if it lands at on a point where λmin(XTX) = 0 but XXT ̸= M⋆. While regular GD struggles to make the smallest eigenvalues of XXT converge to zero, ScaledGD gets in trouble by making these eigenvalues converge quickly. In finding a good mix between these two methods, an intuitive idea is to use the damping parameter η to control the rate at which X becomes singular. More rigorously, we can pick an η ≈ ∥XXT − ZZT ∥F and use Lemma 2 to keep the Lipschitz constant LP bounded. Substituting Lemma 3 into Lemma 2 and using RIP to upper-bound f(X) ≤ (1 + δ)∥XXT −M⋆∥2F and δ ≤ 1 yields
η ≥ Clb∥XXT − ZZT ∥F =⇒ LP (X,D) ≤ 16 + 136/Clb + 256/C2lb. (15) However, the gradient dominance condition (10) will necessarily fail if η is set too large. Our main result in this paper is that keeping η within the same order of magnitude as the error norm ∥XXT − ZZT ∥F is enough to maintain gradient dominance. The following is the noiseless version of this result.
Theorem 4 (Noiseless gradient dominance). Let minX f(X) = 0 for M⋆ ̸= 0. Suppose that X satisfies f(X) ≤ ρ2 · (1 − δ)λ2r⋆(M⋆) with radius ρ > 0 that satisfies ρ2/(1 − ρ2) ≤ (1 − δ2)/2. Then, we have
η ≤ Cub∥XXT − ZZT ∥F =⇒ ∥∇f(X)∥2P∗ ≥ 2µP f(X) where
µP =
(√ 1 + δ2
2 − δ
)2 ·min  ( Cub√ 2− 1 )−1 , ( 1 + 3Cub √ (r − r⋆) 1− δ2 )−1 . (16) The proof of Theorem 4 is involved and we defer the details to the Appendix. In the noiseless case, we get a good estimate of η for free as a consequence of RIP:
η = √ f(X) =⇒ √ 1− δ∥XXT −M⋆∥F ≤ η ≤ √ 1 + δ∥XXT −M⋆∥F .
Repeating (8)-(12) with Lemma 2, (15) and (16) yields our main result below. Corollary 5 (Linear convergence). Let X satisfy the same initial conditions as in Theorem 4. The search direction D = ∇f(X)(XTX + ηI)−1 with damping parameter η = √ f(X) and step-size α ≤ 1/LP yields f(X − αD) ≤ (1− αµP /2) f(X) where LP is as in (15) with Clb = √ 1− δ and µP is as in (16) with Cub = √ 1 + δ.
For a fixed RIP constant δ, Corollary 5 says that PrecGD converges at a linear rate that is independent of the current iterate X , and also independent of possible ill-conditioning in the ground truth. However, it does require an initial point X0 that satisfies ∥A(X0XT0 −M∗)∥2 < ρ2(1− δ)λr∗(M⋆)2 (17) with a radius ρ > 0 satisfying ρ2/(1− ρ2) ≤ (1− δ2)/2. Such an initial point can be found using spectral initialization, even if the measurements are tainted with noise. Concretely, we choose the initial point X0 as
X0 = Pr
( 1
m m∑ i=1 yiAi
) where Pr(M) = arg min
X∈Rn×r ∥XXT −M∥F , (18)
where we recall that y = A(M⋆) + ϵ are the m possibly noisy measurements collected of the ground truth, and that the rank-r projection operator can be efficiently implemented with a singular value decomposition. The proof of the following proposition can be found in the appendix. Proposition 6 (Spectral Initialization). Suppose that δ ≤ (8κ √ r∗)−1 and m ≳ 1+δ1−δ σ2rn logn ρ2λ2
r⋆ (M⋆)
where κ = λ1(M⋆)/λr⋆(M⋆). Then, with high probability, the initial point X0 produced by (18) satisfies the radius condition (17).
However, if the measurements y are noisy, then √ f(X) = ∥A(XXT − M⋆) + ε∥ now gives a biased estimate of our desired damping parameter η. In the next section, we show that a good choice of ηk is available based on an approximation of the noise variance.

heading: 6 Extension to Noisy Setting
text: In this section, we extend our analysis to the matrix sensing with noisy measurements. Our main goal is to show that, with a proper choice of the damping coefficient η, the proposed algorithm converges linearly to an “optimal” estimation error. Theorem 7 (Noisy measurements with optimal η). Suppose that the noise vector ϵ ∈ Rm has subGaussian entries with zero mean and variance σ2 = 1m ∑m i=1 E[ϵ2i ]. Moreover, suppose that ηk = 1√ m ∥A(XkXTk −M∗)∥, for k = 0, 1, . . . ,K, and that the initial point X0 satisfies ∥A(X0XT0 − M∗)∥2 < ρ2(1−δ)λr∗(M⋆)2. Consider k∗ = argmink ηk, and suppose that the step-size α ≤ 1/L, where L > 0 is a constant that only depends on δ. Then, with high probability, we have
∥Xk∗XTk∗ −M⋆∥2F ≲ max { 1 + δ
1− δ
( 1− αµP
2
)K ∥X0XT0 −M∗∥2F , Estat } , (19)
where Estat := σ 2nr logn
µP (1−δ)m .
Assuming fixed parameters for the problem, the above theorem shows that PrecGD outputs a solution with an estimation error of O (Estat) in O (log (1/Estat)) iterations. Moreover, the error O (Estat) is minimax optimal (modulo logarithmic factors), and cannot be improved significantly. In particular, Candes and Plan [14] showed that any estimator X̂ must satisfy ∥X̂X̂T −M∗∥2F ≳ σ2nr/m with non-negligible probability. The classical methods for achieving this minimax rate suffer from computationally-prohibitive per iteration costs [15, 21, 48]. Regular gradient descent alleviates this issue at the expense of a slower convergence rate of O( √ 1/Estat) [8]. Our proposed PrecGD achieves the best of both worlds: it converges to the minimax optimal error with cheap per-iteration complexity of O(nr2 + r3), while benefiting from an exponentially faster convergence rate than regular gradient descent in the over-parameterized regime.
Theorem 7 highlights the critical role of the damping coefficient η in the guaranteed linear convergence of the algorithm. In the noiseless regime, we showed in the previous section that an “optimal” choice η = √ f(X) is available for free. In the noisy setting, however, the same choice of η becomes biased by the noise variance, and is therefore no longer optimal. As is typically the case for regularized estimation methods [49–51], selecting the ideal parameter would amount to some kind of resampling, such as via cross-validation or bootstrapping [52–54], which is generally expensive to implement and use in practice. As an alternative approach, we show in our next theorem that a good choice of η is available based on an approximation of the noise variance σ2. Theorem 8 (Noisy measurements with variance proxy). Suppose that the noise vector ϵ ∈ Rm has sub-Gaussian entries with zero mean and variance σ2 = 1m ∑m i=1 E[ϵ2i ]. Moreover, suppose that
ηk = √ |f(Xk)− σ̂2| for k = 0, 1, . . . ,K, where σ̂2 is an approximation of σ2, and that the initial point X0 satisfies ∥A(X0XT0 −M∗)∥2F < ρ2(1− δ)λr∗(M⋆)2. Consider k∗ = argmink ηk, and suppose that the step-size α ≤ 1/L, where L > 0 is a constant that only depends on δ. Then, with high probability, we have
∥Xk∗XTk∗ −M∗∥2F ≲ max
{ 1 + δ
1− δ
( 1− αµP
2
)K ∥X0XT0 −M∗∥2F , Estat, Edev, Evar } , (20)
where
Estat := σ2nr log n
µP (1− δ)m , Edev :=
σ2
1− δ
√ log n
m , Evar := |σ2 − σ̂2|. (21)
In the above theorem, Edev captures the deviation of the empirical variance 1m ∑m i=1 ϵ 2 i from its expectation σ2. On the other hand, Evar captures the approximation error of the true variance. According to Theorem 8, it is possible to chose the damping factor ηk merely based on f(Xk) and an approximation of σ2, at the expense of a suboptimal estimation error rate. In particular, suppose that the noise variance is known precisely, i.e., σ̂2 = σ2. Then, the above theorem implies that the estimation error is reduced to
∥Xk∗XTk∗ −M∗∥2F ≲ max {Estat, Edev} after O ( log ( 1
max {Estat, Edev}
)) iterations.
If m is not too large, i.e., m ≲ σ2n2r2 log n, the estimation error can be improved to ∥Xk∗XTk∗ − M∗∥2F ≲ Estat, which is again optimal (modulo logarithmic factors). As m increases, the estimation error will become smaller, but the convergence rate will decrease. This suboptimal rate is due to the heavy tail phenomenon arising from the concentration of the noise variance. In particular, one can write
f(X)− σ2 = 1 m ∥A(XXT −M⋆)∥2 + 1
m ∥ϵ∥2 − σ2︸ ︷︷ ︸ variance deviation
+ 2
m ⟨A(ZZT −XXT ), ϵ⟩︸ ︷︷ ︸
cross-term
(22)
Evidently, f(X)− σ2 is in the order of 1m∥A(XX T −M⋆)∥2 if both variance deviation and crossterm are dominated by 1m∥A(XX T −M⋆)∥2. In the proof of Theorem 8, we show that, with high probability, the variance deviation is upper bounded by (1− δ)Edev and it dominates the cross-term. This implies that the choice of η = √ |f(X)− σ2| behaves similar to 1√
m ∥A(XXT −M⋆)∥, and
hence, the result of Theorem 7 can be invoked, so long as 1
m ∥A(XXT −M⋆)∥2 ≥ (1− δ)∥XXT −M⋆∥2F ≳ (1− δ)Edev.

heading: 7 Numerical Experiments
text: Finally, we numerically compare PrecGD on other matrix factorization problems that fall outside of the matrix sensing framework. We consider the ℓp empirical loss fp(X) = ∑m i=1 |⟨Ai, XXT−M⋆⟩|p for 1 ≤ p < 2, in order to gauge the effectiveness of PrecGD for increasing nonsmooth loss functions. Here, we set the damping parameter ηk = [fp(Xk)]1/p as a heuristic for the error ∥XXT −M⋆∥F . The data matrices A1, . . . , Am were taken from [13, Example 12], the ground truth M⋆ = ZZT was constructed by sampling each column of Z ∈ Rn×r⋆ from the standard Gaussian, and then rescaling the last column to achieve a desired condition number.
The recent work of Tong et al. [55] showed that in the exactly-parameterized setting, ScaledGD works well for the ℓ1 loss function. In particular, if the initial point is close to the ground truth, then with a Polyak stepsize αk = f(Xk)/∥∇f(Xk)∥∗P , ScaledGD converges linearly to the ground truth. However, these theoretical guarantees no longer hold in the over-parameterized regime.
When r > r∗, our numerical experiments show that ScaledGD blows up due to singularity near the ground truth while PrecGD continues to converge linearly in this nonsmooth, over-parameterized setting. In Figure 2 we compare GD, ScaledGD and PrecGD in the exact and over-parameterized regimes for the ℓp norm, with p = 1.1, 1.4 and 1.7. For ScaledGD and PrecGD, we used a modified version of the Polyak step-size where αk = f(Xk)p/∥∇f(Xk)∥∗P . For GD we use a decaying stepsize. When r = r∗, we see that both ScaledGD and PrecGD converge linearly, but GD stagnates due to ill-conditioning of the ground truth. When r > r∗, GD still converges slowly and ScaledGD blows up very quickly, while PrecGD continues to converge reliably.
p = 1.1 p = 1.4 p = 1.7

heading: 8 Conclusions
text: In this paper, we propose a preconditioned gradient descent or PrecGD for nonconvex matrix factorization with a comparable per-iteration cost to classical gradient descent. For over-parameterized matrix sensing, gradient descent slows down to a sublinear convergence rate, but PrecGD restores
the convergence rate back to linear, while also making the iterations immune to ill-conditioning in the ground truth. While the thoeretical analysis in our paper uses some properties specific to RIP matrix sensing, our numerical experiments find that PrecGD works well for even for nonsmooth loss functions. We believe that these current results can be extended to similar problems such as matrix completion and robust PCA, where properties like incoherence can be used to select the damping parameter ηk with the desired properties, so that PrecGD converges linearly as well. It remains future work to provide rigorous justification for these observations.

heading: Acknowledgements
text: G.Z. and R.Y.Z are supported by the NSF CAREER Award ECCS-2047462. S.F. is supported by MICDE Catalyst Grant and MIDAS PODS Grant. We also thank an anonymous reviewer who provided a simplified proof of Lemma 14 and made various insightful comments to help us improve an earlier version of this work.

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

This paper proposed a preconditioned gradient descent (PreGD) method for low rank matrix factorization. In practice, the true rank r* of the matrix is unknown. To be conservative, people usually choose an over-estimated r to start with. However, in such case, the convergence of local search algorithms falls from a linear rate (r=r*) to a sublinear rate (r>r*). The PreGD method proposed in this paper restores the linear convergence rate for the r>r* case. Numerically, it is also found that the PreGD method works well in restoring the linear convergence for some variants of the problem.