Source-linked AI summary

LoRA-GA: Low-Rank Adaptation with Gradient Approximation

Shaowen Wang, Linxi Yu, Jian Li

arXiv:2407.05000v2cs.LGcs.CL

TL;DR

Large-model full fine-tuning is costly, while vanilla LoRA can converge substantially more slowly despite reducing per-step costs. The paper introduces LoRA-GA, which initializes low-rank adapters to align their gradients with full fine-tuning, and reports comparable convergence with similar or better performance. The method is evaluated on T5-Base and Llama 2-7B, but not larger models such as Llama 2-70B.

  • Problem

    Full fine-tuning has prohibitive computational and storage costs, while vanilla LoRA often converges more slowly and can require greater overall computation.

  • Method

    LoRA-GA initializes low-rank factors using eigenvectors of the full gradient matrix so the low-rank update approximates full-model gradients.

  • Results

    LoRA-GA achieves convergence comparable to full fine-tuning and similar or superior performance, including up to 2-4 times faster convergence than vanilla LoRA.

  • Takeaways & Limitations

    Because LoRA-GA changes only initialization, it offers an efficient implementation that can also be incorporated with other LoRA variants.

  • Takeaways & Limitations

    The study validates LoRA-GA on T5-Base and Llama 2-7B but not larger pretrained models such as Llama 2-70B.

Abstract

from arXiv · show

Fine-tuning large-scale pretrained models is prohibitively expensive in terms of computational and memory costs. LoRA, as one of the most popular Parameter-Efficient Fine-Tuning (PEFT) methods, offers a cost-effective alternative by fine-tuning an auxiliary low-rank model that has significantly fewer parameters. Although LoRA reduces the computational and memory requirements significantly at each iteration, extensive empirical evidence indicates that it converges at a considerably slower rate compared to full fine-tuning, ultimately leading to increased overall compute and often worse test performance. In our paper, we perform an in-depth investigation of the initialization method of LoRA and show that careful initialization (without any change of the architecture and the training algorithm) can significantly enhance both efficiency and performance. In particular, we introduce a novel initialization method, LoRA-GA (Low Rank Adaptation with Gradient Approximation), which aligns the gradients of low-rank matrix product with those of full fine-tuning at the first step. Our extensive experiments demonstrate that LoRA-GA achieves a convergence rate comparable to that of full fine-tuning (hence being significantly faster than vanilla LoRA as well as various recent improvements) while simultaneously attaining comparable or even better performance. For example, on the subset of the GLUE dataset with T5-Base, LoRA-GA outperforms LoRA by 5.69% on average. On larger models such as Llama 2-7B, LoRA-GA shows performance improvements of 0.34, 11.52%, and 5.05% on MT-bench, GSM8K, and Human-eval, respectively. Additionally, we observe up to 2-4 times convergence speed improvement compared to vanilla LoRA, validating its effectiveness in accelerating convergence and enhancing model performance. Code is available at https://github.com/Outsider565/LoRA-GA.

1 Introduction

LoRA reduces per-step fine-tuning costs through low-rank adapters, but vanilla LoRA converges more slowly than full fine-tuning. LoRA-GA addresses this by initializing adapters to approximate full-model gradients, improving convergence and performance without changing LoRA’s architecture or training algorithm.

  • Motivation: LoRA reduces trainable parameters and per-step computational and memory costs by updating low-rank matrices while keeping the original weights fixed.The modified layer is represented as y = (W + ηBA)x.
  • Motivation: LoRA typically requires 5-6x more iterations and FLOPs than full fine-tuning to reach the same performance under the same learning rate.
  • Motivation: LoRA’s default random initialization is suboptimal, motivating investigation of initialization strategies as a cause of slow convergence.The default initializes A with Kaiming initialization and B to zero.
  • LoRA-GA: LoRA-GA initializes its low-rank factors from eigenvectors of the full gradient matrix so the gradient of BA aligns with the gradient of W.
  • LoRA-GA: LoRA-GA also identifies a scaling factor that keeps adapter-output variance invariant to adapter rank and input dimension.
  • Results: 5.69% improvement on the GLUE subset with T5-Base, alongside 0.34, 11.52%, and 5.05% improvements on MT-bench, GSM8K, and HumanEval with Llama 2-7B.The reported results also include up to 2-4 times faster convergence than vanilla LoRA.

2 Related Work

