The Impact of Initialization on LoRA Finetuning Dynamics

Authors: soufiane  hayou
nikhil  ghosh
bin  yu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 11:28:33

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, one can either initialize $B$ to zero and $A$ to random, 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 (of initializing $B$ to zero and $A$ to random) on average in our experiments 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.

Main paper: 
 introduction: 
 The pretrain-finetune paradigm (e.g., [7, 9] ) has revolutionized deep learning, replacing task-specific models trained from scratch with finetuning of pretrained base models. These base models, trained on generic unsupervised objectives, learn powerful features that can be rapidly adapted to downstream tasks. The most effective models are consistently the largest ones [14, 25] , with state-ofthe-art models reaching hundreds of billions of parameters. While many such models are openly available (e.g., Llama by Touvron et al. [38] ), full finetuning remains computationally prohibitive for most practitioners. This has led to parameter-efficient finetuning methods, including adapters [11] , prompt tuning [20] , and (IA) 3 [24] . Low Rank Adaptation (LoRA) [19] has emerged as a leading parameter-efficient method, training only low-rank adapter matrices added to pretrained weights, typically using Adam [3] . LoRA often matches or exceeds full-finetuning performance [35, 39] , though it may underperform on complex generation tasks. While prior work has examined rank [31] and learning rate [44] hyperparameters, initialization schemes remain understudied. This work provides experimental and theoretical justification for choosing between seemingly equivalent initialization approaches.In standard LoRA training, one of the two LoRA matrices is initialized with random values and the other is initialized to zero (see Section 2.1). Recently, in Meng et al. [48] the authors proposed an alternative initialization scheme to LoRA which uses the top singular vectors of the pretrained weights as opposed to a random initialization and showed improved training on several tasks. To further improve LoRA training with quantization, Li et al. [34] introduced a new method called LoftQ for computing a better initialization for quantized training [27] . However, to the best of our knowledge, there has not been any study concerning the random initialization in vanilla LoRA. Specifically, it is not clear from prior work which of the two LoRA matrices should be initialized to be zero. Empirical results by Zhu et al. [50] suggested that the two initialization schemes mentioned above yield similar performance, but it is not clear if the learning rate was well-tuned for each initialization scheme. Our findings suggest that these two initialization schemes lead to fundamentally different finetuning dynamics, and that one of these schemes generally yields better result compared to the other. LoRA Variations. Beyond altering the LoRA initialization scheme, there have been a series of works which try to address limitations of vanilla LoRA using different variations. To further reduce the number of trainable parameters, LoRA-FA [42] freezes the A matrix which leads to small performance loss while reducing memory consumption by up to 1.4×. The performance of this training scheme is also investigated in Zhu et al. [50] . VeRA [33] freezes random weight tied adapters and learns vector scalings of the internal adapter activations. LoRA-XS [43] initializes the A and B matrices using the SVD of the pretrained weights and trains a low-rank update of the form BRA where R is a trainable r × r matrix and B, A are fixed. NOLA [32] parametrizes the adapter matrices to be linear combinations of frozen random matrices and optimizes the linear coefficients of the mixtures. VB-LORA [46] shares adapter parameters using a global vector bank. In order to improve the learning ability for more challenging finetuning tasks, Kalajdzievski [31] proposes a scaling rule for the scalar adapter multiplier to unlock increased gains with higher adapter ranks. MoRA [45] learns high-rank updates while still preserving parameter efficiency by applying hand-designed compress and decompress operations before and after a trainable adapter matrix. DoRA [47] decomposes the pretrained weight into magnitude and direction components to allow for better training dynamics.Contributions. We study the impact of Initialization in LoRA through a theory of large width for neural networks. The core approach is to take the width of a neural network to infinity and determine how the behavior of the limit depends on the choice of the hyperparameters, such as the learning rate and initialization. This approach allows to derive principled scaling choices for these hyperparameters such that desired properties (e.g. stable feature learning) are achieved as the network size grows (see Appendix A.2 for more details). Examples of the infinite-width limit include works on initialization (e.g. He et al. [4] ), and training dynamics (e.g. [21] ). Examples for the depth limit include initialization strategies [6, 10, 30] , and depth scaling (see e.g. [18, 23, 28, 29, 37, 41] ). A similar strategy was used to derive scaling rules for the LoRA learning rate in Hayou et al. [44] (LoRA+) that concluded that the learning rates for different LoRA matrices should be scaled differently to ensure optimal feature learning. In this work we use the same approach to provide a systematic comparison between two different random initialization schemes for vanilla LoRA finetuning (using the same learning rate for the A and B matrices). Using the notation Init[A] to refer to the case where A is initialized to random and B to zero (as in [19] ) and Init[B] for the opposite, we show that Init[A] and Init[B] lead to fundamentally different training dynamics (as shown in Figure 1 ):  
