论文阅读笔记——LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS

LoRA 论文
传统全面微调，对每个任务学习的参数与原始模型相同：
$$ma{x}_{\Phi }\sum _{(x,y)\in Z}\sum _{t=1}^{|y|}log(P_{\Phi }(y_t|x,y<t))\text{式(1)}$$
LoRA 提出对模型中权重更新部分低秩分解，编码任务特定的参数，大幅减少所需参数规模，同时优化 $\Theta$ 来寻找 $\Delta \Theta$ 。对于 175B 的 GPT-3 参数量只有原来的 0.01%。
$$ma{x}_{\Theta }\sum _{(x,y)\in Z}\sum _{t=1}^{|y|}log(p_{{\Phi }_0+\Delta \Phi (\Theta })(y_t|x,y<t))$$

传统方法不足

• 添加 adapters 的策略虽然参数少，但会在推理阶段引入延迟——增加了模型深度。并且有额外参数和计算，在模型中这些会被放大。
• 直接优化输入层（prefix）在训练参数方面并非单调变化且保留一部分长度进行调整降低了下游任务的序列长度——占用了一部分序列长度，减少了可用的输入序列长度。并且他的优化难度也大。

LoRA

核心思想：对于一个预训练的 $W_0\in R^{d\times k}$ ，训练低秩矩阵 $BA$ 来替代权重更新部分 $\Delta W$，$dim(A)=r\times k,dim(B)=d\times rr<<min(d,k)$。
$$h=W_{0}x+\Delta Wx=W_{0}x+BAx$$
其中，A 采取随机高斯初始化，B 为 0。

LoRA 在适应期间不需要满足满秩的条件，只需要将 r 设置为预训练权重矩阵的秩，大致可恢复完全微调的能力，可以维持原来架构。

LoRA 优势：

1. 参数高效：训练参数减少了数千倍，例如在GPT-3中，训练参数从1750亿减少到数百万甚至更少。
2. 计算资源节省：由于需要计算梯度的参数大大减少，显存占用降低，训练速度加快。
3. 无额外推理延迟：训练完成后，可以将低秩更新融合到预训练权重中，推理时无需额外计算。
4. 任务切换灵活：不同任务只需存储和加载小的低秩矩阵，实现快速切换，减少存储需求。
   将 LoRA 应用于 transformer 架构中，只需要对自注意力模块 $(W_q,W_k,W_v,W_0)$ 中的 $W_q,W_v$ 进行适应，保持 MLP 不变。
   在下游部署时，只需要减去 $BA$ 即可恢复 $W_0$，再根据任务需求加上对应 $B^{{}^{\prime }}{A}^{{}^{\prime }}$ 。最明显的好处在于内存和存储使用量减少。

实验结果

代码实现

class LoRALayer():
    def __init__(
        self, 
        r: int, 
        lora_alpha: int, 
        lora_dropout: float,
        merge_weights: bool,
    ):
        self.r = r
        self.lora_alpha = lora_alpha
        # Optional dropout
        if lora_dropout > 0.:
            self.lora_dropout = nn.Dropout(p=lora_dropout)
        else:
            self.lora_dropout = lambda x: x
        # Mark the weight as unmerged
        self.merged = False
        self.merge_weights = merge_weights

class Embedding(nn.Embedding, LoRALayer):
    def __init__(
        self,
        num_embeddings: int,
        embedding_dim: int,
        r: int = 0,
        lora_alpha: int = 1,
        merge_weights: bool = True,
        **kwargs
    ):
        nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs)
        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=0,
                           merge_weights=merge_weights)

        if r > 0:
            self.lora_A = nn.Parameter(self.weight.new_zeros((r, num_embeddings)))
            self.lora_B = nn.Parameter(self.weight.new_zeros((embedding_dim, r)))
            self.scaling = self.lora_alpha / self.r
            # 冻结预训练权重
            self.weight.requires_grad = False
        self.reset_parameters()

    def reset_parameters(self):
        nn.Embedding.reset_parameters(self)
        if hasattr(self, 'lora_A'):
            nn.init.zeros_(self.lora_A)
            nn.init.normal_(self.lora_B)

    def train(self, mode: bool = True):
        nn.Embedding.train(self, mode)
        if mode:
            if self.merge_weights and self.merged:
                if self.r > 0:
                    self.weight.data -= (self.lora_B @ self.lora_A).transpose(0, 1) * self.scaling
                self.merged = False
        else:
            if self.merge_weights and not self.merged:
                # Merge the weights and mark it
                if self.r > 0:
                    self.weight.data += (self.lora_B @ self.lora_A).transpose(0, 1) * self.scaling
                self.merged = True
                
	    def forward(self, x: torch.Tensor):
	        if self.r > 0 and not self.merged:
	            result = nn.Embedding.forward(self, x)
	            after_A = F.embedding(
	                x, self.lora_A.transpose(0, 1), self.padding_idx, self.max_norm,
	                self.norm_type, self.scale_grad_by_freq, self.sparse
	            )
	            result += (after_A @ self.lora_B.transpose(0, 1)) * self.scaling
	            return result
	        else:
	            return nn.Embedding.forward(self, x)