On the Convergence of Stochastic Low-Rank Adaptation

Abstract

Low-rank adaptation (LoRA) optimizes $J(B,A)=\mathcal{L}(W_{\mathrm{base}}+sBA)$ over two adapters $B\in {\mathbb{R}}^{m\times r}$ and $A\in {\mathbb{R}}^{r\times n}$ that form a low-rank update to a frozen pretrained weight matrix $W_{\mathrm{base}}\in {\mathbb{R}}^{m\times n}$. The prior analysis shows LoRA-GD takes $\exp \{\mathcal{O}(\hat{I}{\mu }^{-2})\}$ oracle calls to find an $\hat{I}\mu$-stationary point such that $\Vert \nabla J(B,A)\Vert \le \hat{I}\mu$ in the deterministic setting. We sharpen the analysis and show that $\mathcal{O}(\hat{I}{\mu }^{-4})$ full-gradient evaluations suffice for the same first-order criterion. We further study stochastic LoRA under unbiased gradient estimates and finite variance. We propose LoRA-NSGDM, which finds an $\hat{I}\mu$-stationary point with $\mathcal{O}(\hat{I}{\mu }^{-8})$ stochastic oracle complexity. Under the additional mean-square smoothness condition, we use variance reduction strategy and propose LoRA-STORM, which improves the stochastic oracle complexity to $\mathcal{O}(\hat{I}{\mu }^{-6})$.