Procrustes Flow (Factored Gradient Descent with $U^\top U = V^\top V$ Balancing Regularization) for Low-Rank Matrix Recovery

A two-stage non-convex Burer-Monteiro factored optimization algorithm and geometric convergence analysis (introduced by Stephen Tu, Ross Boczar, Max Simchowitz, Mahdi Soltanolkotabi, and Benjamin Recht in 2016) for recovering a rank-$r$ matrix $X^*\in\mathbb{R}^{n_1\times n_2}$ from linear measurements $b = \mathcal{A}(X^*)$. For a symmetric positive semidefinite target $X^* = X_0 X_0^\top$ ($X_0 = \Phi_0\Sigma_0^{1/2}$), *Procrustes Flow* minimizes $f(U) = \frac{1}{4}\|\mathcal{A}(U U^\top) - b\|_2^2$ via gradient descent $U_{\tau+1} = U_\tau - \mu_\tau \mathcal{A}^*(\mathcal{A}(U_\tau U_\tau^\top) - b)U_\tau$ (Equation 2.5) measured in the orthogonal-Procrustes distance $d(U,X_0) = \min_{R\in O(r)}\|U - X_0 R\|_F$ (Equation 2.4). For a general rectangular rank-$r$ matrix $X^* = \Phi_0\Sigma_0\Psi_0^\top$ factored as $X = U V^\top$ with lifted stack $W = \begin{bmatrix} U \\ V \end{bmatrix}\in\mathbb{R}^{(n_1+n_2)\times r}$, asymmetric rescaling $U \mapsto U T, V \mapsto V T^{-\top}$ degrades local curvature; rectangular Procrustes Flow eliminates this degeneracy by adding the factor-balancing regularizer $g(U,V) = \frac{1}{16}\|U^\top U - V^\top V\|_F^2$ to form $F(W) = \frac{1}{2}\|\mathcal{A}(U V^\top) - b\|_2^2 + \frac{1}{16}\|U^\top U - V^\top V\|_F^2$ (Equation 5.2), whose balanced gradient step $U_{\tau+1} = U_\tau - \mu_\tau[\nabla_U f + \frac{1}{4}U_\tau(U_\tau^\top U_\tau - V_\tau^\top V_\tau)]$ and $V_{\tau+1} = V_\tau - \mu_\tau[\nabla_V f + \frac{1}{4}V_\tau(V_\tau^\top V_\tau - U_\tau^\top U_\tau)]$ (Equation 5.3) satisfies the local regularity condition $\langle\nabla F(W), W - W_0 R\rangle \ge \frac{\sigma_r(X^*)}{5}\|W - W_0 R\|_F^2 + \frac{1}{100\sigma_1(X^*)}\|\nabla F(W)\|_F^2$ inside $d(W,W_0) \le \frac{1}{4}\sigma_r(X^*)^{1/2}$ under rank-$6r$ Restricted Isometry $\delta_{6r} \le 1/25$ (Lemma 5.10, Theorem 5.2), guaranteeing linear contraction $d(W_\tau,W_0) \le (1 - \frac{\mu\sigma_r(X^*)}{5})^\tau d(W_0^{\mathrm{init}},W_0)$ from $m \ge C r(n_1+n_2)$ projected-gradient initialization.