Stable-LoRA: Stabilizing Feature Learning of Low-Rank Adaptation

TL;DR

Stable-LoRA introduces a weight-shrinkage strategy to improve the stability of Low-Rank Adaptation in fine-tuning large language models, addressing theoretical limitations and enhancing performance without extra memory costs.

Contribution

It proposes Stable-LoRA, a novel method that dynamically stabilizes feature learning in LoRA by shrinking low-rank matrices during early training, backed by theoretical analysis and empirical validation.

Findings

Stable-LoRA effectively eliminates instability in LoRA feature learning.

It outperforms baseline methods across various models and tasks.

The approach requires no additional memory and minimal computation overhead.

Abstract

Low-Rank Adaptation (LoRA) is a widely adopted parameter-efficient method for fine-tuning Large Langauge Models. It updates the weight matrix as $W=W_0+sBA$, where $W_0$ is the original frozen weight, $s$ is a scaling factor and $A$,$B$ are trainable low-rank matrices. Despite its robust empirical effectiveness, the theoretical foundations of LoRA remain insufficiently understood, particularly with respect to feature learning stability. In this paper, we first establish that, LoRA can, in principle, naturally achieve and sustain stable feature learning (i.e., be self-stabilized) under appropriate hyper-parameters and initializations of $A$ and $B$. However, we also uncover a fundamental limitation that the necessary non-zero initialization of $A$ compromises self-stability, leading to suboptimal performances. To address this challenge, we propose Stable-LoRA, a weight-shrinkage optimization strategy that dynamically enhances stability of LoRA feature learning. By progressively shrinking $A$ during the earliest training steps, Stable-LoRA is both theoretically and empirically validated to effectively eliminate instability of LoRA feature learning while preserving the benefits of the non-zero start. Experiments show that Stable-LoRA consistently outperforms other baselines across diverse models and tasks, with no additional memory usage and only negligible computation overheads. The code is available at https://github.com/Yize-Wu/Stable-LoRA.