4. 认识 LoRA：从线性层到注意力机制

如果你有使用过 AI 生图，那你一定对 LoRA 有印象，下图来自Civitai LoRA，上面有很多可供下载的LoRA模型。

你可能也曾疑惑于为什么只导入 LoRA 模型不能生图，读下去，你会解决它。

文章目录

- 为什么需要 LoRA？
- LoRA 的核心思想
- - 低秩分解
  - 应用到神经网络中的线性层
  - - 参数量对比
    - 直观示意图
  - 代码实现：线性层的 LoRA
- LoRA 在注意力机制中的应用
- - 代码实现：带 LoRA 的注意力
- 回到最初的问题：为什么只导入 LoRA 模型不能生图？
- 总结
- 推荐阅读

这篇文章将从基础的线性层开始，带你一步步了解 LoRA 的核心思想，并深入探索它在注意力机制中的应用。

LoRA，全称 Low-Rank Adaptation，是一种用于微调大型预训练模型的技术。它的核心思想是通过 低秩分解（常见的形式是奇异值分解）减少微调时的参数量，而不牺牲模型的性能。

论文原文：LoRA: Low-Rank Adaptation of Large Language Models

为什么需要 LoRA？

大型预训练模型的出现，为我们带来了强大的自然语言处理和计算机视觉能力，这是一个推动时代的成功。但大模型的“大”，不仅体现在其参数量上，更体现在我们无法轻松进行微调 : )，全量微调一个预训练大模型的代价非常高，而且一般的设备根本训练不动。而 LoRA 提供了一种高效的微调方法，使得在小型设备上微调大模型成为可能。

根据论文中的描述：

Compared to GPT-3 175B fine-tuned with Adam, LoRA can reduce the number of trainable parameters by 10,000 times and the GPU memory requirement by 3 times.

相比于对 GPT-3 175B 模型使用全量参数的微调，LoRA 减少了训练参数量的 10,000 倍，GPU 显存需求的 3 倍。

LoRA performs on-par or better than fine-tuning in model quality on RoBERTa, DeBERTa, GPT-2, and GPT-3, despite having fewer trainable parameters, a higher training throughput, and, unlike adapters, no additional inference latency.

LoRA 的可训练参数更少，但在 RoBERTa、DeBERTa、GPT-2 和 GPT-3 上的模型质量与全量微调相当甚至更好，而且不会增加推理延迟。

LoRA 的核心思想

LoRA 的核心在于利用低秩分解来近似模型权重的更新。

低秩分解

在线性代数中，任何矩阵都可以分解为多个低秩矩阵的乘积。例如，一个大的矩阵 $W$ 可以近似表示为两个小矩阵 $B$ 和 $A$ 的乘积：

$\Delta W = BA$

其中：

$in_features A \in \mathbb{R}^{r \times \text{in\_features}}$ ， $r$ 是低秩值， $in_features \text{in\_features}$ 是输入特征维度。
$out_features × r B \in \mathbb{R}^{\text{out\_features} \times r}$ ， $out_features \text{out\_features}$ 是输出特征维度。

通过训练这两个小矩阵，我们可以近似地更新原始权重矩阵 $W$ ，而无需训练整个大的 $W$ 。

应用到神经网络中的线性层

在线性层中，前向传播的计算为：

$y = W x + b$

其中：

$in_features x \in \mathbb{R}^{\text{in\_features}}$ 是输入向量。
$out_features × in_features W \in \mathbb{R}^{\text{out\_features} \times \text{in\_features}}$ 是权重矩阵。
$out_features b \in \mathbb{R}^{\text{out\_features}}$ 是偏置向量。
$out_features y \in \mathbb{R}^{\text{out\_features}}$ 是输出向量。

在微调过程中，通常需要更新 $W$ 和 $b$ 。但在 LoRA 中，我们可以冻结原始的 $W$ ，仅仅在其基础上添加一个可训练的增量 $\Delta W$ ：

$\Delta W)x + b$

其中：

$\Delta W = BA$

通过训练 $A$ 和 $B$ ，我们大大减少了需要更新的参数数量。

参数量对比

假设（回归论文的符号）：

$in_features = d \text{in\_features} = d$
$out_features = k \text{out\_features} = k$
低秩值为 $r$ （通常 $\ll \min(d, k)$ ）

全量微调：

需要训练的参数数量为 $\times d + k$ ，其中：
- $\times d$ 是权重矩阵 $W$ 的参数数量。
- $k$ 是偏置向量 $b$ 的参数数量。

使用 LoRA 微调：

需要训练的参数数量为 $\times d + k \times r + k$ ，其中：
- $\times d$ 是矩阵 $A$ 的参数数量。
- $\times r$ 是矩阵 $B$ 的参数数量。
- $k$ 是偏置向量 $b$ 的参数数量。

参数量减少的比例：

计算：
$\text{减少比例} = \frac{\text{LoRA 参数量}}{\text{全量微调参数量}} = \frac{r d + k r + k}{k d + k}$

为了简化，我们可以将偏置参数忽略（因为它们相对于权重参数来说数量很小），得到：

$\text{减少比例} \approx \frac{r(d + k)}{k d}$

如果假设 $\approx d$ ，则有：

$\text{减少比例} \approx \frac{r(2d)}{d^2} = \frac{2r}{d}$

所以，当 $\approx d$ 时，参数减少比例近似为 $\frac{2r}{d}$ 。
由于 $\ll d$ ，所以参数量大幅减少。

举例说明：

假设：

输入特征维度 $in_features = d = 1024 \text{in\_features} = d = 1024$
输出特征维度 $out_features = k = 1024 \text{out\_features} = k = 1024$
低秩值 $r = 4$

全量微调参数量：

