Efficient Optimization with Orthogonality Constraint: a Randomized Riemannian Submanifold Method

TL;DR

This work tackles optimization on the Stiefel manifold with orthogonality constraints, where traditional retractions incur high computational cost. It introduces Riemannian Random Submanifold Descent (RSDM), which updates on low-dimensional random submanifolds by parameterizing updates via a random orthogonal factor, reducing per-iteration complexity from $O(np^2)$ to $O(nr^2)$ (or $O(npr)$) and, in the subcase $r=2$, recovers Riemannian coordinate descent. The authors provide convergence guarantees for general nonconvex and Polyak–Łojasiewicz (PL) functions, as well as stochastic/finite-sum variants, and extend the framework to quotient manifolds such as Grassmann and flag manifolds. Empirically, RSDM demonstrates fast convergence across Procrustes, PCA, QUAD, and orthogonal neural network tasks, often surpassing standard RGD and related methods, especially on large-scale problems. The work highlights a tunable efficiency–convergence trade-off between orthogonal and permutation sampling and suggests avenues for integrating RSDM with line-search, momentum, and preconditioning techniques to further boost performance in scalable, orthogonality-constrained optimization.

Abstract

Optimization with orthogonality constraints frequently arises in various fields such as machine learning. Riemannian optimization offers a powerful framework for solving these problems by equipping the constraint set with a Riemannian manifold structure and performing optimization intrinsically on the manifold. This approach typically involves computing a search direction in the tangent space and updating variables via a retraction operation. However, as the size of the variables increases, the computational cost of the retraction can become prohibitively high, limiting the applicability of Riemannian optimization to large-scale problems. To address this challenge and enhance scalability, we propose a novel approach that restricts each update on a random submanifold, thereby significantly reducing the per-iteration complexity. We introduce two sampling strategies for selecting the random submanifolds and theoretically analyze the convergence of the proposed methods. We provide convergence results for general nonconvex functions and functions that satisfy Riemannian Polyak-Lojasiewicz condition as well as for stochastic optimization settings. Additionally, we demonstrate how our approach can be generalized to quotient manifolds derived from the orthogonal manifold. Extensive experiments verify the benefits of the proposed method, across a wide variety of problems.

Figures (33)

• Figure 1: Proposed random submanifold method on $2$-sphere. Each iteration restricts the update to a $1$-dimensional randomly selected submanifold, i.e., a circle.
• Figure 2: Experiments on Procrustes problem and PCA problem under various settings. The numbers in brackets represent the size of $n, p$. For the Procrustes problem, we see RSDM converges competitively against the best baselines due to the simplicity of the problem. For PCA problem, we see RSDM converges the fastest.
• Figure 3: Experiment results on PCA ($n=2000, p = 1500$) by (a) varying low-dimension $r$ and (b) random seed with $r = 700$. The results suggest the outperformance of proposed RSDM over RGD is robust to changes in $r$ as well as random seed.
• Figure 4: Experiment results on quadratic assignment problem (QUAD), $n = p = 1000$.
• Figure 5: Test accuracy for training orthogonal neural network (O-FNN) and orthogonal vision transformer (O-ViT) on MNIST dataset and CIFAR10 dataset in five runs.

Theorems & Definitions (78)

• Remark 4.1: Riemannian coordinate descent is a special case
• Lemma 5.2
• Lemma 5.3
• Definition 5.4: Riemannian Polyak-Ł ojasiewicz
• Proposition 5.5
• Remark 5.6: Proof techniques of Proposition \ref{['prop:1']}
• Theorem 5.7
• Remark 5.8: Total complexity of RSDM to RGD
• Theorem 5.9
• Remark 5.10: The trade-off between efficiency and convergence