# Stiefel-AdamW: Geometry-Aware AdamW for Linear Factorization Blocks

arXiv: https://arxiv.org/abs/2609.21039 · HTML: https://mlsanity.com/paper/2609.21039/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 31 of its 70 references that are in the map.
- First posted on arXiv 17 Sep 2026.

## Builds on (most cited first, 30 of 31)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [LoRA+: Efficient Low Rank Adaptation of Large Models](https://mlsanity.com/paper/2402.12354/) (2024)
- [AdaLoRA: Adaptive Budget Allocation for Parameter-Efficient Fine-Tuning](https://mlsanity.com/paper/2303.10512/) (2023)
- [The E2E Dataset: New Challenges For End-to-End Generation](https://mlsanity.com/paper/1706.09254/) (2017)
- [PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization](https://mlsanity.com/paper/1905.13727/) (2019)
- [ReLoRA: High-Rank Training Through Low-Rank Updates](https://mlsanity.com/paper/2307.05695/) (2023)
- [Riemannian Adaptive Optimization Methods](https://mlsanity.com/paper/1810.00760/) (2018)
- [Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform](https://mlsanity.com/paper/2002.01113/) (2020)
- [LoRA-Pro: Are Low-Rank Adapters Properly Optimized?](https://mlsanity.com/paper/2407.18242/) (2024)
- [Riemannian SVRG: Fast Stochastic Optimization on Riemannian Manifolds](https://mlsanity.com/paper/1605.07147/) (2016)
- [Pufferfish: Communication-efficient Models At No Extra Cost](https://mlsanity.com/paper/2103.03936/) (2021)
- [Initialization and Regularization of Factorized Neural Layers](https://mlsanity.com/paper/2105.01029/) (2021)
- [Riemannian Preconditioned LoRA for Fine-Tuning Foundation Models](https://mlsanity.com/paper/2402.02347/) (2024)
- [Low-rank lottery tickets: finding efficient low-rank neural networks via matrix differential equations](https://mlsanity.com/paper/2205.13571/) (2022)
- [LoRA Done RITE: Robust Invariant Transformation Equilibration for LoRA Optimization](https://mlsanity.com/paper/2410.20625/) (2024)
- [Matrix Compression via Randomized Low Rank and Low Precision Factorization](https://mlsanity.com/paper/2310.11028/) (2023)
- [Riemannian conjugate gradient methods: General framework and specific algorithms with convergence analyses](https://mlsanity.com/paper/2112.02572/) (2021)
- [Geometry-aware training of factorized layers in tensor Tucker format](https://mlsanity.com/paper/2305.19059/) (2023)
- [Riemannian Adaptive Optimization Algorithm and Its Application to Natural Language Processing](https://mlsanity.com/paper/2004.00897/) (2020)
- [A new, globally convergent Riemannian conjugate gradient method](https://mlsanity.com/paper/1302.0125/) (2013)
- [A general framework of Riemannian adaptive optimization methods with a convergence analysis](https://mlsanity.com/paper/2409.00859/) (2024)
- [GeoLoRA: Geometric integration for parameter efficient fine-tuning](https://mlsanity.com/paper/2410.18720/) (2024)
- [A geometric framework for momentum-based optimizers for low-rank training](https://mlsanity.com/paper/2506.17475/) (2025)

## Abstract (arXiv)

A pervasive structural pattern in modern deep learning is the linear factorization block: a submodule of the form W = BA in which two parameter matrices are multiplied directly, with no intervening nonlinearity. Such blocks appear in LoRA adapters, low-rank compressed layers, query-key products of self-attention, and share a common pathology: the factorization is non-unique, which can destabilize training and limit usable learning rates. Despite this, factorization blocks are typically optimized with standard Euclidean methods that ignore the underlying geometry. We introduce Stiefel-AdamW, a near drop-in replacement for AdamW for use wherever such blocks appear. By constraining one factor on the Stiefel manifold while leaving the other Euclidean, Stiefel-AdamW relaxes the full GL(ℝr) gauge symmetry to a compact orthogonal symmetry, ruling out factor blow-up while retaining the coordinate-wise diagonal preconditioning that gives AdamW its practical strength. Moment estimation is performed in the ambient Euclidean space, with geometry entering only through a tangent-space projection and a manifold retraction. The implementation overhead over AdamW is minimal, and we show that the resulting optimizer inherits both the stability benefits of Riemannian methods and standard convergence guarantees. We validate Stiefel-AdamW on LoRA-style fine-tuning of GPT2, ViT, and Mistral 7B and on full pretraining of GPT2 on OpenWebText, showing consistent improvements over strong baselines at essentially no additional cost over AdamW.