setup and definitions: 
 We consider a general neural network model of the formEQUATIONwherex ∈ R d is the input, L ≥ 1 is the network depth, (F l ) l∈[L]are mappings that define the layers, and W l ∈ R n×n are the hidden weights, where n is the network width, and W in , W out are input and output embedding weights. 1 This model will represent the pretrained model that will later be finetuned on some new task.To finetune a (large) pretrained model with a limited amount of computational resources, a popular resource efficient approach is to use the LoRA finetuning method defined below. Definition 1 (Low Rank Adapters (LoRA) from [19] ). To apply LoRA to a weight matrix W ∈ R n1×n2 in the model, we constrain its update in the fine-tuning process by representing the latter with a low-rank decomposition W = W * + α r BA. Here, only the weight matrices B ∈ R n1×r , A ∈ R r×n2 are trainable and the original pretrained weights W * remain frozen. The rank r ≪ min(n 1 , n 2 ) and α ∈ R are tunable constants.As the width n grows,2 the network initialization scheme and the learning rate should be adapted to avoid numerical instabilities and ensure efficient learning. For instance, the variance of the initialization weights (in hidden layers) should scale like 1/n to prevent the pre-activations from blowing up as we increase model width n (e.g., He initialization [4] ). To derive proper scaling rules, a principled approach consist of analyzing the statistical properties of key quantities in the model (e.g. second moment of the pre-activations) as n grows and then adjust the initialization variance, the learning rate, and the architecture to achieve desirable properties in the limit n → ∞ [5, 10, 13, 40] . We use this approach to study the effect of initialization on the feature learning dynamics of LoRA in the infinite-width limit. For more details about the theory of scaling of neural networks, see Appendix A.2.Throughout the paper, we will be using asymptotic notation to describe the behaviour of several quantities as the width n grows. Note that the width n will be the only scaling dimension of neural network training which grows and all other scaling dimensions such as the LoRA rank r, number of layers L, sequence length, number of training steps, etc., will be considered as fixed. We use the following notation for the asymptotic analysis.Notation. Given sequences c n ∈ R and d n ∈ R + , we write c n = O(d n ), resp. c n = Ω(d n ), to refer to c n < κd n , resp. c n > κd n , for some constant κ > 0. We write c n = Θ(d n ) if both c n = O(d n ) and c n = Ω(d n ) are satisfied. For vector sequences c n = (c i n ) 1≤i≤k ∈ R k (for some k > 0), we write c n = O(d n ) when c i n = O(d i n ) for all i ∈ [k],and same holds for other asymptotic notations. Finally, when the sequence c n is a vector of random variables, convergence is understood to be convergence in second moment (L 2 norm).The standard way to initialize trainable weights is to take an iid initialization of the entriesA ij ∼ N (0, σ 2 A ), B ij ∼ N (0, σ 2 B) for some σ A , σ B ≥ 0 (this includes initialization with zeros if σ B or σ A are set to 0). 3 . Due to the additive update structure of LoRA, we want to initialize the product BA to be 0 so that finetuning starts from the pretrained model [19] . This can be achieved by initializing one of the weights A and B to 0. If both are initialized to 0, no learning occurs in this case since this is a saddle point and the parameter gradients will remain zero. Thus, we should initialize one of the parameters A and B to be non-zero and the other to be zero. If we choose a non-zero initialization for A, then following standard initialization schemes (e.g., He Init [4] , LeCun Init [1] ), one should set σ 2 A = Θ(n -1 ) to ensure Ax does not explode for large n. This is justified by the Central Limit Theorem (CLT). On the other hand, if we choose a non-zero initialization for B, one should make sure that σ 2 b = Θ(r -1 ) = Θ(1). This leaves us with two possible initialization schemes:• Init[A]: σ 2 B = 0, σ 2 A = Θ(n -1 ) (default initialization in LoRA [19]). • Init[B]: σ 2 B = Θ(r -1 ) = Θ(1), σ 2 A = 0. 4These two initialization achieve the goal of starting finetuning from the pretrained model. A priori, it is unclear if there is a material difference between the two initialization schemes. Surprisingly, as we will show later in this paper, these two initialization schemes lead to fundamentally different training dynamics when model width is large.Notation. For a given LoRA layer in the network, we use Z to denote the input to that layer and Z for the output after adding the pretrained weights. More precisely, we can write the layer operation as Z = W * Z + α r BA Z. Our main analysis relies on a careful estimation of the magnitude of several quantities involving LoRA features. Let us first give a formal definition.Definition 2 (LoRA Features). Given a general neural architecture and a LoRA layer (Definition 1), we define LoRA features (Z A , Z B ) as Z A = AZ, and Z B = BZ A = BAZ. At fine-tuning step t, we use the superscript t to denote the value of LoRA features Z t A , Z t B , and the subscript t to denote the weights A t , B t . 
