Manifold constrained steepest descent

TL;DR

MCSD introduces a single-loop manifold optimization framework that chooses descent directions via a linear minimization oracle on the Riemannian gradient and retracts back to the manifold. The SPEL specialization on the Stiefel manifold uses the spectral-norm LMO and enables scalable updates through fast matrix sign computations, with convergence guarantees under standard smoothness assumptions and a stochastic momentum variant. Empirical results on PCA, orthogonality-constrained CNNs, and LLM adapters show improved stability and competitive or superior performance relative to Riemannian baselines and existing tangent-space LMO methods. The approach provides a practical, memory-efficient alternative for large-scale problems with orthogonality or tangent-space constraints, bridging Euclidean and manifold optimization with broad applicability.

Abstract

Norm-constrained linear minimization oracle (LMO)-based optimizers such as spectral gradient descent and Muon are attractive in large-scale learning, but extending them to manifold-constrained problems is nontrivial and often leads to nested-loop schemes that solve tangent-space subproblems iteratively. We propose \emph{Manifold Constrained Steepest Descent} (MCSD), a single-loop framework for optimization over manifolds that selects a norm-induced steepest-descent direction via an LMO applied to the Riemannian gradient, and then returns to the manifold via projection. Under standard smoothness assumptions, we establish convergence guarantees for MCSD and a stochastic momentum variant. We further introduce \emph{SPEL}, the spectral-norm specialization of MCSD on the Stiefel manifold, which admits scalable implementations via fast matrix sign computations. Experiments on PCA, orthogonality-constrained CNNs, and manifold-constrained LLM adapter tuning demonstrate improved stability and competitive performance relative to standard Riemannian baselines and existing manifold-aware LMO methods.

Figures (7)

• Figure 1: One iteration of Riemannian Gradient Descent, Manifold Muon, and SPEL on the unit sphere.
• Figure 2: Convergence of manifold optimizers on PCA.
• Figure 3: Training Wide ResNet-28 on CIFAR-100. (Left) Test accuracy versus epoch. (Middle) Training loss versus epoch. (Right) Learning-rate sensitivity of SPEL and Muon. Curves are averaged over 3 runs; shaded regions indicate run-to-run variability.
• Figure 4: PCA with RGD under different constant step sizes ($p=5$, $d=1000$). The figure shows the convergence behavior of RGD with varying constant step sizes, where $\alpha=0.001$ provides the best convergence rate, which is used in PCA experiments (Figure \ref{['fig:pca_dis']}).
• Figure 5: Orthogonality constraint violation $\|WW^T - I\|_F$ under the $\mathrm{msign}$ operation using the Polar Express algorithm amsel2025polar with eight iterations. Numerical precision is set to float64. The experiment shows that all algorithms maintain high precision in satisfying the orthogonality constraint at each iteration, validating the feasibility of using the Polar Express approximation as a projection operator onto the Stiefel manifold without relying on SVD.

Theorems & Definitions (17)

• Remark 3.1: On efficient implementation of MCSD
• Proposition 4.2
• Proposition 4.3
• Lemma 4.4
• Theorem 4.5
• Remark 4.6: On the convergence Manifold Muon
• Theorem 4.8
• Theorem 1.1
• Lemma 1.2
• proof