Related work reduces PEFT costs through adapters and structured LoRA variants, while LoRA-GA differs by approximating gradients rather than pretrained weights. Existing improvements target expressiveness, scaling, merging, or learning rates without changing LoRA’s basic role as a low-rank weight-update method.

  • Initialization stability: Variance-stability analyses distinguish LoRA-GA’s initialization concerns from established Xavier, He, and LSUV initialization approaches.
  • Adapter-based methods: Adapter-based PEFT inserts trainable layers into existing model layers, reducing resource consumption but introducing additional latency.
  • LoRA variants: AdaLoRA dynamically prunes insignificant weights for rank allocation, while DoRA adds learnable magnitudes to low-rank direction adjustments.
  • LoRA improvements: ReLoRA periodically merges adapters into weights, LoRA+ uses different learning rates for the two matrices, and rsLoRA makes output scale invariant to rank.
  • Initialization methods: PiSSA initializes factors to approximate the original weight matrix W, whereas LoRA-GA approximates W’s gradient using sampled gradients and scaling.

3 Methods

LoRA-GA initializes LoRA adapters to approximate the first-step full-finetuning update while preserving stable adapter scales. It combines gradient-direction approximation, rank and scale stability, and memory-efficient layerwise gradient computation.

  • 3.1 Review of Vanilla LoRA: LoRA represents weight changes with trainable low-rank matrices while keeping pretrained weights frozen.The adapter uses A and B with rank r much smaller than the layer dimensions.
  • 3.2 Gradient Approximation: LoRA-GA initializes A and B so the first-step adapter update approximates a scaled full-finetuning weight update.The approximation is formulated by minimizing the Frobenius-norm difference between the adapter and full-model updates.
  • 3.2 Gradient Approximation: The adapter gradients are linear mappings of the effective-layer gradient, which equals the full-finetuning gradient at training initialization.This equality lets LoRA-GA use the initial full-model gradient direction to construct the adapter initialization.
  • 3.3 Scale Stability: The scaling factor ζ is selected to provide forward and backward scale stability across rank and layer dimensions.The method combines this stable scaling with gradient approximation to form the complete LoRA-GA initialization.
  • 3.4 LoRA-GA Initialization: LoRA-GA uses the SVD of the full gradient matrix to initialize A and B from selected singular vectors.Theorem 3.1 provides the initialization for a chosen scaling factor ζ.
  • Implementation: Layerwise backward computation and immediate gradient deletion reduce initialization memory from O(L) to O(1).The method can also use gradient accumulation when the sampled batch is large.

4 Experiments

Experiments across GLUE, MT-Bench, GSM8K, and Human-eval show that LoRA-GA improves performance and convergence over vanilla LoRA while approaching full fine-tuning, with stable gains from its components and low initialization cost.

  • Natural Language Understanding: LoRA-GA consistently outperforms vanilla LoRA and other baselines on T5-Base GLUE experiments, achieving performance comparable to full fine-tuning.
  • Natural Language Understanding: LoRA-GA excels on smaller datasets such as CoLA and MRPC, demonstrating faster convergence and effective use of limited training data.
  • Large Language Model: LoRA-GA outperforms or matches competing methods on Llama 2-7B, surpassing them on GSM8K and Human-eval while remaining competitive on MT-Bench.It requires fewer parameters and approximately 70% of DoRA’s training time.
  • Ablation Study: Higher ranks produce faster loss reduction and curves closer to full fine-tuning, while LoRA-GA remains stable across rank settings.The reported rank experiments include rank=32 and rank=128.
  • Ablation Study: Stable output and gradient approximation each improve performance and convergence, while their combination produces the strongest results and curves closest to full fine-tuning.The ablation compares non-zero initialization, stable output, gradient approximation, and their combination.
  • Memory Costs and Running Time: LoRA-GA requires no extra memory beyond LoRA, and its initialization cost is negligible relative to fine-tuning.In Code-Feedback, initialization took about 1 minute versus approximately 10 hours for training.

5 Conclusions

The paper introduces LoRA-GA, an initialization scheme that approximates full fine-tuning gradients from the first step. Experiments show comparable convergence and similar or superior performance without changing LoRA’s architecture or training algorithm.

  • LoRA-GA initializes LoRA so the low-rank product approximates full fine-tuning gradients from the first training step.
  • LoRA-GA achieves convergence comparable to full fine-tuning with similar or superior performance while modifying only initialization.The method can also be incorporated with other LoRA variants.

A.1 Proof of Theorem 3.1

