在当前大规模语言模型（LLM）的参数高效微调方法中，LoRA（Low-Rank Adaptation）无疑是一个重要的手段。此前，我们在《梯度视角下的LoRA：简介、分析、猜测及推广》中曾简单讨论过LoRA。今天，我们来学习LoRA的一个新结论：给LoRA的两个矩阵分配不同的学习率，LoRA的效果还能进一步提升。这一结论出自最近的论文《LoRA+: Efficient Low Rank Adaptation of Large Models》（下称“LoRA+”）。
友情链接：ACEJoy
为什么调整学习率显得如此重要？
乍一看，给不同的矩阵分配不同的学习率似乎并没有什么特别之处，因为配置不同的学习率相当于引入了新的超参数，而通常来说，只要引入并精调超参数，模型性能都会有所提升。然而，“LoRA+”的特别之处在于，它从理论角度肯定了这一必要性，并且明确指出最优解必然是右矩阵的学习率大于左矩阵的学习率。
LoRA基本原理回顾
假设预训练参数为
[latex]( W_0 \in \mathbb{R}^{n \times m} )[/latex]
，如果使用全量参数微调，那么增量也是一个 ( n \times m ) 矩阵。为了降低参数量，LoRA将更新量约束为低秩矩阵，即设 ( W = W_0 + AB )，其中
[latex]( A \in \mathbb{R}^{n \times r}, B \in \mathbb{R}^{r \times m}, r \ll \min(n, m) )[/latex]
。在训练时，只更新 ( A ) 和 ( B )。
为了便于理解，我们假设层的输入是
[latex]( X \in \mathbb{R}^{b \times n} )[/latex]
，层的运算为
[latex]( XW = X(W_0 + AB) )[/latex]
。由于“LoRA+”的结论跟预训练权重无关，不失一般性可以设
[latex] ( W_0 = 0 )[/latex]
，那么层运算简化为
[latex]( Y = XAB \in \mathbb{R}^{b \times m} )[/latex]
。
结论简析
“LoRA+”的结论是：为了使LoRA的效果尽可能接近最优，权重 ( B ) 的学习率应该要大于权重 ( A ) 的学习率。
这个结论的核心在于，虽然 ( A ) 和 ( B ) 在表面上是对称的，但实际上它们有着固有的不对称性。即使将 ( A ) 或 ( B ) 全零初始化，结论依然成立。
数值稳定性与贡献相当假设
1. 数值稳定性
数值稳定性意味着 ( X )、 ( XA )、 ( XAB ) 的每个分量都应该是 ( O(1) ) 级别的，而不依赖于网络宽度 ( n, m )。虽然 ( XA ) 是中间变量，但它的数值稳定性对梯度稳定性至关重要。
为了实现数值稳定性，假设 ( X ) 是 ( O(1) )，那么 ( A ) 和 ( B ) 应该分别用 ( 1/n ) 和 ( 1/r ) 的方差初始化。
2. 贡献相当
为了使LoRA最优， ( A ) 和 ( B ) 两个矩阵对效果应该有同等程度的贡献。我们可以通过损失函数 ( L ) 的变化来衡量 ( A ) 和 ( B ) 的贡献：
[latex][ L(A + \Delta A, B + \Delta B) – L(A, B) \approx \langle \frac{\partial L}{\partial A}, \Delta A \rangle + \langle \frac{\partial L}{\partial B}, \Delta B \rangle ][/latex]
在Adam优化器中，如果两个参数的学习率之比是 ( \lambda )，那么经过长期训练后，这两个参数的数量级之比也是
[latex]( \lambda )[/latex]
。因此，如果 [latex]( \eta_B / \eta_A = O(\sqrt{n}) )[/latex]，那么 ( B ) 和 ( A ) 的数量级之比也是 [latex]( O(\sqrt{n}) )[/latex]，与我们的期望一致。
快速推导
通过矩阵求导，我们可以得到：
[latex][ \frac{\partial L}{\partial A} = X^\top \frac{\partial L}{\partial Y} B^\top, \quad \frac{\partial L}{\partial B} = A^\top X^\top \frac{\partial L}{\partial Y} ][/latex]
结合数值稳定性和贡献相当的假设，可以推导出：
[latex][ \eta_B / \eta_A = O(\sqrt{n}) ][/latex]
总结
在这篇文章中，我们介绍并推导了“LoRA+”的结果，它支持LoRA的两个低秩矩阵 ( A ) 和 ( B ) 存在固有的不对称性。不管将哪个矩阵全零初始化，都应该将 ( B ) 的学习率设置得大于 ( A )，以达到更优的效果。这一理论不仅在实践中有效，也为我们提供了一个经典的理论指导训练的范例。