Can LoRA still be improved by using different learning rates?

配置不同的学习率，LoRA还能再涨一点？ 苏剑林 · 2024-02-27 · Machine Learning gradient optimizer low-rank LoRA

LoRA (Low-Rank Adaptation) is one of the current parameter-efficient fine-tuning methods for LLMs. We previously discussed it briefly in LoRA from a Gradient Perspective: Introduction, Analysis, Speculation, and Generalization. In this article, let's take a look at a new result about LoRA:

Assigning different learning rates to the two matrices of LoRA can further improve LoRA's performance.

This result comes from the recent paper LoRA+: Efficient Low Rank Adaptation of Large Models (hereafter "LoRA+"). At first glance, this conclusion doesn't seem particularly special, since setting different learning rates is essentially introducing a new hyperparameter, and generally introducing and tuning hyperparameters tends to yield some improvement. What makes "LoRA+" special is that it confirms this necessity theoretically, and asserts that the optimal solution must have the learning rate of the right-hand matrix greater than that of the left-hand matrix. In short, "LoRA+" is a fine example of theory guiding training that also proves effective in practice, and it's well worth studying carefully.

A Brief Look at the Conclusion

Suppose the pretrained parameter is $W_0 \in \mathbb{R}^{n\times m}$. If we use full-parameter fine-tuning, the update is also a $n\times m$ matrix. To reduce the parameter count, LoRA constrains the update to be a low-rank matrix, i.e., letting $W=W_0 + AB$, where $A\in\mathbb{R}^{n\times r},B\in\mathbb{R}^{r\times m}$ and $r\ll \min(n,m)$. We replace the original model parameter with the new $W$, then keep $W_0$ fixed and only update $A,B$ during training, as shown in the figure below:

$$\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}}$$more

Note that LoRA is usually applied to Dense layers, but the analysis in the original paper is based on weights left-multiplying the input, whereas in practice the input almost always right-multiplies the weight. To avoid confusion, the notation in this article is aligned with the implementation, i.e., we assume the layer input is $X\in\mathbb{R}^{b\times n}$ and the layer operation is $XW = X(W_0 + AB)$. Since the conclusion of "LoRA+" is independent of the pretrained weight, without loss of generality we can set $W_0=0$, so the layer operation simplifies to $Y=XAB\in\mathbb{R}^{b\times m}$.

The conclusion of "LoRA+" is:

To make LoRA's performance as close to optimal as possible, the learning rate for weight $B$ should be larger than the learning rate for weight $A$.

Note that, to make the initial model equivalent to the original pretrained model, LoRA usually initializes one of $A,B$ to all zeros. The author initially thought this conclusion resulted from this all-zero initialization, and thus should depend on which matrix is zero-initialized. However, after careful reading, it turns out that the conclusion claimed by "LoRA+" has nothing to do with the all-zero initialization. In other words, although $A,B$ appear symmetric on the surface, they in fact have an inherent asymmetry, such that regardless of whether $A$ or $B$ is zero-initialized, the conclusion is always that the learning rate of $B$ should be greater than that of $A$. This is where it gets interesting.

That said, it must be noted that the explanation in the original "LoRA+" paper is rather hard to follow, so what follows below is the author's own attempt at a simplified derivation. Broadly speaking, it rests on two assumptions:

1. Numerical stability: The output value of every layer of the model should be numerically stable, independent of the network width;
2. Equal contribution: For LoRA to be optimal, the two matrices $A,B$ should contribute to the outcome to an equal degree.

Let's analyze and quantify these two assumptions one by one.

Numerical Stability

First, numerical stability means that each component of $X,XA,XAB$ should be of order $\mathcal{O}(1)$, independent of the network width $n,m$. Here $\mathcal{O}(1)$ mainly describes that its order with respect to network width is zero, not that its absolute value is necessarily close to 1. This assumption should not be controversial — it's hard to imagine a numerically unstable network achieving good predictive performance. However, some readers might question the necessity of "$XA$ is $\mathcal{O}(1)$," since $X$ is the input and $XAB$ is the output, and requiring numerical stability for both of them is reasonable — but $XA$ is merely an intermediate variable, so does it also need to be numerically stable?

Looking purely at the forward pass, numerical stability of $XA$ is indeed not strictly necessary. But if $XA$ is numerically unstable while $XAB$ is numerically stable, there are two cases: either $XA$ is too large and $B$ is too small, in which case, according to the chain rule for derivatives, the gradient of $A$ will be too small and the gradient of $B$ will be too large; or conversely, $XA$ is too small and $B$ is too large, which will make the gradient of $A$ too large and the gradient of $B$ too small. In short, numerical instability of $XA$ leads to gradient instability for $A,B$, which increases optimization difficulty — so it's better to include the numerical stability of $XA$ as a condition.

