LoRA from a Gradient Perspective: Introduction, Analysis, Speculation, and Extensions

梯度视角下的LoRA：简介、分析、猜测及推广 苏剑林 · 2023-04-17 · Mathematics gradient optimizer low-rank LoRA

With the popularity of ChatGPT and its various imitators, parameter-efficient fine-tuning methods have also been on the rise, and one of the most popular schemes is the protagonist of this post, LoRA, from the paper LoRA: Low-Rank Adaptation of Large Language Models. LoRA is fairly simple and direct in its methodology, and there are quite a few off-the-shelf implementations already available, so it's easy to get up to speed with it whether for understanding or for practical use — there isn't really that much to elaborate on about the method itself.

However, implementing LoRA directly requires modifying the network architecture, which is a bit of a hassle. At the same time, LoRA has always struck me as quite similar to the earlier optimizer AdaFactor, so the question I want to pursue is: can we analyze and implement LoRA from the perspective of the optimizer instead? This post is built around that theme.

Method Overview

Some earlier results (e.g., Exploring Universal Intrinsic Task Subspace via Prompt Tuning) show that although a pretrained model has a huge number of parameters, the intrinsic dimension corresponding to each downstream task is actually quite small. In other words, in theory we should be able to fine-tune a very small number of parameters and still achieve good results on a downstream task.

Building on this observation, LoRA proposes that for a pretrained parameter matrix $W_0\in\mathbb{R}^{n\times m}$, instead of directly fine-tuning $W_0$, we make a low-rank decomposition assumption on the update:

\begin{equation}W = W_0 + A B,\qquad A\in\mathbb{R}^{n\times r},B\in\mathbb{R}^{r\times m}\end{equation}more

Here one of $A,B$ is initialized to all zeros, $W_0$ is kept fixed, and the optimizer only optimizes $A,B$. Since the intrinsic dimension is small, we can make $r$ quite small — a common choice is $r=8$, and in extreme cases we can even take $1$. So LoRA is a parameter-efficient fine-tuning method, at least in the sense that the number of parameters actually being optimized is drastically reduced.

Here's a schematic diagram drawn directly with MathJax:

