The power of preconditioning in overparameterized low-rank matrix sensing

Abstract We propose $\mathsf{ScaledGD(\lambda )}$, a preconditioned gradient descent ($\mathsf{GD}$) 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, $\mathsf{ScaledGD(\lambda )}$ starts from a small random initialization, and proceeds by GD with a specific form of damped preconditioning to combat bad curvatures induced by overparameterization and ill-conditioning. $\mathsf{ScaledGD(\lambda )}$ is remarkably robust to ill-conditioning compared to vanilla $\mathsf{GD}$ even with overparameterization. Specifically, we show that, under the restricted isometry property of the sensing operator, $\mathsf{ScaledGD(\lambda )}$ 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 $\mathsf{GD}$ which suffers from a polynomial dependency on the condition number. Furthermore, we show that in the presence of measurement noise, $\mathsf{ScaledGD(\lambda )}$ 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.