LoRA的学习率 LoRA+: Efficient Low Rank Adaptation

LoRA 用 low-rank update 限制下游适配需要训练的参数；AdaLoRA 继续讨论有限的 rank 应该分配给哪些 matrix；QLoRA 则通过量化冻结的基础模型，进一步降低训练时的显存成本。

LoRA+ 没有再改变 low-rank update 的结构，而是回到一个更基础、也更容易被忽略的问题：

LoRA 中的两个 matrix 扮演不同角色，初始化方式也不同，为什么训练时却默认使用相同的学习率？

它给出的答案非常简单： $\underline{\underline{A}}$ 和 $\underline{\underline{B}}$ 不应该被当成完全对称的参数。LoRA+ 保留 LoRA 的参数量、计算图和推理方式，只让 $\underline{\underline{B}}$ 使用比 $\underline{\underline{A}}$ 更大的学习率。

LoRA+: Efficient Low Rank Adaptation of Large Models

Soufiane Hayou, Nikhil Ghosh, Bin Yu

ICML 2024

arXiv PDF Code

LoRA+ 改了什么#

LoRA 将一个权重 matrix 的更新写成 $\underline{\underline{\Delta W}}=\frac{\alpha}{r}\underline{\underline{B}}\underline{\underline{A}}$ 。其中 $\underline{\underline{A}}$ 先把输入投影到 $r$ 维空间， $\underline{\underline{B}}$ 再把它映射回输出空间。

标准 LoRA 通常对两者使用同一个学习率，即 $\eta_A=\eta_B=\eta$ 。LoRA+ 将它改成 $\eta_B=\lambda\eta_A$ ，其中 $\lambda>1$ 。

实际使用时，只需要让 optimizer 为两组参数设置不同的学习率。 $\underline{\underline{A}}$ 和 $\underline{\underline{B}}$ 的形状没有改变，可训练参数数量没有改变，forward 也仍然计算 $\underline{\underline{B}}\underline{\underline{A}}\underline{x}$ 。因此 LoRA+ 不是一种新的 low-rank parameterization，而是对 LoRA 优化方式的修正。

为什么相同的学习率不一定合理#

在常见的 LoRA 初始化中， $\underline{\underline{A}}$ 被随机初始化，而 $\underline{\underline{B}}$ 从零开始。这样可以保证训练开始时 $\underline{\underline{B}}\underline{\underline{A}}=\underline{\underline{0}}$ ，模型的初始输出仍然等于预训练模型的输出。

但这个初始化也让两者从第一步开始就具有不同的训练行为。因为 $\underline{\underline{B}}=\underline{\underline{0}}$ ，初始时 $\frac{\partial\mathcal{L}}{\partial\underline{\underline{A}}}=\underline{\underline{0}}$ ；而 $\frac{\partial\mathcal{L}}{\partial\underline{\underline{B}}}$ 仍然可以通过非零的 $\underline{\underline{A}}$ 得到 gradient。也就是说， $\underline{\underline{B}}$ 会先开始变化， $\underline{\underline{A}}$ 要等 $\underline{\underline{B}}$ 离开零点后才真正参与训练。

这种不对称并不只存在于第一步。两个 matrix 的形状相反，在 $\underline{\underline{B}}\underline{\underline{A}}\underline{x}$ 中承担的作用也不同： $\underline{\underline{A}}$ 决定保留哪些输入方向， $\underline{\underline{B}}$ 决定这些方向怎样作用到输出。给它们相同的学习率，并不代表它们会以相近的速度改变模型输出。

理论分析最重要的结论#

论文从模型宽度 $n$ 不断增大的极限出发，要求由 $\underline{\underline{A}}$ 和 $\underline{\underline{B}}$ 更新产生的输出变化都保持在 $\Theta(1)$ 。

这里的 $\Theta(1)$ 不需要理解成某个固定数值。它表达的是：当模型越来越宽时，这一部分变化仍然具有实际作用，既不会趋近于零，也不会不断放大。论文把这种状态称为 efficient feature learning。

分析得到的关键结论是：

