On the Convergence of Stochastic Low-Rank Adaptation
Abstract Overview
This paper studies the optimization theory of simultaneous two-factor LoRA, where a frozen pretrained weight matrix is adapted through a low-rank update parameterized by factors B and A. The authors analyze both deterministic and stochastic optimization for the nonconvex factored objective, emphasizing that standard smooth nonconvex theory does not directly apply because the gradient in factor space is generally not globally Lipschitz. They refine prior analysis for deterministic LoRA-GD and introduce two stochastic methods, LoRA-NSGDM and LoRA-STORM, under standard unbiased finite-variance assumptions, with LoRA-STORM additionally using mean-square smoothness. The paper also includes experiments on logistic regression over pretrained features, ResNet-18 training on CIFAR-10, and TinyLlama-1.1B fine-tuning to examine practical convergence behavior.
Novelty
The main novelty is a finite-time convergence analysis for simultaneous stochastic LoRA under relatively standard stochastic-oracle assumptions, including a proof that plain LoRA-SGD can fail under finite variance alone. The work also sharpens the known deterministic guarantee for LoRA-GD from an exponential dependence to a polynomial complexity bound, and proposes normalized-momentum and variance-reduced stochastic LoRA variants with proved oracle complexities.
Results
The paper shows that deterministic LoRA-GD can reach an epsilon-stationary point in O(epsilon^-4) full-gradient evaluations, improving on a previously reported exponential bound. For stochastic optimization, LoRA-NSGDM attains O(epsilon^-8) stochastic oracle complexity under unbiased finite-variance gradients, while LoRA-STORM improves this to O(epsilon^-6) under an added mean-square smoothness condition. Experimentally, LoRA-NSGDM is reported to converge faster and more steadily than LoRA-GD baselines on CIFAR-10 logistic regression, ResNet-18 training, and TinyLlama fine-tuning, and LoRA-STORM achieves the lowest mean loss in the logistic-regression setting over the final updates shown.
Key Points
- The paper establishes polynomial convergence guarantees for simultaneous two-factor LoRA in both deterministic and stochastic settings.
- It proves that plain simultaneous LoRA-SGD may diverge in expectation under finite-variance assumptions, motivating the normalized momentum design of LoRA-NSGDM.
- A variance-reduced extension, LoRA-STORM, further improves stochastic oracle complexity under mean-square smoothness and is supported by convergence experiments on vision and language tasks.