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

Emanuele Zangrando; Marco Sutti; Francesco Tudisco

cs.LGSep 17, 2026

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

Authors: Emanuele Zangrando, Marco Sutti, Francesco Tudisco

Organizations: Gran Sasso Science Institute, L’Aquila, Italy · University of Edinburgh, Edinburgh, UK

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.

Figures & tables


Method	BLEU	NIST	MET	ROUGE-L	CIDEr
AdamW Loshchilov and Hutter (2019)	68.41 $\pm$ 0.4950	8.65 $\pm$ 0.04	46.38 $\pm$ 0.12	71.13 $\pm$ 0.17	2.51 $\pm$ 0.001
Scaled AdamW Zhang and Pilanci (2024)	69.17 $\pm$ 0.43	8.72 $\pm$ 0.058	46.44 $\pm$ 0.16	71.57 $\pm$ 0.235	2.51 $\pm$ 0.005
Stiefel-AdamW	69.20 $\pm$ 0.964	8.74 $\pm$ 0.102	46.48 $\pm$ 0.259	71.40 $\pm$ 0.44	2.51 $\pm$ 0.02
GeoLoRA Schotthöfer et al. (2025a)	68.11 $\pm$ 0.271	8.60 $\pm$ 0.07	45.82 $\pm$ 0.24	70.47 $\pm$ 0.388	2.41 $\pm$ 0.02
LoRA-RITE Yen et al. (2025)	68.98 $\pm$ 1.03	8.69 $\pm$ 0.11	46.35 $\pm$ 0.219	71.17 $\pm$ 0.410	2.48 $\pm$ 0.05

Table 1: Fine-tuning performance with low-rank adapters. Best results highlighted in bold. Left: GPT2 on the E2E Natural Language Generation challenge, with rank = 4. AdamW and Scaled AdamW are reported from Zhang and Pilanci (2024) . Right: ViT-Base on CIFAR-10, for three different choices of adapter rank. We report with

\pm\sigma

the standard deviation over

5

random initializations.

Figure 1: Left and center: Loss function descent for ViT Base on CIFAR-10 LoRA fine-tuning (LoRA rank

64

and learning rate

10^{-3}

). Right: Learning rate versus best training loss for different optimizers when fine-tuning ViT Base on CIFAR-10. In this experiment, we fixed the LoRA rank at

32

and trained all models for

50

epochs.

Method	GLUE
Method	MNLI	SST-2	MRPC	CoLA	QNLI	QQP	RTE	STS-B	WNLI	Avg.
SGD	88.15	96.10	70.10	55.89	94.22	88.59	50.90	47.64	49.30	71.21
SGD	(-4.14%)	(-1.18%)	(-22.07%)	(-22.23%)	(-1.21%)	(-3.94%)	(-44.27%)	(-48.37%)	(-43.54%)	71.21
Scaled GD Zhang and Pilanci (2024)	90.21	96.90	81.62	68.17	94.40	91.15	54.15	90.31	56.34	80.36
Scaled GD Zhang and Pilanci (2024)	(-1.89%)	(-0.36%)	(-9.26%)	(-5.14%)	(-1.01%)	(-1.16%)	(-40.72%)	(-2.13%)	(-35.48%)	80.36
AdamW	91.64	97.25	87.01	71.87	94.79	91.81	90.25	90.51	85.91	89.00

Table 2: Scores for rank

16

LoRA fine-tuning of the 4-bit quantized Mistral 7B model on the GLUE benchmark for Natural Language Understanding (NLU) challenges with different optimizers (best results highlighted in bold). In parentheses, we report, for each task, the percentage deviation from the best performer. SGD, Scaled GD, and Scaled AdamW results reported from ( Zhang and Pilanci, 2024 , Table 2) .

Method	Test Loss $\pm\sigma$	Peak Memory
Method	Test Loss $\pm\sigma$	(GB)
AdamW	3.260 $\pm$ 0.05	13.8
Stiefel-AdamW	3.236 $\pm$ 0.07	13.8

Table 3: Left : Full GPT2 pretraining on OpenWeb Text for 6000 iterations. Right: Ablation over different retractions on ViT Base.

Method	Loss $\pm\sigma$
AdamW	2.66 $\pm$ 0.05
Stiefel-AdamW	2.59 $\pm$ 0.004

Table 4: Qwen2 pretraining on Wikitext-103, standard deviation reported over

5

random seeds.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Comparison of different numerical retraction methods.