如果始终令 $\eta_A=\eta_B$ ，就无法让两个分支对模型输出的影响同时保持在 $\Theta(1)$ ；其中一个分支会在宽模型中变得相对不重要。
将学习率分开后，这个冲突可以消失。在论文分析的尺度下，合适的关系是 $\eta_A=\Theta(n^{-1})$ 、 $\eta_B=\Theta(1)$ 。
因此模型宽度越大， $\underline{\underline{B}}$ 的学习率相对 $\underline{\underline{A}}$ 越大。这也是实际算法使用 $\eta_B=\lambda\eta_A$ 且 $\lambda>1$ 的来源。

这个理论不是要给每个有限大小的模型算出唯一的 $\lambda$ ，而是说明一个方向性的结论：两个 matrix 需要不同的更新尺度，而且 $\underline{\underline{B}}$ 应该更新得更快。

实验告诉我们什么#

论文在 RoBERTa-base、GPT-2 和 Llama-7B 上进行了实验，任务包括 MNLI、QQP、SST-2、QNLI，以及用 flan-v2 训练后在 MMLU 上评估。

最值得保留的不是某一组孤立的分数，而是以下几个整体结论：

最优结果通常出现在 $\eta_B>\eta_A$ 的区域，而不是两者相等的位置。
论文整体报告了约 $1\%\text{--}2\%$ 的性能提升，以及最高约 $2\times$ 的 fine-tuning speedup。
这里的 speedup 主要表示模型用更少的训练 step 达到相同或更高的性能。LoRA+ 的单步计算量与 LoRA 基本相同，并不是每一个 step 本身快了两倍。
任务需要的适配越明显，区分两个学习率通常越有价值；对于已经很容易适配的任务，相同学习率也可能接近最优。

这说明 LoRA+ 并不会在所有任务上带来同样大小的提升。它修正的是 LoRA 的训练动力学，而一个任务是否真正需要这部分额外的 feature learning，仍然取决于模型与数据。

$\lambda$ 应该怎样选择#

$\lambda$ 不是一个对所有模型都通用的常数。论文发现，它会受到模型、任务和初始化方式影响。

在论文的实验中，RoBERTa 使用常见的“ $\underline{\underline{A}}$ 随机初始化、 $\underline{\underline{B}}$ 初始化为零”时，较好的比例大约落在 $4\text{--}8$ ；使用相反初始化时，比例可能增大到约 $16$ 。Llama 实验中的较优比例则大约在 $2\text{--}4$ 。

因此，更稳妥的理解不是“LoRA+ 必须令 $\lambda=16$ ”，而是：

先确定实现中哪一个 matrix 是 down projection、哪一个是 up projection；
保持 $\eta_B>\eta_A$ ，可以从 $\{2,4,8,16\}$ 这样的倍数开始搜索；
同时观察最终效果和收敛速度，不要只根据训练 loss 判断比例。

不同库可能对两个 adapter matrix 使用不同的命名或乘法顺序。本文沿用 $\underline{\underline{\Delta W}}=\underline{\underline{B}}\underline{\underline{A}}$ ： $\underline{\underline{A}}$ 是 down projection， $\underline{\underline{B}}$ 是 up projection。应用 LoRA+ 时，应该先确认实现中的对应关系，而不是只按变量名机械地放大学习率。

LoRA+ 在整个发展路径中的位置#

LoRA 讨论的是 update 能否具有低 rank；AdaLoRA 讨论的是 rank budget 应该放在哪里；QLoRA 讨论的是冻结的基础模型应该怎样保存；LoRA+ 讨论的则是 low-rank factor 应该怎样被 optimizer 更新。

所以 LoRA+ 的贡献并不是增加一种更复杂的 adapter，而是指出：即使 parameterization 已经选对了，不合适的优化尺度仍然可能浪费它的训练能力。

从结构上看，它也可以和 QLoRA、AdaLoRA 等方法组合，因为它只改变 optimizer 对不同参数组使用的学习率。不过组合后最合适的 $\lambda$ 仍然需要重新验证，不能直接假设原论文中的比例保持不变。

LoRA+ 最终留下的 whole picture 很清楚：

low-rank update 的两个 factor 并不对称。让 $\underline{\underline{B}}$ 比 $\underline{\underline{A}}$ 学得更快，可以在不增加参数量和单步计算成本的前提下，更充分地利用 LoRA 已经拥有的训练能力。