LoRA

Low-Rank Adaptation of Large Language Models

https://arxiv.org/abs/2106.09685

背景

以往的一些结果（比如Exploring Universal Intrinsic Task Subspace via Prompt Tuning (opens new window)）显示，尽管预训练模型的参数量很大，但每个下游任务对应的本征维度（Intrinsic Dimension）并不大，换句话说，理论上我们可以微调非常小的参数量，就能在下游任务取得不错的效果。因此，论文的作者认为权重更新的那部分参数矩阵尽管随机投影到较小的子空间，仍然可以有效的学习，可以理解为针对特定的下游任务这些权重矩阵就不要求满秩。

Intrinsic Dimension: 优化下游任务达到一定水平的子空间

现有微调大模型方法的不足：

推理延迟（inference latency）: Adapter
无法与 fine-tuning 媲美

技术原理

原理

LoRA: LoRA 通过向预训练模型的每一层的权重矩阵中添加低秩矩阵来实现微调。即固定原始模型权重，然后定义两个低秩矩阵作为新增 weight 参与运算，并将两条链路的结果求和后作为本层的输出，而在微调时，只梯度下降新增的两个低秩矩阵。

以单个 Linear 层( $h = W X$ )举例，用公式表达就是假设原始预训练模型的权重为 $W_{0} \in R^{d \times d}$ , 两个低秩矩阵的 shape 为 $B \in R^{d \times r}$ 和 $A \in R^{r \times d}$ ，其中 $r << d$ ，并将 Linear 层的计算过程有原始 $h = W_{0} X$ 调整为

h = W_{0} x + Δ W x = W_{0} x + B A x

对矩阵 $A$ 使用随机高斯初始化，对 $B$ 使用零初始化, 保证在微调开始之前 $Δ W x = 0$ 。也可以使用 $\frac{α}{r}$ 作为缩放参数来调节 $Δ W$ , 通过调节缩放比例可以调节预训练模型与 LoRA 的加权占比。

相对于 Adapter，LoRA 有个好处就是不会引入额外的推理延迟，因为前者相当于在原始模型结构上新增了一些结构(或者说 FLOPs) ，而对于 LoRA 而言 $B \times A$ 的结果的 shape 与 $W_{0}$ 的 shape 是一致的，也就是说微调完成后是可以将结果直接累加到原始 weight 中的，即 $W = W_{0} + \frac{α}{r} B A$ ，这样一来，在推理时微调前后的 FLOPs 的是一样的。

batch inputs 选择不同 $A$ 、 $B$ 的延迟问题：类增量学习中，不同的图片选择不同的专家，即选择不同的 $A$ 、 $B$

梯度分析

多参数高效的微调实际上只是降低了显存需求，并没有降低计算量，LoRA 其实也不例外。为了认识到这一点，我们只需要观察 $B$ , $A$ 的梯度：

\frac{\partial L}{\partial B} = \frac{\partial L}{\partial h} A^{T}, \frac{\partial L}{\partial A} = B^{T} \frac{\partial L}{\partial h}

$L$ 是损失函数，在训练过程中，求模型梯度是主要的计算量，如果全量更新，那么所用的梯度是 $\frac{\partial L}{\partial h}$ ，而 LoRA 所用的梯度则是 $\frac{\partial L}{\partial B}$ 和 $\frac{\partial L}{\partial A}$ ，它们是建立在全量更新的梯度 $\frac{\partial L}{\partial h}$ 基础上的，所以理论上 LoRA 的计算量比全量更新还大。

	前向	梯度
全量微调	$W_{0} x$	$\frac{\partial L}{\partial W_{0}} = \frac{\partial L}{\partial h} x^{T} \in R^{d \times d}$
LoRA	$(W_{0} + B A) x$	$\frac{\partial L}{\partial B} = \frac{\partial L}{\partial h} (A x)^{T} \in R^{d \times r} \frac{\partial L}{\partial A} = B^{T} \frac{\partial L}{\partial h} x^{T} \in R^{r \times d}$

LoRA 显存占用少在哪里

预训练权重 $W_{0}$ 的梯度存储开销，实际就是 LoRA 能大大减少了显存占用的关键。在 LoRA 训练时， $W_{0}$ 仍然会参与前向传播和反向传播，但是不会计算其对应梯度 $\frac{\partial L}{\partial W_{0}}$ ，更不会更新其参数。因此，这一步不再需要计算和保存梯度 $\frac{\partial L}{\partial W_{0}}$ ，以及更新 $W_{0}$ 。以 $d = 4096, r = 16$ 为例，这部分减少的梯度显存占用粗略估计为: $d * d - 2 * d * r = 1 - \frac{2 r}{d}$ , 减少了 99.2187% 。

相对于全局微调，这些高效微调技术为什么会使得训练的速度变快呢？

只更新了部分参数：比如 LoRA 原论文就选择只更新 Self Attention 的参数，实际使用时我们还可以选择只更新部分层的参数；
减少了通信时间：由于更新的参数量变少了，所以（尤其是多卡训练时）要传输的数据量也变少了，从而减少了传输时间；
采用了各种低精度加速技术，如 FP16、FP8 或者 INT8 量化等。

所以总的来说，LoRA 在反向传播阶段计算复杂度还要多一些，只是需要梯度下降的参数少，所以节省显存，梯度下降的也快。

最后，关于 $B$ 和 $A$ 这两个低秩矩阵的初始化问题，首先我们是需要 $B A$ 的结果初始是 0，这样就能保证微调开始时新引入的低秩矩阵不会对最终结果造成影响，那么最直接的方式就是令其中一个低秩矩阵初始阶段为全零，另一个为非全零即可，两者不能都为全零，通过带上 LoRA 的梯度计算公式(2)如果和这两个低秩矩阵初始化都是全 0 那么两个矩阵的梯度都是 0，也就训不起来。