The proof establishes an initialization whose low-rank update optimally approximates the full gradient, making LoRA’s initial effective-weight gradient equal to full fine-tuning. The supporting analysis connects approximation quality to gradient-matrix rank and observed convergence speed.

  • Gradient equivalence: LoRA adapter gradients are linear mappings of the effective weight gradient.The effective layer is represented as W′ = W0 + ηBA.
  • Gradient equivalence: At initialization, LoRA’s effective-weight gradient equals the full fine-tuning gradient.Both methods have identical outputs and inputs at the beginning of training.
  • Optimal initialization: The proposed initialization selects orthogonal components for Ainit and Binit from the gradient decomposition, with rank governed by the chosen adapter rank.The construction partitions index sets of the singular directions between the two adapter matrices.
  • Optimal initialization: SVD of the full gradient provides the optimal rank-constrained approximation under the Frobenius norm.The theorem uses the singular vectors of ∇W L to construct the initialization.
  • Empirical convergence: LoRA-GA converges significantly faster than vanilla LoRA and ablations, approaching full fine-tuning in the reported training-loss curves.The comparison is based on Figures 3 and 4.
  • Rank and approximation quality: When the gradient matrix is close to rank 2r, its low-rank approximation is expected to be more accurate.Figure 5 is described as showing rapidly declining singular values and concentration of squared singular values in a few ranks.

B.3 Detailed Ablation Study Result of GLUE

The GLUE ablation study finds that LoRA-GA outperforms the alternative ablations, with gains attributable to both its scaling and gradient-approximation components. Additional experiments report robustness across learning rates.

  • Ablation results: LoRA-GA outperforms all other ablation models on the GLUE subset.The full ablation results are reported in Table 6, with average scores summarized in Table 4.
  • Component contributions: Both +SO and +GA improve over vanilla LoRA, indicating positive contributions from both LoRA-GA components.The passage attributes the improvement to the two components jointly evaluated in the ablation.
  • Learning-rate robustness: LoRA-GA maintains strong performance under learning rates 1e-5 and 5e-5.Tables 7 and 8 report the corresponding comparisons across MT-Bench, GSM8K, and Human-eval.

D.1 Experiments on Natural Language Understanding

The T5-Base experiments evaluate LoRA-GA on a GLUE subset under a standardized AdamW training setup. The configuration fixes the adapter rank, target modules, sequence length, batch size, and one training epoch.

  • Training setup: T5-Base experiments use AdamW with learning rate 1e-4, 0.03 warmup, cosine decay, and zero weight decay.The optimizer uses β1 = 0.9, β2 = 0.999, and ϵ = 1e-8.
  • LoRA configuration: The LoRA configuration uses rank r = 8 and α = 16, targeting all linear modules except embeddings, layer normalization, and the language-model head.These settings define the adapter placement and scale for the T5-Base experiments.
  • Training setup: Training uses sequence length 128, batch size 32, one epoch, and FP32 precision.These settings apply to the reported T5-Base experiments.

D.2 Experiment on Large Language Model

The Llama 2-7B experiments evaluate dialogue, mathematical reasoning, and coding with fixed LoRA-GA and generation settings. The related-work discussion distinguishes LoRA-GA’s gradient-based initialization from weight-based SVD initialization in PiSSA.

  • Training setup: Llama 2-7B experiments use AdamW with learning rate 2e-5, 0.03 warmup, cosine decay, and zero weight decay.The optimizer uses β1 = 0.9, β2 = 0.999, and ϵ = 1e-8.
  • Precision: The backbone uses bf16, while LoRA adapter matrices A and B use fp32 during training.This precision setup follows the cited PEFT implementation.
  • Generation: Generation uses top_p = 0.95 and temperature T = 0.8.These parameters govern the reported generated outputs.
  • Related initialization methods: LoRA-GA applies SVD to sampled gradients, whereas PiSSA applies SVD to pretrained weights.The methods therefore use different initialization motivations and scaling strategies.

F Limitations

The paper identifies limits in validation scope, compatibility with other LoRA variants, and checkpoint storage. These constraints bound how broadly the reported results can be generalized and increase storage relative to vanilla LoRA.

  • LoRA-GA was validated on T5-Base and Llama 2-7B, but not on larger models such as Llama 2-70B.The authors attribute this unvalidated scale range to computational resource constraints.
  • The evaluation covered MTBench, GSM8K, and Human-eval, so consistency across all benchmarks is not guaranteed.
  • LoRA-GA was not implemented with orthogonal LoRA variants such as ReLoRA, leaving its compatibility with those architectures untested.
  • LoRA-GA requires double the checkpoint storage of original LoRA because it stores both initial and final adapter checkpoints.

G Compute Resources

The experiments used single-GPU setups, with hardware varying by model and baseline. The paper also connects lower fine-tuning costs to reduced energy use and notes a potential misuse risk.

  • T5-Base experiments used one RTX 3090, while Llama 2-7B full fine-tuning and DoRA used one A100 and other variants used one RTX 3090.
  • LoRA-GA is presented as reducing parameters and computational costs while maintaining comparable performance, and as converging faster than vanilla LoRA.
  • Lower fine-tuning costs could reduce electricity consumption, while easier access to large-model training could broaden participation.
  • The authors caution that easier language-model training could also facilitate models generating fake news or misleading information.
Loading 2407.05000v2…