LoRA-GA:梯度近似的低秩参数高效微调.

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$对模型结果没有影响。

如果两个矩阵$A,B$都采用非全零初始化,并且希望训练开始时新增的通路没有影响,则可以把整个模块的等效权重视为:

\[W = (W_0-B_0A_0) + BA\]

若采用随机梯度下降训练可学习矩阵$A,B$,则权重更新过程为:

\[\begin{aligned} A_1 &= A_0 - \eta \frac{\partial \mathcal{L}}{\partial A}= A_0 - \eta B_0^T\frac{\partial \mathcal{L}}{\partial W} \\ B_1 &= B_0 - \eta \frac{\partial \mathcal{L}}{\partial B}= B_0 - \eta \frac{\partial \mathcal{L}}{\partial W}A_0^T \end{aligned}\]

这等效于新的权重:

\[\begin{aligned} W_1 &= (W_0-B_0A_0) + B_1A_1 \\ &= W_0-B_0A_0+\left(B_0 - \eta \frac{\partial \mathcal{L}}{\partial W}A_0^T\right)\left(A_0 - \eta B_0^T\frac{\partial \mathcal{L}}{\partial W}\right) \\ & \approx W_0- \eta\left(\frac{\partial \mathcal{L}}{\partial W}A_0^TA_0 +B_0 B_0^T\frac{\partial \mathcal{L}}{\partial W}\right) \\ \end{aligned}\]

考虑全量微调,此时新的权重应该更新为:

\[W_1 = W_0 - \eta\frac{\partial \mathcal{L}}{\partial W}\]

若使得参数高效微调接近全量微调的效果,则可建立最小化目标函数:

\[\mathop{\arg\min}_{A_0,B_0} \left\| \frac{\partial \mathcal{L}}{\partial W}A_0^TA_0 +B_0 B_0^T\frac{\partial \mathcal{L}}{\partial W} - \frac{\partial \mathcal{L}}{\partial W} \right\|_F^2\]

对梯度\(\frac{\partial \mathcal{L}}{\partial W}\)执行奇异值分解:

\[\frac{\partial \mathcal{L}}{\partial W} = U\Sigma V, U\in R^{n\times n}, \Sigma\in R^{n\times m}, V\in R^{m\times m}\]

注意到$U,V$是正交矩阵,$\Sigma$是对角矩阵,则目标函数可重写为:

\[\begin{aligned} & \left\| U\Sigma VA_0^TA_0 +B_0 B_0^TU\Sigma V - U\Sigma V \right\|_F^2 \\ = & \left\| U\Sigma (A_0V^T)^TA_0V^TV +UU^TB_0 (U^TB_0)^T\Sigma V - U\Sigma V \right\|_F^2 \\ = & \left\| U\left[\Sigma (A_0V^T)^TA_0V^T +U^TB_0 (U^TB_0)^T\Sigma - \Sigma \right] V \right\|_F^2 \\ = & \left\| \Sigma (A_0V^T)^T(A_0V^T) +(U^TB_0) (U^TB_0)^T\Sigma - \Sigma \right\|_F^2 \\ \end{aligned}\]

不妨记$X=A_0V^T,Y=U^TB_0$,根据Eckart-Young定理,$\Sigma X^TX+YY^T\Sigma = \Sigma$的最优近似是:

\[Y=(I_n)_{[:,:r]}, X=(I_m)_{[r:2r,:]}\]

因此得到$A_0,B_0$的值:

\[A_0=(I_m)_{[r:2r,:]}V=V_{[r:2r,:]} \\ B_0=U(I_n)_{[:,:r]}=U_{[:,:r]}\]

因此本文提出的LoRA-GA做法是,采样一批样本计算初始梯度\(\frac{\partial \mathcal{L}}{\partial W}\),并进行奇异值分解\(\frac{\partial \mathcal{L}}{\partial W}=U\Sigma V\);取$U$的前$r$列初始化 $B$,取$V$的第$r+1\sim 2r$行初始化$A$。

实验结果表明LoRA-GA策略在GLUE的一个子集上取得了最接近全量微调的效果;并且训练数据量越少,相对提升的幅度越大,可以用更少的训练步数就能达到更优的效果。