Stiefel-AdamW: Geometry-Aware AdamW for Linear Factorization Blocks
Emanuele Zangrando, Marco Sutti, Francesco Tudisco
cs.LG · math.NA
Quick verdict
Proposes Stiefel-AdamW, a drop-in Riemannian AdamW for W=BA blocks with regret guarantees; gains are real but modest and significance is untested.
Summary
The paper identifies linear factorization blocks W=BA — ubiquitous in LoRA adapters, low-rank compressed layers, and attention query-key products — as sharing a non-uniqueness (gauge) pathology under GL(R^r) that can cause factor blow-up and training instability. It proposes Stiefel-AdamW: constrain one factor to the Stiefel manifold (row-orthonormal), keep the other Euclidean, reducing the symmetry to compact O(r) while retaining AdamW's coordinate-wise diagonal preconditioning. All moment estimation happens in the ambient Euclidean space; geometry enters only via a tangent projection and a Stiefel retraction (Cayley fixed-point by default, with QR/polar/Newton–Schulz alternatives). Theoretical results include a convex regret bound R(T)=o(T) under AMSGrad-style assumptions, robustness to approximate retractions (adding a linear δ term), and gradient boundedness on compact fibers (Proposition 4.3). Experiments cover LoRA fine-tuning of GPT2 (E2E), ViT-Base (CIFAR-10), Mistral 7B (GLUE), plus GPT2 pretraining on OpenWebText and Qwen2 on WikiText-103, against AdamW, Scaled AdamW, GeoLoRA, LoRA-RITE, LoRA-Pro, and Cayley Adam. Gains are consistent but often small (E2E BLEU 69.20±0.96 vs 69.17±0.43 for Scaled AdamW; GLUE average 89.57 vs 89.00); pretraining is comparable to AdamW with identical hyperparameters and equal peak memory. Limitations: no significance testing, theory assumes convexity/AMSGrad max/alignment not met in practice, tuning asymmetries, and no code availability statement.
Review complete
Assessment at a glance
6
risk flags identified
● high 0
● medium 3
● low 3
5 strengths · 5 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- Constraining one factor of W=BA to the Stiefel manifold relaxes the GL(R^r) gauge symmetry to compact O(r), ruling out factor blow-up while uniquely preserving AdamW's coordinate-wise diagonal preconditioning (Sections 1, 3.1).
- Stiefel-AdamW attains zero-regret convergence R(T)=o(T) under AMSGrad-style assumptions and remains convergent with approximate retractions, adding only a linear δ term (Theorem 4.2, Appendix A.2).
- Optimizing on compact fibers yields bounded gradients and larger stable learning rates than the unconstrained factorization (Proposition 4.3, Section 5.4).
- Across LoRA fine-tuning (GPT2, ViT, Mistral 7B) and pretraining (GPT2, Qwen2), Stiefel-AdamW consistently improves over six baselines at essentially AdamW's cost and memory (Sections 5.1-5.3, Tables 1-4).
Strengths
Well-motivated geometric relaxation uniquely compatible with Adam-style adaptivity
Section 1 and 3.1 articulate the key trade-off: full GL(R^r) quotient invariance forces scalar preconditioners in prior Riemannian Adams (RAdam/BG19, GeoLoRA), while relaxing to O(r) on St(n,r)×R^{m×r} keeps coordinate-wise diagonal preconditioning well-defined. This is the paper's sharpest conceptual contribution.
Minimal-overhead design with validated cost claims
Algorithm 1 computes moments in ambient space and touches geometry only via projection P_X and retraction. Table 3 confirms identical 13.8 GB peak memory vs AdamW in GPT2 pretraining; retraction costs O(nr²+r³), and Table 3-right plus Figure 2 ablate six retraction implementations (wall-clock and orthogonality error ‖YᵀY−I‖_F).
Inexact-retraction theory feeding practice
Broad and current baseline set across five settings
Honest untuned pretraining comparison and a direct stability experiment
Weaknesses
Headline gains within noise on the flagship benchmark
Theory-practice gap in the convergence guarantee
Tuning-budget asymmetry in the Mistral/GLUE comparison
Reproducibility friction
Citation slip and residual invariance limitation
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
This is a method paper at the intersection of Riemannian optimization and deep-learning practice. The target object is the linear factorization block W=BA — LoRA adapters, low-rank compressed layers, attention query-key products — whose GL(R^r) gauge symmetry makes the parametrization non-unique and permits factor blow-up (Section 1; concrete divergence example in Section 4.2, where GD on rank-r recovery diverges once σ > √(3/η)). The central claim: constraining one factor to the Stiefel manifold relaxes the symmetry to compact O(r) while *preserving* AdamW's coordinate-wise diagonal preconditioning — precisely the trade-off that prior Riemannian adaptive methods (RAdam/Cayley Adam with scalar preconditioners, GeoLoRA with full-quotient geometry) forfeit (Section 2, contribution 2). Secondary claims: a zero-regret bound R(T)=o(T) (Theorem 4.2), gradient boundedness on compact fibers (Proposition 4.3, Appendix C), and consistent empirical gains at AdamW-like cost (Sections 5.1–5.3).
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 5 strengths & 5 weaknesses, analyzed in detail
- 6 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is Stiefel-AdamW in simple terms?
It is a modified version of the AdamW optimizer for neural-network blocks built as a product of two matrices, W=BA (as in LoRA adapters or attention query-key products). Because such products can be written many ways, training can blow up one factor. Stiefel-AdamW constrains one factor to have orthonormal rows (it lives on the Stiefel manifold), which removes that instability while keeping AdamW's usual per-coordinate scaling. Geometry is applied only at the final update step via a projection and a retraction.
How much overhead does Stiefel-AdamW add over AdamW?
Very little. Moments are computed in the ordinary Euclidean space, and only the constrained factor needs a tangent projection plus a retraction costing O(nr²+r³). In the paper's GPT2 pretraining test, peak GPU memory was identical to AdamW (13.8 GB, Table 3), and the default Cayley retraction is computed with a cheap fixed-point iteration whose accuracy is covered by the convergence theory.
Does Stiefel-AdamW actually beat AdamW?
The improvements are consistent but often small. On GPT2 LoRA fine-tuning (E2E benchmark, rank 4) it reaches BLEU 69.20±0.96 vs 68.41 for AdamW but essentially ties Scaled AdamW (69.17±0.43) and loses on ROUGE-L. On Mistral 7B GLUE the average score is 89.57 vs AdamW's 89.00. On GPT2 pretraining with identical hyperparameters it is comparable (test loss 3.236±0.07 vs 3.260±0.05). Its clearest advantage is stability at large learning rates (Section 5.4).
Is code for Stiefel-AdamW available?
The paper text reviewed contains no code or data availability statement, so availability could not be confirmed. However, the algorithm is short (one projection and one retraction on top of AdamW), the hyperparameters for the main benchmarks are listed (Table 5), and the same authors previously released GeoLoRA, which could serve as a starting codebase.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures 1 and 2 are available only as text fragments; plots (loss curves, learning-rate sweep shape, retraction timing) could not be visually analyzed.
- The ViT column of Table 1 lost its method labels in extraction; row-to-method attribution was inferred from the ordering of the adjacent GPT2 table (AdamW, Scaled AdamW, Stiefel-AdamW, GeoLoRA, LoRA-RITE, LoRA-Pro, Cayley Adam) and should be re-verified against the PDF.
- Appendix proofs were checked only as extracted text; LaTeX flattening may hide subtle notation issues.
- No code availability statement was found in the extracted text; if the arXiv paper has one outside the extracted body, reproducibility ratings would improve.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.