PaperlayerView exact revision on arXiv

Source-linked AI summary

LoRA+: Efficient Low Rank Adaptation of Large Models

Soufiane Hayou, Nikhil Ghosh, Bin Yu

arXiv:2402.12354v2cs.LGcs.AIcs.CLstat.ML

TL;DR

LoRA uses equal learning rates for its adapter matrices, which the paper finds inefficient for feature learning in wide models. Using scaling arguments, the paper proposes LoRA+ with unequal rates for A and B, and reports improved performance and finetuning speed at the same computational cost.

  • Problem

    Equal learning rates for LoRA’s A and B matrices are suboptimal for feature learning in large-width models.

  • Method

    The paper analyzes infinite-width LoRA dynamics and introduces LoRA+ by assigning different learning rates to A and B.

  • Results

    LoRA+ improves training speed and performance in extensive empirical evaluations at the same computational cost as LoRA.

  • Takeaways & Limitations

    Efficient LoRA finetuning requires decoupling the adapter learning rates, with B generally receiving the larger rate.

  • Takeaways & Limitations

    The optimal ηB/ηA ratio is sensitive to the model and task, which the paper’s analysis does not fully estimate.

Abstract

from arXiv · show

In this paper, we show that Low Rank Adaptation (LoRA) as originally introduced in Hu et al. (2021) leads to suboptimal finetuning of models with large width (embedding dimension). This is due to the fact that adapter matrices A and B in LoRA are updated with the same learning rate. Using scaling arguments for large width networks, we demonstrate that using the same learning rate for A and B does not allow efficient feature learning. We then show that this suboptimality of LoRA can be corrected simply by setting different learning rates for the LoRA adapter matrices A and B with a well-chosen ratio. We call this proposed algorithm LoRA$+$. In our extensive experiments, LoRA$+$ improves performance (1-2 $\%$ improvements) and finetuning speed (up to $\sim$ 2X SpeedUp), at the same computational cost as LoRA.

1. Introduction

Large pretrained models make full finetuning computationally infeasible, motivating parameter-efficient methods such as LoRA. The paper identifies equal adapter learning rates as suboptimal at large width and proposes LoRA+ with theoretically guided, unequal rates.

  • Motivation: Models with tens or hundreds of billions of parameters make full finetuning computationally infeasible for many practitioners.Resource-efficient methods freeze pretrained weights and train only a small set of inserted parameters.
  • Background: LoRA adapts pretrained weights by training only a low-rank matrix update while keeping the original weights fixed.The trainable update uses matrices A and B with rank r much smaller than the matrix dimensions.
  • Problem: Principled guidelines exist for LoRA rank selection, but not for learning-rate selection beyond common choices around 1e-4.The paper frames learning-rate selection as a missing hyperparameter guideline.
  • Contribution: LoRA+ sets different learning rates for the A and B adapter matrices, with the B rate scaled by λ≫1 relative to A.The paper validates this design with empirical results across language models and tasks.
  • Approach: The paper studies infinite-width LoRA dynamics using neural-network scaling arguments to derive learning-rate rules for feature learning.This extends scaling analysis traditionally used to choose initialization and learning-rate parameterizations.

2. Setup and Definitions

The paper formalizes resource-efficient finetuning and LoRA within a general width-n neural-network setup. LoRA freezes pretrained weights and represents each trainable update with a low-rank factorization whose scaling is analyzed as width grows.

  • Network setup: The model consists of input and output embeddings surrounding L hidden layers of width n.The hidden weights are W_l∈R^{n×n}, while the layer maps are F_l.
  • Finetuning setup: Pretrained models can be adapted to downstream tasks by training low-rank weight matrices instead of performing full finetuning.This reduces computational cost for finetuning on devices with limited GPUs.
  • LoRA definition: LoRA represents a pretrained weight update as α/r BA, training only B∈R^{n1×r} and A∈R^{r×n2}.The rank satisfies r≪min(n1,n2), while α and r are tunable constants.
  • Scaling analysis: Scaling analysis studies how initialization, learning rates, and network quantities change as width n approaches infinity.The paper uses asymptotic orders such as O, Ω, and Θ to characterize these dependencies.

3. An Intuitive Analysis of LoRA

