The Power of Preconditioning in Overparameterized Low-Rank Matrix Sensing

We propose ScaledGD(λ), a preconditioned gradient descent method to tackle the low-rank matrix sensing problem when the true rank is unknown, and when the matrix is possibly ill-conditioned. Using overparameterized factor representations, ScaledGD(λ) starts from a small random initialization, and proceeds by gradient descent with a specific form of damped preconditioning to combat bad curvatures induced by overparameterization and ill-conditioning. ScaledGD(λ) is remarkably robust to ill-conditioning compared to vanilla gradient descent (GD) even with overparameterization. Specifically, we show that, under the restricted isometry property (RIP) of the sensing operator, ScaledGD(λ) converges to the true low-rank matrix at a constant linear rate after a small number of iterations that scales only logarithmically with respect to the condition number and the problem dimension. This significantly improves over the convergence rate of vanilla GD which suffers from a polynomial dependency on the condition number. Furthermore, we show that in the presence of measurement noise, ScaledGD(λ) converges to the minimax optimal error up to a multiplicative factor of the condition number at the same rate as in the noiseless setting, which is the first nearly minimax-optimal overparameterized gradient method for low-rank matrix sensing scaling with the true rank rather than the (possibly much larger) overparameterized rank. Our results also extend to the setting when the matrix is only approximately low-rank under the Gaussian design. Our work provides evidence on the power of preconditioning in accelerating the convergence without hurting generalization in overparameterized learning.