Recursive Gradient Retraction Methods for Stochastic Nonconvex Optimization on Stiefel Manifolds

Abstract

Nonconvex optimization on Stiefel manifolds underlies many problems in machine learning and deep learning. However, existing stochastic Riemannian retraction methods often rely on stepsizes that require prior knowledge of the gradient Lipschitz constant, a prescribed iteration horizon, or restrictive batch-size choices. We develop {\sf RGR}, a recursive gradient retraction method for stochastic nonconvex optimization on Stiefel manifolds. Instead of recursively updating gradient estimators in changing tangent spaces, {\sf RGR} constructs them in the Euclidean space and subsequently maps them to tangent space, thereby avoiding vector transport. The method comprises three variants: {\sf RGR-SPIDER} for finite-sum optimization and {\sf RGR-SRM} and {\sf RGR-EMA} for online stochastic optimization. {\sf RGR-SPIDER} is Lipschitz-free and horizon-free, as its adaptive stepsize requires neither prior knowledge of the gradient Lipschitz constant nor a prescribed iteration horizon. {\sf RGR-SRM} and {\sf RGR-EMA} are mini-batch and horizon-free methods whose convergence guarantees do not rely on prescribed large batch sizes and whose parameters can be chosen without prior knowledge of the total iteration budget. We show that {\sf RGR-SPIDER} reaches an $\varepsilon$-critical point using $\mathcal{O}(\varepsilon^{-2})$ iterations and $\mathcal{O}(N+\sqrt{N}\varepsilon^{-2})$ component-gradient evaluations, where $N$ denotes the number of component functions. For the online setting, {\sf RGR-SRM} and {\sf RGR-EMA} achieve iteration complexities of $\tilde{\mathcal{O}}(\varepsilon^{-3})$ and $\tilde{\mathcal{O}}(\varepsilon^{-4})$, respectively. Experiments on principal component analysis, low-rank matrix completion, and orthogonal convolutional neural networks demonstrate faster empirical convergence than methods that construct recursive gradient estimators directly in Riemannian tangent spaces.