arXiv preprint ·
Stiefel-AdamW: Geometry-Aware AdamW for Linear Factorization Blocks
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
Abstract
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 $\mathrm{GL}(\mathbb{R}^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.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Machine Learning
Beyond PINNs: A Unified Gauss--Newton and Petrov--Galerkin Framework for Neural and Hybrid PDE Solvers
FREESIA: Covariance-Aware Posterior Transport for Expressive and Scalable Data Assimilation
MGRD: Compact morphology-gated residual diffusion for variance-aware cross-domain neurite forecasting
Domain-decomposed Evolutional Deep Neural Network with Random Features for Transient Pressure Diffusion with Discontinuous and High-Contrast Coefficients
Cost-Accuracy Trade-offs: Neural Operator vs Classical Numerical Solver
Get the few papers worth your time in Machine Learning every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.