梯度视角下的LoRA：简介、分析、猜测及推广

随着ChatGPT及其平替的火热，各种参数高效（Parameter-Efficient）的微调方法也“水涨船高”，其中最流行的方案之一就是本文的主角LoRA了，它出自论文《LoRA: Low-Rank Adaptation of Large Language Models》。LoRA方法上比较简单直接，而且也有不少现成实现，不管是理解还是使用都很容易上手，所以本身也没太多值得细写的地方了。

然而，直接实现LoRA需要修改网络结构，这略微麻烦了些，同时LoRA给笔者的感觉是很像之前的优化器AdaFactor，所以笔者的问题是：能否从优化器角度来分析和实现LoRA呢？本文就围绕此主题展开讨论。

方法简介 #

以往的一些结果（比如《Exploring Universal Intrinsic Task Subspace via Prompt Tuning》）显示，尽管预训练模型的参数量很大，但每个下游任务对应的本征维度（Intrinsic Dimension）并不大，换句话说，理论上我们可以微调非常小的参数量，就能在下游任务取得不错的效果。

LoRA借鉴了上述结果，提出对于预训练的参数矩阵$W_0\in\mathbb{R}^{m\times n}$，我们不去直接微调$W_0$，而是对增量做低秩分解假设：
\begin{equation}W = W_0 + U V,\qquad U\in\mathbb{R}^{m\times r},V\in\mathbb{R}^{r\times n}\end{equation}
其中$U,V$之一用全零初始化，$W_0$固定不变，优化器只优化$U,V$。由于本征维度很小的结论，所以$r$我们可以取得很小，很多时候我们甚至可以直接取$1$。所以说，LoRA是一种参数高效的微调方法，至少被优化的参数量大大降低了。

梯度分析 #

正如《Ladder Side-Tuning：预训练模型的“过墙梯”》所提到的，很多参数高效的微调实际上只是降低了显存需求，并没有降低计算量，LoRA其实也不例外。为了认识到这一点，我们只需要观察$U,V$的梯度：
\begin{equation}\frac{\partial \mathcal{L}}{\partial U} = \frac{\partial \mathcal{L}}{\partial W} V^{\top},\quad \frac{\partial \mathcal{L}}{\partial V} = U^{\top}\frac{\partial \mathcal{L}}{\partial W}\label{eq:grad}\end{equation}
这里$\mathcal{L}$是损失函数，并且约定参数的shape跟其梯度的shape一致。在训练过程中，求模型梯度是主要的计算量，如果全量更新，那么所用的梯度是$\frac{\partial \mathcal{L}}{\partial W}$，而LoRA所用的梯度则是$\frac{\partial \mathcal{L}}{\partial U}$和$\frac{\partial \mathcal{L}}{\partial V}$，它们是建立在全量更新的梯度$\frac{\partial \mathcal{L}}{\partial W}$基础上的，所以理论上LoRA的计算量比全量更新还大。

那为什么实际使用时LoRA的训练速度也会变快呢？有如下几个原因：

1、只更新了部分参数：比如LoRA原论文就选择只更新Self Attention的参数，实际使用时我们还可以选择只更新部分层的参数；

2、减少了通信时间：由于更新的参数量变少了，所以（尤其是多卡训练时）要传输的数据量也变少了，从而减少了传输时间；

3、采用了各种低精度加速技术，如FP16、FP8或者INT8量化等。

这三部分原因确实能加快训练速度，然而它们并不是LoRA所独有的，事实上几乎都有参数高效方法都具有这些特点。LoRA的优点是它的低秩分解很直观，在不少场景下跟全量微调的效果一致，以及在预测阶段可以直接把$W_0,U,V$合并成单个矩阵从而不增加推理成本。

优化视角 #

事实上，梯度$\eqref{eq:grad}$也告诉了我们如何从优化器角度来实现LoRA。优化器可以直接获取到全量梯度$\frac{\partial \mathcal{L}}{\partial W}$，然后我们只需要按照公式$\eqref{eq:grad}$对梯度进行投影，就得到$U,V$的梯度，接着就可以按照常规的优化器实现$U,V$的更新了。

假如优化器是SGD，那么就是
\begin{equation}\begin{aligned}
U_{t+1} =&\, U_t - \eta\frac{\partial \mathcal{L}}{\partial W_t} V_t^{\top},\quad V_{t+1} = V_t - \eta U_t^{\top}\frac{\partial \mathcal{L}}{\partial W_t}\\[5pt]
W_{t+1} =&\, W_0 + U_{t+1} V_{t+1} = W_t + (U_{t+1} V_{t+1} - U_t V_t)
\end{aligned}\end{equation}
如果是Adam之类的带滑动变量的优化器，则只需要滑动投影后的梯度，因此是降低了优化器的参数量，节省了一定的显存。模型越大，这部分参数所占的显存比例也就越大。

LoRA约定$U$或$V$之一使用全零初始化，这是为了保证初始状态模型跟预训练一致，但同时也带来了不对称问题（一个全零，一个非全零）。事实上，$U,V$都使用非全零初始化也是可以的，只需要事先将预训练权重减去$U_0 V_0$就行了，或者等价地说，将$W$参数化为
\begin{equation}W = W_0 - U_0 V_0 + U V\end{equation}
这样同时保持了初始状态一致，同时允许$U,V$都用非全零初始化，增强了对称性。

