改进LoRA：从零实现权重分解低秩适配DoRA
低秩适配（Low-Rank Adaptation，简称LoRA）是一种机器学习技术，它仅调整模型参数中一个小型的低秩子集，即可修改预训练模型（例如大语言模型或视觉Transformer），使其更好地适配特定的（通常规模较小的）数据集。
该方法的重要意义在于，它支持在任务特定数据上对大模型进行高效微调，大幅降低了微调所需的计算成本(与)时间。
上周，研究人员提出了LoRA的一种全新替代方案：https://arxiv.org/abs/2402.09353，其性能可能大幅超越LoRA。
DoRA是标准LoRA的一种极具潜力的替代方案（图注来自DoRA论文：https://arxiv.org/abs/2402.09353）
为了理解这些方法的工作原理，本文将从零开始，用PyTorch实现LoRA(与)DoRA两种算法！
LoRA 回顾
在深入讲解DoRA之前，我们先简要回顾一下https://arxiv.org/abs/2106.09685中LoRA的工作原理。
由于大语言模型体量庞大，受GPU显存限制，训练时更新全部模型权重的成本很高。假设某一层有一个大型权重矩阵$W$，在反向传播过程中，我们会得到一个$\Delta W$矩阵，它包含了为最小化训练损失函数，需要对原始权重进行的更新幅度信息。
在常规训练(与)微调中，权重更新的公式如下：
$$W_{\text{updated}} = W + \Delta W$$
https://arxiv.org/abs/2106.09685 提出的LoRA方法提供了一种更高效的替代方案：它不去直接计算权重更新量$\Delta W$，而是学习$\Delta W$的近似值，即$\Delta W \approx AB$。换言之，在LoRA中，权重更新公式如下，其中$A$和$B$是两个小型权重矩阵：
$$W_{\text{updated}} = W + A \cdot B$$
（“$A \cdot B$”中的“$\cdot$”代表矩阵乘法。）
下图并排展示了全量微调(与)LoRA微调的公式原理。
图：常规微调（左）(与)LoRA微调（右）示意图
LoRA是如何节省GPU显存的？假设预训练权重矩阵$W$是$1000 \times 1000$(的LoRA矩阵)，那么常规微调中的权重更新矩阵$\Delta W$同样是$1000 \times 1000$(的LoRA矩阵)，此时$\Delta W$包含1,000,000个参数。如果我们取LoRA的秩为2，那么$A$是$1000 \times 2$(的LoRA矩阵)，$B$是$2 \times 1000$(的LoRA矩阵)，使用LoRA时仅需要更新$2 \times 2 \times 1000 = 4000$个参数。在这个例子中，秩为2的LoRA参数量仅为全量更新的1/250。
当然，$A$和$B$无法捕捉$\Delta W$所能包含的全部信息，但这是刻意设计的。使用LoRA时，我们基于这样一个假设：模型需要满秩的大型矩阵$W$来捕获预训练数据集中的全部知识；但在微调大语言模型时，我们不需要更新所有权重，用比$\Delta W$少得多的参数就能捕获适配任务的核心信息，因此我们通过$AB$实现低秩更新。
如果你仔细观察，上图中全量微调(与)LoRA的图示和我之前给出的公式看起来略有不同。这是因为矩阵乘法的分配律：我们不必将权重(与)更新量合并，而是可以分开保存。例如，设$x$为输入数据，那么常规微调可以写成：
$$x \cdot (W+\Delta W) = x \cdot W + x \cdot \Delta W$$
同理，LoRA可以写成：
$$x \cdot (W+A \cdot B) = x \cdot W + x \cdot A \cdot B$$
LoRA权重矩阵可以分开保存的特性，让LoRA极具实用价值。在实际应用中，这意味着我们完全不需要修改预训练模型的权重，只需在推理时动态叠加LoRA矩阵即可。如果你需要为多个客户提供模型托管服务，这一点尤其有用：你不必为每个客户保存一份庞大的更新后模型，只需在原始预训练模型之外，保存一组小型的LoRA权重即可。
为了让这个概念不那么抽象，并提供更直观的理解，我们将在下一节从零开始用代码实现LoRA。
LoRA层的代码实现
我们首先初始化一个LoRALayer层，它会创建矩阵$A$和$B$，以及alpha缩放超参数和秩超参数。该层可以接收输入并计算对应的输出，如下图所示。
秩为r的LoRA矩阵A与B示意图
在代码中，上图所示的LoRA层实现如下：
import torch.nn as nn
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
在上述代码中，rank（秩）是控制矩阵$A$和$B$内部维度的超参数。换言之，该参数决定了LoRA引入的额外参数量，是平衡模型适配能力(与)参数效率的关键因素。
第二个超参数alpha是作用于低秩适配输出的缩放超参数。它本质上控制了适配层的输出对被适配层原始输出的影响程度，可以看作是调节低秩适配对层输出影响大小的手段。
到目前为止，我们实现的LoRALayer类可以对层输入$x$进行变换。但在LoRA的实际应用中，我们通常需要替换已有的线性（Linear）层，从而将权重更新作用到现有的预训练权重上，如下图所示：
应用于现有线性层的LoRA
为了融入上图所示的原始线性层权重，我们将实现一个LinearWithLoRA层，它复用之前实现的LoRALayer，可用于替换神经网络中已有的线性层，例如大语言模型中的自注意力模块或前馈模块：
class LinearWithLoRA(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):
return self.linear(x) + self.lora(x)
注意，由于我们在LoRA层中将权重矩阵$B$（LoRALayer中的self.B）初始化为零值，$A$(与)$B$(的LoRA矩阵)乘法结果是一个全零矩阵，因此不会影响原始权重（因为给原始权重加0不会改变其数值）。
我们用一个仅含单个线性层的小型神经网络层来测试LoRA：
输入：
torch.manual_seed(123)
layer = nn.Linear(10, 2)
x = torch.randn((1, 10))
print("Original output:", layer(x))
输出：
Original output: tensor([[0.6639, 0.4487]], grad_fn=<AddmmBackward0>)
现在，我们给这个线性层加上LoRA，可以看到输出结果和原来完全一致，因为我们还没有训练LoRA的权重。也就是说，一切都符合预期：
输入：
layer_lora_1 = LinearWithLoRA(layer, rank=2, alpha=4)
print("LoRA output:", layer_lora_1(x))
输出：
LoRA output: tensor([[0.6639, 0.4487]], grad_fn=<AddmmBackward0>)
之前我提到了矩阵乘法的分配律：
$$x \cdot (W+A \cdot B) = x \cdot W + x \cdot A \cdot B$$
这意味着我们也可以将LoRA矩阵(与)原始权重合并，得到等效的实现方式。在代码中，LinearWithLoRA层的这种替代实现方式如下：
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)
简而言之，LinearWithLoRAMerged计算的是等式$x \cdot (W+A \cdot B) = x \cdot W + x \cdot A \cdot B$的左侧，而LinearWithLoRA计算的是右侧——二者是等价的。
我们可以通过以下代码验证它的输出和之前完全一致：
输入：
layer_lora_2 = LinearWithLoRAMerged(layer, rank=2, alpha=4)
print("LoRA output:", layer_lora_2(x))
输出：
LoRA output: tensor([[0.6639, 0.4487]], grad_fn=<AddmmBackward0>)
现在我们已经有了可运行的LoRA实现，下一节将介绍如何将它应用到神经网络中。
《Ahead of AI》是一份读者支持的出版物。想要接收新文章并支持我的创作，欢迎成为免费或付费订阅者。
应用LoRA层
我们为什么要用PyTorch模块以这种方式实现LoRA？因为这种方法可以让我们轻松地将现有神经网络中的线性层（例如大语言模型的前馈模块或注意力模块）替换为我们新的LinearWithLoRA（或LinearWithLoRAMerged）层。
为简单起见，我们暂时不用大语言模型，而是聚焦于一个小型的3层多层感知机，如下图所示：
一个简单的3层多层感知机
在代码中，我们可以这样实现上图的多层感知机：
输入：
class MultilayerPerceptron(nn.Module):
def __init__(self, num_features, num_hidden_1, num_hidden_2, num_classes):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(num_features, num_hidden_1),
nn.ReLU(),
nn.Linear(num_hidden_1, num_hidden_2),
nn.ReLU(),
nn.Linear(num_hidden_2, num_classes)
)
def forward(self, x):
x = self.layers(x)
return x
model = MultilayerPerceptron(
num_features=num_features,
num_hidden_1=num_hidden_1,
num_hidden_2=num_hidden_2,
num_classes=num_classes
)
print(model)
输出：
MultilayerPerceptron(
(layers): Sequential(
(0): Linear(in_features=784, out_features=128, bias=True)
(1): ReLU()
(2): Linear(in_features=128, out_features=256, bias=True)
(3): ReLU()
(4): Linear(in_features=256, out_features=10, bias=True)
)
)
使用LinearWithLoRA，我们可以替换多层感知机模型中的原始线性层，从而加入LoRA层：
输入：
model.layers[0] = LinearWithLoRA(model.layers[0], rank=4, alpha=8) model.layers[2] = LinearWithLoRA(model.layers[2], rank=4, alpha=8) model.layers[4] = LinearWithLoRA(model.layers[4], rank=4, alpha=8) print(model)
输出：
MultilayerPerceptron(
(layers): Sequential(
(0): LinearWithLoRA(
(linear): Linear(in_features=784, out_features=128, bias=True)
(lora): LoRALayer()
)
(1): ReLU()
(2): LinearWithLoRA(
(linear): Linear(in_features=128, out_features=256, bias=True)
(lora): LoRALayer()
)
(3): ReLU()
(4): LinearWithLoRA(
(linear): Linear(in_features=256, out_features=10, bias=True)
(lora): LoRALayer()
)
)
)
接下来，我们可以冻结原始线性层，仅让LoRALayer层可训练，方法如下：
输入：
def freeze_linear_layers(model):
for child in model.children():
if isinstance(child, nn.Linear):
for param in child.parameters():
param.requires_grad = False
else:
# Recursively freeze linear layers in children modules
freeze_linear_layers(child)
freeze_linear_layers(model)
for name, param in model.named_parameters():
print(f"{name}: {param.requires_grad}")
输出：
layers.0.linear.weight: False layers.0.linear.bias: False layers.0.lora.A: True layers.0.lora.B: True layers.2.linear.weight: False layers.2.linear.bias: False layers.2.lora.A: True layers.2.lora.B: True layers.4.linear.weight: False layers.4.linear.bias: False layers.4.lora.A: True layers.4.lora.B: True
通过上面的True和False值，我们可以直观地确认现在只有LoRA层是可训练的（True代表可训练，False代表冻结）。在实际应用中，我们就可以用这个LoRA配置的网络在新数据集或新任务上进行训练。
为了不让文章过于冗长，我省略了训练该模型的样板代码。如果你对完整代码感兴趣，可以在这里找到独立的代码笔记本：https://github.com/rasbt/dora-from-scratch。
此外，如果你想了解从零开始的LoRA讲解以及在大语言模型上的应用，也可以查看我的Lightning Studio：https://lightning.ai/lightning-ai/studios/code-lora-from-scratch 。
理解权重分解低秩适配（DoRA）
你可能已经注意到，我们花了很多时间来实现和讲解LoRA。这是因为DoRA（https://arxiv.org/abs/2402.09353）可以看作是在LoRA基础上的改进(与)扩展，我们现在可以很方便地复用之前的部分代码来实现DoRA。
DoRA可以分为两个步骤：第一步是将预训练权重矩阵分解为幅值向量（$m$）和方向矩阵（$V$）；第二步是对方向矩阵$V$应用LoRA，同时单独训练幅值向量$m$。
这种将向量分解为幅值和方向分量的思路，源于一个数学原理：任意向量都可以表示为其幅值（表示长度的标量）(与)方向（表示空间朝向的单位向量）的乘积。
单个向量的方向(与)幅值示意图。例如，对于二维向量[1, 2]，我们可以将其分解为幅值2.24和方向向量[0.447, 0.894]，而$2.24 \times [0.447, 0.894] = [1, 2]$。
在DoRA中，我们将这种幅值-方向分解应用到整个预训练权重矩阵$W$上，而不是单个向量。权重矩阵的每一列（向量）对应连接所有输入到某个输出神经元的权重。
因此，分解$W$后得到的幅值向量$m$，代表了权重矩阵中每一列向量的尺度或长度，如下图所示。
DoRA中权重矩阵分解示意图
随后，DoRA对方向矩阵$V$应用标准LoRA，例如：
$$W’ = m \cdot (V + \Delta V)/\text{norm} = m \cdot (W + AB)/\text{norm}$$
这里的归一化（为了不让概述过于复杂，我简称为“norm”）基于Salimans和Kingma在2016年提出的权重归一化方法，参见https://arxiv.org/abs/1602.07868 。
DoRA的两步流程（分解预训练权重矩阵、对方向矩阵应用LoRA）在下面DoRA论文的图中有进一步说明。
来自DoRA论文的带注释示意图（https://arxiv.org/abs/2402.09353）
研发DoRA的动机，源于对LoRA(与)全量微调学习模式的分析对比。DoRA的作者发现，LoRA对幅值和方向的更新是成比例增减的，无法像全量微调那样仅对方向做细微调整。因此，研究人员提出将幅值分量(与)方向分量解耦。
换言之，DoRA方法的目标是仅对方向分量$V$应用LoRA，同时让幅值分量$m$可以单独训练。
引入幅值向量$m$后，DoRA的参数量仅比LoRA多0.01%。但研究人员发现，在大语言模型和视觉Transformer的基准测试中，即使DoRA的秩减半（即参数量仅为常规LoRA的一半），其性能依然优于LoRA，如下方性能对比图所示。
DoRA论文中LoRA(与)DoRA的对比（https://arxiv.org/abs/2402.09353）
正如我几个月前在另一篇文章中所写，LoRA需要仔细调整秩才能优化性能：https://magazine.sebastianraschka.com/p/practical-tips-for-finetuning-llms。但如下图的对比所示，DoRA对秩的变化鲁棒性要强得多。
DoRA对秩超参数的鲁棒性优于LoRA（图注来自DoRA论文：https://arxiv.org/abs/2402.09353）
DoRA可以在秩相对较小的情况下取得良好效果，这让该方法比LoRA的参数效率更高。
总的来说，这些结果让我印象深刻，将LoRA实现升级为DoRA的工作量并不大，我们将在下一节完成这项工作。
用PyTorch实现DoRA层
在本节中，我们将看到DoRA的代码实现形式。之前我们提到，可以将预训练权重$W_0$初始化为幅值$m$和方向分量$V$。例如，我们有如下公式：
其中$|V|_c$是$V$的逐列向量范数。然后我们可以写出包含LoRA权重更新$BA$的DoRA公式，如下所示：
在DoRA论文中，作者给出的DoRA公式如下：他们直接将初始预训练权重$W_0$作为方向分量，在训练中学习幅值向量$m$：
其中$\Delta V$是方向分量矩阵$V$的更新量。
虽然原作者尚未发布官方实现，但你可以找到一个第三方实现：https://github.com/catid/dora/blob/main/dora.py，我下面的实现也大致参考了它。
基于我们之前的LinearWithLoRAMerged实现，我们可以将其升级为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))
# Code loosely inspired by
# https://github.com/catid/dora/blob/main/dora.py
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)
LinearWithDoRAMerged类(与)之前的LinearWithLoRAMerged类有几个关键区别，主要体现在修改和应用线性层权重的方式上。不过两个类都集成了LoRALayer来增强原始线性层的权重，而DoRA额外加入了权重归一化(与)调整。
下图并排展示了两个类的代码差异对比：
LinearWithLoRAMerged(与)LinearWithDoRAMerged的代码差异对比
从上图可以看到，LinearWithDoRAMerged额外加入了一步对增强后权重的动态归一化。
在将原始权重(与)LoRA调整后的权重合并（self.linear.weight + self.lora.alpha*lora.T）之后，它会计算这些合并后权重的列范数（column_norm），然后通过除以各自的范数来归一化合并权重（$V = \text{合并权重} / \text{列范数}$）。这一步确保了合并权重矩阵的每一列都具有单位范数，能够通过维持权重更新的尺度来稳定训练过程。
DoRA还引入了可学习向量self.m，它代表归一化后权重矩阵每一列的幅值。该参数让模型可以在训练中动态调整合并权重矩阵中每个权重向量的尺度，这种额外的灵活性有助于模型更好地捕捉不同特征的重要性。
总而言之，LinearWithDoRAMerged在LinearWithLoRAMerged的概念基础上进行了扩展，加入了动态权重归一化(与)缩放，以提升训练性能。
在实际应用中，对于之前的多层感知机，我们可以直接将现有线性层替换为LinearWithDoRAMerged层，方法如下：
输入：
model.layers[0] = LinearWithDoRAMerged(model.layers[0], rank=4, alpha=8) model.layers[2] = LinearWithDoRAMerged(model.layers[2], rank=4, alpha=8) model.layers[4] = LinearWithDoRAMerged(model.layers[4], rank=4, alpha=8) print(model)
输出：
MultilayerPerceptron(
(layers): Sequential(
(0): LinearWithDoRAMerged(
(linear): Linear(in_features=784, out_features=128, bias=True)
(lora): LoRALayer()
)
(1): ReLU()
(2): LinearWithDoRAMerged(
(linear): Linear(in_features=128, out_features=256, bias=True)
(lora): LoRALayer()
)
(3): ReLU()
(4): LinearWithDoRAMerged(
(linear): Linear(in_features=256, out_features=10, bias=True)
(lora): LoRALayer()
)
)
)
在微调模型之前，我们可以复用之前实现的freeze_linear_layers函数，仅让LoRA权重和幅值向量可训练：
输入：
freeze_linear_layers(model)
for name, param in model.named_parameters():
print(f"{name}: {param.requires_grad}")
输出：
layers.0.m: True layers.0.linear.weight: False layers.0.linear.bias: False layers.0.lora.A: True layers.0.lora.B: True layers.2.m: True layers.2.linear.weight: False layers.2.linear.bias: False layers.2.lora.A: True layers.2.lora.B: True layers.4.m: True layers.4.linear.weight: False layers.4.linear.bias: False layers.4.lora.A: True layers.4.lora.B: True
包含模型训练在内的完整代码示例，可以在我的GitHub仓库中找到：https://github.com/rasbt/dora-from-scratch。
DoRA代码笔记本：https://github.com/rasbt/dora-from-scratch
结语
在我看来，DoRA是LoRA一种合理、有效且极具潜力的扩展，我很期待在实际的大语言模型微调场景中应用它。
同时，我也将上述DoRA实现加入到了Lightning Studio中（https://lightning.ai/lightning-ai/studios/code-lora-from-scratch），用于微调DistilBERT语言模型（参见bonus_02_finetune-with-dora.ipynb）。即使没有进行超参数调优，我也已经观察到它的预测准确率比LoRA高出1%以上。
这本杂志是我的个人兴趣项目。如果你愿意支持我的创作，欢迎购买我的著作：https://amzn.to/4fqvn0D 。（我相信你会从这本书中收获颇丰，因为它对大语言模型工作原理的讲解深度是其他地方找不到的。）
《从零构建大语言模型》现已发售：https://amzn.to/4fqvn0D
如果你读过这本书，并且能抽出几分钟时间，我会非常感谢你留下评价：https://www.amazon.com/Build-Large-Language-Model-Scratch/dp/1633437167 。这对我们作者帮助很大！
另外，我最近也在Substack上开通了付费订阅选项，可以直接支持这本杂志。
《Ahead of AI》是一份读者支持的出版物。想要接收新文章并支持我的创作，欢迎成为免费或付费订阅者。