The Impact of Initialization on LoRA Finetuning Dynamics

We appreciate the reviewer's feedback, but we respectfully disagree with several points and believe there are significant misunderstandings in their assessment. We address these below:

1. Initialization schemes: The reviewer suggests investigating "different initialization schemes aside from init[A] and init[B]," citing Gaussian, Kaiming, and principal component initializations. We believe this is a fundamental misunderstanding of our work. Init[A] and Init[B] in our setup already use Kaiming (Gaussian) initialization. The distinction lies in whether A or B is set to zero to ensure $BA=0$ at the start. More generally, our method covers (as stated in Footnote 2 in page 3) general random distributions (not only Gaussian).

2. Connection to prior work: The reviewer states, "In [1] only different learning rates for A and B have been investigated." This is precisely our point - [1] focused on learning rates, while our work addresses the unexplored impact of initialization schemes in vanilla LoRA. While our paper uses similar tools to those used in [1], to study the learning dynamics in the large width limit, our results are orthogonal to those of [1]. In this paper, we study the impact of the initialization scheme in vanilla LoRA as introduced in [2] (same learning rate for A and B), while in [1], the authors study the impact of setting different learning rates for A and B. Moreover, note that the tools used in our paper and also in [1] are not new and are based on the theory of infinite-width networks. For instance, the same machinery was used in the Tensor Programs series (see e.g. [10]). We decided to use the $\gamma$ notation introduced in [1] because it simplifies the message for practitioners. However, as we explained, our results are orthogonal to those of [1]. We will add this discussion to the revised version of the paper.

3. Contribution: The reviewer claims our findings have "marginal" impact because init[A] is already commonly used. We respectfully agree with this argument for the following reasons: a) It ignores the scientific value of providing the first rigorous explanation for an empirical practice. b) It overlooks the importance of validating existing methods, which is crucial in scientific research. c) It fails to recognize that our work prevents potential future misapplications of init[B], which could lead to suboptimal results. d) it judges our paper as a 'method' paper where the main contribution is the introduction of a novel method. This is not the case. Our paper studies theoretically and empirically an existing method, and shows which initialization is better.

4. Scope: Multiple layer analysis presents challenges. With two or more layers, changes in LoRA features $Z_B = BAz$ are not only due to updates of $A$ and $B$ but also changes in $z$ via previous LoRA layers. This requires a more refined, layerwise definition of efficient feature learning. While the fundamental analysis should remain similar, this adds unnecessary complexity to the setup. Importantly, our empirical results with multiple LoRA layers confirm that our findings hold in these more complex scenarios.

[1] Hayou et al., LoRA+: Efficient low rank adaptation of large models., ICML 2024

[2] Hu et al., LoRA: Low-Rank Adaptation of Large Language Models, ICLR 2022

[10] Feature Learning in Infinite-Width Neural Networks (2021). Yang and Hu.

We believe our work provides valuable, novel insights into LoRA dynamics, filling a gap in the literature. We are surprised by the low score assigned to our paper, and we hope this clarification addresses the misunderstandings raised above.

We thank the reviewer for their positive and constructive comments. We address their main questions below:

1. Sensitivity to LoRA rank: Our results hold for different rank values. For LoRA rank $r$, we used two primary values: $r=8$ (for RoBERTa) and $r=64$ (for LLama). We also conducted limited experiments with $r=4$ for RoBERTa. However, we chose to allocate more compute resources to $r=8$ to increase the number of random seeds. We designed these experiments to maximize the amount of useful empirical results within our compute budget constraints. All our experiments were conducted using the LoRA+ codebase (available on GitHub) and can be easily reproduced. To replicate vanilla LoRA, one can set the lora_plus_ratio to 1 (ensuring the same learning rate for both A and B) and choose between Init[A] or Init[B] using the use_original_lora_init argument.

2. Interaction with LoRA variants: We appreciate this interesting question. While we haven't conducted a theoretical analysis of initialization's impact on advanced LoRA variants (e.g., QLoRA and LoRA+), we can provide some intuition: For QLoRA, we expect similar results since the only difference from LoRA is the quantization step. Our analysis should remain fundamentally the same. For LoRA+, the situation is more complex as it involves choosing different learning rates for $A$ and $B$. While we can't definitively state the outcomes, our preliminary results suggest that the optimal ratio in LoRA+ is affected by the choice of initialization (init[A] vs init[B]).

3. Instability of Init[A] in practice: As specified in the paper, with Init[A], LoRA "internal" features ($Az$) grow at most as $\Theta(n^{1/2})$ with respect to width (embedding dim). This growth, although potentially problematic, is slow (in $n$) and should only become an issue for extremely large widths. In practical settings (e.g., $n \approx 10^3$ for LLaMA 3.1 405B), this is not significantly problematic. The constant in the $\Theta(.)$ growth term depends on the model and downstream task, which affects this growth term.