随机投影 #

如果我们将SGD场景下的更新量$U_{t+1} V_{t+1} - U_t V_t$展开，结果将是
\begin{equation}- \eta\left(\frac{\partial \mathcal{L}}{\partial W_t} V_t^{\top} V_t + U_t U_t^{\top}\frac{\partial \mathcal{L}}{\partial W_t}\right) + \eta^2 \frac{\partial \mathcal{L}}{\partial W_t} V_t^{\top} U_t^{\top}\frac{\partial \mathcal{L}}{\partial W_t}\end{equation}
假设$\eta^2$项是可以忽略的高阶项，那么就剩下
\begin{equation}- \eta\left(\frac{\partial \mathcal{L}}{\partial W_t} V_t^{\top} V_t + U_t U_t^{\top}\frac{\partial \mathcal{L}}{\partial W_t}\right)\end{equation}
从这个角度来看，相比全量微调的SGD，LoRA就是用括号中的结果替代了全量的梯度$\frac{\partial \mathcal{L}}{\partial W_t}$。

简单起见，接下来我们只关心$r=1$的情形，留意到在上式中，$t$时刻的投影向量$U_t,V_t$是依赖于$t$的，如果我们将它们换成不依赖于$t$的随机向量（每步训练都重新随机生成），那么会发生什么呢？我们考虑$u,v\sim\mathcal{N}(0,1)$，其中$u\in\mathbb{R}^{m\times 1}, v\in\mathbb{R}^{1\times n}$，那么更新量就变为
\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}
可以证明的是
\begin{equation}\mathbb{E}_{u\sim \mathcal{N}(0,1)}[u u^{\top}] = I_{m\times m},\quad \mathbb{E}_{v\sim \mathcal{N}(0,1)}[v^{\top} v] = I_{n\times n}\end{equation}
这里的$I_{m\times m},I_{n\times n}$分别指$m\times m,n\times n$的单位矩阵。因此，跟“零阶梯度”类似，在平均意义下，这种每步都重新初始化的LoRA事实上等价于满秩的SGD。然而，真要按照这个方式实现的话，其速度甚至可能比满秩的SGD都要慢，所以它的目的不是提速，而是希望能缓解灾难遗忘问题——通过对单个（batch）样本使用低秩矩阵（而不是满秩）更新量的方式，减少对整个模型权重的影响。当然，这只是猜测，实际效果如何，笔者还没有实验过。

一个变体 #

同样还是先只考虑$r=1$的情形，LoRA相当于假设了$\Delta w_{i,j} = u_i v_j$，我们能不能做其他低秩分解假设呢？比如$\Delta w_{i,j} = u_i + v_j$？写成矩阵形式就是
\begin{equation}W = W_0 + U \mathbb{1}_{1\times n} + \mathbb{1}_{m\times 1} V,\qquad U\in\mathbb{R}^{m\times 1},V\in\mathbb{R}^{1\times n}\end{equation}
其中$\mathbb{1}_{1\times n},\mathbb{1}_{m\times 1}$分别指$1\times n,m\times 1$的全1矩阵。容易求出它的梯度是：
\begin{equation}\frac{\partial \mathcal{L}}{\partial U} = \frac{\partial \mathcal{L}}{\partial W} \mathbb{1}_{n\times 1},\quad \frac{\partial \mathcal{L}}{\partial V} = \mathbb{1}_{1\times m}\frac{\partial \mathcal{L}}{\partial W}\end{equation}
其实就是原本梯度的行求和与列求和。相比原版LoRA，这个加性分解有两个优点：1、加比乘计算量更低，梯度形式也更简单；2、$UV$的秩一定是1，但是$U \mathbb{1}_{1\times n} + \mathbb{1}_{m\times 1} V$的秩可能是2，如果秩代表了模型能力的话，那也就是说同样的参数量，加性的表达能力可能还更强。至于具体效果如何，后面笔者用到LoRA的时候，再做对比实验吧。

那么，加性分解能不能推广到$r > 1$的情形呢？自然是可以的，但稍微有些技巧。这里约定$m,n$都能被$r$整除，那么我们只需要将参数化方式改为
\begin{equation}W = W_0 + U I_{r(1\times n/r)} + I_{r(m/r\times 1)} V,\qquad U\in\mathbb{R}^{m\times r},V\in\mathbb{R}^{r\times n}\end{equation}
这里的$I_{r(1\times n/r)}$、$I_{r(m/r\times 1)}$分别指$1\times n/r$、$m/r\times 1$的分块矩阵，每一块则是$r\times r$的单位阵。这个形式说白了，就是分别将$U$、$V$看成是$m/r\times 1$、$1\times n/r$的分块矩阵，然后套用$r=1$的思路来操作。

文章小结 #

本文介绍了从梯度角度来理解LoRA，除了基本的介绍外，还包含了笔者的一些猜测和推广，供读者参考。

转载到请包括本文地址：https://kexue.fm/archives/9590

更详细的转载事宜请参考：《科学空间FAQ》

打赏

微信打赏

支付宝打赏

因为网站后台对打赏并无记录，因此欢迎在打赏时候备注留言。
你还可以点击这里或在下方评论区留言来告知你的建议或需求。