The toy-model analysis shows that equal learning rates cannot produce efficient LoRA feature learning at increasing width. Decoupling the rates resolves this scaling problem when A uses Θ(n^-1) and B uses Θ(1).

  • Motivation: A and B have transposed shapes, motivating separate learning-rate analysis for the two LoRA matrices.The analysis focuses on how their updates behave as the model width grows.
  • Toy model: The toy model uses a single-sample linear setup with frozen pretrained weights and LoRA weights a and b.Its dimensions correspond to a rank-one LoRA update, and the analysis assumes input coordinates remain order one with width.
  • Initialization: Nonzero initialization of a requires variance Θ(n^-1) so that a^T x remains stable as width increases.The paper considers Gaussian initialization schemes in which one LoRA parameter is initialized to zero.
  • Scaling diagnosis: η=Θ(n^-1/2) leaves model outputs at order Θ(n^-1/2), contradicting the desired Θ(1) feature updates.The resulting updates to b and a scale as Θ(n^-1/2) and Θ(n^-1), respectively.
  • Equal-rate inefficiency: With a shared learning rate η=Θ(n^c), the analysis proves that both linear update contributions cannot remain Θ(1) for all training steps.Thus equal-rate LoRA cannot efficiently update both adapter factors in the infinite-width limit.
  • Decoupled rates: ηa=Θ(n^-1) and ηb=Θ(1) ensure stable and efficient LoRA finetuning, with all three update terms remaining Θ(1).The larger rate for b remains valid for general rank r and is expected to extend to general neural networks.

4. Stability and Feature Learning with LoRA in the Infinite Width Limit

The paper analyzes LoRA stability and feature learning as model width grows, showing that equal learning rates for A and B are suboptimal. Efficient finetuning instead requires width-dependent, unequal rates, implemented in LoRA+ with a fixed ratio.

  • General architectures: The analysis extends from linear models to general neural architectures, where equal learning rates for A and B yield suboptimal feature learning.The same conclusion holds across general neural networks with LoRA layers.
  • Analysis setup: The analysis treats LoRA rank r as fixed while studying finetuning dynamics in the infinite-width limit.This setup reflects practical cases where r is much smaller than model width n.
  • LoRA features: LoRA features are defined as ZA = AZ and ZB = BZA for a LoRA layer.These quantities support the paper’s analysis of how adapter updates affect network features.
  • Stability: Stability requires Z, ZA, and ZB to remain O(1) as width n approaches infinity.This prevents network quantities from exploding as model width grows.
  • Learning-rate scaling: Efficient LoRA is impossible with ηA = ηB, but is achieved with ηA = Θ(n^-1) and ηB = Θ(1).The result implies an efficient rate ratio ηB/ηA = Θ(n).
  • LoRA+: LoRA+ fixes ηB/ηA and tunes only one learning rate, reducing tuning to a one-dimensional grid search at standard LoRA’s computational cost.This avoids jointly tuning both learning rates as a two-dimensional search.

5. Experiments with Language Models

Across Roberta, GPT-2, and Llama experiments, performance generally improves when ηB is much larger than ηA, especially on harder tasks. The optimal ratio varies by model, task, and initialization, while easy tasks may show little benefit.

  • GLUE tasks with GPT-2 and RoBERTa: ηB ≫ηA consistently maximizes Roberta-base test accuracy relative to learning rates with ηA ≈ηB.The advantage is more pronounced on MNLI and QQP than on SST2 and QNLI.
  • GLUE tasks with GPT-2 and RoBERTa: ηB ≫ηA also yields maximal GPT-2 test accuracies, with larger gaps over ηA ≈ηB on harder tasks.
  • Llama: 1.3% gain: Llama-7b MMLU accuracy improves with ηB ≫ηA compared with the optimal ηB = ηA.The MMLU experiment trains on a 100k flan-v2 subset for one epoch.
  • Llama: For Llama-7b on MNLI, ηB = ηA is nearly optimal, consistent with the task requiring less efficient feature learning.The model is trained for one epoch because MNLI is relatively easy for Llama.
  • How to set LoRA+ Ratio?: The best ηB/ηA ratio is model- and task-sensitive and exhibits significant variance across top-performing runs.The ratio distribution summarizes the interquartile range and median of the top four learning-rate pairs.
  • How to set LoRA+ Ratio?: Initialization changes the preferred ratio: approximately 24 for Roberta with Init[2], but roughly 22-23 with Init[1], while Llama favors roughly 21-22.

6. Conclusion and Limitations

