🤖 machine learning

Riemannian Gradient Descent for Low-Rank Architectures

This paper investigates Riemannian gradient descent across ten algorithmic designs for rank-factored matrix parameters in deep learning, but finds that, despite tuning, these methods do not conclusively outperform an AdamW baseline when applied to multihead attention in small language models.

Original authors: Nicholas Knight

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Nicholas Knight

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Navigating a Mountain vs. Driving a Car

Imagine you are trying to find the lowest point in a vast, foggy valley (this represents training a machine learning model to make fewer mistakes). Usually, computer scientists use a standard method called AdamW, which is like driving a car with a very smart GPS. The GPS tells you which way is "down," and you take a step. If you hit a bump, the car's suspension (momentum) helps you keep moving. This works great, but it treats every part of the car's engine as a separate, independent part.

This paper asks a different question: What if we treat the engine parts as a single, connected system?

In modern AI, many parts of the brain (the model) are built using "low-rank" matrices. Think of these not as a giant, solid block of data, but as a folded piece of paper. You can describe the whole shape of the paper by just knowing the coordinates of the fold lines (two smaller matrices, $A$ and $B$ ).

The author argues that standard methods treat the fold lines as if they are independent, which is like trying to smooth out a crumpled piece of paper by pulling on the corners without realizing the paper is connected. Instead, the author proposes using Riemannian Gradient Descent.

The Analogy:

Standard Method (Euclidean): Imagine walking on a flat floor. You take a step in the direction the GPS says. If you are on a crumpled paper, you might accidentally step off the paper entirely, forcing you to jump back on.
Riemannian Method: Imagine you are a surfer riding a wave. You are constrained to stay on the surface of the wave. You don't just walk; you glide along the curve of the wave. This method ensures you never leave the "shape" of the low-rank matrix, keeping the math clean and theoretically sound.

The Experiment: Testing New Surfboards

The author built ten different "surfboards" (algorithms) based on this wave-riding idea. He tested them on a small language model (a brain that learns to predict the next word in a sentence).

He tested two main types of waves:

Fixed-Rank Waves: Just keeping the paper folded at a specific tightness.
Partial Isometry Waves: A stricter rule where the paper must be perfectly orthogonal (like a perfect right angle), which is like a more rigid, structured wave.

He also tested a "Grid" version, where multiple parts of the model share the same fold lines (like a group of people holding the same rope).

The Results: A "Good" but Not "Great" Outcome

Here is the honest conclusion from the paper, stripped of hype:

It Works, But It's Not a Magic Bullet: The new methods successfully trained the model. They didn't crash, and they learned the task. This proves the math is sound and the code works.
No Clear Victory: After carefully tuning the "speed" (learning rate) of the new surfboards, they did not consistently beat the standard AdamW car.
- In some cases, they were slightly better.
- In other cases, they were slightly worse.
- Overall, the difference was so small it could be blamed on random noise (like a slight breeze changing the path of a leaf).
The Cost: The new methods are more computationally expensive. It's like driving a high-tech electric car that requires a special charging station, whereas the standard method is a reliable gas car that works everywhere. Since the new car didn't get you to the destination significantly faster or better, the author concludes it's not yet ready to replace the standard method for everyday use.

Key Takeaways for the General Audience

The Theory is Beautiful: The idea of treating matrix parameters as a curved surface (a manifold) rather than a flat grid is mathematically elegant and logically consistent. It respects the hidden relationships between the numbers.
The Practice is Tricky: While the theory says "you should get there faster," the reality of deep learning (with its noise, massive data, and complex architectures) means the standard, simpler methods (AdamW) are still the champions for now.
Future Potential: The author remains optimistic. Just because a new engine doesn't win a race on a small track doesn't mean it won't win on a massive highway. The author suggests that as models get bigger (scaling up), these geometric methods might finally show their true power.

What the Paper Does Not Claim

It does not claim this method will cure diseases or solve climate change.
It does not claim this is the future of AI right now.
It does not claim the method is faster or cheaper (in fact, it admits it is likely slower and more complex).

In summary: The author built a sophisticated, geometric way to train AI models that respects the internal structure of the data. It works perfectly fine, but for now, the old, simple way (AdamW) is still the best choice for getting the job done efficiently. The new method is a promising "long shot" that might pay off when we build much larger models in the future.

Technical Summary: Riemannian Gradient Descent for Low-Rank Architectures

Problem Statement
In modern deep learning, model parameters are often matrices constrained by low rank to reduce computational cost and parameter count (e.g., in multihead attention mechanisms where $W = AB^T$ ). Standard training updates the factors $A$ and $B$ independently using Euclidean gradients, ignoring the geometric structure of the set of rank- $r$ matrices. This paper investigates whether treating the rank-constrained matrix $W$ as a point on a smooth submanifold and applying Riemannian optimization can yield better convergence or solutions. The core challenge is that the factorization $W=AB^T$ is not unique (invariant under $A \to AS, B \to BS^{-T}$ ), and the set of rank- $r$ matrices is a manifold, not a vector space. The authors aim to restrict modifications to the optimizer only, avoiding changes to model architecture, while maintaining per-iteration complexity comparable to standard methods ( $O(mr^2 + nr^2)$ ).

