Research map / Architectures and efficiency / Pruning and efficient tuning
Riemannian Preconditioned LoRA for Fine-Tuning Foundation Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes an r×r Riemannian preconditioner (scaled GD/AdamW) for LoRA fine-tuning that stabilizes feature learning under infinite width and improves SGD/AdamW on GPT-2, Mistral 7B and diffusion models.
Problem
LoRA training with standard optimizers is sensitive to learning rates for the two low-rank factors A and B, and infinite-width analyses show unpreconditioned training needs separately tuned learning rates for stable feature learning.
Approach
Precondition LoRA gradients with r×r matrices (B^TB)^{-1} for A and (AA^T)^{-1} for B, derived from a quotient-manifold Riemannian metric (scaled GD); apply this to AdamW by scaling each raw gradient, which needs only a few lines of code change and yields stable feature learning with a single shared learning rate.
Key results
- On E2E with GPT-2 medium (r=4), scaled AdamW improves over AdamW on all five metrics (BLEU 69.6 vs 68.9), and scaled GD (69.2 BLEU) closes most of the gap to AdamW from SGD (66.6).
- On GLUE with 4-bit Mistral 7B (r=16), scaled AdamW averages 89.01 vs 88.28 for AdamW and scaled GD averages 80.36 vs 71.21 for SGD.
- Theory: under infinite-width analysis, preconditioned Adam achieves stable feature learning with a single learning rate eta=Theta(1), whereas unpreconditioned Adam needs eta_A=Theta(n^-1) and eta_B=Theta(1).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply LoRA gradient preconditioning with scaled AdamW | GPT-2 medium LoRA fine-tuning · E2E | BLEU: 68.9 -> 69.6 | Helped |
| Apply LoRA gradient preconditioning with scaled AdamW | GPT-2 medium LoRA fine-tuning · E2E | NIST: 8.69 -> 8.77 | Helped |
| Apply LoRA gradient preconditioning with scaled AdamW | GPT-2 medium LoRA fine-tuning · E2E | MET: 46.5 -> 46.6 | Helped |
| Apply LoRA gradient preconditioning with scaled AdamW | GPT-2 medium LoRA fine-tuning · E2E | ROUGE-L: 71.3 -> 71.8 | Helped |
| Apply LoRA gradient preconditioning with scaled AdamW | GPT-2 medium LoRA fine-tuning · E2E | CIDEr: 2.51 -> 2.52 | Helped |
| Apply LoRA gradient preconditioning with scaled GD | GPT-2 medium LoRA fine-tuning · E2E | BLEU: 66.6 -> 69.2 | Helped |
| Apply LoRA gradient preconditioning with scaled AdamW | 4-bit Mistral 7B LoRA fine-tuning · GLUE MNLI | score: 89.86 -> 90.68 | Helped |
| Apply LoRA gradient preconditioning with scaled AdamW | 4-bit Mistral 7B LoRA fine-tuning · GLUE SST-2 | score: 96.79 -> 97.25 | Helped |
| Apply LoRA gradient preconditioning with scaled AdamW | 4-bit Mistral 7B LoRA fine-tuning · GLUE MRPC | score: 88.48 -> 89.46 | Helped |
| Apply LoRA gradient preconditioning with scaled AdamW | 4-bit Mistral 7B LoRA fine-tuning · GLUE CoLA | score: 71.05 -> 71.30 | Helped |
| Apply LoRA gradient preconditioning with scaled AdamW | 4-bit Mistral 7B LoRA fine-tuning · GLUE QQP | score: 91.24 -> 92.22 | Helped |
| Apply LoRA gradient preconditioning with scaled AdamW | 4-bit Mistral 7B LoRA fine-tuning · GLUE RTE | score: 90.61 -> 91.34 | Helped |
Surprising findings
- Scaling each single gradient inside AdamW worked better empirically than scaling the processed gradient, which differs from the version analyzed in the proof.
- Scaled GD nearly matches AdamW performance while using less optimizer storage.
Limitations
- Stable feature learning theory relies on an Adam sign-like gradient assumption (Assumption on g_A x = Theta(n)) and vector-output setting; scalar output needs learning-rate schedules.
- Empirical gains are reported for a limited set of models, tasks and ranks, with grid-searched learning rates per method.
Experimental setup
Datasets: E2E NLG Challenge GLUE
Benchmarks: E2E NLG Challenge GLUE
Metrics: BLEU NIST METEOR ROUGE-L CIDEr accuracy Matthews correlation F1 Pearson correlation
Baselines: SGD AdamW LoRA+
Builds on: LoRA LoRA+ scaled GD (Tong et al.) AdamW Mix-of-Show
Closest papers by meaning
- A Stronger Mixture of Low-Rank Experts for Fine-Tuning Foundation Models (ICML 2025 · similarity 0.92)
- AdaPreLoRA: Adafactor Preconditioned Low-Rank Adaptation (NeurIPS 2026 · similarity 0.91)
- SDS-LoRA: Overcoming Anisotropic Gradient Scaling in Low-Rank Adaptation (NeurIPS 2026 · similarity 0.89)
- LoRA-GA: Low-Rank Adaptation with Gradient Approximation (NeurIPS 2024 · similarity 0.89)
- LoRA-Pro: Are Low-Rank Adapters Properly Optimized? (ICLR 2025 · similarity 0.89)
- Flat-LoRA: Low-Rank Adaptation over a Flat Loss Landscape (ICML 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.