权重参数： $1024 \times 1024 = 1,048,576$
偏置参数： $1024$
总参数量： $1, 048, 576 + 1024 = 1, 049, 600$

使用 LoRA 微调参数量：

矩阵 $A$ 参数： $\times 1024 = 4,096$
矩阵 $B$ 参数： $1024 \times 4 = 4,096$
偏置参数： $1024$
总参数量： $4, 096 + 4, 096 + 1024 = 9, 216$

参数量对比：

全量微调： $1, 049, 600$ 参数
LoRA 微调： $9, 216$ 参数
参数减少比例： $\frac{9,216}{1,049,600} \approx 0.0088$

也就是说，使用 LoRA 后，参数量减少了约 $114$ 倍，即参数量仅为原来的 $0.88\%$ 。

直观示意图

论文中的这张图直观地展示了这一点，为了更符合直觉，我们使用在后续继续使用 $in_features \text{in\_features}$ 和 $out_features \text{out\_features}$ 替代 $d$ 和 $k$ 进行描述：

LoRA 原理示意图

代码实现：线性层的 LoRA

下面我们来实现一个带有 LoRA 的线性层。

import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, in_features, out_features, r):
        super(LoRALinear, self).__init__()
        self.in_features = in_features  # 对应 d
        self.out_features = out_features  # 对应 k
        self.r = r  # 低秩值

        # 原始权重矩阵，冻结
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.weight.requires_grad = False  # 冻结

        # LoRA 部分的参数，初始化为零
        self.A = nn.Parameter(torch.zeros(r, in_features))  # 形状为 (r, d)
        self.B = nn.Parameter(torch.zeros(out_features, r))  # 形状为 (k, r)

        # 偏置项，可选
        self.bias = nn.Parameter(torch.zeros(out_features))

    def forward(self, x):
        # 原始部分
        original_output = torch.nn.functional.linear(x, self.weight, self.bias)
        # LoRA 增量部分
        delta_W = torch.matmul(self.B, self.A)  # 形状为 (k, d)
        lora_output = torch.nn.functional.linear(x, delta_W)
        # 总输出
        return original_output + lora_output

在这个实现中，self.weight 是原始的权重矩阵，被冻结不参与训练。self.A 和 self.B 是可训练的低秩矩阵。

LoRA 在注意力机制中的应用

Transformer 模型的核心是注意力机制，其中涉及到 Query、Key、Value 的计算，这些都是线性变换。

在标准的注意力机制中，计算公式为：

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

其中 $Q$ 、 $K$ 、 $V$ 的计算为：

$X_Q W_Q, \quad K = X_K W_K, \quad V = X_V W_V$

$X_Q$ 、 $X_K$ 、 $X_V$ 的输入可以相同，也可以不同。例如，在 Cross-Attention 中，解码器的隐藏状态作为 $X_Q$ ，编码器的输出作为 $X_K$ 和 $X_V$ 。

LoRA 可以应用到 $W_Q$ 、 $W_K$ 、 $W_V$ 上，采用与线性层类似的方式。

代码实现：带 LoRA 的注意力

下面我们实现一个带有 LoRA 的单头注意力层。

import torch
import torch.nn as nn

class LoRAAttention(nn.Module):
    def __init__(self, embed_dim, r):
        super(LoRAAttention, self).__init__()
        self.embed_dim = embed_dim  # 对应 d_model
        self.r = r  # 低秩值

        # 原始的 QKV 权重，冻结
        self.W_Q = nn.Linear(embed_dim, embed_dim)
        self.W_K = nn.Linear(embed_dim, embed_dim)
        self.W_V = nn.Linear(embed_dim, embed_dim)
        self.W_O = nn.Linear(embed_dim, embed_dim)

        for param in self.W_Q.parameters():
            param.requires_grad = False
        for param in self.W_K.parameters():
            param.requires_grad = False
        for param in self.W_V.parameters():
            param.requires_grad = False

        # LoRA 的 Q 部分
        self.A_Q = nn.Parameter(torch.zeros(r, embed_dim))  # 形状为 (r, d_model)
        self.B_Q = nn.Parameter(torch.zeros(embed_dim, r))  # 形状为 (d_model, r)

        # LoRA 的 K 部分
        self.A_K = nn.Parameter(torch.zeros(r, embed_dim))
        self.B_K = nn.Parameter(torch.zeros(embed_dim, r))

        # LoRA 的 V 部分
        self.A_V = nn.Parameter(torch.zeros(r, embed_dim))
        self.B_V = nn.Parameter(torch.zeros(embed_dim, r))

    def forward(self, query, key, value):
        """
        query, key, value: 形状为 (batch_size, seq_length, embed_dim)
        """
        # 计算原始的 Q、K、V
        Q = self.W_Q(query)  # (batch_size, seq_length, embed_dim)
        K = self.W_K(key)
        V = self.W_V(value)

        # 计算 LoRA 增量部分
        delta_Q = torch.matmul(query, self.A_Q.t())  # (batch_size, seq_length, r)
        delta_Q = torch.matmul(delta_Q, self.B_Q.t())  # (batch_size, seq_length, embed_dim)
        delta_K = torch.matmul(key, self.A_K.t())
        delta_K = torch.matmul(delta_K, self.B_K.t())
        delta_V = torch.matmul(value, self.A_V.t())
        delta_V = torch.matmul(delta_V, self.B_V.t())

        # 更新后的 Q、K、V
        Q = Q + delta_Q
        K = K + delta_K
        V = V + delta_V

        # 计算注意力得分
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.embed_dim ** 0.5)
        attn_weights = torch.nn.functional.softmax(scores, dim=-1)
        context = torch.matmul(attn_weights, V)

        # 输出层
        output = self.W_O(context)

        return output