lora finetuning dynamics in the large width limit: 
 We fix the LoRA rank r throughout the analysis and examine the finetuning dynamics in the limit of large width. This setup aligns well with practical scenarios where the rank is much smaller than the width (i.e., r ≪ n ). Typically, for Llama models the rank r is generally of order 2 k for k ∈ {2, . . . , 6}, and model width n is generally larger than 2 12 . We will refer to a layer of the network to which LoRA is applied (see Definition 1) as a LoRA layer. For the theoretical analysis, we adopt a simplified setting that facilitates a rigorous yet intuitive derivations of the results.The following simplified setup was considered in Hayou et al. [44] to derive asymptotic results concerning the learning rates in LoRA. We use the same setup in our analysis to investigate the impact of initialization.Finetuning Dataset. We assume that the dataset used for finetuning consists of a single datapoint (x, y), 5 and the goal is to minimize the loss calculated with the model with adjusted weights W * + BA for all LoRA layers (here θ = {A, B, for all LoRA layers in the model}). Z t is the input to the LoRA layer, computed with data input x. Similarly, we write d Zt to denote the gradient of the loss function with respect to the layer output features Z evaluated at data point (x, y).Single LoRA Module. Given a LoRA layer, LoRA feature updates are not only driven by the change in the A, B weights, but also the changes in Z, d Z which are updated as we finetune the model (assuming there are multiple LoRA layers). To isolate the contribution of individual LoRA layers to feature learning, we assume that only a single LoRA layer is trainable and all other LoRA layers are frozen. 6 For this LoRA layer the layer input Z is fixed and does not change with t, whereas d Z changes with step t (because Zt = (W * + α r B t A t )Z). After step t, Z B is updated as followsEQUATIONAs discussed in Hayou et al. [44] , the terms δ 1 t , δ 2 t represent 'linear' feature updates that we obtain if we fix one weight matrix and only train the other. The third term δ 3 t represents the 'multiplicative' feature update which captures the compounded update due to updating both A and B.Hayou et al. [44] introduced the notion of stability of LoRA features as width grows. We introduce here a slightly more relaxed notion of stability.Definition 3 (Feature Stability). We say that LoRA finetuning is stable if for all LoRA layers in the model, and all training steps t, we have Z, Z B = O(1), as the width n goes to infinity.Here, feature stability implies that LoRA output Z B remains bounded (in L 2 norm) as width grows. To achieve such stability, hyperparameters (initialization, learning rate) should be scaled as n grows. We will show that the dependence of the optimal learning rate on n is highly sensitive to the choice ofNote that feature stability also requires that Z = O(1) which is directly related to pretraining dynamics since it depends on some pretrained weights W * . We assume that pretraining parameterization (how initialization and learning rate are parametrized w.r.t width) ensures this kind of stability (see Appendix A for more details). 7As discussed above, feature updates are driven by the terms (δ i t ) i∈{1,2,3,} . As n grows, these feature updates might become trivial (i.e. vanish as n → ∞) or unstable (i.e. grows unbounded). To avoid such scenarios, we want to ensure that ∆Z B = Θ(1). Such conditions are the main ideas behind µP [26] and Depth-µP [41] , which are network parametrizations that ensure stability and feature learning in the large width and depth limits for pretraining. We recall this definition from [44] .Definition 4 (Feature Learning). We say that LoRA finetuning induces stable feature learning in the limit of large width if the dynamics are stable (Definition 3), and for all finetuning steps t, we have∆Z t B def = Z t+1 B -Z t B = Θ(1).∆Z B is the sum of the terms δ i t 's (Equation ( 2)). To achieve optimal feature learning, we want to ensure that δ 1 t = Θ(1) and δ 2 t = Θ(1) which means that both weight matrices A and B are efficiently updated and contribute to the update in Z B . An intuitive explanation is provided in Appendix A.1. This leads us to the following definition of efficient learning with LoRA.Definition 5 (Efficient Learning with LoRA). We say that LoRA fine-tuning is efficient if it is stable (Definition 3), and for all LoRA layers in the model, and all fine-tuning steps t > 1, we haveδ i t = Θ(1), i ∈ {1, 2}.Next, we introduce the γ-operator, an essential tool in our analysis of the large width dynamics.In the theory of scaling, one usually tracks the asymptotic behavior of key quantities as we scale some model ingredient. For instance, if we scale the width n of a neural network, we are interested in quantifying how certain quantities in the network behave as n grows. This is a standard approach for (principled) model scaling and it has so far been used to derive scaling rules for initialization [5] , activation function [10] , network parametrization [41] , amongst other things.With Init[A] and Init[B], initialization weights are of order Θ(n -β ) for some β ≥ 0. Assuming that the learning rate also scales polynomialy with n, it is straightforward that preactivations, gradients, and weight updates are all asymptotically polynomial in n. Note that this is only possible because all neural computations consists of sums of Θ(n α ) terms, where typically α ∈ {0, 1}. For instance, when calculating the features AZ, each entry is a sum of n terms, while when calculating BZ A , each entry is a sum of r terms (r fixed as n goes to infinity). This is true for general neural computation that can be expressed as Tensor Programs [15] .Consequently, for some quantity v in the computation graph, it is natural to track the exponent that determines the asymptotic behavior of v with respect to n. We write v = Θ(γ[v]) to capture this polynomial dependence. Elementary operations with the γ-operator include:8 Zero.When v = 0, we write γ[v] = -∞ (as a limit of γ[n -β ] when β → ∞). Multiplication. Given two real-valued variables v, v ′ , we have γ[v × v ′ ] = γ[v] + γ[v ′ ].Addition. Given two real-valued variables v, v ′ , wegenerally have γ[v + v ′ ] = max(γ[v], γ[v ′ ]).The only case where this is violated is when v ′ = -v. This is generally a zero probability event if v and v ′ are random variables that are not perfectly (negatively) correlated, which is the case in most situations where we make use of this formula. See Appendix A.3 for discussion.We have now introduced all required notions for the subsequent analysis. For better readability, we defer all the proofs to the appendix.Using the γ-operator, we can track the asymptotic behavior of the finetuning dynamics as model width n grows. At finetuning step t, the weights are updated as followsA t = A t-1 -ηg t-1 A , B t = B t-1 -ηg t-1B , where g A , g B are processed gradients (e.g. normalized gradients with momentum as in AdamW). We assume that the gradients are processed in a way that makes their entries Θ(1). This is generally satisfied in practice (with Adam for instance) and has been considered in [40] to derive the µ-parametrization for general gradient processing functions. From this, we obtain the following recursive formulas for γ[Z t A ] and γ[B t ], which characterizes their behavior in the large width limit. Lemma 1 (Informal). For t fixed, the asymptotic dynamics of Z t A and B t follow the recursive formulaEQUATIONThe formal proof of Lemma 1 is provided in Appendix A and relies on Assumption 1 which fairly represents practical scenarios (see Appendix A for a detailed discussion). Lemma 1 captures the change in asymptotic behavior of quantities Z t A and B t as width grows. Naturally, the dynamics depend on the the initialization scheme which lead to completely different behaviors as we show in the next two results.Next, we provide a precise characterization of stability and feature learning when using Init[A]. Theorem 1 (Informal). For t fixed, with Init[A] and learning rate η, we have• Stability: Z t B = O(1) if and only if γ[η] ≤ -1/2.• Feature Learning:∆Z t B = Θ(1) if and only if γ[η] = -1/2.In this case, we also have δ 1 t , δ 2 t = Θ(1) (efficient feature learning, Definition 5).Moreover, "internal" instability (Z t A = Ω(1)) occurs when γ[η] ∈ (-1, 1/2].With Init[A], the maximal learning rate 9 that does not lead to instability in Z B scales as Θ(n -1/2 ). This can be seen as an asymptotic form of the edge of stability phenomenon [17] where if we increase the learning rate beyond some level, instability occurs. Interestingly, in this case (i.e. with Θ(n -1/2 ) learning rate) the features are efficiently updated (Definition 5). However, this comes with caveat: the features Z t A grow as Θ(n 1/2 ) which can potentially cause numerical instabilities. We call this phenomenon internal instability: only the features Z A (internal LoRA features) grows, LoRA output Z B remains Θ(1) in this case.The fact that Θ(n -1/2 ) is the maximal learning rate that does not cause instability in Z B does not mean it is the optimal learning rate. As the width n grows, this internal instability in Z A will become more and more problematic. Intuitively, we expect that a trade-off appears in this case: the optimal learning rate (found by grid search) to be larger than Θ(n -1 ) but smaller than Θ(n -1/2 ), i.e. the network will try to achieve a balance between optimal feature learning (γ[η] = -1/2) and internal stability Z t A = Θ(1) (γ[η] = -1). We verify this empirically in the next section.In the next result, we show that the maximal learning rate allowed with Init[B] is different from that with Init[A], leading to completely different dynamics. Theorem 2 (Informal). For t fixed, with Init[B], we have• Stability: Z t B = O(1) if and only if γ[η] ≤ -1.• Feature Learning:∆Z t B = Θ(1) if and only if γ[η] = -1.Moreover, efficient feature learning cannot be achieved with Init[B] for any choice of learning rate scaling γ[η] (that does not violate the stability condition). More precisely, with Θ(n -1 ) learning rate, the limiting dynamics (when n → ∞) are the same if B was not trained and A is trained.With Init[B], the maximal learning rate (that does not violate stability) scales as Θ(n -1 ) (for any ϵ > 0, a learning rate of Θ(n -1+ϵ ) leads to Z B = Ω(1)).Because of this bound on the maximal learning rate, no internal instability occurs with Init[B]. In this case, feature learning is suboptimal since the B weight matrix is undertrained in the large width limit (δ 2 t → 0).Conclusions from Sections 3.5 and 3.6. The results of Theorem 1 and Theorem 2 suggest that Init[A] allows the use of larger learning rates compared to Init[B], which might lead to better feature learning and hence better performance at the expense of some internal instability. Here, 'larger' learning rate should be interpreted in asymptotic terms: with Init[A] the maximal learning rate that does not cause instability satisfies γ[η] = -1/2. With Init[B], we have γ[η] = -1 instead.Note that because of the constants in Θ(n β ) learning rates (for some β) , the optimal learning rate with Init[A] is not systematically larger than Init[B] for finite width. However, as width grows, we will see that it is case.Another important insight from this analysis is that with both initializations, the dynamics are suboptimal in the limit: internal instability with Init[A] and undertraining of B with Init[B]. 10 We will later discuss possible solutions to this behavior. -------------------------------- -------------------------------- -------------------------------- -------------------------------- To validate our theory in a controlled setting, we consider the following simple model:EQUATIONwhere Optimal Learning Rate. We finetune model (4) on synthetic data generated from the teacher model. In Figure 2 , we show the optimal learning rate when using either Init[A] or Init[B] to initialize the finetuning, as a function of width n. For n ≫ 1 (typically n ≥ 2 9 ), the optimal learning rate with Init[A] is larger than the optimal learning rate with Init[B]. This is in agreement with the theoretical results obtained in Theorem 1 and Theorem 2 which predict asymptotic maximal learning rates (that satisfy the stability condition) of Θ(n -1/2 ) and Θ(n -1 ) respectively.W in ∈ R n×d , W h ∈ R n×n , W out ∈ RWith Init[A], we observe the stability/feature learning trade-off for large n. The optimal learning rate with Init[A] in this regime (e.g. n = 2 13 ) is smaller than the maximal theoretical learning 11 Here, the pretrained model is effectively given by Yout = W teacher out ϕ(W teacher in x), and the finetuning dataset is simulated by injecting the LoRA weights A teacher , B teacher . 12 In this setup, a student model can have larger width n than the teacher model.rate n -1/2 that achieves optimal feature learning (Theorem 1). Here, the model seems to balance the internal instability that occurs in the Z A features with feature learning and thus favors smaller learning rates: the optimal learning rates is smaller than Θ(n -1/2 ) and larger than Θ(n -1 ).Internal Instability and Feature Learning. ), which implies that Z t B = o(1), i.e. the Z B features vanish as width grows. While this might problematic, it only becomes an issue when the width is extremely large: if the optimal learning rate scales as Θ(n -β ) for some β ∈ (1/2, 1) (so that the learning rate is between Θ(n -1 ) and Θ(n -1/2 ), balancing internal instability and efficient feature learning), the LoRA output feature scales as Z B = B t A t Z = Θ(n -β+1 ). Therefore, if β ≈ 0.7 for instance, the vanishing rate of LoRA output feature is Z B ≈ Θ(n -0.3 ) which is slow given the order of magnitude of width in practice (for n = 2 12 , we have n -0.3 ≈ 0.08). 
