Riemannian Gradient Descent for Low-Rank Architectures

2026-06-01

AI Summary

This work investigates efficient optimization of matrix parameters in deep learning—such as multi-head attention weights—under low-rank structural constraints. The authors systematically explore ten Riemannian optimization strategies, encompassing various low-rank matrix manifolds, including the rank‑r partial isometry manifold and its block-shared factor variants, and compare them against an AdamW baseline. Experiments are conducted on small-scale language models with thorough learning rate tuning. While the proposed approaches do not significantly outperform AdamW, this study provides the first systematic evaluation of the suitability of different Riemannian geometric structures for optimizing low-rank neural networks. The accompanying code has been made publicly available.

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 language models. After tuning learning rates, our methods do not conclusively outperform an AdamW baseline. Our implementations are available online.

Problem

• Riemannian optimization
• low-rank matrices
• deep learning
• multihead attention
• rank-factored parameters

Innovation

• Riemannian optimization
• low-rank matrix
• rank factorization
• partial isometry
• multihead attention