One email each morning — yesterday's AI, sortedGet it in your inbox

METAL LAB

LoRA-GA$^2$: Low Rank Adaptation with Multi-step Gradient Adaptive Alignment

arXiv:2608.198002026-08-21

Peeking at a few early training gradients before fine-tuning starts to set up LoRA smarter

LoRA lets you fine-tune huge models with far less memory than full fine-tuning, but it still tends to underperform full fine-tuning. This paper introduces LoRA-GA2, which briefly runs a lightweight probe over a handful of early training steps to gather gradient information, then uses it to decide how much capacity each layer gets and how to initialize LoRA's weights. This consistently beat existing LoRA variants on GLUE, GSM8K, and HumanEval benchmarks.

What they did

  1. Prior gradient-guided LoRA methods used only a single snapshot of the gradient at the very start of training, which fails to capture how the actual optimization path unfolds.
  2. An alternative that tracks gradients at every training step is more faithful but requires changing the optimizer itself, which sharply increases GPU memory and training time.
  3. LoRA-GA2 instead runs a temporary, lightweight optimizer called AdaLomo for just a few steps, accumulates the gradients on CPU, then restores the original pretrained weights before using that accumulated signal to allocate ranks and initialize the LoRA matrices — adding no extra GPU memory and only a small time cost.
  4. For rank allocation, it multiplies a sensitivity score (how much a layer matters for the task) with an effective-rank score (how spread out that layer's gradient directions are), avoiding the waste that comes from relying on either signal alone.
  5. On T5-Base with GLUE (rank 8), LoRA-GA2 beat the best baseline by 0.66 points on average; on Llama3.1-8B-Base it improved GSM8K by 1.03 points and HumanEval by 0.87 points over the strongest baseline, and it also outperformed the best baseline by 0.63 points on CLIP-ViT-B/16 image classification.
Figure 1: Illustration of (a) LoRA; (b) One-step gradients based LoRA Variants; (c) Multi-step gradients based LoRA Variants; and (d) Ours (LoRA-GA2), which introduces an efficient probe for estimating multi-step gradients of pre-trained weights and utilizing the directions and intrinsic dimensionalities of the gradients to initialize low-rank weights.
Figure 1: Illustration of (a) LoRA; (b) One-step gradients based LoRA Variants; (c) Multi-step gradients based LoRA Variants; and (d) Ours (LoRA-GA2), which introduces an efficient probe for estimating multi-step gradients of pre-trained weights and utilizing the directions and intrinsic dimensionalities of the gradients to initialize low-rank weights.
Table 1: Performance of fine-tuning T5-Base on 5 sub-tasks of the GLUE benchmark. Bold and underline indicate the highest and second-highest scores of low-rank methods with r=8 or rref=8.
MethodMNLISST-2CoLAQNLIMRPCAverage
Full86.33±0.0094.75±0.2180.70±0.2493.19±0.2284.56±0.7387.91
LoRA (13)85.30±0.0494.04±0.1169.35±0.0593.19±0.2284.56±0.7385.29
Convergence Optimization Methods for LoRA
rsLoRA (15)85.73±0.1094.19±0.2372.32±1.1293.12±0.0952.86±2.2779.64
DoRA (19)85.67±0.0994.04±0.5372.04±0.9493.04±0.0668.08±0.5182.57
LoRA+ (9)85.81±0.0993.85±0.2477.53±0.2093.14±0.0374.43±1.3984.95
Initialization Optimization Methods for LoRA
PiSSA (22)85.75±0.0794.07±0.0674.27±0.3993.15±0.1476.31±0.5184.71
LoRA-GA (30)85.70±0.0994.11±0.1880.57±0.2093.18±0.0685.29±0.2487.77
Adaptive Methods for LoRA
AdaLoRA (37)85.45±0.1193.69±0.2069.16±0.2491.66±0.0568.14±0.2881.62
RaLoRA (34)85.76±0.0394.22±0.2978.11±0.4593.36±0.1484.74±0.2787.24
GoRA (10)85.91±0.2294.68±0.4379.86±0.3593.27±0.0886.10±0.2087.96
LoRA-GA2 (Ours)85.91±0.0194.72±0.3782.39±0.2493.19±0.0586.88±0.1488.62
Figure 2: Framework of LoRA-GA2. A temporary multi-step probe collects trajectory gradients, restores the pretrained weights, and uses the accumulated signal for both rank allocation and SVD initialization before standard LoRA training.
Figure 2: Framework of LoRA-GA2. A temporary multi-step probe collects trajectory gradients, restores the pretrained weights, and uses the accumulated signal for both rank allocation and SVD initialization before standard LoRA training.
Table 2: Performance of fine-tuning Llama3.1-8B-Base on GSM8K and HumanEval.
MethodGSM8KHumanEval
Full73.69±0.2851.63±1.27
LoRA (13)67.78±1.2543.09±0.35
rsLoRA (15)68.36±0.7445.78±2.80
DoRA (19)69.17±1.0043.70±1.54
LoRA+ (9)71.29±0.9344.51±2.11
OLoRA (2)68.54±0.4243.29±2.44
PiSSA (22)68.56±1.0344.10±1.54
LoRA-GA (30)71.39±0.9043.29±0.61
AdaLoRA (37)70.63±0.7741.46±3.66
RaLoRA (34)72.25±0.5948.78±1.61
GoRA (10)72.91±0.7648.98±2.14
LoRA-GA2 (Ours)73.94±0.4849.85±0.33
Figure 3: Loss comparison of fine-tuning Llama3.1-8B-Base on MetamathQA using LoRA and LoRA-GA2
Figure 3: Loss comparison of fine-tuning Llama3.1-8B-Base on MetamathQA using LoRA and LoRA-GA2
Table 3: Performance of fine-tuning CLIP-ViT-B/16 on seven image classification tasks.
MethodCarsDTDEuroSATGTSRBRESISC45SUN397SVHNAverage
Zero-shot63.7544.3942.2235.2256.4662.5615.5345.73
LoRA (13)82.31±0.0876.97±0.5198.38±0.2097.10±0.0694.99±0.1177.19±0.1996.62±0.0689.08±0.10
MELoRA (26)82.65±0.3875.16±0.5998.64±0.0598.88±0.0595.78±0.1674.69±0.2296.95±0.0988.96±0.15
MoRA (14)84.61±0.2177.34±0.1498.65±0.1698.68±0.1896.33±0.1978.12±0.0697.17±0.1590.13±0.16
AdaLoRA (37)73.58±0.0973.79±0.4896.96±0.1258.87±0.3889.07±0.6072.00±0.1094.26±0.1379.79±0.27
DoRA (19)82.44±0.2676.86±0.8498.43±0.1797.25±0.1295.10±0.1677.30±0.1796.63±0.0489.14±0.07
rsLoRA (15)83.94±0.2277.64±0.3398.51±0.1798.69±0.1795.90±0.2077.96±0.2196.94±0.0689.94±0.06
LoRA+ (9)86.61±0.3673.33±1.3098.54±0.1498.99±0.2096.06±0.3876.80±0.3496.98±0.0889.62±0.19
PiSSA (22)83.36±0.3877.38±0.5798.54±0.0998.32±0.0995.92±0.4077.46±0.1397.00±0.0989.71±0.25
OLoRA (2)83.85±0.1378.60±0.2598.62±0.0398.49±0.1496.01±0.2877.30±0.0897.15±0.1490.00±0.15
RaLoRA (34)86.63±0.3077.75±0.2098.66±0.2798.98±0.1196.62±0.2877.86±0.0597.24±0.1190.53±0.03
LoRA-GA2 (Ours)87.82±0.1379.49±0.1198.83±0.1199.01±0.0896.48±0.2479.05±0.1497.42±0.0791.16±0.05
Figure 4: Rank allocation results of fine-tuning Llama3.1-8B-Base-base on MetamathQA.
Figure 4: Rank allocation results of fine-tuning Llama3.1-8B-Base-base on MetamathQA.
Table 4: Key ablations with standard deviations on Llama3.1-8B-Base.
VariantGSM8KHumanEval
One-step gradient72.13±0.3449.65±1.01
No SVD init.68.92±0.1248.37±2.50
No rank alloc.72.40±1.3248.88±0.63
Sensitivity only72.98±0.6548.17±2.28
Effective rank only73.51±0.5148.78±1.00
Full LoRA-GA273.94±0.4849.85±0.33
Figure 5: Ablation of the hyperparameters N and γ.
Figure 5: Ablation of the hyperparameters N and γ.

Why it matters

Since LoRA is already widely used for fine-tuning large models, an improvement like this means better model quality for the same computational budget. It's practically appealing because it adds essentially no extra GPU memory and only marginal time overhead, making it easy to slot into existing LoRA workflows.

Terms in this paper

  • LoRA (Low-Rank Adaptation) · A fine-tuning method that trains only two small matrices instead of the full model weights, saving memory
  • one-step vs. multi-step gradients · A gradient shows which direction reduces loss; one-step means measuring it only at the start, multi-step means tracking it across several early training steps
  • SVD (singular value decomposition) initialization · Breaking down the accumulated gradient into its main directions and magnitudes, then setting LoRA's starting weights along those directions
  • effective rank · A measure of how many independent directions a layer's gradient is spread across
  • sensitivity · A score showing how much changing a layer's weights affects the final loss, i.e. how important that layer is to the task

Original abstract (English)

Low-Rank Adaptation (LoRA) is a prominent fine-tuning method for large models, achieving competitive performance with reduced memory overhead. However, a persistent performance gap remains between LoRA and full fine-tuning. Recent studies have sought to narrow this gap by employing one-step gradient approximations of pretrained weights to align LoRA updates with the principal directions or intrinsic dimensionalities of full fine-tuning updates. Nevertheless, these approaches fail to capture the full dynamics of the gradients. In this paper, we propose LoRA-GA$^2$, an effective fine-tuning algorithm that fully leverages multi-step gradient information. Specifically, we introduce a lightweight probe for multi-step gradients of pretrained weights that incurs no additional GPU memory cost and only marginal time overhead. We further employ a spectrum-aware, importance-based rank allocation and optimal initialization derived from multi-step gradients. Extensive experimental results demonstrate that LoRA-GA$^2$ consistently outperforms existing LoRA variants while preserving the efficiency advantages of vanilla LoRA. For instance, LoRA-GA$^2$ surpasses the leading baseline by an average of 0.66 points on the GLUE benchmark, and outperforms the strongest baseline by 1.03 points on GSM8K and 0.87 points on HumanEval, respectively.

Authors · Haonan He, Xinyue Fan

Read on arXiv

Latest papers

All papers →

Latest from METAL LAB

Figures: Haonan He et al., arXiv:2608.19800, CC BY 4.0