experiments with language models: 
 Our theoretical results from earlier provides a detailed asymptotic analysis of the finetuning dynamics when LoRA modules are initialized with Init[A] or Init [B] . The main conclusions are that Init[A] generally leads to more efficient feature learning (which can be justified by the fact that optimal learning rate is larger when using Init[A] compared to when using Init[B]). To provide evidence of this claim on real-world tasks, we use LoRA to finetune a set of language models on different benchmarks. Details about the experimental setup and more empirical results are provided in Appendix B. We use LoRA+ code [44] for our experiments (available at https://github.com/nikhil-ghosh-berkeley/loraplus).The GLUE benchmark (General Language Understanding Evaluation) consists of several language tasks that evaluate the understanding capabilities of langugage models [8] . Using LoRA, we finetune Roberta-large from the RoBERTa family [12] on MNLI, SST2, and QNLI tasks with varying learning rates η and initialization schemes (Init[A] or Init[B]). We use the same experimental setup of [19] for Roberta-Large to compare our results with theirs (see Appendix B for more details). ------------------------------- ------------------------------- ------------------------------- The results in Figure 4 are aligned with our theory: we observe that Init[A] generally leads to better performance, and the optimal learning rate with Init[A] is generally larger than with Init[B]. Models initialized with Init[A] match the performances reported in [19] , while those initialized with Init[B] generally underperform that baseline. For MNLI task (the hardest one amongst the three tasks), we observe a significant difference in the best test accuracy (over 3 random seeds) with 90.69 with Init[A] and 89.47 with Init[B]. We also observe that for MNLI, the optimal learning rate with Init[A] (η * = 8e-5) is much larger than the optimal learning rate with Init[B] (η * = 1e-5), which aligns with our theoretical predictions. However, note that for QNLI for instance (an easier task), while the optimal test accuracy is significantly better with Init[A], the optimal learning rate (from the grid search) is the same for Init[A] and Init [B] . There are many possible explanations for this: 1) the width is not large enough in this case to see the gap between optimal learning rates (for RoBERTa-Large, the width is n = 2 10 ) 2) The constants in Θ(n -1 ) are Θ(n -1/2 ) are significantly different in magnitude due to dependence on finetuning task. We notice similar behaviour with LLama experiments below. A precise analysis of this observation is beyond the scope of this paper, we leave it for future work. ------------------------------------------------- To further validate our theoretical findings on more modern models and datasets, we report the results of finetuning the Llama-7b model [38] on the Flan-v2 dataset [36] and the GSM8k dataset [16] , and finetuning the TinyLlama model [49] on WikiText-2 using LoRA. Each trial is averaged over two seeds and the shaded region indicates one standard error. In the left panel of Figure 5 we see that when finetuning TinyLlama using LoRA the optimal learning rate using Init[A] is larger than with Init[B] and the corresponding test perplexity is lower. Similarly, for the center panel of Figure 5 , when finetuning the Llama-7b model on Flan-v2, the optimal learning rates for Init[A] and Init[B] are the same (for the learning rate grid we used), but the the optimal MMLU accuracy for Init[A] is slightly higher than for Init[B]. For learning rates close to the optimal choice, the accuracy using Init[A] is generally higher than for Init[B]. An analagous result holds for the GSM8k dataset as shown in the rightmost panel of Figure 5 . More details about this setting are provided in Appendix B. 
conclusion and limitations: 
 We showed that LoRA dynamics are highly sensitive to initialization. Init[A] is associated with larger optimal learning rates, compared to Init[B]. Larger learning rates typically result in better performance, as confirmed by our empirical results. Note that this is a zero-cost adjustment with LoRA finetuning: we simply recommend usingInit[A] instead of Init[B].One limitation of our work is that we only define feature learning via the magnitude of feature updates in the limit of large width. In this way, our definition of feature learning is data-agnostic and therefore no conclusion about generalization can be obtained with this analysis. The constants in Θ(.) asymptotic notation naturally depend on the data (the finetuning task) and therefore such data-agnostic approach does not allow us to infer any information about the impact of the data on the finetuning dynamics.More importantly, our results indicate that both initialization schemes lead to suboptimal scenarios, although Init[A] allows more efficient feature learning. In both cases, instability and/or suboptimal feature learning present fundamental issues, which can potentially be mitigated by approaches such as LoRA+ [44] . Understanding the interaction of LoRA+ and related efficiency methods with the initialization scheme is an important question for future work. 
 
 Appendix: Recall the feature update decompositionEQUATIONTo achieve optimal feature learning, we want to ensure that δ 1 t = Θ(1) and δ 2 t = Θ(1) which means that both weight matrices A and B are efficiently updated and contribute to the update in Z B . To justify why this is a desirable property, let us analyze how changes in matrices A and B affect LoRA feature Z B = BA Z. Let (B :,i ) 1≤i≤r denote the columns of B. We have the following decomposition of Z B :Z B = r i=1 (A Z) i B :,i ,where (AZ) i is the i th coordinate of AZ. This decomposition suggests that the direction of Z B is a weighted sum of the columns of B, and A modulates the weights. With this, we can also write   δ 1 t = r i=1 (∆A t Z) i (B :,i ) t-1 δ 2 t = r i=1 (A t-1 Z) i (∆B :,i ) t-1 ,where (B :,i ) t refers to the columns of B at time step t. Having both δ 1 t and δ 2 t of order Θ(1) means that both A and B are 'sufficiently' updated to induce a change in weights (AZ) i and directions B :,i . If one of the matrices A, B is not efficiently updated, we might end up with suboptimal finetuning, leading to either non updated directions B or direction weights (A t-1 Z). For instance, assuming that the model is initialized with Init[B], and that B is not efficiently updated, the direction of Z B will be mostly determined by the vector (sub)space of dimension r generated by the columns of B at initialization. This intuition was discussed in details in [44] .Scaling refers to the process of increasing the size of one of the ingredients in the model to improve performance (see e.g. [22] ). This includes model capacity which can be increased via width (embedding dimension) or depth (number of layers) or both, compute (training data), number of training steps etc. In this paper, we are interested in scaling model capacity via the width n. This is motivated by the fact that most state-of-the-art language and vision models have large width.It is well known that as the width n grows, the network initialization scheme and the learning should be adapted to avoid numerical instabilities and ensure efficient learning. For instance, the initialization variance should scale 1/n to prevent arbitrarily large pre-activations as we increase model width n (e.g. He init [4] ). To derive such scaling rules, a principled approach consist of analyzing statistical properties of key quantities in the model (e.g. pre-activations) as n grows and then adjust the initialization, the learning rate, and the architecture itself to achieve desirable properties in the limit n → ∞ [5, 10, 13] .In this context, Yang et al. [26] introduces the Maximal Update Parameterization (or µP), a set of scaling rules for the initialization scheme, the learning rate, and the network architecture that ensure stability and maximal feature learning in the infinite width limit. Stability is defined by Y i l = Θ(1) for all l and i where the asymptotic notation 'Θ(.)' is with respect to width n (see next paragraph for a formal definition), and feature learning is defined by ∆Y l = Θ (1) , where ∆ refers to the feature update after taking a gradient step. µP guarantees that these two conditions are satisfied at any training step t. Roughly speaking, µP specifies that hidden weights should be initialized with Θ(n -1/2 ) random weights, and weight updates should be of order Θ(n -1 ). Input weights should be initialized Θ(1) and the weights update should be Θ(1) as well. While the output weights should be initialized Θ(n -1 ) and updated with Θ(n -1 ). These rules ensure both stability and feature learning in the infinite-width limit, in contrast to standard parameterization (exploding features if the learning rate is well tuned), and kernel parameterizations (e.g. Neural Tangent Kernel parameterization where ∆Y l = Θ(n -1/2 ), i.e. no feature learning in the limit).A.3 When does γ-Operator fail to capture asymptotic behavior?When non-polynomial dependencies (in terms of n) appear in neural computations, then the γ operator cannot capture asymptotic behavior of the learning dynamics. For instance, if one of the layers has embedding dimension e n or n×log(n), polynomial exponents are no longer sufficient to capture the asymptotic dynamics. Fortunately, such cases are generally not considered in practice.In this section, we provide the formal proof of Lemma 1. The proof relies on the following assumption on the processed gradient g A . This assumption was used in [44] to derive scaling rules for the optimal learning rates for A and B weight matrices. Here, we use it to study the sensitivity of LoRA dynamics to initialization. We provide an intuitive discussion that shows why this assumption is realistic. Assumption 1. With the same setup of Section 3, at training step t, we have Z, d Z = Θ(1) andg t A Z = Θ(n).Assumption 1 consists of two parts: that 1) Z, d Z = Θ(1) and 2) g t A Z = Θ(n). The first condition is mainly related to pretraining paramterization which we assume satisfied such conditions. 13 The second condition is less intuitive, so let us provide an argument to justify why it is sound in practice. Let us study the product g t A Z in the simple case of Adam with no momentum, a.k.a SignSGD which is given byg A = sign ∂L ∂A ,where the sign function is applied element-wise. At training step t, we have∂L t ∂A = α r B ⊤ t-1 d Zt-1 ⊗ Z, Let S t = α r B ⊤ t-1 d Zt-1 . Therefore we have g A = sign(S t ⊗ Z) = (sign(S t i Z j )) 1≤i,j≤n .However, note that we also have sign(S t i Z j ) = sign(S t i )sign(Z j ), and as a result g t A = sign(S t ) ⊗ sign(Z).Hence, we obtain g t A Z = (sign(Z) ⊤ Z)sign(S t ) = Θ(n), where we used the fact that sign(Z) ⊤ Z = Θ(n).This intuition should in-principle hold for the general variant of Adam with momentum as long as the gradient processing function (a notion introduced in [2] ) roughly preserves the sign(Z) direction. This reasoning can be made rigorous for general gradient processing function using the Tensor Program framework and taking the infinite-width limit where the components of g A , Z, d Z all become iid. However this necessitates an intricate treatment of several quantities in the process, which we believe is an unnecessary complication and does not serve the main purpose of this paper. Lemma 1. Under Assumption 1, the asymptotic behaviour of Z t A and B t follow the recursive formula γ[Z t A ] = max(γ[Z t-1 A ], γ[η] + 1) γ[B t ] = max(γ[B t-1] ], γ[η]).Proof. At finetuning step t, the weights are updated as followsA t = A t-1 -ηg t-1 A , B t = B t-1 -ηg t-1 B .Using the elementary operations with the γ-operator, we obtainγ[Z t A ] = max(γ[Z t-1 A ], γ[ηg t-1 A Z]) = max(γ[Z t-1 A ], γ[η] + γ[g t-1 A Z]).We conclude for Z t A using Assumption 1. The formula for γ[B t ] follows using the same techniques.A. Z] = max(0, γ[η] + 1) γ[B t ] = γ[η]To achieve Z B = O(1), we should therefore haveγ[η] + max(0, γ[η] + 1) ≤ 0, which is equivalent to γ[η] ≤ -1/2.This implies that the maximum learning rate that does not cause instability is Θ(n -1/2 ). Such learning rate causes internal instability, i.e. the feature Z A explodes with width. Why? Because, with this learning rate, we have γ[A t Z] = 1/2, i.e. A t Z = Θ(n 1/2 ) which diverges as n grows. However, this growth is compensated with the fact that γ[B t ] = -1/2, i.e. B t = Θ(n -1/2 ). This analysis is valid for any γ[η] ∈ (-1, 1/2].In this case, feature learning is efficient in the sense of Definition 5: δ 1 t = Θ(1) and δ 2 t = Θ(1). To see this, recall that δ Proof. Here, we show that maximal learning rate that does not cause instability in LoRA output features Z B is Θ(n -1 ) and no internal instability occurs in this scenario.1 t = B t-1 ∆Z 1 A which yields γ[δ 1 t ] = γ[B t-1 ] + γ[∆Z t A ] = γ[η] + γ[η] + 1 = 0 and γ[δ 2 t ] = γ[∆B t ] + γ[Z t-1 A ] = γ[η] + max(γ[η] + 1, 0) = 0.WithInit[B], we have that γ[B 0 ] = 0 and γ[A 0 Z] = -∞. From Equation (3), we obtain that γ[A t Z] = γ[η] + 1 γ[B t ] = max(0, γ[η]).As a result, LoRA output stability is achieved if and only ifγ[η] + 1 + max(0, γ[η]) ≤ 0, which is equivalent to having γ[η] ≤ -1.Moreover, with η = Θ(n -1 ) we have that γ[δ 1 t ] = γ[B t-1 ] + γ[∆Z t A ] = 0 + γ[η] + 1 = 0 and γ[δ 2 t ] = γ[∆B t ] + γ[Z t-1 A ] = γ[η] + 0 = -1.As a result, feature learning is not efficient in this case, and the learning dynamics are asymptotically equivalent to not training matrix B (because δ 2 t → 0).This section complements the empirical results reported in the main text. We provide the details of our experimental setup, and show the acc/loss heatmaps for several configurations.In Figure 2 , we trained a simple model with LoRA layers to verify the results of the analysis in ??.Here we provide the empirical details for these experiments.Model. We consider a simple model given byf (x) = W out ϕ(W in x + (W h + BA)ϕ(W in x)),where W in ∈ R n×d , W out ∈ R 1×n , A ∈ R r×n , B ∈ R n×r are the weights, and ϕ is the ReLU activation function. Training. We train the model with AdamW with β 1 = 0.9 and β 2 = 0.99 for a range for values of η. The weights are initialized as follows: W in ∼ N (0, 1/d), W h ∼ N (0, 1/n), W out ∼ N (0, 1/n) and fixed. Only the weight matrices A, B are trainable.For our experiments with RoBERTa models, finetuned on GLUE tasks, we use the following setup:Training Alg Details Model Roberta-Large Learning Rates {2 k × 10 -5 , for k = 0, 1, 2, . . . , 10} β 1 0.9 ------------------------------- ------------------------------- ------------------------------- ------------------------------- Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: The paper does not involve crowdsourcing nor research with human subjects.β38th Conference on Neural Information Processing Systems (NeurIPS 2024).We use the same notation from Hayou et al.[44].The width in SOTA models is typically large, i.e. of width n > 10.Gaussianity is not important and can be replaced by any zero-mean distribution with finite-variance for our purposes.Here, we assumed that r = Θ(1) (in width), i.e. it doesn't grow with width. In general, the right scaling for Init[B] is σ 2 B = Θ(r -1 ).Although this a simplifying assumption for our analysis, the results can be extended to mini-batched gradients without affecting the conclusions. Such results will require additional assumptions to be fully rigorous.This is equivalent to having only a single LoRA layer in the model since LoRA layers are initialized to zero.When taking the infinite width limit, we can for instance assume that pretraining parameterization is µP[26]. This is a technicality for the infinite-width limit and does not have any implications on practical scenarios where the width is finite. The most important implications of this assumption is that in the pretrained network (before introducing LoRA layers), we have Z = Θ(1), Z = Θ(1), which holds for a general input-output pair (x, y).The γ-operator is a mapping from the set {v, s.t.v = Θ(n β ) for β ∈ R ∪ {-∞}} to the set R ∪ {-∞}.Maximal γ[η] that does not cause instability in ZBMore precisely, one can show that with Init[B], for fixed t, in the limit n → ∞, Bt converges to B0, i.e. B is untrained in this limit.There is a technical intricacy on this point. While Z depends only on pretraining, the Jacobian d Z depends on finetuning. However, under the stability conditions mentioned in Definition 3, if d Z = Θ(1), it should remain so during finetuning as well. 


