深入解析 LoRA+:如何通过调整学习率进一步优化LLM

在当前大规模语言模型(LLM)的参数高效微调方法中,LoRA(Low-Rank Adaptation)无疑是一个重要的手段。此前,我们在《梯度视角下的LoRA:简介、分析、猜测及推广》中曾简单讨论过LoRA。今天,我们来学习LoRA的一个新结论:给LoRA的两个矩阵分配不同的学习率,LoRA的效果还能进一步提升。这一结论出自最近的论文《LoRA+: Efficient Low Rank Adaptation of Large Models》(下称“LoRA+”)。

为什么调整学习率显得如此重要?

乍一看,给不同的矩阵分配不同的学习率似乎并没有什么特别之处,因为配置不同的学习率相当于引入了新的超参数,而通常来说,只要引入并精调超参数,模型性能都会有所提升。然而,“LoRA+”的特别之处在于,它从理论角度肯定了这一必要性,并且明确指出最优解必然是右矩阵的学习率大于左矩阵的学习率。

LoRA基本原理回顾

假设预训练参数为

\(( W_0 \in \mathbb{R}^{n \times m} )\)

,如果使用全量参数微调,那么增量也是一个 ( n \times m ) 矩阵。为了降低参数量,LoRA将更新量约束为低秩矩阵,即设 ( W = W_0 + AB ),其中

\(( A \in \mathbb{R}^{n \times r}, B \in \mathbb{R}^{r \times m}, r \ll \min(n, m) )\)

。在训练时,只更新 ( A ) 和 ( B )。

为了便于理解,我们假设层的输入是

\(( X \in \mathbb{R}^{b \times n} )\)

,层的运算为

\(( XW = X(W_0 + AB) )\)

。由于“LoRA+”的结论跟预训练权重无关,不失一般性可以设

\( ( W_0 = 0 )\)

,那么层运算简化为

\(( Y = XAB \in \mathbb{R}^{b \times m} )\)

结论简析

“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 ) 的贡献:

\([ 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 ]\)

在Adam优化器中,如果两个参数的学习率之比是 ( \lambda ),那么经过长期训练后,这两个参数的数量级之比也是

\(( \lambda )\)

。因此,如果 \(( \eta_B / \eta_A = O(\sqrt{n}) )\),那么 ( B ) 和 ( A ) 的数量级之比也是 \(( O(\sqrt{n}) )\),与我们的期望一致。

快速推导

通过矩阵求导,我们可以得到:

\([ \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} ]\)

结合数值稳定性和贡献相当的假设,可以推导出:

\([ \eta_B / \eta_A = O(\sqrt{n}) ]\)

总结

在这篇文章中,我们介绍并推导了“LoRA+”的结果,它支持LoRA的两个低秩矩阵 ( A ) 和 ( B ) 存在固有的不对称性。不管将哪个矩阵全零初始化,都应该将 ( B ) 的学习率设置得大于 ( A ),以达到更优的效果。这一理论不仅在实践中有效,也为我们提供了一个经典的理论指导训练的范例。

0 0 投票数
Article Rating
订阅评论
提醒
0 评论
最旧
最新 最多投票
内联反馈
查看所有评论
0
希望看到您的想法,请您发表评论x