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

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

方法简介 #

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

LoRA借鉴了上述结果，提出对于预训练的参数矩阵W0∈Rn×m，我们不去直接微调W0，而是对增量做低秩分解假设：
W=W0+AB,A∈Rn×r,B∈Rr×m
其中A,B之一用全零初始化，W0固定不变，优化器只优化A,B。由于本征维度很小的结论，所以r我们可以取得很小，常见的是r=8，极端情况下我们甚至可以取1。所以说，LoRA是一种参数高效的微调方法，至少被优化的参数量大大降低了。

用MathJax直接画了个示意图：
W0∈Rn×m+A∈Rn×r×B∈Rr×m

梯度分析 #

正如《Ladder Side-Tuning：预训练模型的“过墙梯”》所提到的，很多参数高效的微调实际上只是降低了显存需求，并没有降低计算量。那么LoRA是否例外呢？它在显存和计算量方面的效率如何呢？下面我们来分析一下。

首先，我们知道训练模型所消耗的显存来源包括模型参数、模型梯度、模型激活值、优化器状态四部份，LoRA通过低秩分解降低了模型参数量，那么梯度和优化器状态也会随之降低，因此节省的显存是很明显的。那它能否节省计算量呢？

这取决于LoRA的实现方式，不同的实现方式计算梯度的复杂度不一样。LoRA的两种等效实现如下：
Y=XW=X(W0+AB)Y=XW0+XAB=XW0+ZB
其中X∈Rb×n是模型输入，Z=XA∈Rb×r是中间输出。针对实现(2)，我们有
∂L∂A=∂L∂WB⊤=(X⊤∂L∂Y)B⊤,∂L∂B=A⊤∂L∂W=A⊤(X⊤∂L∂Y)
L是损失函数。很明显，这种实现导致的后果是需要算完整梯度∂L∂W∈Rn×m，然后才能算A,B的梯度，这意味着它比不LoRA还慢，也费显存。对于实现(3)，我们则有
∂L∂A=X⊤∂L∂Z=X⊤(∂L∂YB⊤),∂L∂B=Z⊤∂L∂Y=(XA)⊤∂L∂Y
此时的Z,∂L∂Z∈Rb×r，相比完整的梯度显然省了不少，计算复杂度也明显降低。所以，LoRA想要节省显存和计算最大化，关键是按照(3)而不是(2)来实现。

（注：关于矩阵计算梯度，我们可以根据链式法则和输出形状来“凑”，比如∂L∂A，根据链式法则我们知道它必然是∂L∂W和B以某种方式相乘，我们约定∂L∂A的形状跟A一致，即n×r，想要用∂L∂W和B凑出一个n×r的结果来，那就只有∂L∂WB⊤了。）

其他原因 #

除了低秩分解带来的好处外，如下几点也是LoRA能节省显存和提速的原因：

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

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

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

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

优化视角 #

梯度(4)还告诉了我们如何从优化器角度来实现LoRA。优化器可以直接获取到全量梯度∂L∂W，然后我们只需要按照公式(4)对梯度进行投影，就得到A,B的梯度，接着就可以按照常规的优化器实现A,B的更新了。

假如优化器是SGD，那么就是
At+1=At−η∂L∂WtB⊤t,Bt+1=Bt−ηA⊤t∂L∂WtWt+1=W0+At+1Bt+1=Wt+(At+1Bt+1−AtBt)
如果是Adam之类的带滑动变量的优化器，则只需要滑动投影后的梯度，因此是降低了优化器的参数量，节省了一定的显存。模型越大，这部分参数所占的显存比例也就越大。

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

随机投影 #

如果我们将SGD场景下的更新量At+1Bt+1−AtBt展开，结果将是
−η(∂L∂WtB⊤tBt+AtA⊤t∂L∂Wt)+η2∂L∂WtB⊤tA⊤t∂L∂Wt
假设η2项是可以忽略的高阶项，那么就剩下
−η(∂L∂WtB⊤tBt+AtA⊤t∂L∂Wt)
从这个角度来看，相比全量微调的SGD，LoRA就是用括号中的结果替代了全量的梯度∂L∂Wt。

简单起见，接下来我们只关心r=1的情形，留意到在上式中，t时刻的投影向量At,Bt是依赖于t的，如果我们将它们换成不依赖于t的随机向量（每步训练都重新随机生成），那么会发生什么呢？我们考虑u,v∼N(0,1)，其中u∈Rm×1,v∈R1×n，那么更新量就变为
−η(∂L∂Wtv⊤v+uu⊤∂L∂Wt)
可以证明的是
Eu∼N(0,1)[uu⊤]=In×n,Ev∼N(0,1)[v⊤v]=Im×m
这里的In×n,Im×m分别指n×n,m×m的单位矩阵。因此，跟“零阶梯度”类似，在平均意义下，这种每步都重新初始化的LoRA事实上等价于满秩的SGD。然而，真要按照这个方式实现的话，其速度甚至可能比满秩的SGD都要慢，所以它的目的不是提速，而是希望能缓解灾难遗忘问题——通过对单个（batch）样本使用低秩矩阵（而不是满秩）更新量的方式，减少对整个模型权重的影响。当然，这只是猜测，实际效果如何，笔者还没有实验过。

一个变体 #

同样还是先只考虑r=1的情形，LoRA相当于假设了Δwi,j=uivj，我们能不能做其他低秩分解假设呢？比如Δwi,j=ui+vj？写成矩阵形式就是
W=W0+A11×m+1n×1B,A∈Rn×1,B∈R1×m
其中11×m,1n×1分别指1×m,n×1的全1矩阵。容易求出它的梯度是：
∂L∂A=∂L∂W1m×1,∂L∂B=11×n∂L∂W
其实就是原本梯度的行求和与列求和。相比原版LoRA，这个加性分解有两个优点：1、加比乘计算量更低，梯度形式也更简单；2、AB的秩一定是1，但是A11×m+1n×1B的秩可能是2，如果秩代表了模型能力的话，那也就是说同样的参数量，加性的表达能力可能还更强。至于具体效果如何，后面笔者用到LoRA的时候，再做对比实验吧。

那么，加性分解能不能推广到r>1的情形呢？自然是可以的，但稍微有些技巧。这里约定m,n都能被r整除，那么我们只需要将参数化方式改为
W=W0+AIr(1×m/r)+Ir(n/r×1)B,A∈Rn×r,B∈Rr×m
这里的Ir(1×m/r)、Ir(n/r×1)分别指1×m/r、n/r×1的分块矩阵，每一块则是r×r的单位阵。这个形式说白了，就是分别将A、B看成是n/r×1、1×m/r的分块矩阵，然后套用r=1的思路来操作。

文章小结 #

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

转载到请包括本文地址：https://spaces.ac.cn/archives/9590

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

如果您还有什么疑惑或建议，欢迎在下方评论区继续讨论。

如果您觉得本文还不错，欢迎分享/打赏本文。打赏并非要从中获得收益，而是希望知道科学空间获得了多少读者的真心关注。当然，如果你无视它，也不会影响你的阅读。再次表示欢迎和感谢！

如果您需要引用本文，请参考：

苏剑林. (Apr. 17, 2023). 《梯度视角下的LoRA：简介、分析、猜测及推广 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/9590

@online{kexuefm-9590,
        title={梯度视角下的LoRA：简介、分析、猜测及推广},
        author={苏剑林},
        year={2023},
        month={Apr},
        url={\url{https://spaces.ac.cn/archives/9590}},
}