配置不同的学习率,LoRA还能再涨一点?
LoRA(Low-Rank Adaptation)是当前 LLM 的参数高效微调手段之一,此前我们在《梯度视角下的LoRA:简介、分析、猜测及推广》也有过简单讨论。这篇文章我们来学习 LoRA 的一个新结论:
给 LoRA 的两个矩阵分配不同的学习率,LoRA 的效果还能进一步提升。
该结论出自最近的论文《LoRA+: Efficient Low Rank Adaptation of Large Models》[1](下称“LoRA+”)。咋看之下,该结论似乎没有什么特别的,因为配置不同的学习率相当于引入了新的超参数,通常来说只要引入并精调超参数都会有提升。
“LoRA+”的特别之处在于,它从理论角度肯定了这个必要性,并且断定最优解必然是右矩阵的学习率大于左矩阵的学习率。简而言之,“LoRA+”称得上是理论指导训练并且在实践中确实有效的经典例子,值得仔细学习一番。
结论简析
“LoRA+”的结论是:为了使 LoRA 的效果尽可能接近最优,权重 B 的学习率应该要大于权重A的学习率。
注意,为了使初始模型等价于原始预训练模型,LoRA 通常会将 A,B 之一全零初始化。笔者一开始以为,该结论是由于全零初始化导致的,所以应该依赖于全零初始化的位置,但仔细阅读后发现,“LoRA+”所声称的结论跟全零初始化无关,也就是说,表面上 A,B 是对称的,但实际上它们有着固有的不对称性,以至于不管选择 A 还是 B 来全零初始化,结论都是 B 的学习率要大于 A。这就有意思起来了。
不得不说。“LoRA+”原文的推导写得相当让人费解,下面笔者尽量用自己的思路来完成推导。大体上,它基于两点假设:1. 数值稳定:模型每一层的输出值都应该是数值稳定的,跟网络宽度无关;2. 贡献相当:为了使 LoRA 最优,A,B 两个矩阵对效果应该有同等程度的贡献。接下来我们逐一分析并量化这两点假设。
数值稳定
单看前向传播来说,XA 的数值稳定性确实不是必要的。但如果 XA 数值不稳定同时 XAB 数值稳定的话,那么有两种情况:XA 数值偏大、B 数值偏小,根据求导公式,这将导致 A 的梯度偏小、B 的梯度偏大;反过来,XA 数值偏小、B 数值偏大,这将导致 A 的梯度偏大、B 的梯度偏小。
总而言之,XA 的数值不稳定会导致 A,B 的梯度不稳定,从而增加优化难度,所以还是加上 XA 的数值稳定性为条件比较好。
贡献相当
这里使用了一阶线性近似,其中 是 A,B 的梯度, 是(Frobenius)内积运算,右端两项就可以理解为 A,B 对效果的分别贡献。但注意线性近似的有效性取决于增量 是小量,但对于训练好的权重,它对于原始权重的增量还真未必是小量。
所以退而求其次,我们将“贡献相当”假设改为“A,B 在每一步更新中应该对效果有同等程度上的贡献”,由于单步更新的量通常很小,因此线性近似能比较好地满足。
既然要考虑每一步的更新量,那么就引导我们到了优化器的方向上。当前预训练和微调的主流优化器都是 Adam,那么我们就以 Adam 为主要分析对象。
其中 是各自的学习率,“LoRA+”的结论就是 。
将 SignSGD 的增量(2)代回式(1),那么就得到
快速推导
进一步的分析需要求出梯度的具体形式。再次设 Y=XAB,那么可以求出:
但是还没完,我们要检查一下结果是否自洽,因为我们用到的条件之一是“前向的数值稳定性”,至今为止还只是一个理想的假设。如何让假设尽可能成立呢?战胜一个假设的方法是引入另一个假设:
根据 Adam 的近似式(2),每步增量的数量级确实正比于学习率,但总的更新结果又不完全是每一步简单叠加,所以这个假设给人的感觉就是“看上去有点道理,但又不完全有道理”。但不要紧,假设通常都是这样子的,有点道理就行,剩下的就只能靠信仰了。
文章小结
这篇文章中,我们介绍并推导了一个名为“LoRA+”的结果,它支持 LoRA 的两个低秩矩阵 A,B 存在固有的不对称性,不管将哪个矩阵全零初始化,都应该将 B 的学习率设置得大于 A,以达到更优的效果。
本文转载自 PaperWeekly,作者: 苏剑林