Provable accelerated gradient method for nonconvex low rank optimization

Optimization over low rank matrices has broad applications in machine learning. For large-scale problems, an attractive heuristic is to factorize the low rank matrix to a product of two much smaller matrices. In this paper, we study the nonconvex problem $$\min _{\mathbf {U}\in \mathbb {R}^{n\times r}} g(\mathbf {U})=f(\mathbf {U}\mathbf {U}^T)$$ min U ∈ R n × r g ( U ) = f ( U U T ) under the assumptions that $$f(\mathbf {X})$$ f ( X ) is restricted $$\mu $$ μ -strongly convex and L -smooth on the set $$\{\mathbf {X}:\mathbf {X}\succeq 0,\text{ rank }(\mathbf {X})\le r\}$$ { X : X ⪰ 0 , rank ( X ) ≤ r } . We propose an accelerated gradient method with alternating constraint that operates directly on the $$\mathbf {U}$$ U factors and show that the method has local linear convergence rate with the optimal dependence on the condition number of $$\sqrt{L/\mu }$$ L / μ . Globally, our method converges to the critical point with zero gradient from any initializer. Our method also applies to the problem with the asymmetric factorization of $$\mathbf {X}={\widetilde{\mathbf {U}}}{\widetilde{\mathbf {V}}}^T$$ X = U ~ V ~ T and the same convergence result can be obtained. Extensive experimental results verify the advantage of our method.