ManifoldFlow: SPD-Relaxed Stiefel Layers with Learnable Singular Spectrum

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

• The paper introduces ManifoldFlow, which relaxes fixed-spectrum orthonormal constraints by learning a bounded SPD singular spectrum to enhance model expressivity.
• It employs a two-step update on a product manifold, combining Riemannian Stiefel optimization with affine-invariant SPD retraction to control the singular spectrum.
• Empirical and theoretical analyses demonstrate that ManifoldFlow outperforms fixed-spectrum layers by improving signal propagation, convergence, and accuracy across various architectures.

ManifoldFlow: Relaxing Fixed-Spectrum Stiefel Layers with Learnable SPD Singular Spectrum

Motivation and Problem Formulation

Fixed-spectrum Stiefel layers have emerged as a reliable approach for exact spectral norm control, using weights of the form $W = Q$ with $Q \in St(p, r)$, guaranteeing every nonzero singular value of $W$ is exactly one. While this orthonormal constraint stabilizes signal propagation, especially in recurrent and deep architectures, it severely restricts expressivity: all represented directions are of unit gain, so no anisotropic attenuation or amplification is possible. The fundamental modeling question is whether, in cases where the Stiefel basis serves as the correct inductive bias, one should relax the fixed spectrum to improve flexibility and optimization without sacrificing orthonormality.

ManifoldFlow addresses this via a minimal relaxation that maintains the orthonormal basis while introducing a learnable, bounded SPD spectrum, parameterizing weights as $W = Q S^{1/2}$ with $Q \in St(p, r)$ and $S \succ 0$. The spectrum $\sigma_i^2(W) = \lambda_i(S)$ is directly controlled via eigenvalue clipping, and $W^\top W = S$ is an exact identity, not an approximation or regularizer. As illustrated in (Figure 1), this allows ManifoldFlow to interpolate between strict Stiefel layers and more expressive, spectrum-adaptive linear maps without losing the mathematical tractability of Riemannian optimization.

Figure 1: ManifoldFlow architecture: fixed Stiefel constrains all singular values to one, whereas ManifoldFlow learns a bounded SPD spectrum.

Methodology and Optimization on the Product Manifold

The layer update is decomposed into orthogonal "basis" and SPD "spectrum" steps. The $Q$ update proceeds as in Riemannian Stiefel optimization, employing tangent projections and Stiefel retractions. The $S$ update is an affine-invariant SPD retraction, enforcing positive-definiteness and spectral box constraints. The rejected normal component from the projected Euclidean gradient—termed the "pressure" matrix—is used as the default update direction for $Q \in St(p, r)$0, justified by first-order geometry: at $Q \in St(p, r)$1, the tangent-projected gradient gives the natural update for $Q \in St(p, r)$2, and the symmetric normal component is the leading-order signal for updating $Q \in St(p, r)$3.

A gate mechanism damps anisotropic or inconsistent updates to $Q \in St(p, r)$4 and, along with spectral clipping, ensures layer stability and boundedness. The ManifoldFlow update is illustrated in detail in (Figure 2):

Figure 2: Detailed schematic of the ManifoldFlow update, showing separation of tangent and normal update branches.

The parametric form $Q \in St(p, r)$5 guarantees, via theory, that $Q \in St(p, r)$6 is not an auxiliary or regularization-only variable: it is the true Gram matrix of $Q \in St(p, r)$7, so spectrum control is exact and transparent.

Theoretical Guarantees

ManifoldFlow preserves all spectral-norm control guarantees of fixed Stiefel weights, while strictly enlarging the model class to all full-column-rank matrices. Each weight matrix has condition number $Q \in St(p, r)$8, with eigenvalues clipped to $Q \in St(p, r)$9. Nonconvex stochastic optimization on the product manifold $W$0 retains standard stationarity guarantees, and generalized product-manifold stability and generalization bounds follow.

Notable properties:

• The pressure update direction for $W$1 is mathematically grounded as the symmetric normal component of the projected Euclidean gradient.
• Any symmetric update direction in $W$2 with equal Frobenius norm induces the same affine-invariant movement, making the precise choice of pressure less determinant near isotropic $W$3.
• The theory establishes that ManifoldFlow is strictly more expressive than fixed Stiefel, and behaves as a spectrum-bounded, exact singular-value controlling variant. Condition numbers and all propagation bounds for signal and gradient are explicit in the learned spectrum.

Empirical Analysis: Convergence, Dynamics, and Boundary

Diagnostics and Initial Motivation