Configuration	Stiefel-AdamW	AdamW	Scaled AdamW	GeoLoRA
Learning rate	$8\times 10^{-3}$	$8\times 10^{-3}$	$8\times 10^{-3}$	$8\times 10^{-3}$
Weight decay	$10^{-4}$	$10^{-2}$	$10^{-2}$	$10^{-4}$
Learning rate schedule	Linear	Linear	Linear	Linear
$(\beta_{1},\beta_{2})$	$(0.98,0.98)$	$(0.9,0.999)$	$(0.7,0.8)$	$(0.98,0.98)$

Appendix

Table 5: Hyperparameters for GPT2 LoRA fine-tuning on E2E.

Explore similar work

Sep 16, 2026

Directions That Don't Drift: Stiefel Manifold Routing for Transformer Attention

The query and key projections

\WQ,\WK

in attention are almost always trained by Euclidean optimizers with no geometric constraint. We constrain them to the Stiefel manifold and optimize with a Riemannian Adam carrying one scalar second moment per frame---the form of \citet{becigneul2019}, here extended to the compact, non-Hadamard

\St(d,r)

with a tangent projector, step-norm cap, and polar retraction. Four propositions prove steepest descent in the embedded metric, gradient-scale independence, well-conditioning, and exact

\mathrm{O}(d)

-equivariance. A fifth records that weight decay has identically zero Riemannian gradient on

\St(d,r)

(

W{=}WI_r

lies in the normal space), so decay cannot act on the constrained frames. On a CIFAR-10 patch benchmark at

n{=}10\mathrm{k}

this rule gains

\mathbf{+6.79}

,pp over AdamW across 12 paired starts (

t{=}38.33

12/12

); earlier fixed-step Riemannian SGD gains

+1.97

,pp, of which

+1.69

,pp comes from frozen orthonormal initialization alone. The corrected Adam's lead grows with data:

+1.9

,pp at

n{=}1\mathrm{k}

+6.7

,pp at

n{=}50\mathrm{k}

. A 12-seed ablation credits all gain to the scale-free step (

+4.63

,pp,

12/12

), nothing to the projector or equivariance; a targeted

\varepsilon

-sweep causally confirms the mechanism (

-2.6

,pp at

\varepsilon{=}0.1

p{<}0.001

). Two five-seed grokking studies confirm the constrained arm does not grok better than the baseline (

p{=}0.019

, A2 wins): the weight-decay exemption has no grokking consequence. A single-seed pilot exploiting this localization achieves the first stable grokking under slingshot conditions---Stiefel + targeted circuit regularization keeps routing-frame isometry error

10^6\times

lower than the unconstrained ablation through every collapse.

Rubén Darío Guerrero

NeuroTechNet S.A.S., 1108831, Bogotá, Colombia

Stiefel Manifold Optimization Riemannian Optimization Transformer Attention Attention Mechanisms

Ask PDF

Jul 28, 2026

Retraction-Free Optimization over the Stiefel Manifold for the LoRA Fine-Tuning

Optimization over the Stiefel manifold plays a significant role in various machine learning tasks. Existing methods either use the retraction operators, requiring costly orthonormalization for large-scale matrices, or employ landing methods that rely on careful step size selection and penalty parameter tuning. To address these challenges, we propose a retraction-free and penalty parameter-free algorithm that directly lands on the manifold. By leveraging the strongly-convex-like property of the quadratic penalty function and the proximal smoothness of the Stiefel manifold, we establish global convergence guarantees with the best-known iteration complexities under both constant and diminishing step sizes. Then, we reformulate the low-rank adaptation (LoRA) fine-tuning problem for large language models as a manifold optimization problem, introducing Manifold-LoRA for geometry-accelerated adaptation. This approach employs the proposed landing technique and a carefully designed step size strategy to accelerate the training process. Numerical experiments on benchmark datasets demonstrate the efficiency and strong downstream performance of the proposed method.

Yuan Zhang, Jiang Hu, Zhijian Lai +2

Center for Data Science, Peking University, Beijing 100871, China · Yau Mathematical Sciences Center, Tsinghua University, Beijing 100084, China · Beijing International Center for Mathematical Research, Peking University, Beijing 100871, China +2

Stiefel Manifold Optimization Riemannian Optimization LLM Fine-Tuning Low-Rank Adaptation

Ask PDF

Jun 1, 2026

Riemannian Gradient Descent for Low-Rank Architectures

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.

Nicholas Knight

NVIDIA

Riemannian Optimization Low-Rank Approximation Matrix Optimization Riemannian Geometry Multi-Head Attention

Ask PDF