Computing Matrix Squareroot via Non Convex Local Search

We consider the problem of computing the squareroot of a positive semidefinite (PSD) matrix. Several fast algorithms (some based on eigenvalue decomposition and some based on Taylor expansion) are known to solve this problem. In this paper, we propose another way to solve this problem: a natural algorithm performing gradient descent on a non-convex formulation of the matrix squareroot problem. We show that on an $n\times n$ input PSD matrix ${M}$, if the initial point is well conditioned, then the algorithm finds an $\epsilon$-accurate solution in $O\left(\kappa^{3/2} \log \frac{\left\|{M}\right\|_F}{\epsilon}\right)$ iterations, where $\kappa$ is the condition number of $M$. Each iteration involves three matrix multiplications (and does not use either matrix inversions or solutions of linear system), giving a total run time of $O\left(n^{\omega}\kappa^{3/2}\log\frac{\left\|{M}\right\|_F}{\epsilon}\right)$, where $\omega$ is the matrix multiplication exponent. Furthermore we show that our algorithm is robust to errors in each iteration. We also show a lower bound of $\Omega(\kappa)$ iterations for our algorithm demonstrating that the dependence of our result on $\kappa$ is necessary. Existing analyses of similar algorithms (e.g., Newton's method) require commutativity of the input matrix with each iterate of the algorithm which is ensured by choosing the starting iterate carefully. Our analysis, on the other hand, is much more general and does not require each iterate to commute with the input matrix. Consequently, our result guarantees convergence from a wide range of starting points. More generally, our result demonstrates that non-convex optimization can be a viable approach to obtaining fast and robust algorithms. Our argument is quite general and we believe it will find application in designing such algorithms for other problems in numerical linear algebra.