PaperlayerView exact revision on arXiv

Source-linked AI summary

A Rank Stabilization Scaling Factor for Fine-Tuning with LoRA

Damjan Kalajdzievski

arXiv:2312.03732v1cs.CLcs.LG

TL;DR

Large language model fine-tuning is computationally demanding, and conventional LoRA scaling can cause higher-rank adapters to learn slowly or collapse. The paper analyzes this scaling choice and introduces rsLoRA, showing that square-root rank scaling stabilizes learning and enables higher-rank performance gains without increasing inference cost.

  • Problem

    Conventional LoRA’s rank scaling causes gradient collapse and limits the benefits of higher-rank adapters during fine-tuning.

  • Method

    The paper analyzes LoRA scaling through learning-trajectory theory and introduces rsLoRA, which scales adapters by the inverse square root of rank.

  • Results

    rsLoRA remains stable and non-collapsing at very large ranks, allowing better fine-tuning performance while preserving LoRA’s inference-cost profile.

  • Takeaways & Limitations

    Higher ranks can use additional training resources to improve fine-tuning performance, while the merged adapters keep inference computation unchanged.

  • Takeaways & Limitations

    The theorem does not establish how feature quality varies across ranks when learning is stable, and its practical implications require experimental validation.

Abstract

from arXiv · show

As large language models (LLMs) have become increasingly compute and memory intensive, parameter-efficient fine-tuning (PEFT) methods are now a common strategy to fine-tune LLMs. A popular PEFT method is Low-Rank Adapters (LoRA), which adds trainable low-rank "adapters" to selected layers. Each adapter consists of a low-rank matrix product, multiplicatively scaled by a rank-dependent factor. This scaling factor, which divides adapters by a factor of the rank, results in slowed learning and stunted performance for LoRA with higher-rank adapters. Consequently, the use of LoRA in practice has generally been limited to very low ranks. In this work, we study the impact of the scaling factor on the learning process and prove that LoRA adapters should be divided by a factor of the square root of the rank. Modifying LoRA with the appropriate scaling factor, which we call the rank-stabilized LoRA (rsLoRA) method, easily provides for a fine-tuning compute/performance trade-off, where larger ranks can be used to trade off increased computational resources during training for better fine-tuning performance, with no change in inference computing cost.

1. Introduction

Parameter-efficient fine-tuning reduces the trainable parameter set for computational efficiency, with LoRA adding scaled low-rank adapters to selected pretrained layers. This paper argues that conventional rank scaling limits higher-rank LoRA, while rsLoRA stabilizes learning and enables better performance without extra inference cost.

  • Fine-tuning large language models often requires substantial computational resources despite their broad capabilities.
  • LoRA adds trainable low-rank matrix-product adapters, multiplied by a scaling factor, to selected pretrained parameter matrices.
  • Conventional LoRA divides adapters by rank, causing gradient collapse and slowed learning as rank increases.
  • rsLoRA divides adapters by the square root of rank, preserving gradients and allowing higher ranks to trade increased training compute for better fine-tuning performance.
  • Because rsLoRA adapters retain LoRA’s form and can be merged into pretrained weights, changing rank does not change inference computational cost.

2. Background and Relevant Works

The background introduces parameter-efficient adapters and LoRA’s parallel low-rank formulation, which can be merged into pretrained weights for efficient inference. It identifies rank-dependent scaling as a problem because conventional scaling makes larger ranks learn no better than small ones, and motivates an analytical scaling study.

  • 2.1. Low-Rank Adapters (LoRA): LoRA computes low-rank adapters in parallel with transformer submodules so they can be combined with pretrained parameters after fine-tuning.
  • 2.1. Low-Rank Adapters (LoRA): A LoRA adapter augments a pretrained linear module with low-rank matrices A and B and a rank-dependent scaling factor γr.
  • 2.1. Low-Rank Adapters (LoRA): After fine-tuning, W + γrBA replaces W, so the merged model incurs no additional inference compute cost.
  • 2.1. Low-Rank Adapters (LoRA): AdaLoRA dynamically prunes singular values to adjust adapter rank during training, while retaining LoRA’s scaling factor.
  • 2.1. Low-Rank Adapters (LoRA): Using γr = α/r makes larger-rank LoRA learning little different from low-rank learning, despite the extra training cost.
  • 2.2. Scaling-Initialization-Update Schemes: Scaling-initialization-update schemes analyze how parameter initialization, learning rates, and scaling choices jointly affect learning dynamics.

