The Impact of Initialization on LoRA Finetuning Dynamics

Abstract: In this paper, we study the role of initialization in Low Rank Adaptation (LoRA) as originally introduced in Hu et al. (2021). Essentially, to start from the pretrained model as initialization for finetuning, one can either initialize B to zero and A to random (default initialization in PEFT package), or vice-versa. In both cases, the product BA is equal to zero at initialization, which makes finetuning starts from the pretrained model. These two initialization schemes are seemingly similar. They should in-principle yield the same performance and share the same optimal learning rate. We demonstrate that this is an incorrect intuition and that the first scheme (initializing B to zero and A to random) on average yields better performance compared to the other scheme. Our theoretical analysis shows that the reason behind this might be that the first initialization allows the use of larger learning rates (without causing output instability) compared to the second initialization, resulting in more efficient learning of the first scheme. We validate our results with extensive experiments on LLMs.

• The paper demonstrates that the initialization scheme Init[#1]{A} significantly outperforms Init[#1]{B} by allowing larger learning rates and more effective feature learning.
• The authors combine rigorous theoretical analysis in the infinite-width limit with empirical validations on synthetic models, GLUE benchmarks, and LLMs.
• The findings offer practical insights for applying optimal low-rank adaptation strategies to enhance finetuning efficiency in large-scale neural networks.

The Impact of Initialization on LoRA Finetuning Dynamics

Introduction

The pretrain-finetune paradigm has become a fundamental approach in deep learning, particularly with the emergence of large-scale models. LLMs such as GPT and BERT are pretrained on vast amounts of data to learn general features, which can then be quickly adapted to specific tasks using finetuning. The Low Rank Adaptation (LoRA) method is a crucial player in this context, enabling efficient parameter finetuning by introducing low-rank updates instead of updating the entire model, addressing computational feasibility issues. However, the role of initialization in LoRA finetuning remains underexplored. This paper investigates this aspect rigorously.

LoRA and Initialization Schemes

LoRA, by design, updates a pretrained model using low-rank matrices $A$ and $B$. The typical strategy initializes one of these matrices to zero and the other randomly, theoretically ensuring that the product $BA$ is initially zero, thereby retaining the pretrained model's performance. This paper focuses on two initialization schemes:

• Init[#1]{A}: Matrix $B$ initialized to zero and $A$ to random.
• Init[#1]{B}: Matrix $A$ initialized to zero and $B$ to random.

The research challenges the intuition that both schemes should perform similarly, by demonstrating, both theoretically and experimentally, that Init[#1]{A} outperforms Init[#1]{B}.

Theoretical Analysis

Stability and Learning Dynamics

The paper grounds its analysis in the infinite-width limit, a principle that allows studying neural network behavior at large scales. Using the $\gamma$-operator to track the polynomial dependence of key training variables on the network width $n$, the study derives critical insights:

• Init[#1]{A}: This scheme allows for a larger learning rate, scaling as $\Theta(n^{-1/2})$. It achieves efficient feature learning but induces internal instability in the matrix $B$0. This trade-off implies that while $B$1 becomes large, the product $B$2 remains stable.
• Init[#1]{B}: This configuration limits the learning rate to $B$3, resulting in suboptimal feature learning since the weight matrix $B$4 is undertrained. Although this avoids internal instability, it constrains performance improvements.

These outcomes point to a fundamental difference in the finetuning dynamics between the two initialization strategies, with Init[#1]{A} offering a balance between feature learning efficiency and numerical stability.

Empirical Validation

The theoretical findings are substantiated through experiments on synthetic models and real-world datasets. For instance:

• Synthetic Model: By training a simple model with varying widths, the study observes that Init[#1]{A} allows for higher learning rates and results in lower training loss compared to Init[#1]{B}.
• GLUE Benchmark: Finetuning RoBERTa-large on MNLI, SST2, and QNLI tasks, Init[#1]{A} consistently shows better test accuracy and benefits from higher optimal learning rates.
• LLMs: Finetuning TinyLlama and Llama-7b models on datasets like WikiText-2, Flan-v2, and GSM8k, verifies that Init[#1]{A} outperforms Init[#1]{B} in terms of perplexity and accuracy, affirming the theoretical predictions.

Practical Implications

The findings have significant practical implications. Deploying Init[#1]{A} in LoRA finetuning can enhance model performance efficiently by leveraging larger learning rates, vital for handling modern LLMs. This zero-cost adjustment can be immediately implemented, leading to notable improvements in various NLP tasks.

Future Directions

The research opens avenues for future work to address the identified suboptimal limitations of current LoRA methods. Approaches such as LoRA$B$5, which calibrates learning rates for $B$6 and $B$7 separately, could be coupled with optimal initialization strategies to further boost performance. Additionally, extending the analysis to consider the interaction between data characteristics and finetuning dynamics could provide deeper insights into model generalization and robustness.

Conclusion

This paper illuminates the nuanced role of initialization in LoRA finetuning, revealing that properly chosen schemes like Init[#1]{A} can drive significant improvements in model performance. These insights not only enhance our understanding of parameter-efficient finetuning methods but also provide actionable strategies to optimize training dynamics in large-scale neural networks.