$$\style{display: inline-block; width: 24ex; padding: 10ex 0; border: 1px solid #6C8EBF; background-color: #DAE8FC}{W_0\in\mathbb{R}^{n\times m}} \quad + \quad \style{display: inline-block; width: 8ex; padding: 10ex 0; border: 1px solid #D79B00; background-color: #FFE6CC}{A\in\mathbb{R}^{n\times r}}\quad\times\quad \style{display: inline-block; width: 24ex; padding: 3ex 0; border: 1px solid #D79B00; background-color: #FFE6CC}{B\in\mathbb{R}^{r\times m}}$$

Gradient Analysis

As mentioned in Ladder Side-Tuning: A "Ladder over the Wall" for Pretrained Models, many parameter-efficient fine-tuning methods actually only reduce GPU memory requirements without reducing the amount of computation. So is LoRA an exception? How efficient is it really in terms of memory and compute? Let's analyze this below.

First, we know that the GPU memory consumed during training comes from four parts: model parameters, model gradients, model activations, and optimizer states. LoRA reduces the number of model parameters via low-rank decomposition, so the gradients and optimizer states shrink accordingly — the memory savings here are obvious. But can it also save on computation?

That depends on how LoRA is implemented — different implementations lead to different complexities when computing gradients. There are two equivalent implementations of LoRA:

\begin{align}Y =&\, XW = X(W_0 + AB) \label{eq:lora-1}\\[5pt] Y =&\, XW_0 + XAB = XW_0 + ZB \label{eq:lora-2}\end{align}

Here $X\in\mathbb{R}^{b\times n}$ is the model input and $Z=XA\in\mathbb{R}^{b\times r}$ is the intermediate output. For implementation $\eqref{eq:lora-1}$, we have

\begin{equation}\frac{\partial \mathcal{L}}{\partial A} = \frac{\partial \mathcal{L}}{\partial W} B^{\top} = \left(X^{\top}\frac{\partial \mathcal{L}}{\partial Y}\right) B^{\top},\quad \frac{\partial \mathcal{L}}{\partial B} = A^{\top}\frac{\partial \mathcal{L}}{\partial W} = A^{\top}\left(X^{\top}\frac{\partial \mathcal{L}}{\partial Y}\right)\label{eq:grad-1}\end{equation}

where $\mathcal{L}$ is the loss function. Clearly, the consequence of this implementation is that we need to compute the full gradient $\frac{\partial \mathcal{L}}{\partial W}\in\mathbb{R}^{n\times m}$ first before we can compute the gradient of $A,B$, which means it's actually slower than not using LoRA at all, and it's memory-hungry too. For implementation $\eqref{eq:lora-2}$, on the other hand, we have

\begin{equation}\frac{\partial \mathcal{L}}{\partial A} = X^{\top}\frac{\partial \mathcal{L}}{\partial Z} = X^{\top}\left(\frac{\partial \mathcal{L}}{\partial Y} B^{\top}\right),\quad \frac{\partial \mathcal{L}}{\partial B} = Z^{\top}\frac{\partial \mathcal{L}}{\partial Y} = (XA)^{\top}\frac{\partial \mathcal{L}}{\partial Y}\label{eq:grad-2}\end{equation}

Here $Z,\frac{\partial \mathcal{L}}{\partial Z}\in\mathbb{R}^{b\times r}$ is clearly much cheaper than the full gradient, and the computational complexity is significantly lower. So, in order for LoRA to maximize memory and compute savings, the key is to implement it as $\eqref{eq:lora-2}$ rather than $\eqref{eq:lora-1}$.

(Note: for computing gradients with respect to matrices, we can often "work it out" using the chain rule together with the required output shape. For example, for $\frac{\partial \mathcal{L}}{\partial A}$, the chain rule tells us it must be some product of $\frac{\partial \mathcal{L}}{\partial W}$ and $B$. If we stipulate that the shape of $\frac{\partial \mathcal{L}}{\partial A}$ matches that of $A$, i.e., $n\times r$, then to combine $\frac{\partial \mathcal{L}}{\partial W}$ and $B$ into a result of shape $n\times r$, the only option is $\frac{\partial \mathcal{L}}{\partial W} B^{\top}$.)

Other Reasons

Besides the benefits brought by low-rank decomposition, the following points also contribute to LoRA's memory savings and speedup:

1. Only part of the parameters are updated: for instance, the original LoRA paper chose to update only the Self-Attention parameters, and in practice we can also choose to update only some of the layers;
2. Reduced communication time: since fewer parameters are being updated, (especially in multi-GPU training) less data needs to be transmitted, reducing transfer time;
3. Various low-precision acceleration techniques are used, such as FP16, FP8, or INT8 quantization.

Of course, these three factors do speed up training, but they are not unique to LoRA — in fact, almost all parameter-efficient methods share these characteristics. LoRA's standout advantage is that its low-rank decomposition is very intuitive, matches the performance of full fine-tuning in many scenarios, and at inference time we can directly merge $W_0,A,B$ into a single matrix so that there is no extra inference cost.

An Optimization Perspective

The gradient $\eqref{eq:grad-1}$ also tells us how to implement LoRA from the optimizer's perspective. The optimizer can directly obtain the full gradient $\frac{\partial \mathcal{L}}{\partial W}$, and then we simply need to project the gradient according to formula $\eqref{eq:grad-1}$ to get the gradient of $A,B$, after which we can update $A,B$ using a standard optimizer implementation.

If the optimizer is SGD, this becomes

\begin{equation}\begin{aligned} A_{t+1} =&\, A_t - \eta\frac{\partial \mathcal{L}}{\partial W_t} B_t^{\top},\quad B_{t+1} = B_t - \eta A_t^{\top}\frac{\partial \mathcal{L}}{\partial W_t}\\[5pt] W_{t+1} =&\, W_0 + A_{t+1} B_{t+1} = W_t + (A_{t+1} B_{t+1} - A_t B_t) \end{aligned}\end{equation}

If it's an optimizer with sliding/moving variables like Adam, then we only need to keep a running average of the projected gradient, which reduces the parameter count of the optimizer states and saves some memory. The larger the model, the larger the proportion of memory taken up by this part.

LoRA stipulates that one of $A$ or $B$ is initialized to all zeros, which ensures that the initial state of the model matches the pretrained model, but it also introduces an asymmetry (one is all-zero, the other is not). In fact, it's also fine to initialize both $A,B$ with non-zero values — we just need to first subtract $A_0 B_0$ from the pretrained weights beforehand, or equivalently, parameterize $W$ as

\begin{equation}W = W_0 - A_0 B_0 + A B\end{equation}

This keeps the initial state consistent while allowing both $A,B$ to be initialized with non-zero values, which enhances the symmetry.

Random Projection

If we expand the update $A_{t+1} B_{t+1} - A_t B_t$ in the SGD setting, the result is

\begin{equation}- \eta\left(\frac{\partial \mathcal{L}}{\partial W_t} B_t^{\top} B_t + A_t A_t^{\top}\frac{\partial \mathcal{L}}{\partial W_t}\right) + \eta^2 \frac{\partial \mathcal{L}}{\partial W_t} B_t^{\top} A_t^{\top}\frac{\partial \mathcal{L}}{\partial W_t}\end{equation}

Assuming the $\eta^2$ term is a negligible higher-order term, what remains is

\begin{equation}- \eta\left(\frac{\partial \mathcal{L}}{\partial W_t} B_t^{\top} B_t + A_t A_t^{\top}\frac{\partial \mathcal{L}}{\partial W_t}\right)\end{equation}

From this perspective, compared to full fine-tuning with SGD, LoRA is essentially replacing the full gradient $\frac{\partial \mathcal{L}}{\partial W_t}$ with the term in the parentheses.

For simplicity, let's focus only on the case $r=1$. Notice that in the expression above, the projection vector $A_t,B_t$ at time step $t$ depends on $t$. What happens if we replace it with a random vector that doesn't depend on $t$ (re-generated randomly at every training step)? Consider $u,v\sim\mathcal{N}(0,1)$, where $u\in\mathbb{R}^{m\times 1}, v\in\mathbb{R}^{1\times n}$, so the update becomes

\begin{equation}- \eta\left(\frac{\partial \mathcal{L}}{\partial W_t} v^{\top} v + u u^{\top}\frac{\partial \mathcal{L}}{\partial W_t}\right)\end{equation}

It can be shown that

\begin{equation}\mathbb{E}_{u\sim \mathcal{N}(0,1)}[u u^{\top}] = I_{n\times n},\quad \mathbb{E}_{v\sim \mathcal{N}(0,1)}[v^{\top} v] = I_{m\times m}\end{equation}

Here $I_{n\times n},I_{m\times m}$ refer respectively to the identity matrices of size $n\times n,m\times m$. So, similar to "zeroth-order gradients", in an averaged sense, this version of LoRA — where the projection is re-initialized at every step — is in fact equivalent to full-rank SGD. However, if we were to actually implement it this way, it might well be even slower than full-rank SGD, so the point of doing this isn't speed — it's the hope that it might alleviate catastrophic forgetting: by using a low-rank (rather than full-rank) update for each individual (batch of) sample(s), the impact on the overall model weights is reduced. Of course, this is just speculation on my part — I haven't actually experimented with it to see how it performs in practice.

A Variant

Again, let's first consider only the case $r=1$. LoRA is equivalent to assuming $\Delta w_{i,j} = u_i v_j$ — can we make some other low-rank decomposition assumption instead? For example, $\Delta w_{i,j} = u_i + v_j$? Written in matrix form, this becomes

\begin{equation}W = W_0 + A \mathbb{1}_{1\times m} + \mathbb{1}_{n\times 1} B,\qquad A\in\mathbb{R}^{n\times 1},B\in\mathbb{R}^{1\times m}\end{equation}

where $\mathbb{1}_{1\times m},\mathbb{1}_{n\times 1}$ refer respectively to the all-ones matrices of shape $1\times m,n\times 1$. It's easy to work out that the gradient is:

\begin{equation}\frac{\partial \mathcal{L}}{\partial A} = \frac{\partial \mathcal{L}}{\partial W} \mathbb{1}_{m\times 1},\quad \frac{\partial \mathcal{L}}{\partial B} = \mathbb{1}_{1\times n}\frac{\partial \mathcal{L}}{\partial W}\end{equation}

which is simply the row-sum and column-sum of the original gradient. Compared to the original LoRA, this additive decomposition has two advantages: 1) addition is cheaper than multiplication, and the gradient form is simpler too; 2) the rank of $AB$ is necessarily 1, but the rank of $A \mathbb{1}_{1\times m} + \mathbb{1}_{n\times 1} B$ can be 2 — if rank represents model capacity, this means that with the same number of parameters, the additive form may actually have stronger expressive power. As for how well it actually performs, I'll leave that to a future comparative experiment once I get around to using LoRA myself.

So, can this additive decomposition be extended to the case $r > 1$? Naturally, yes, though it requires a bit of a trick. Here we stipulate that $m,n$ is divisible by $r$, and then we just need to change the parameterization to

\begin{equation}W = W_0 + A I_{r(1\times m/r)} + I_{r(n/r\times 1)} B,\qquad A\in\mathbb{R}^{n\times r},B\in\mathbb{R}^{r\times m}\end{equation}

Here $I_{r(1\times m/r)}$ and $I_{r(n/r\times 1)}$ refer respectively to block matrices of shape $1\times m/r$ and $n/r\times 1$, where each block is the identity matrix of size $r\times r$. In plain terms, this amounts to treating $A$ and $B$ as block matrices with block structure $n/r\times 1$ and $1\times m/r$, respectively, and then applying the same idea as $r=1$.

Summary

This post introduced how to understand LoRA from a gradient perspective. Besides the basic introduction, it also includes some speculation and extensions of my own, for readers' reference.