Diagnostics on standard benchmarks (Adult MLP) demonstrate that pressure vectors are non-random and increasingly coherent in deeper layers. Training trajectories show that $W$4 grows over training (eigenvalue adaptation) and that accuracy increases with spectrum adaptation.

Figure 3: Diagnostics on Adult MLP: Layerwise pressure coherence, eigenvalue interval dynamics, and accuracy as a function of spectral bound.

Main Results: Sequence, Tabular, Transformer, and Image Models

Strong paired comparisons are provided between fixed-spectrum (FS) Stiefel and ManifoldFlow (MF) layers, controlling for all other factors.

Recurrent LLMs

In LSTM/WikiText-2 and across multiple recurrent settings (LSTM, GRU, WikiText-2/103), MF universally improves perplexity, and the separation is most pronounced under SGD (e.g., FS: 397.05, MF: 212.54, $W$5 +184.51 PPL).

Figure 4: Main convergence: MF consistently outperforms FS in both sequence and tabular settings.

Tabular and MLPs

In MLPs on Adult, Covertype, and image-derived data (Fashion-MNIST, CIFAR-10), MF achieves consistent accuracy gains, up to +11.4% on Covertype MLP under SGD. In some image-MLP configurations, the improvement is optimizer-dependent, with MF favored by adaptive methods.

Mini-Transformer and Convolutional Heads

Wrapping only the feedforward layers of a Mini-Transformer shows smaller but consistent gains. In convolutional classifier-heads, the sign becomes boundary-dependent: ReLU-ResNet classifier heads sometimes favor FS or are optimizer dependent, clarifying that ManifoldFlow is not a universal replacement but a focused relaxation for orthonormal-prior settings.

Figure 5: Validation curves for Mini-Transformer FFN blocks: MF consistently achieves lower perplexity.

Spectral Adaptation and Trajectories

Eigenvalue traces during training directly demonstrate that $W$6 is not dormant, but actively moves away from the identity and in some cases saturates near the spectral boundary. Layerwise traces in deeper architectures show architecture-dependent spectrum adaptation.

Figure 6: MF-Adam induces persistent movement in the singular spectrum; spectrum learning is active and layer-specific.

Ablations

Ablation studies replacing the pressure update with random symmetric directions, or removing the EMA/gate, rarely degrade performance by more than 0.1%, supporting the claim that the driving mechanism is spectrum relaxation and not any particular update heuristic.

Figure 7: Pressure persists over time, showing layerwise structure rather than acting as random noise.

Summary and Scope of Claims

A summary heatmap across all paired cells confirms that the most robust gains are in recurrent projections and MLP hidden layers, with boundary and negative cases only in image classifier heads under specific optimizer-task pairings.

Figure 8: Heatmap of paired gains (perplexity reduction, accuracy improvement): MF dominates in recurrent/MLP settings, mixed in boundary cases.

Practical and Theoretical Implications

ManifoldFlow creates a new design space between fixed-spectrum, strictly orthogonal layers and unconstrained dense layers, allowing explicit control over both basis and spectrum. In regimes where orthonormal basis is a useful inductive bias but the unit spectrum constraint is unnecessarily restrictive, MF provides a practical mechanism to improve convergence, signal propagation, and final performance—especially under optimizers without external diagonal scaling (e.g., SGD). Theoretical analysis shows that all spectral generalization bounds, Rademacher complexity, and capacity control measures become explicit functions of the learned and bounded $W$7.

In large-scale or structured settings ($W$8), the need for efficient block-diagonal or low-rank variants is highlighted. Moreover, the mechanism directly interacts with adaptive optimizers, shifting the "preconditioning" role from the optimizer's state to the intrinsic layer structure.

Future directions include scaling the SPD update to very large dimensions, joint spectrum-basis coupling, and extension to GNNs and message-passing layers. The precise mathematical interpretation of product-manifold optimization on $W$9 also opens exploration of geometric priors in spectral regularization and non-Euclidean overparameterization.

Conclusion

ManifoldFlow is a mathematically principled, spectrum-learnable relaxation of fixed-spectrum Stiefel layers. It exposes the fixed-spectrum bottleneck, allows for architecture-dependent spectrum adaptation, and provides generalization and optimization guarantees equivalent or superior to existing orthonormal parametrizations. Empirically, MF yields the strongest improvements in recurrent projections and MLPs, clarifying that orthonormality need not entail spectral rigidity. ManifoldFlow is thus of direct practical interest in architectures where spectral control is desirable, but spectrum flexibility is needed for expressivity and learnability.

Reference: "ManifoldFlow: SPD-Relaxed Stiefel Layers with Learnable Singular Spectrum" (2607.04535)