ManifoldFlow: Spectrum-Learnable Stiefel Optimization via SPD Relaxation

Abstract

Orthogonal and Stiefel constraints are widely used in neural network optimization to control weight geometry, improve conditioning, and stabilize signal propagation. Yet standard Stiefel optimization imposes a hidden restriction: while learning the directions of a linear map, it fixes every singular value to one, preventing the model from learning how strongly different directions should be amplified or suppressed. We introduce ManifoldFlow, a spectrum-learnable extension of Stiefel optimization that separates directional geometry from spectral scaling through $W=Q{S}^{1/2}$, where $Q^{\top }Q=I$ and $S\succ 0$. Beyond simply making the spectrum trainable, ManifoldFlow uses the gradient component rejected by the Stiefel constraint as a signal for how the spectrum should change, accumulates this signal across aligned tangent spaces, and jointly updates the orthonormal basis and a bounded SPD spectrum. This construction preserves exact spectral interpretability because $W^{\top }W=S$, while recovering standard Stiefel optimization when $S=I$. We establish the geometry of the rejected gradient, characterize optimization over the resulting Stiefel--SPD product space, and derive how the learned spectrum affects conditioning, signal propagation, and spectral complexity. Across diverse tasks, ManifoldFlow consistently outperforms matched fixed-spectrum Stiefel baselines, with particularly strong gains in recurrent language modeling. Our code is available at https://github.com/Hik289/manifold_flow.git.

arXiv comment37 pages