LoRA Initialization Strategies

As we move into more sophisticated applications of LoRA, understanding how to initialize the trainable parameters becomes significant. The way we set the initial values for the low-rank matrices $A$ and $B$ can influence training stability, convergence speed, and ultimately, the quality of the fine-tuned model. This section examines common initialization strategies and their implications.

Recall that the LoRA update $\Delta W$ for a weight matrix $W_0$ is calculated as $\Delta W = BA$, where $B \in \mathbb{R}^{d \times r}$ and $A \in \mathbb{R}^{r \times k}$ for a weight matrix $W_0 \in \mathbb{R}^{d \times k}$. The modified weight matrix during fine-tuning is $W = W_0 + \alpha \frac{BA}{r}$, although for simplicity in discussing initialization, we often focus on the core $BA$ term, keeping in mind the scaling factor $\alpha/r$ applies later. The goal of initialization is to set the starting values for $A$ and $B$ appropriately.

Default Initialization: Zero for B, Gaussian for A

The most widely adopted and often default strategy for LoRA initialization is:

1. Initialize matrix $A$ using a standard random initialization method, typically Kaiming uniform or Gaussian distribution ($\mathcal{N}(0, \sigma^2)$) with a small variance. This ensures that the initial representation projected into the lower-rank space has some variability.
2. Initialize matrix $B$ with all zeros.

The core idea behind this approach is to ensure that at the very beginning of the fine-tuning process (step $t=0$), the adaptation term $\Delta W = BA$ is exactly zero.

$$W_{t=0} = W_0 + BA = W_0 + 0 \cdot A = W_0$$

This means the model with LoRA layers initially behaves identically to the pre-trained base model. The fine-tuning process then gradually learns non-zero values for $B$, allowing the adaptation $\Delta W$ to emerge and adjust the model's behavior based on the task-specific data.

Advantages:

• Stability: Starting with $\Delta W = 0$ prevents any initial disruption to the carefully learned representations of the pre-trained model. This generally leads to more stable training dynamics, especially at the beginning.
• Preserves Pre-trained Knowledge: Ensures that the fine-tuning process begins exactly from the state of the base model, without any initial random perturbation from the LoRA layers.

Disadvantages:

• Potentially Slower Initial Learning: Since $B$ starts at zero, the gradient information needs to flow back to update $B$ to non-zero values before meaningful adaptation can occur. This might slightly slow down the initial phase of convergence compared to methods where $\Delta W$ is non-zero from the start.

This strategy is implemented by default in popular libraries like Hugging Face's PEFT (peft). For instance, the LoraLayer often initializes lora_B weights to zeros and lora_A weights using Kaiming uniform initialization.

Gaussian Initialization for A and B

An alternative approach is to initialize both matrices $A$ and $B$ using a random distribution, typically Gaussian ($\mathcal{N}(0, \sigma^2)$) with a carefully chosen (usually small) variance $\sigma^2$.

In this case, at $t=0$, the adaptation term $\Delta W = BA$ will be a non-zero matrix, albeit likely one with small entries if $\sigma$ is small.

$$W_{t=0} = W_0 + BA \neq W_0 \quad (\text{unless } BA \text{ happens to be zero})$$

Rationale:

The motivation here is that starting with a small, non-zero random adaptation might allow the model to start learning the required adjustments more quickly, potentially accelerating convergence. The initial random $\Delta W$ provides an immediate, albeit noisy, direction for adaptation.

Considerations:

• Choice of Variance ($\sigma^2$): This is a critical hyperparameter. If $\sigma$ is too large, the initial random $\Delta W$ could significantly perturb the pre-trained weights $W_0$, destabilizing training or immediately degrading performance. If $\sigma$ is too small, the effect might be negligible and similar to zero-initializing $B$. The optimal variance often requires tuning.
• Interaction with Learning Rate: A non-zero initial $\Delta W$ might necessitate a smaller initial learning rate compared to the zero-initialization strategy for $B$ to maintain stability.

Advantages:

• Potential for Faster Initial Convergence: The model doesn't need to "break the symmetry" of $B=0$ and might adapt more quickly in the initial steps if the random initialization provides a useful signal.

Disadvantages:

• Risk of Instability: The initial random $\Delta W$ can disrupt the pre-trained model's functionality if not scaled appropriately.
• Increased Sensitivity to Hyperparameters: Requires more careful tuning of the initialization variance ($\sigma^2$) and potentially the learning rate.

This chart illustrates how the magnitude of the LoRA update term $BA$ might evolve. Zero-initialization for $B$ starts the update at zero, while Gaussian initialization for both $A$ and $B$ starts with a small, non-zero update.

Practical Considerations and Recommendations

For most practical applications, starting with the default strategy (zero-initialize $B$, random-initialize $A$) is recommended. It provides a stable and reliable baseline that preserves the integrity of the pre-trained model at the onset of fine-tuning. This approach is less sensitive to hyperparameter choices compared to initializing both matrices randomly.

Consider experimenting with random initialization for both $A$ and $B$ only if:

• You observe particularly slow initial convergence with the default method.
• You have the computational resources to perform careful hyperparameter tuning for the initialization variance and potentially the learning rate.
• You are exploring variations for research purposes.

Remember that the effective initial update magnitude is also influenced by the LoRA rank $r$ and the scaling factor $\alpha$. The update is $W_0 + \alpha \frac{BA}{r}$. Even with Gaussian initialization for $A$ and $B$, a very small $\alpha$ or a very large $r$ will diminish the initial perturbation caused by $BA$. Conversely, a large $\alpha$ or small $r$ will amplify it. These factors interact, and tuning them together is part of optimizing LoRA performance.

While other initialization schemes could be devised, perhaps drawing inspiration from matrix factorization techniques or incorporating prior knowledge about the task, the Kaiming/Gaussian for $A$ and zero for $B$ remains the prevalent and effective starting point for applying LoRA.