LoRA Mathematics › LoRA Gradient Computation

Gradient flow through LoRA reveals the training dynamics and connects to the optimization concepts from Part X. The gradient formulas explain why the zero initialization of $B$ does not prevent learning, how $A$ and $B$ co-evolve during training, and why the optimization objective makes LoRA training well-behaved.

Understanding gradients has practical value for LoRA. The specific structure of the gradient formulas determines which parameters learn first, how quickly the two matrices co-adapt, and what happens when you change the scaling factor or rank. Practitioners who understand gradient flow can diagnose training problems more easily and tune LoRA more effectively.

Consider the forward pass for a single layer with LoRA adaptation:

$$h = W_0 x + \frac{\alpha}{r} BA x$$

where:

• $h$: the output vector of the adapted layer
• $W_0$: the pre-trained weight matrix (frozen, no gradients computed)
• $x$: the input vector to the layer
• $\alpha$: the scaling factor
• $r$: the rank
• $B, A$: the LoRA up-projection and down-projection matrices (both trainable)

Let $\mathcal{L}$ be the loss function we are minimizing. To update $A$ and $B$ via gradient descent, we need $\frac{\partial \mathcal{L}}{\partial A}$ and $\frac{\partial \mathcal{L}}{\partial B}$. The chain rule traces how parameter changes propagate through the computation to affect the loss: a small change in $A$ changes the intermediate activation $Ax$, which changes the LoRA output $BAx$, which changes the layer output $h$, which eventually changes the loss $\mathcal{L}$.

First, let us define the upstream gradient. Let $\frac{\partial \mathcal{L}}{\partial h} = g \in \mathbb{R}^d$ be the gradient of the loss with respect to the layer's output. This gradient comes from the layers above in the network and tells us how changes in $h$ affect the loss. It is computed by standard backpropagation through the rest of the network and is the "signal" that drives adaptation in both $A$ and $B$.

For the gradient with respect to $B$, we apply the chain rule. The LoRA contribution to $h$ is $\frac{\alpha}{r} B(Ax)$. Treating the product $Ax$ as a fixed vector (since we are differentiating with respect to $B$, not $A$), the derivative of $B(Ax)$ with respect to $B$ has a clean outer product form:

$$\begin{aligned} \frac{\partial \mathcal{L}}{\partial B} &= \frac{\alpha}{r} \frac{\partial \mathcal{L}}{\partial h} \cdot \frac{\partial h}{\partial B} && \text{(chain rule)} \\ &= \frac{\alpha}{r} g (Ax)^T && \text{(derivative of linear map)} \end{aligned}$$

where:

• $\mathcal{L}$: the loss function
• $g = \frac{\partial \mathcal{L}}{\partial h}$: the upstream gradient vector from layers above
• $Ax$: the intermediate vector (input projected down to rank $r$)
• $(Ax)^T$: the transpose of the intermediate vector
• $\alpha$: the scaling factor
• $r$: the rank

The gradient of $B$ is an outer product between the upstream gradient $g$ and the intermediate representation $Ax$. Each entry of the gradient $(B)_{il}$ measures how much increasing that entry would affect the loss, which depends on how much gradient signal arrives at output dimension $i$ (captured by $g_i$) and how active the corresponding bottleneck dimension $l$ was (captured by $(Ax)_l$). Dimensions that were heavily activated and carry strong gradient signal receive large updates, while inactive or gradient-free dimensions receive small updates.

For processing batches efficiently, we extend this to matrix form:

$$\frac{\partial \mathcal{L}}{\partial B} = \frac{\alpha}{r} G (AX)^T$$

where:

• $\mathcal{L}$: loss function
• $\alpha$: scaling factor
• $r$: rank
• $G \in \mathbb{R}^{d \times n}$: matrix of upstream gradients for the batch, with one column per example
• $X \in \mathbb{R}^{k \times n}$: matrix of input vectors for the batch, with one column per example
• $A$: down-projection matrix

For the gradient with respect to $A$, we again apply the chain rule, but now we trace through both the multiplication by $B$ and the multiplication by $x$:

$$\begin{aligned} \frac{\partial \mathcal{L}}{\partial A} &= \frac{\alpha}{r} B^T \cdot \frac{\partial \mathcal{L}}{\partial h} \cdot x^T && \text{(chain rule through B)} \\ &= \frac{\alpha}{r} B^T g x^T && \text{(substitute } g \text{)} \end{aligned}$$

where:

