Riemannian Gradient Descent for Low-Rank Architectures
2026-06-01 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study different ways to optimize matrices that are broken down into smaller parts, a technique useful in deep learning. They test ten variations involving different mathematical shapes (geometries) of these smaller parts and ways to share factors across blocks. Their focus is on improving the part of language models called multihead attention. After adjusting settings, the authors found their new methods did not clearly beat a common existing optimizer called AdamW. They have shared their work so others can try it too.
Riemannian optimizationrank-factored matrixpartial isometrymultihead attentiondeep learningAdamW optimizermatrix factorizationgeometry in optimizationlanguage models
Authors
Nicholas Knight
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.