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

2023-04-20 04:04

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

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

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

方法简介

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

LoRA 借鉴了上述结果，提出对于预训练的参数矩阵，我们不去直接微调，而是对增量做低秩分解假设：

其中之一用全零初始化，固定不变，优化器只优化。由于本征维度很小的结论，所以我们可以取得很小，很多时候我们甚至可以直接取 1。所以说，LoRA 是一种参数高效的微调方法，至少被优化的参数量大大降低了。

梯度分析

正如《Ladder Side-Tuning：预训练模型的“过墙梯”》所提到的，很多参数高效的微调实际上只是降低了显存需求，并没有降低计算量，LoRA 其实也不例外。为了认识到这一点，我们只需要观察的梯度：

这里是损失函数，并且约定参数的 shape 跟其梯度的 shape 一致。在训练过程中，求模型梯度是主要的计算量，如果全量更新，那么所用的梯度是，而 LoRA 所用的梯度则是和，它们是建立在全量更新的梯度基础上的，所以理论上 LoRA 的计算量比全量更新还大。

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

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

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

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

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

优化视角

事实上，梯度（2）也告诉了我们如何从优化器角度来实现 LoRA。优化器可以直接获取到全量梯度，然后我们只需要按照公式（2）对梯度进行投影，就得到的梯度，接着就可以按照常规的优化器实现的更新了。

假如优化器是 SGD，那么就是

如果是 Adam 之类的带滑动变量的优化器，则只需要滑动投影后的梯度，因此是降低了优化器的参数量，节省了一定的显存。模型越大，这部分参数所占的显存比例也就越大。

LoRA 约定或之一使用全零初始化，这是为了保证初始状态模型跟预训练一致，但同时也带来了不对称问题（一个全零，一个非全零）。事实上，都使用非全零初始化也是可以的，只需要事先将预训练权重减去就行了，或者等价地说，将参数化为：

这样同时保持了初始状态一致，同时允许都用非全零初始化，增强了对称性。

随机投影

如果我们将 SGD 场景下的更新量展开，结果将是

假设项是可以忽略的高阶项，那么就剩下

从这个角度来看，相比全量微调的 SGD，LoRA 就是用括号中的结果替代了全量的梯度。

简单起见，接下来我们只关心的情形，留意到在上式中，时刻的投影向量是依赖于的，如果我们将它们换成不依赖于的随机向量（每步训练都重新随机生成），那么会发生什么呢？我们考虑，其中，那么更新量就变为

可以证明的是

这里的分别指的单位矩阵。因此，跟“零阶梯度”类似，在平均意义下，这种每步都重新初始化的 LoRA 事实上等价于满秩的 SGD。然而，真要按照这个方式实现的话，其速度甚至可能比满秩的 SGD 都要慢，所以它的目的不是提速，而是希望能缓解灾难遗忘问题——通过对单个（batch）样本使用低秩矩阵（而不是满秩）更新量的方式，减少对整个模型权重的影响。当然，这只是猜测，实际效果如何，笔者还没有实验过。

一个变体

同样还是先只考虑的情形，LoRA 相当于假设了，我们能不能做其他低秩分解假设呢？比如？写成矩阵形式就是

其中分别指的全 1 矩阵。容易求出它的梯度是：

其实就是原本梯度的行求和与列求和。相比原版 LoRA，这个加性分解有两个优点：1）加比乘计算量更低，梯度形式也更简单；2）UV 的秩一定是 1，但是的秩可能是 2，如果秩代表了模型能力的话，那也就是说同样的参数量，加性的表达能力可能还更强。至于具体效果如何，后面笔者用到 LoRA 的时候，再做对比实验吧。