Riemannian Gradient Descent for Low-Rank Architectures

Abstract: We explore Riemannian optimization techniques for rank-factored matrix parameters, targeting contemporary deep learning applications. We examine ten points in the algorithm design space: two geometries for rank-$r$ matrices, three geometries for rank-$r$ partial isometries, and block-matrix variants of these five, where factors are shared across block-rows and block-columns. We apply our methods to the multihead attention parameters in small LLMs. After tuning learning rates, our methods do not conclusively outperform an AdamW baseline. Our implementations are available online.

• The paper presents a manifold-aware optimizer that leverages Riemannian geometry to update low-rank network weights accurately.
• It introduces retraction schemes and gradient projection methods that maintain tangent-space coherence on fixed-rank and partial isometry manifolds.
• Experimental results show performance comparable to AdamW, suggesting future gains for larger-scale or more structured architectures.

Riemannian Gradient Descent for Low-Rank Architectures: An Expert Summary

Motivation and Problem Setting

The paper "Riemannian Gradient Descent for Low-Rank Architectures" (2606.02328) systematically investigates Riemannian optimization methods for neural network weights constrained to low-rank and structured manifolds. The motivation is to strengthen modern deep learning systems, where rank-factored representations like $W=AB^T$ offer substantial parameter efficiency, especially in high-dimensional multihead attention modules. While standard approaches optimize $A$ and $B$ independently in flat (Euclidean) space, this work prioritizes the product $W$ as the true parameter and leverages its geometrical structure.

The central question is whether manifold-aware optimization, with explicit control of the rank and geometric equivalence classes (since $AS, S^{-1}B^T$ yield identical $W$), can yield practical benefits in convergence or solution quality versus adaptive Euclidean optimizers like AdamW. The technical scope includes both strict rank-$r$ matrices and further submanifolds such as partial isometries—where singular values are fixed at $1$, enforcing orthonormal factors.

Riemannian Optimization Framework

The authors adopt an embedded optimization strategy that minimally perturbs the existing architecture. Riemannian gradient descent, parameterized by various geometries, replaces only the optimizer step. Each iteration comprises:

1. Gradient Recovery: Given back-propagated gradients for $A$ and $B$, reconstruct the "true" Euclidean gradient with respect to $A$0 (only unique if the objective depends strictly on $A$1).
2. Riemannian Gradient Computation: Project the Euclidean gradient onto the tangent space of the target manifold (e.g., fixed-rank matrices, partial isometries).
3. Step and Momentum Update: Use exponential smoothing with a user-controlled parameter $A$2 (or optionally heavy-ball momentum). A novel normalization scales step size according to the manifold-induced metric.
4. Retraction: Instead of stepping in Euclidean space, move along the manifold using a retraction (SVD truncation, QR/polar projections) that approximates the exponential map.
5. Parallel Transport: After moving to the new point, transport momentum to the tangent space at the new location.

Algorithmic and implementation details are optimized for $A$3 per step, thus tractable for large matrices with moderate rank.

Geometries and Manifolds

A core contribution is the systematic exploration of the optimizer design space, spanning:

• Fixed-Rank Matrices ($A$4): Both embedded (ambient metric; tangent-space projections, SVD retractions) and quotient (factor space with a metric invariant to linear transformations of the factor space) geometries are described.
• Partial Isometries: Manifolds where $A$5 with $A$6 column-orthonormal (Stiefel), enforcing all nonzero singular values to $A$7. Three geometries are instantiated: embedded, quotient (with respect to $A$8 gauge transformations), and canonical (using metrics natural to the Stiefel manifold).
• Block-Structured Grids: For architectures with weight-sharing (e.g., group-query attention), block matrices of the form $A$9 are treated. Product manifold structures for both fixed-rank and partial isometry constraints are enforced via appropriately designed metrics and retractions.
• Extensions: The method is discussed in the context of tensor generalizations, symmetry constraints (e.g., to projections/Grassmannians), and possible Finsler geometry settings for further generality.

Each algorithm is carefully engineered for computational efficiency, and technical appendices provide explicit procedures for tangent space computations, gradient projections, and efficient retraction schemes.

Numerical Experiments

Experiments are conducted on small-scale, GPT-style decoder transformers with both regular and grouped attention (MHA and GQA), using FineWeb data. Ten different low-rank manifold-aware optimizers are benchmarked—five each for singleton (MHA) and shared-grid (GQA) architectures—against AdamW.

Key observations:

• Tuning: With appropriate learning rate tuning and normalization, all Riemannian methods attain validation and training loss comparable to AdamW, with negligible variance attributable to manifold geometry, momentum strategy, or normalization technique.
• No Superiority: No configuration convincingly outperformed AdamW. In some settings, manifold-aware methods slightly outperformed, but always within the standard deviation across random seeds.
• Spikes and Instabilities: Sharp, recoverable loss spikes are observed for both AdamW and some Riemannian methods, attributed to mini-batch content (e.g., structured web documents).
• Algorithmic Complexity: Riemannian approaches are substantially more expensive per iteration than AdamW ($B$0 vs. $B$1), with no compensatory reduction in iterations needed.
• Implementation Robustness: Explicit mechanisms are required for maintaining tangent-space coherence (e.g., orthogonality of factors), and handling gradient inconsistencies when the loss is not strictly a function of $B$2.

Theoretical and Practical Implications

The results highlight the following points:

• Manifold Optimization Soundness: The presented Riemannian optimizers are functionally correct and produce stable learning trajectories for deep neural architectures. This validates manifold methods as theoretically solid for modern deep learning settings.
• Expressivity vs. Cost: While more expressivity (e.g., handling orthogonal projections, weight sharing) is supported, in practice, the overhead may not justify the cost for standard transformer architectures—at least for small model and data regimes studied.
• Scaling Hypothesis: Potential advantages of geometrically informed optimization may surface at larger model and data scales, or with architectures where parameter efficiency and structure alignment are crucial (e.g., large MoEs, tensorized modules).
• Generalization Potential: The techniques generalize naturally to tensor factorizations and to settings where regularization or architectural constraints (e.g., symmetry, normed descent) are prominent. Applications outside attention—such as low-rank compression or structured pruning—may benefit.
• Future Prospects: The work surfaces several research directions, notably:
  • Riemannian variants of non-Euclidean (Finsler) descent methods to match the success of adaptive optimizers like AdamW in practice.
  • Theorizing and managing stratified or singular geometry (lower-rank subspaces, degeneration events).
  • Alternative retraction and vector transport schemes, especially retraction-free approaches.
  • Better theoretical and engineering integration of momentum and weight decay.
  • Extending to objectives only "nearly" separable in $B$3, or with regularizations coupling factor spaces.

Conclusion

This work offers a comprehensive design and implementation guide for Riemannian gradient descent on fixed-rank and structured matrix manifolds as applied to deep learning. Despite careful optimization and hyperparameter tuning, no substantial performance gains over AdamW were observed in small LLM experiments. The absence of practical improvement does not discount potential future utility; larger-scale, more structured neural systems may eventually expose benefits unique to manifold-aware optimization. The authors' contributions include robust algorithmic treatments, explicit tradeoff analysis, and a roadmap for future developments in structured optimization for deep learning.