Decision: This paper investigates the impact of initialization schemes for LoRA. The reviewers appreciated, among others, the importance of the topic, theoretical analysis, and experimental contributions. The discussion that followed convinced two skeptical reviewers to increase their ratings to 5. Consequently, the reviewers unanimously recommend acceptance; hence, I recommend accepting.

Reviews: review: paper_summary: paper_summary: Finetuning large language models has become a common technique among practitioners, however due to the memory footprint of practically viable LLMs, there is a dire need of techniques that allow memory-lightweight finetuning (PEFT). Among those techniques LoRA has gained immense popularity. In LoRA a dense matrix is modified via addition of a low-rank update consisting of matrices A and B. In order for this update to initially conserve the model’s outputs while preserving the ability to learn, exactly one of the matrices needs to be initialized to zeros.
“The Impact of Initialization on the Finetuning Dynamics in LoRA” proposes that the B matrix should be initialized with zeros. It hypothesizes that the reason for that is that this initialization (Init[A]) allows using greater learning rate without sacrificing stability.
The main points of the paper are supported by theoretical derivations stemming from the perspective of infinite-width networks. Furthermore, experimental evaluations serve as proof of the soundness of the theoretical approach and as practical justification of the method.
strength_weakness: strength_weakness: - There is virtually no discussion of why the training hyperparameters are set to those exact values, in particular with regard to the number of epochs and weight decay = 0. I suspect changing those HPs might mitigate the need to choose the initialization schema so carefully. Also the r parameter seems to be chosen arbitrarily.
- The gains provided by the method seem to be small and would be better visualized if they were put in the perspective of comparison with other methods such as differing r or even, if possible, full finetuning.
- The Toy Model setup is not explained clearly.
questions: questions: - There are other works that try to improve LoRA’s performance (LoRA+) or are inspired by it (VeRA, AdaLoRA, etc.). Do you believe your work could be used with those techniques?
- Do you feel that there is some simple, informal intuition suggesting why the Init[A] outperforms Init[B]?
- AdamW seems to be among the most popular optimizers, however the paper focuses on rigorous derivation in case of SignSGD claiming that the result should extend to the typical choice of the Adam optimizer. Even though the authors claim they are using the AdamW optimizer in the experiments, they set weight decay = 0, effectively replacing AdamW with Adam. Would the method perform as well while using weight decay? Would it be justified theoretically the same way Adam is?
- Can you justify the fixed hyperparameters’ values?
- Can you explain the toy model setup more clearly? As I understand it now, the same model is being adapted and is used to create the dataset which does not seem to make sense.
limitations: limitations: - I do not believe the authors have sufficiently explained the limitations of the fixed hyperparameters in the evaluations
paper_summary: paper_summary: The paper analyzes the impact of different initialization techniques for A and B matrices in LoRA adapters. Typically, either A or B matrix is initialized with zero while the other is initialized form a Gaussian distribution. This is done so that fine-tuning starts with LoRA adapters ($A \times B = 0$) having zero effect. However, the paper argues that it matters which of A and B initialize to 0. As per the results, $B = 0$ init allows higher optimal learning rates which leads to better overall performance. However, there is still instability during training which could be mitigated with a method like LoRA+.
strength_weakness: strength_weakness: Perhaps I am not the right person to review this paper, but I don't fully understand the impact of this paper. It shows that there are two ways of initializing AB matrices and the default one that we have been using until now is the better option. Additionally, it says that the default initialization is not good enough but doesn't propose anything new. Maybe the theory part is important, but I couldn't understand it fully in the limited time I had. Having said that, I still feel that the paper has something worthwhile to contribute to the community. So, I vote to accept it for now.
questions: questions: None.
limitations: limitations: The limitations section in the paper is adequate.
paper_summary: paper_summary: This paper investigates the impact of initialization schemes on the finetuning dynamics of Low Rank Adaptation (LoRA) in large language models. The authors compare two initialization approaches: Init[A], where A is initialized randomly and B to zero, versus Init[B], where B is initialized randomly and A to zero. Through theoretical analysis and empirical experiments, they demonstrate that Init[A] allows for larger learning rates and leads to more efficient feature learning compared to Init[B], albeit at the cost of some internal instability. The paper provides a rigorous theoretical framework based on large width limits and supports the findings with experiments on toy models and real-world language tasks.
strength_weakness: strength_weakness: W1: While the theoretical analysis is sound, the paper could benefit from a more intuitive explanation of why Init[A] allows for larger learning rates and more efficient feature learning. This would make the insights more accessible to some practitioners.