3. rsLoRA: Rank-Stabilized Adapters

The paper defines rank stabilization through rank-independent adapter outputs and gradients, then analyzes the infinite-rank limit to derive the required scaling. It concludes that γr must scale as 1/√r for stable learning, while acknowledging that the theorem alone does not establish rank-dependent feature quality.

  • A rank-stabilized adapter preserves Θr(1) output moments for iid inputs and Θr(1) backpropagated gradients across ranks.
  • The analysis studies the infinite-width limit r →∞ to derive a scaling factor that stabilizes LoRA learning.
  • The theorem identifies γr ∈ Θr(1/√r) as the scaling required for rank-stabilized adapters; other asymptotic scalings produce unstable or collapsing learning at sufficiently large ranks.
  • The analysis applies along the learning trajectory under its assumptions, with gradient magnitudes maintained at Θr(1) through successive adapters.
  • The theorem does not determine whether features learned at different ranks vary in quality when learning remains stable.

4. Experimental Results

Experiments show that rsLoRA improves both rank-scaled fine-tuning performance and learning stability, whereas conventional LoRA exhibits little rank benefit and collapsing gradients at higher ranks. These patterns persist across optimizers, models, datasets, and adapter configurations.

  • The experiments fine-tune Llama 2 on 20,000 OpenOrca examples using AdamW and adapters in all linear attention and feed-forward MLP sub-modules.
  • rsLoRA unlocks improved fine-tuning performance with higher ranks, while LoRA shows little performance difference across ranks.LoRA can even perform slightly worse at larger ranks, whereas rsLoRA benefits from increased rank.
  • LoRA has collapsing gradients with higher ranks, whereas rsLoRA maintains comparable gradient norms across ranks during training.For rsLoRA, gradient norms are equal across ranks at training onset and remain approximately within the same order of magnitude.
  • The same performance pattern appears with a different pretrained model, dataset, and optimizer.
  • The same gradient-norm stability pattern appears with SGD, controlling for stabilization effects from adaptive optimizers.
  • Rank-4 LoRA cannot match high-rank rsLoRA at the default learning rate, indicating that rsLoRA’s gains are not merely a learning-rate boost.
  • Attention-only adapters give similar results, while rank-correcting initialization without reparameterization becomes unstable at larger ranks.

5. Conclusion

The paper derives and validates a rank-correcting scaling factor for LoRA adapters. The resulting rsLoRA method remains stable at large ranks and enables better fine-tuning performance, while its application to AdaLoRA remains future work.

  • The proposed rank-stabilized adapters remain stable and non-collapsing at very large ranks, enabling better performance with larger adapters.
  • Using the maximum rank allowed by memory can provide the best fine-tuning performance under the paper’s supported setting.
  • Applying the correction factor to AdaLoRA is identified as future work because AdaLoRA currently uses conventional LoRA scaling.

A. Proof of theorem 3.2

The theorem identifies γ_r ∈ Θ_r(1/√r) as the scaling required to prevent unstable or collapsing learning at large rank. The proof analyzes adapter outputs and gradients across the learning trajectory under stated initialization and input assumptions.

  • Theorem setup: The theorem assumes B is initialized to zero, A has rank-independent iid variance, and adapter scaling is γ_rBA.These assumptions define the initialization and parameterization used in the theorem.
  • Forward pass: Adapter outputs have Θ_r(1) moments exactly when the scaling preserves their magnitude under the stated iid input assumptions.The forward-pass analysis evaluates moments of adapter outputs over initialization.
  • Backward pass: Loss gradients entering the adapter input have Θ_r(1) moments exactly when the scaling preserves gradient magnitude under the stated assumptions.The backward-pass analysis applies the same rank-scaling criterion to loss gradients.
  • Conclusion: γ_r ∈ Θ_r(1/√r) is necessary to avoid unstable or collapsing learning as r becomes large.The result applies at any point in the learning trajectory when the theorem’s assumptions hold.
  • Proof strategy: The proof tracks SGD updates of A and B through f(x)=γ_rBAx and the resulting loss gradients.The derivation starts from the zero initialization of B and follows the parameters after updates.

B. Ablations and additional experiments