Methodology
The authors propose a Riemannian gradient descent framework adapted for rank-factored parameters. The general iteration involves:

Recovering the ambient Euclidean gradient $G_W$ from the factor gradients $G_A$ and $G_B$ .
Projecting this gradient onto the tangent space of the specific manifold to obtain the Riemannian gradient $G_W^R$ .
Updating momentum and normalizing the step size.
Moving along the geodesic (using a retraction) to the next iterate.
Parallel transporting the momentum vector to the new tangent space.

The paper instantiates this template across ten specific algorithmic designs categorized by geometry and constraints:

Geometries for Fixed-Rank Matrices ( $R_{r}^{m \times n}$ ):
- Embedded Geometry: Uses the ambient Frobenius metric. Retraction is performed via rank- $r$ truncated SVD of $W + \Xi$ . Parallel transport uses metric projection.
- Quotient Geometry: Treats the manifold as a quotient of the product space of factors $(A, B)$ . The metric is defined "upstairs" on the factors, inducing a metric on the quotient. Retraction is additive on factors; transport involves horizontal projection.
Geometries for Rank- $r$ Partial Isometries ( $P_{r}^{m \times n}$ ):
- Here, $W$ is constrained such that its non-zero singular values are all 1 (equivalent to $A$ and $B$ being column-orthonormal).
- Embedded Geometry: Restricts the Frobenius metric. Retraction uses truncated SVD with specific orthogonality constraints.
- Quotient Geometry: Uses the product of Stiefel manifolds ( $St_{m,r} \times St_{n,r}$ ) as the "upstairs" space. Retractions use QR or Polar decompositions.
- Canonical Geometry: Uses the canonical metric on the Stiefel factors, inducing a product metric on the quotient.
Weight Sharing (Grid Manifolds):
- Extends the above to architectures where factors are shared across block-rows and block-columns (e.g., Group-Query Attention).
- For Fixed-Rank, the authors show that stacking factors into larger matrices allows the use of singleton algorithms, provided rank preservation is checked blockwise.
- For Partial Isometries, simple stacking fails due to orthogonality constraints. The authors develop more complex algorithms, including a "hybrid" approach for the embedded case where they lift to the quotient geometry to perform retractions and transport, then project back down.

Key Contributions

Systematic Exploration: The paper provides a comprehensive survey and implementation of ten distinct Riemannian optimization variants for low-rank matrix parameters, covering both fixed-rank and partial isometry constraints, as well as weight-sharing scenarios.
Implementation Details: It details the specific retractions (SVD, QR, Polar) and vector transports required for each geometry, addressing the non-uniqueness of factorizations and the preservation of manifold constraints (rank and orthogonality).
Empirical Evaluation: The methods are applied to the Q, K, V, O matrices in small-scale language models (6 layers, 512 embedding dimension) using the FineWeb dataset.

Experimental Results

Performance: After tuning learning rates, the Riemannian methods did not conclusively outperform the AdamW baseline. In Multi-Head Attention (MHA) and Group-Query Attention (GQA) settings, the loss differences were often within the variance of random seeds.
Hyperparameters: Metric normalization provided slight improvements (0.01–0.05 nats). Heavy momentum ( $\nu=7/8$ ) generally hurt performance or provided negligible gains compared to the baseline.
Stability: Some runs exhibited loss spikes (e.g., around step 6660) which recovered, often correlated with structured data in the training set.
Complexity: The authors acknowledge that while the methods are mathematically sound, their per-iteration computational complexity is higher than AdamW, making them less attractive for the specific small-scale experiments conducted.

Significance and Claims
The paper maintains a modest stance regarding its novelty and immediate impact:

Rediscovery: The authors explicitly state that any apparent novelty is "almost surely a rediscovery," building heavily on established textbooks (Absil et al., Boumal) and prior literature.
Heuristic Motivation: The work is motivated by the heuristic that better accounting for parameter relationships (via manifold geometry) could lead to faster convergence or better solutions, though this was not definitively proven in the experiments.
Future Potential: The authors argue that the lack of outperformance in small-scale models does not rule out merit at larger scales. They suggest the approach is mathematically sound and that the "cost-accuracy tradeoff" might shift favorably with larger models or different architectures (e.g., tensor-based models like MLA, MoLAE).
Limitations: The study admits significant limitations, including small model size, lack of learning rate annealing, and the use of a single baseline (AdamW) without tuning its hyperparameters.

In conclusion, the paper serves as a rigorous technical demonstration of applying Riemannian optimization to low-rank deep learning components, providing a library of algorithms and a baseline for future research, while acknowledging that current empirical results do not yet justify replacing standard optimizers like AdamW.