• $B^T$: transpose of the up-projection matrix, which back-projects the upstream gradient into the bottleneck space
• $\mathcal{L}$: the loss function
• $\alpha$: the scaling factor
• $r$: the rank
• $g$: the upstream gradient vector
• $x$: the input vector to the layer
• $x^T$: the transpose of the input vector

This formula shows that the gradient with respect to $A$ involves two stages: first, the upstream gradient $g$ is projected back through $B$ using $B^T g$, which maps the output-space gradient into the bottleneck space. Then, this bottleneck-space gradient is combined with the input $x$ to form the outer product that updates $A$. In matrix form for batch processing:

$$\frac{\partial \mathcal{L}}{\partial A} = \frac{\alpha}{r} B^T G X^T$$

where:

• $B^T$: transpose of the up-projection matrix
• $G$: upstream gradient matrix (one column per batch example)
• $X^T$: transpose of the input batch matrix

Several properties of these gradients are noteworthy and reveal important aspects of LoRA's training dynamics. These properties have direct practical implications for how training progresses and what kinds of convergence behavior you should expect.

Independence from $W_0$: The gradients $\frac{\partial \mathcal{L}}{\partial A}$ and $\frac{\partial \mathcal{L}}{\partial B}$ are independent of $W_0$. Frozen weights affect the upstream gradient $g$ through their contribution to the forward pass, but they never appear explicitly in the LoRA gradient formulas. This separation enables efficient training where $W_0$ is never updated. You do not need to compute or store gradients for the much larger frozen weight matrices, which is the key source of LoRA's memory savings during training.

Coupled dynamics: Although $A$ and $B$ are separate parameters, their gradients are intimately coupled:

• $\frac{\partial \mathcal{L}}{\partial B}$ depends on $A$ through the term $Ax$
• $\frac{\partial \mathcal{L}}{\partial A}$ depends on $B$ through the term $B^T g$

This coupling means $A$ and $B$ co-evolve during training, with $A$'s current value determining $B$'s gradient and vice versa. This is similar to the dynamics in other factorized parameterizations like matrix factorization for recommendation systems, leading to interesting optimization behavior including potential saddle points and multiple stable solutions that represent the same effective weight update through different factorizations.

Zero initialization dynamics: At initialization ($B = 0$), the gradients behave asymmetrically:

• $\frac{\partial \mathcal{L}}{\partial B} = \frac{\alpha}{r} g (Ax)^T \neq 0$ (generally non-zero, since $A$ is random and $g$ is non-zero)
• $\frac{\partial \mathcal{L}}{\partial A} = \frac{\alpha}{r} B^T g x^T = 0$ (exactly zero because $B = 0$)

This asymmetry is important: $B$ updates immediately from the first gradient step, while $A$ has zero gradient initially. After the first update step, $B \neq 0$, and both matrices begin receiving non-zero gradients. The initial zero-gradient phase for $A$ is brief, lasting only a single step in theory, though in practice the gradients for $A$ remain small until $B$ has moved appreciably away from zero. This "bootstrap" phase, where $B$ adapts alone before both matrices co-evolve, is a direct consequence of the asymmetric initialization.

LoRA's gradient computations are efficient, requiring operations proportional to the parameter count rather than the full weight matrix size. This efficiency is not accidental; it is a direct consequence of the factorized structure and is the key reason LoRA training is so much faster than full fine-tuning.

Computing $\frac{\partial \mathcal{L}}{\partial B} = \frac{\alpha}{r} g z^T$ where $z = Ax$ involves the following costs: during the forward pass, we must store $z \in \mathbb{R}^r$, adding minimal memory overhead of just $r$ numbers per training example. The gradient computation itself is an outer product between vectors of dimension $d$ and $r$, requiring $O(dr)$ operations and creating a $d \times r$ matrix that is directly usable as the gradient for $B$.

Computing $\frac{\partial \mathcal{L}}{\partial A} = \frac{\alpha}{r} B^T g x^T$ requires first computing $B^T g$, which is a matrix-vector product costing $O(dr)$ operations (transforming the $d$-dimensional upstream gradient into an $r$-dimensional signal using $B^T$), followed by forming the outer product with $x^T$, which costs $O(rk)$ operations and produces an $r \times k$ matrix that is the gradient for $A$.

The total gradient computation is therefore $O(r(d + k))$, matching the parameter count. This is essentially optimal: we need at least one operation per parameter to compute a gradient, and LoRA achieves this lower bound up to constant factors. This efficiency extends to the memory required for gradient storage, which also scales with $r(d+k)$ rather than $dk$, explaining why LoRA reduces GPU memory requirements by over 99% compared to full fine-tuning for large models.