ManifoldFlow: SPD-Relaxed Stiefel Layers with Learnable Singular Spectrum

2026-07-05

AI Summary

This work proposes ManifoldFlow, a novel approach that overcomes the limitations of existing Stiefel-constrained layers, which enforce all singular values to be exactly one and thus cannot capture direction-dependent scaling. ManifoldFlow decomposes the weight matrix as \( Q S^{1/2} \), where \( Q \) lies on the Stiefel manifold to preserve orthogonality of basis vectors, and \( S \) is a learnable positive-definite matrix that models a bounded singular spectrum. This formulation is the first to introduce learnable singular values while retaining the Stiefel structure, thereby breaking free from the restrictive fixed-spectrum assumption of prior methods. By integrating manifold optimization, positive-definite parameterization, and eigenvalue clipping, ManifoldFlow enables direct and controllable learning of the singular spectrum. Experiments demonstrate consistent improvements over fixed-spectrum Stiefel layers across sequence modeling, tabular data, and image tasks, with particularly notable gains in recurrent language model projections.

Abstract

Orthogonal and Stiefel layers give neural weights exact spectral control, but they also impose a strong modeling constraint: all represented singular values are fixed at one. Many settings that benefit from an orthonormal basis still need direction-dependent attenuation or amplification. We introduce ManifoldFlow, a minimal relaxation of a fixed-spectrum Stiefel layer that keeps the basis on the Stiefel manifold while learning a bounded positive spectrum through W = Q S^{1/2}, with Q^T Q = I and S positive definite. Since W^T W = S, the eigenvalues of S are exactly the squared singular values of the realized weight, making eigenvalue clipping a direct singular-value control mechanism. Across paired sequence, tabular, and image experiments, the learnable SPD spectrum improves the fixed-spectrum Stiefel counterpart in the reported settings where the Stiefel prior is useful, with the largest gains in recurrent language-model projections. Boundary cases in convolutional classifier heads clarify the intended scope: ManifoldFlow is not a universal dense-layer replacement, but a spectrum-learnable Stiefel relaxation for settings where an orthonormal basis is a useful prior. When the basis should be orthonormal, its spectrum need not be frozen. Code available at https://github.com/Hik289/manifold_flow

Problem

• Stiefel manifold
• singular values
• orthogonal layers
• spectral control
• learnable spectrum

Innovation

• ManifoldFlow
• Stiefel manifold
• learnable singular spectrum
• orthogonal layers
• SPD relaxation