W2: The experiments focus primarily on language models and NLP tasks. This is definitely acceptable and overall a very sound choice, however the inclusion of experiments from other domains (e.g., finetuning an image classification model, or some other architecture that relates to a vision transformer) would strengthen the experimental section significantly, and increase the overall impact of the paper in the broader NeurIPS community.
questions: questions: Q1: How sensitive are the results to the choice of LoRA rank? Do the theoretical predictions and empirical findings hold across different rank values?

Q2: The paper focuses on vanilla LoRA. How do the authors expect these initialization schemes to interact with variants like QLoRA or LoRA+? Would the trade-offs and recommendations change?

Q3: Given the internal instability observed with Init[A], are there any potential negative consequences for downstream task performance or generalization ability? How might this instability manifest in practice (and would there be a difference depending on input domain and/or other parameters of the model + problem space)?

Q4: The theoretical analysis seems to assume a single LoRA layer, for simplicity. How well do the insights generalize to more realistic scenarios with multiple LoRA layers throughout a network?
limitations: limitations: L1: The theoretical analysis seems to rely on assumptions about the asymptotic behavior in the large width limit. The paper could benefit from a more detailed discussion of how these assumptions may or may not hold in practical scenarios with finite-width networks.

L2: The empirical evaluation, while comprehensive, is limited to a subset of NLP tasks and model architectures. Expanding the evaluation to a broader range of tasks, model sizes, and architectures would strengthen the generalizability of the findings.
paper_summary: paper_summary: The paper investigates the impact on training dynamics of two initialization schemes for LoRA. For this purpose the authors investigate the asymptotic behaviour of activations and weights for LoRA adapters. The authors find that init[A] where A is intialized randomly and B is initalized with zeros leads to more efficient learning dynamics but leads to some numerical instabilities. This initialzation method is the one most commonly used for LoRA-style methods. Further, the authors found that init[B] leads to suboptimal learning dynamics but is less prone to numerical instabilities. Finally, init[A] performs better with larger learning rates than init[B], which is supported by experiments on NLU as well as NLG finetuning tasks.
strength_weakness: strength_weakness: **Novelty**

