DoRA: Weight-Decomposed Low-Rank Adaptation阅读

Paper: DoRA: Weight-Decomposed Low-Rank Adaptation
Refer to: Improving LoRA: Implementing Weight-Decomposed Low-Rank Adaptation (DoRA) from Scratch

随着大规模预训练语言（LLM）模型的兴起，如何在有限算力资源条件下对语言模型进行微调（高效微调，PEFT）受到了越来越多的关注，其中LoRA(Low-rank Adaptation) 便是高效微调的代表之一，它通过调整模型参数的部分低秩子集来实现对模型参数的微调从而更好地适应特定的数据集，大大降低了微调LLM的计算成本和时间。
近期，相关学者提出了DoRA: Weight-Decomposed Low-Rank Adaptation，一种LoRA的变体并且在性能上大幅优于LoRA。

LoRA回顾：LoRA: Low-Rank Adaptation of Large Language Models

对于一个LLM f, 假设其中某一层参数为W，通过反向传播，我们会学习到一个$\Delta W$ 矩阵来更新参数W：$$W_{new}=W+\Delta W$$

当W维度较大时，对于一般的微调方法$\Delta W$的维度与W相同，因此在计算时需要更多的资源，随着语言模型的规模不断变大，微调模型的成本也越来越大，难以实现。
针对该问题，LoRA提供了一种高效替代的方法来计算$\Delta W$，它通过学习两个矩阵A和B来近似得到$\Delta W$:
$$\Delta W=A\ast B$$
其中，A和B是两个低秩矩阵，如$W,\Delta W\in R^{m\ast n}$, $A\in R^{m\ast h},B\in R^{h\ast n},h<<m,n$，$\ast$表示矩阵乘法。

LoRA是如何节省GPU 显存的呢？

如上图所示，假设W的维度为100*100，那么通过微调计算的$\Delta W$同样也是100*100，即10000个参数需要计算。而通过LoRA计算的话，假设A是100*2，B是2*100，那么$\Delta W$同样是100*100，但由于A和B都是低秩矩阵，因此实际需要计算的仅有2*100*2=400个参数，相比于微调节省了25倍的计算资源。

然而，通过LoRA计算得到的$\Delta W$仅是一个近似值，难免会造成一定的信息丢失。但是对于参数量较大的模型来说，即使通过微调计算得到$\Delta W$，其中的所有参数也并不是都会被使用到，因此，通过LoRA来近似是完全合理可行的。

另外，LoRA是可以在保持模型f原参数不变的条件的实现对模型输出的更新，当通过LoRA计算得到$\Delta W$后，对于一条输入x，模型计算过程为：$$x\ast (W+\Delta W)=x\ast W+x\ast \Delta W=x\ast W+x\ast A\ast B$$因此，在训练LoRA时，我们可以在冻结f的条件下进行，更好的节省了模型参数。

Weight-Decomposed Low-Rank Adaptation (DoRA)是什么？

DoRA其实是LoRA的一种改进版本。研究人员发现LoRA 可以按比例增加或减少幅度和方向更新，但似乎缺乏像微调那样只进行微妙方向变化的能力。因此，研究人员建议将幅度和方向部分解耦。

DoRA主要包含2步，
第一步：将预训练语言模型中的参数W分解为magnitude vector (m) 和 directional matrix (V)。
第二步：在directional matrix (V)应用LoRA并通过微调训练magnitude vector (m)。

我们首先对第一步分解过程进行介绍：

在矩阵论中，任何vertor可以被其magnitude和direction的乘积来表示，如下图中，在二维空间中，假设$v_i$为[1,2],其magnitude $m=\sqrt{1^2+2^2}=2.24$，因此其directional $v=[0.447,0.894]$，因为$[1,2]=2.24\ast [0.447,0.894]$。

在DoRA中，参数更新方式为：$$W_{new}=m(V+\Delta V)/norm=m(W+AB)/norm$$, norm 表示权重归一化。

DoRA具体流程见下图，它对模型参数$W_0\in R^{d\ast k}$进行分解，从而得到magnitude $m\in R^{1\ast k}$和directional $V\in R^{d\ast k}$，其中分解过程为：$$W_0=m\frac{V}{||V|{|}_c}=||W_0|{|}_c\frac{W_0}{||W_0|{|}_c}$$ $||\cdot |{|}_c$ 表示矩阵在每一列上的的正则化。

第二步更新过程即对图中绿色部分进行训练，作者使用LoRA在更新V并通过微调来更新m。

我们同样对DoRA参数进行分析，假设W的维度为100*100，那么通过微调计算的$\Delta W$同样也是100*100，即10000个参数需要计算。而通过LoRA计算的话，假设A是100*2，B是2*100，那么$\Delta W$同样是100*100，但由于A和B都是低秩矩阵，因此实际需要计算的仅有2*100*2=400个参数，相比于微调节省了25倍的计算资源。
对于DoRA，W维度为100100,那么 m的维度为1100，V的维度为100100，需要更新的参数为：m和LoRA，即1100+21002=500。

LoRA、 DoRA代码对比

最后我们通过pytroch代码来对比LoRA和DoRA的区别：

LoRA类：

class LoRALayer(nn.Module):
    def __init__(self, in_dim, out_dim, rank, alpha):
        super().__init__()
        std_dev = 1 / torch.sqrt(torch.tensor(rank).float())
        self.A = nn.Parameter(torch.randn(in_dim, rank) * std_dev)
        self.B = nn.Parameter(torch.zeros(rank, out_dim))
        self.alpha = alpha

    def forward(self, x):
        x = self.alpha * (x @ self.A @ self.B)
        return x
        
class LinearWithLoRAMerged(nn.Module):
    def __init__(self, linear, rank, alpha):
        super().__init__()
        self.linear = linear
        self.lora = LoRALayer(
            linear.in_features, linear.out_features, rank, alpha
        )

    def forward(self, x):
        lora = self.lora.A @ self.lora.B # Combine LoRA matrices
        # Then combine LoRA with orig. weights
        combined_weight = self.linear.weight + self.lora.alpha*lora.T 
        return F.linear(x, combined_weight, self.linear.bias)

DoRA类：

class LinearWithDoRAMerged(nn.Module):

    def __init__(self, linear, rank, alpha):
        super().__init__()
        self.linear = linear
        self.lora = LoRALayer(
            linear.in_features, linear.out_features, rank, alpha
        )
        self.m = nn.Parameter(
            self.linear.weight.norm(p=2, dim=0, keepdim=True))

    def forward(self, x):
        lora = self.lora.A @ self.lora.B
        numerator = self.linear.weight + self.lora.alpha*lora.T
        denominator = numerator.norm(p=2, dim=0, keepdim=True)
        directional_component = numerator / denominator
        new_weight = self.m * directional_component
        return F.linear(x, new_weight, self.linear.bias)

通过代码可以看到，DoRA相比于LoRA额外增加了一个self.m ，并且在计算时不再是直接将lora结果与Linear参数拼接，而是通过计算magnitude 和directional 乘积来作为新的结果。

总结

本文对DoRA进行了简单分析介绍，并对比了与LoRA的区别，在原文中，作者还从梯度角度对DoRA进行更加详细的分析介绍，具体参加原文。