This numerical stability condition naturally reminds us of "LeCun initialization," which says that if $W\in\mathbb{R}^{n\times m}$ are sampled i.i.d. from a distribution with "mean 0, variance $1/n$," then the order of magnitude of each component of $XW$ is roughly the same as that of the components of $X$. Following the same strategy, if the input $X$ is already $\mathcal{O}(1)$, then in order for the components of $XA,XAB$ to also be of order $\mathcal{O}(1)$, $A,B$ should each be initialized with variance $1/n,1/r$ (with mean defaulting to 0, which I won't repeat below).

Of course, as mentioned earlier, to preserve the identity of the initialization, LoRA chooses to zero-initialize one of $A,B$, but this is not particularly important here. We just need to recognize that the variance of $1/n,1/r$ keeps $XA,XAB$ numerically stable, so we can conjecture that after training completes, $A,B$ will likely also approximately have variance $1/n,1/r$. Given $r \ll n$, this is equivalent to saying that the absolute values of the components of $A$ will be noticeably smaller than those of $B$ — this is the origin of the asymmetry between $A,B$.

Equal Contribution

Next, let's look at the second assumption: $A,B$ should contribute to the outcome to an equal degree. This assumption also seems reasonable, since in the LLM + LoRA setting we usually have $m=n$, i.e., $A,B$ have the same number of parameters, so it's reasonable that they should contribute equally to the outcome. If $m\neq n$, we could further generalize this assumption to say that the contribution is proportional to the parameter count. The most fundamental measure of "outcome" is naturally the loss function, denoted here as $\mathcal{L}$.

We want to measure the change in the loss function when $A\to A+\Delta A,B\to B + \Delta B$:

\begin{equation}\mathcal{L}(A+\Delta A,B+\Delta B) - \mathcal{L}(A,B)\approx \left\langle \frac{\partial\mathcal{L}}{\partial A},\Delta A\right\rangle + \left\langle \frac{\partial\mathcal{L}}{\partial B},\Delta B\right\rangle\label{eq:delta-loss}\end{equation}

Here we've used a first-order linear approximation, where $\frac{\partial\mathcal{L}}{\partial A},\frac{\partial\mathcal{L}}{\partial B}$ is the gradient of $A,B$, and $\langle\cdot,\cdot\rangle$ is the (Frobenius) inner product operation. The two terms on the right-hand side can then be understood as the respective contributions of $A,B$ to the outcome. But note that the validity of this linear approximation depends on the increment $\Delta A,\Delta B$ being small — yet for a fully trained weight, its increment relative to the original weight may not actually be small at all. So instead, we relax the "equal contribution" assumption to "$A,B$ should contribute to the outcome to an equal degree at every single update step," since the amount of change in a single update step is usually small, so the linear approximation is much better satisfied there.

Since we now need to consider the update at each step, this naturally leads us toward the optimizer. Since the mainstream optimizer for both pretraining and fine-tuning today is Adam, we'll take Adam as our main object of analysis. As we know, the Adam optimizer has two sets of moving-average states along with corresponding hyperparameters $\beta_1,\beta_2$, which makes a precise analysis rather difficult. But for the purposes of this article, we only need an order-of-magnitude estimate, so we try to consider only one extreme case and assume it gives the same order-of-magnitude estimate as the general case. This extreme case is $\beta_1=\beta_2=0$, where Adam degenerates into SignSGD:

\begin{equation}\Delta A = -\eta_A\,\text{sign}\left(\frac{\partial\mathcal{L}}{\partial A}\right),\quad\Delta B = -\eta_B\,\text{sign}\left(\frac{\partial\mathcal{L}}{\partial B}\right)\label{eq:sign-sgd}\end{equation}

where $\eta_A,\eta_B$ are the respective learning rates, and the conclusion of "LoRA+" is that $\eta_B \gg \eta_A$.

Substituting the SignSGD increment $\eqref{eq:sign-sgd}$ back into equation $\eqref{eq:delta-loss}$, we get

\begin{equation}\mathcal{L}(A+\Delta A,B+\Delta B) - \mathcal{L}(A,B)\approx \underbrace{-\,\eta_A \left\Vert\frac{\partial\mathcal{L}}{\partial A}\right\Vert_1}_{\Delta \mathcal{L}_A}\,\underbrace{-\,\eta_B \left\Vert \frac{\partial\mathcal{L}}{\partial B}\right\Vert_1}_{\Delta \mathcal{L}_B}\end{equation}

Here $\Vert\cdot\Vert_1$ is the $L_1$ norm, i.e., the sum of the absolute values of all components. "Equal contribution" means we want the two terms on the right, $\Delta \mathcal{L}_A,\Delta \mathcal{L}_B$, to be of the same order of magnitude.

A Quick Derivation

Further analysis requires the explicit form of the gradient. Again setting $Y=XAB$, we can derive:

\begin{equation}\frac{\partial \mathcal{L}}{\partial A} = X^{\top}\frac{\partial \mathcal{L}}{\partial Y}B^{\top},\quad \frac{\partial \mathcal{L}}{\partial B} = A^{\top} X^{\top}\frac{\partial \mathcal{L}}{\partial Y}\end{equation}

Readers unfamiliar with matrix calculus may be puzzled by how the above results are derived — to be honest, the author isn't entirely comfortable with it either, but there's a simple trick we can use. For instance, $\frac{\partial \mathcal{L}}{\partial A}$ — we know it's a $n\times r$ matrix (same shape as $A$); likewise, $\frac{\partial \mathcal{L}}{\partial Y}$ is a $b\times m$ matrix. According to the chain rule, $\frac{\partial \mathcal{L}}{\partial A}$ should be some product of $\frac{\partial \mathcal{L}}{\partial Y}$, $X$, and $B$, so we just need to figure out, according to the rules of matrix multiplication, how these three matrices should be multiplied together to produce a matrix of shape $n\times r$.

Once we have the explicit form of $\frac{\partial \mathcal{L}}{\partial A},\frac{\partial \mathcal{L}}{\partial B}$, there's a quick way to understand LoRA+. First, $\Delta \mathcal{L}_A$ is proportional to $\left\Vert\frac{\partial\mathcal{L}}{\partial A}\right\Vert_1$, which is the sum of the absolute values of $nr$ components; assuming each component is comparable in magnitude, this means $\Delta \mathcal{L}_A$ is roughly proportional to $nr$. Next, $\frac{\partial\mathcal{L}}{\partial A}$ is linear in $B$, so we can roughly say that each component of $\frac{\partial\mathcal{L}}{\partial A}$ is proportional in magnitude to the components of $B$; combining these, $\Delta \mathcal{L}_A$ is proportional to the magnitudes of both $nr$ and $B$ simultaneously. Similarly, $\Delta \mathcal{L}_B$ is also roughly proportional to the magnitudes of both $mr$ and $A$. Earlier, in the "Numerical Stability" section, we said that for forward numerical stability, the magnitude of $B$ should be greater than that of $A$ (proportional to their approximate standard deviations $\sqrt{1/r},\sqrt{1/n}$). So for $\Delta \mathcal{L}_A$ and $\Delta \mathcal{L}_B$ to be comparable in magnitude, we should have approximately:

\begin{equation}\eta_A \times nr \times \sqrt{1/r} \approx \eta_B \times mr \times \sqrt{1/n}\quad\Rightarrow\quad \frac{\eta_B}{\eta_A} \approx \frac{n}{m}\sqrt{\frac{n}{r}}\end{equation}

Given that in practice we often have $m=n$ and $r=\mathcal{O}(1)$, this can be simply written as

\begin{equation}\frac{\eta_B}{\eta_A} = \mathcal{O}(\sqrt{n})\end{equation}

But we're not done yet — we need to check whether this result is self-consistent, because one of the conditions we used, "forward numerical stability," has so far only been an idealized assumption. How can we make this assumption hold as closely as possible? The way to reinforce one assumption is to introduce another:

In the Adam optimizer, if the ratio of the learning rates of two parameters is $\lambda$, then after prolonged training, the ratio of the magnitudes of those two parameters will also be $\lambda$.

According to Adam's approximate update rule $\eqref{eq:sign-sgd}$, the magnitude of each step's increment is indeed proportional to the learning rate, but the overall training result is not simply the sum of each step in a straightforward way, so this assumption feels like "it kind of makes sense, but not entirely." Still, that's typical of assumptions — as long as they make some sense, that's enough, and the rest is a matter of faith. Under this assumption, if we train with the learning rate ratio $\frac{\eta_B}{\eta_A} = \mathcal{O}(\sqrt{n})$, then the magnitude ratio of the two parameters $B,A$ will also be $\mathcal{O}(\sqrt{n})$; and since we previously expected them to have comparable standard deviation ratio $\sqrt{1/r},\sqrt{1/n}$, these two ratios turn out to be exactly $\mathcal{O}(\sqrt{n})$ — the result is fully self-consistent!

The result in the original paper differs slightly from the one above: it gives the answer $\mathcal{O}(n)$, because the original paper considers $\Delta A,\Delta B$ contributing equally to $Y$, but $Y$ is just the output of the model layer, which does not represent the final outcome — so this is not quite appropriate. Although the original paper does attempt to relate the increment of $Y$ to the increment of $\mathcal{L}$, it doesn't carry out the calculation carefully, leading to a biased result. Moreover, the derivation in the original paper, strictly speaking, only applies to the special case of $b=1,r=1,m=n$; the general case of $b > 1, r > 1$ is adopted directly without justification, which means the analysis is actually not fully general.

Of course, whether it's exactly $\mathcal{O}(n)$ or $\mathcal{O}(\sqrt{n})$ isn't really that important in practice — you still need to tune it. But LoRA+ ran experiments across models of various sizes, with $r$ generally set to 8 and $n$ ranging from 768 to 4096, and ultimately arrived at a recommended default learning rate ratio of $2^4 = 16$, which happens to be close to $\sqrt{n/r}$ — so the optimal value is closer to $\mathcal{O}(\sqrt{n})$ rather than $\mathcal{O}(n)$.

Summary

In this article, we introduced and derived a result called "LoRA+," which shows that the two low-rank matrices $A,B$ in LoRA have an inherent asymmetry: regardless of which matrix is zero-initialized, the learning rate for $B$ should be set larger than that for $A$ in order to achieve better performance.