The paper concludes that equal learning rates make LoRA inefficient and that LoRA+ improves speed and performance by using different adapter learning rates. The optimal ratio remains task- and model-dependent.

  • LoRA+ sets different learning rates for the adapter matrices and improves both training speed and performance over standard LoRA.The benefits are reported as more significant on harder tasks such as MNLI and MMLU.
  • The optimal ηB/ηA ratio should account for task and model dependence, which the paper leaves for future work.

Impact Statement

The impact statement frames LoRA+ as a way to reduce the resources needed to customize pretrained large language models while improving finetuned-model performance.

  • LoRA+ aims to speed LoRA fine-tuning while improving the performance of finetuned pretrained large language models.
  • The reported speed-up can save computation resources when pretrained models are customized for downstream tasks.

A.3. Proof of Proposition 1

The proposition proves that standard LoRA cannot achieve efficient feature learning under a shared learning rate as width-scaled training dynamics are analyzed.

  • The contradiction implies that LoRA fine-tuning is inefficient under the proposition’s initialization and shared-learning-rate setup.
  • Efficient LoRA training would require a learning rate scaling as η = Θ(n−1/2), but this requirement leads to a contradiction.

A.4. Proof of Proposition 2

Proposition 2 identifies unequal learning-rate scaling that keeps LoRA fine-tuning efficient: ηa = Θ(n−1) and ηb = Θ(1).

  • ηa = Θ(n−1) and ηb = Θ(1) ensure efficient LoRA fine-tuning in the toy model.The proposition states this holds for all t > 1 and i ∈ {1, 2, 3}.
  • The proof derives the necessary scaling relation γ[ηa] + γ[ηb] = −1.
  • The recursive formulas track how the learning-rate exponents determine the scaling of b_t and a_t⊤x across updates.
  • Setting γ[ηb] = 0 and γ[ηa] = −1 satisfies the relation and preserves the required Θ(1) feature-update scaling.

A.5. Proof of Theorem 1

Theorem 1 applies infinite-width scaling analysis to Adam-trained LoRA and concludes that equal learning rates are inefficient, while ηA = Θ(n−1) and ηB = Θ(1) are efficient.

  • The analysis assumes a condition on the processed gradient gA and studies LoRA dynamics as width n grows.
  • Theorem 1 states that ηA = ηB cannot achieve efficient LoRA fine-tuning under the stated Adam assumption.
  • ηA = Θ(n−1) and ηB = Θ(1) yield efficient LoRA fine-tuning under the same theorem assumptions.
  • The proof derives γ[ηA] + γ[ηB] = −1 from scaling constraints on the adapter updates.
  • Equal learning rates imply γ[η] = −1/2 and violate the stability condition, whereas the asymmetric choice preserves the target scaling.
  • The loss analysis links Θ(1) feature changes to more efficient loss reduction because feature-update directions correlate with d̄Zt−1.
  • The experiments section reports supplementary setup details for empirical scaling checks, including MLP and pretrained-model configurations.

C.2. Results of Roberta-base Finetuning on all Tasks

The Roberta-base appendix complements the main finetuning results with test/train accuracy, showing a small divergence between the learning rates that optimize each.

  • Figure 8 reports test/train accuracy for all GLUE tasks with Roberta-base.
  • The learning-rate choice maximizing test accuracy differs slightly from the choice maximizing train accuracy.
  • The passage attributes this small difference to mild overfitting during finetuning.

C.3. Results of GPT2 Finetuning on all Tasks

The GPT2 appendix extends finetuning comparisons across tasks, precisions, losses, ranks, and initializations, while showing that optimal learning-rate ratios depend on initialization.

  • Figure 9 extends GPT2 finetuning results beyond MNLI and QQP to additional tasks using the same setup as Figure 4.
  • Figures 11–13 vary precision and evaluation target by using full precision or 100× test/train loss instead of accuracy.
  • Figures 14–16 examine rank settings r = 4 and r = 16 for Roberta-base and r = 4 for GPT2.
  • The optimal ratio λ = ηB/ηA is generally smaller with Init[1] than with Init[2].
  • Llama-7b flan-v2 results report possible one-epoch overfitting and no observed ηB > ηA optimum for train loss in the tested grid.
  • Figures 18–19 examine Llama-7b initialization effects and train/test loss on MNLI.
Loading 2402.12354v2…