LoRA Factor Gradients and Initialization

Low-Rank Adaptation (LoRA) is a parameter-efficient fine-tuning technique that freezes pre-trained model weights and injects trainable rank decomposition matrices into each layer of the Transformer architecture. For a frozen weight matrix W₀, the update is represented as ΔW = B · A. To ensure the model begins training exactly as the original pre-trained model, the matrix B is typically initialized to zero, forcing the initial update ΔW to be exactly zero.

Worked Example

Let's verify the gradients analytically for a small 3x2 example where the A vector is "Difference" a = [1, -1], the B vector is "Tilt" b = [1, -1, 0.5], the input is "First" x = [1, 0], target is "Zero" y = [0, 0, 0], and scale is s = 1. We are given the frozen base matrix:

W0 = [
  [1, 0],
  [0, 1],
  [1, 1]
]

First, calculate the low-rank adaptation matrix ΔW = s · b · a^T:

ΔW = 1 * [1, -1, 0.5]^T * [1, -1]
   = [
       [ 1, -1],
       [-1,  1],
       [0.5,-0.5]
     ]

Next, find the new model weights W = W0 + ΔW:

W = [
  [2, -1],
  [-1, 2],
  [1.5, 0.5]
]

The prediction is W · x. With input x = [1, 0]:

prediction = [2(1) - 1(0), -1(1) + 2(0), 1.5(1) + 0.5(0)]
           = [2, -1, 1.5]

The target y = [0, 0, 0], so the residual error e = prediction - y = [2, -1, 1.5]. The loss is 0.5 * ||e||² = 0.5 * (4 + 1 + 2.25) = 3.625.

Now we compute the gradients for factors a and b. Using the chain rule, ∇a = s * x * (b · e) and ∇b = s * e * (a · x).

b &middot; e = 1(2) + (-1)(-1) + 0.5(1.5) = 2 + 1 + 0.75 = 3.75
a &middot; x = 1(1) + (-1)(0) = 1

∇a = 1 * [1, 0] * 3.75 = [3.75, 0]
∇b = 1 * [2, -1, 1.5] * 1 = [2, -1, 1.5]

The gradients are non-zero because the factor B is non-zero, initiating adaptation.

Experiment: Zero Initialization Dynamics

The interactive visualization at https://riano.app/module/lora-factor-gradient-initialization illustrates why one factor (typically B) is initialized to zero. Follow these steps to explore the initialization dynamics:

Set the B vector to "Zero", A vector to "Difference", and Scale to 1. Notice the prediction matches the frozen base model behavior, and the total update matrix is exactly zero.
Observe the gradients: ∇A is exactly zero, but ∇B is non-zero. Even though factor B is entirely zero, the non-zero residual allows B to receive a gradient and update.
Now change the A vector to "Zero" while keeping the B vector at "Zero". You'll see that both ∇A and ∇B are frozen at zero. This happens because each gradient depends on the other factor; if both start at zero, neither factor ever updates.
Reset to A vector "Difference" and B vector "Zero", then change the Input to "Sum" and the Target to "Zero". Notice that ∇B becomes zero. The input `x = [1, 1]` is orthogonal to `a = [1, -1]`, making the scalar projection `a · x = 0`, freezing B's gradient for that specific data point.

Limitations

This module uses a simplified 3x2 toy matrix to show the exact arithmetic. A true LLM LoRA adaptation operates on massive matrices (e.g., 4096x4096). The parameter saving ratio here is modest (5 trainable parameters vs 6 for full fine-tuning), but in large matrices, a rank r adaptation requires r(m+n) parameters compared to m · n for full rank, which becomes highly efficient for r « m, n.

Additionally, this model assumes simple MSE loss on vectors and computes the instantaneous gradient. It does not run an optimization loop (like Adam) or model multi-layer dependencies found in a true Transformer architecture.

References

For more information on the motivation and design of this technique, refer to the primary literature below:

LoRA: Low-Rank Adaptation of Large Language Models (Hu et al., 2021). This foundational paper introduces the method of freezing pre-trained model weights and training rank decomposition matrices. It specifies the critical zero initialization strategy for matrix B to ensure parity with the base model at initialization.

LoRA Factor Gradients and Initialization