With SGD, rsLoRA maintains comparable gradient norms across ranks, whereas conventional LoRA exhibits gradient collapse as rank increases. The experiment uses ranks 4 through 2048 and a stable learning rate of 0.0001.

  • Stability with SGD: Ranks 4, 8, 32, 128, 512, and 2048 are compared using average parameter gradient norms for LoRA and rsLoRA.The experiment uses SGD with a high but stable learning rate of 0.0001.
  • Stability with SGD: LoRA gradients collapse as rank increases, producing rank-dependent gradient magnitudes.This is the instability pattern measured in the gradient-norm comparison.
  • Stability with SGD: rsLoRA has the same gradient norm across ranks at initialization, with norms remaining within the same order of magnitude during training.With SGD, all ranks remain near the same norm for about 100 training steps.

B.2. Change of Model/Optimizer/Dataset

The findings generalize to GPT-J fine-tuned on GSM8k with Adafactor: increasing rank leaves LoRA performance essentially unchanged but unlocks additional fine-tuning performance with rsLoRA.

  • Change of Model/Optimizer/Dataset: The experiment trains the 6-billion-parameter GPT-J on GSM8k using the adaptive optimizer Adafactor.This changes the model, optimizer, and fine-tuning dataset from the primary setting.
  • Change of Model/Optimizer/Dataset: For ranks 4, 8, 32, 128, 512, and 2048, larger-rank LoRA training does not change performance, with curves overlapping into a single run.The comparison uses the same rank range for LoRA and rsLoRA.
  • Change of Model/Optimizer/Dataset: The same rank range unlocks additional fine-tuning performance with rsLoRA.The figure reports this pattern in the GPT-J, GSM8k, and Adafactor setting.
  • Change of Model/Optimizer/Dataset: rsLoRA produces similar parameter gradient norms across ranks, while LoRA produces norms differing by orders of magnitude.The LoRA gradient disparity is associated in the passage with collapsed learning at higher ranks.

B.3. Scaling at Initialization Only

Scaling only the initialization does not reproduce the benefits of rank-stabilized LoRA. Without adapter scaling, high-rank training becomes unstable; retaining LoRA scaling remains sub-optimal.

  • Scaling at Initialization Only: The ablation compares ranks 4 and 2048 while scaling only matrix A’s initialization by 1/√r.It tests whether rank-correct initialization alone is sufficient.
  • Scaling at Initialization Only: With no adapter scaling factor, rank-2048 training becomes unstable despite matching gradient-norm magnitudes at initialization.The instability appears during training rather than at initialization.
  • Scaling at Initialization Only: Adding the conventional LoRA adapter scaling factor to the scaled initialization produces the same sub-optimal high-rank learning as LoRA.The result shows initialization alone cannot correct the learning behavior.

B.4. Learning Rate for Low Rank is Not Enough

A learning-rate sweep shows that low-rank LoRA cannot match high-rank rsLoRA, while attention-only adapters reproduce the broader rank-dependent behavior across methods.

  • 1.863 perplexity is the best result for rank 4 LoRA, versus 1.836 perplexity for rank 2048 rsLoRA at the default learning rate.The rank-4 LoRA result comes from sweeping learning rates, so the gap is not explained by an unfavorable low-rank learning rate.
  • Higher-rank rsLoRA trains better, whereas higher-rank LoRA does not, and LoRA gradient norms still collapse at large ranks.
  • The attention-only adapter ablation produces effects similar to experiments using adapters in all MLP and attention modules.
  • Figure 7 varies ranks from 4 to 2048, using blue-green colors for rsLoRA and copper colors for LoRA, with brighter shades indicating higher ranks.The gradient-norm plot shades the range between each method’s minimum and maximum rank.
  • The experiments also examine scaling factors larger or smaller than the rsLoRA and LoRA choices.

B.7. Activations

Averaged activation moments remain small and stable for both LoRA and rsLoRA, although high-rank rsLoRA moments level off later during training.

  • Both LoRA and rsLoRA exhibit relatively stable, non-collapsing averaged first and second activation moments.The moments are measured for post-adapter, pre-LayerNorm activations and averaged over adapters.
  • High-rank rsLoRA at ranks 512 and 2048 shows both activation moments leveling off later in training.
  • The averaged moments may conceal effects because of LayerNorm, small activation magnitudes, or averaging across activations and layers.The paper does not examine these possible explanations further.
Loading 2312.03732v1…