Most of the theoretical Analysis of the LoRA finetuning dynamics in this paper have already been presented in [1]. 
The authors should have made it more clear how their work differs from the Analysis done in [1], especially since most formulas and notation is identical. 
Further, both of the methods evaluated in this paper have also been investigated in [1].
The paper would greatly benefit from a clear explanation on how the analysis in this work differs from results published in [1].
Finally, it is not clear if the contribution of individual LoRA layers to feature learning can be derived from a setting where only a single LoRA layer is trainable while all others are frozen (as claimed in chapter 3.1).

**Experimental results**

Since the main gist of the paper is to investigate the learning dynamics of LoRA, it would be good to also investigate different initialization schemes aside from init[A] and init [B] prevalent in the litearture, e.g. gaussian init [2], kaiming init [3], principal components [4], etc.
In [1] only different learning rates for A and B have been investigated.

**Significance of results**

While the main finding that init[A] has a better optimal performance is interesting, this initialization scheme is already the common initialization scheme used by LoRA.
Therefore, the impact of the main findings is marginal.

[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

[3] He et al., & Sun, J. Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification., ICCV 2015

[4] Meng et al., PiSSA: Principal Singular Values and Singular Vectors Adaptation of Large Language Models., arXiv 2024
questions: questions: - Did the authors investigate other commonly used initialization schemes?
- Will the authors provide code that can be used to reproduce the results in the experiments section?
limitations: limitations: The authors have properly addressed limitations.

Comments: comment: I thank the authors for the detailed response and the clarifications.

There was indeed a misunderstanding, I would recommend to make the fact that the analysis comprises general random distributions more explicit in the main text and not only mention it in a footnote.
Is it possible to extend this analysis to other initialization schemes than general random distrbutions, as in e.g. [1]?
I believe this would strengthen the paper.

Further, after carefully double-checking the contributions of [2], I agree that the paper presents a valid contribution, which is why I decided to raise my score. However, since I did not carefully check the math, I will also decrease my confidence.

[1] Meng, F., Wang, Z., and Zhang, M. Pissa: Principal singular values and singular vectors adaptation 507 of large language models, 2024.

[2] S. Hayou, N. Ghosh, and B. Yu. LoRA+: Efficient Low Rank Adaptation of Large Models. 2024.
I thank the authors for the replies. I raise my score to 5.