LoRA+:大模型的高效低秩微调.

LoRA是一种低秩分解的模型参数微调方法。在涉及到矩阵相乘的模块(如自注意力中计算QKV),在原始矩阵$W\in R^{n\times m}$旁边增加一个新的通路,通过两个可学习矩阵$A\in R^{r\times m},B\in R^{n\times r}$相乘,中间层维度$r ≪ \min(n,m)$。

在下游任务训练时,固定模型的其他参数,只优化新增的矩阵权重参数,并将预训练模型与新增的通路两部分的结果加起来作为最终的结果(两边通路的输入跟输出维度是一致的)。

\[h = W_0x+\Delta Wx = W_0x+BAx\]

第一个矩阵$A$的权重参数会通过高斯函数初始化,而第二个矩阵$B$的权重参数则会初始化为零矩阵,这样能保证训练开始时新增的通路$BA=0$对模型结果没有影响。

本文指出,为了使LoRA的效果尽可能接近最优,权重$B$的学习率应该要大于权重$A$的学习率。

作者首先假设模型每一层的输出值都应该是数值稳定的,跟网络宽度无关。这意味着输入$x$、中间值$Ax$和输出值$BAx$都应该具有相同的方差。假设输入$x$的方差是$O(1)$,则$A$的方差应为$1/m$,$B$的方差应为$1/r$,才能保证中间值$Ax$和输出值$BAx$方差都是$O(1)$。注意到$r ≪ m$,即$A$的分量绝对值会明显小于$B$的分量绝对值。

作者进一步假设为了使LoRA最优,$A,B$两个矩阵对效果应该有同等程度的贡献。分析梯度: \(\frac{\partial \mathcal{L}}{\partial A} = B^T\frac{\partial \mathcal{L}}{\partial h}x^T,\frac{\partial \mathcal{L}}{\partial B} =\frac{\partial \mathcal{L}}{\partial h}x^TA^T\)

若$A,B$两个矩阵的贡献相当,则两个梯度近似有相同的量级,近似估计梯度分量:

\[\left\|\frac{\partial \mathcal{L}}{\partial A} \right\|_1 \propto \left\|B \right\|_1 \propto nr\sqrt{1/r} \\ \left\|\frac{\partial \mathcal{L}}{\partial B} \right\|_1 \propto \left\|A \right\|_1 \propto mr \sqrt{1/m}\\\]

为了调整$A,B$两个矩阵的贡献,设置学习率$\eta_A,\eta_B$:

\[\eta_A \frac{\partial \mathcal{L}}{\partial A} \approx \eta_B \frac{\partial \mathcal{L}}{\partial B} \quad \to \quad \frac{\eta_A}{\eta_B} \propto \frac{\sqrt{rm}}{n}\]

考虑到实际使用时常有$m=n$且$r=O(1)$,则有结论:

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

因此得到结论,矩阵$B$的学习率设置为矩阵$A$的$\sqrt{n}$倍是比较好的。