4. Generalization to multiple layers: Multiple layer analysis presents challenges. With two or more layers, changes in LoRA features $Z_B = BAz$ are not only due to updates of $A$ and $B$ but also changes in $z$ via previous LoRA layers. This requires a more refined, layerwise definition of efficient feature learning. While the fundamental analysis should remain similar, this adds unnecessary complexity to the setup. Importantly, our empirical results with multiple LoRA layers confirm that our findings hold in these more complex scenarios.

We thank the reviewer for their positive comments. We provide some details that explain the scope and contributions of our paper.

1. Scope and contributions: As correctly noted, in low-rank adaptation (LoRA), we generally aim to initialize the model such that $BA=0$. This naturally leads to two possible initialization schemes: either initialize $A$ randomly and $B$ to zero (init[A] in our paper), or vice versa (init[B]). Prior to our study, practitioners had no clear guidance on which scheme to use, often defaulting to the one implemented in their chosen package (e.g., init[A] in the PEFT package). Our paper provides the first definitive answer to this question: init[A] is generally superior to init[B] because it enables finetuning with larger learning rates, resulting in more efficient learning.

2. Issue at extreme width: The reviewer correctly points out that both init[A] and init[B] become problematic in extreme scenarios where the width is exceptionally large (large enough that a value of order $n^{1/2}$ causes numerical instability). However, this is an extreme regime, and current state-of-the-art models have not yet reached this threshold.

We hope this clarifies the scope and significance of our paper.

We thank the reviewer for their constructive feedback. We believe there are some misunderstandings and we take this oppotunity to address them.

1. Connection to other LoRA variants: We emphasize that we do not introduce a new method in this paper; rather, we study the finetuning dynamics of the original LoRA method [1] for two natural initializations: init[A] (set A to random and B to zero) and init[B] (set B to random and A to zero). The impact of initialization on other LoRA variants (e.g., LoRA+, VeRA, AdaLoRA) is an interesting question but outside the scope of this paper, as stated in our conclusion and limitations section. While we haven't conducted a theoretical analysis of initialization's impact on advanced LoRA variants (e.g., QLoRA, VeRA, LoRA+, AdaLoRA), we can provide some intuition: For QLoRA, we expect similar results since the only difference from LoRA is the quantization step. Our analysis should remain fundamentally the same. For LoRA+, the situation is more complex as it involves choosing different learning rates for A and B. While we can't definitively state the outcomes, our preliminary results suggest that the optimal ratio in LoRA+ is affected by the choice of initialization (init[A] vs init[B]). As the reviewer might guess from this intuition, the impact of initialization on advanced LoRA methods will depend on the method itself. We will add this discussion in the revised version.

2. Intuition on why Init[A] outperforms init[B]: Assume LoRA rank is $r$ and the embedding dimension is $n$. LoRA features are given by $BA z$ for input vector $z \in \mathbb{R}^n$ (LoRA input). With init[A], B is set to zero, allowing the magnitude of $Az$ to increase significantly without causing a blow-up in LoRA features $BAz$, as the small magnitude of B compensates for the growth in $Az$. We show in Thm 1 that the maximal learning rate (that does not cause features $BAz$ to explode) scales as $\Theta(n^{-1/2})$ in width. This is not possible with init[B], where B's magnitude is always $\Theta(1)$ due to random initialization with scaled variance, in which case the maximal learning scales as $\Theta(n^{-1})$ (Thm 2). We will add this discussion to the revised version.

3. AdamW vs Adam (impact of weight decay): We ran sweeps with weight decay $wd \in (0, 0.05, 0.1)$ (0.1 used in the original LoRA paper). This hyperparameter did not significantly impact the results for common $r$ values (8, 64), and our conclusion that Init[A] outperforms Init[B] holds for different $wd$ values. With $wd=0$, we reproduced the results of [1] within a statistically sound confidence interval (recall that in [1], the authors used wd=0.1). We expect $wd$'s impact to be more significant for large (though generally impractical) $r$ values due to overparameterization. We will add these details to the revised version.

4. Fixed hyperparameters' values: Our results are insensitive to $wd$ choices within the tested range. We used LoRA ranks $r=8$ (RoBERTa) and $r=64$ (LLama). For train batch size, we tested 4 and 16, selecting the best (4 for RoBERTa, 16 for LLama), aligning with Hu et al. (2021). This setup maximizes useful empirical results within our compute budget. While larger sweeps and more models would be ideal, our experiments confirm our theoretical results beyond a trivial statistical threshold.

5. Toy model setup: In statistics, the teacher-student model is a simplified framework for studying learning. It involves a "teacher" model generating data and a "student" model learning from it. Both models are of the same nature, making learning feasible. This setup allows analytical study of learning dynamics and generalization, providing insights into more complex machine learning scenarios. It's particularly useful for understanding phenomena like overfitting, learning curves, and the relationship between model complexity and sample size.

[1] Hu et al. (2021). LoRA: Low-Rank Adaptation of Large Language Models

We hope these answers clarify any misunderstandings. We're happy to address further questions.