菜单

低秩适应 (LoRA)

相关源文件

低秩适应 (LoRA) 是一种高效微调大型语言模型的技术,它通过冻结预训练模型权重,并在transformer的每一层中注入可训练的秩分解矩阵。这种方法显著减少了可训练参数的数量,同时保持了模型质量,使得在有限计算资源下微调大型模型更加实用。

本页基于论文 LoRA: 大型语言模型的低秩适应,介绍了 labml_nn 库中 LoRA 的实现。

核心概念

LoRA 的工作原理是冻结预训练权重,并添加可高效训练的低秩分解。

  1. 对于预训练权重矩阵 $W_0 \in \mathbb{R}^{d \times k}$,LoRA 引入:

    • $W_0 + \Delta W = W_0 + BA$
    • 其中 $B \in \mathbb{R}^{d \times r}$,$A \in \mathbb{R}^{r \times k}$,且秩 $r \ll \min(d, k)$。
  2. 在微调期间

    • 原始权重 $W_0$ 被冻结(不更新)
    • 仅训练 $A$ 和 $B$ 矩阵(参数显著减少)
    • 在训练开始时,$\Delta W$ 被初始化为零。
  3. 输出乘以一个缩放因子 $\frac{\alpha}{r}$,其中 $\alpha$ 是一个超参数。

来源:labml_nn/lora/__init__.py28-98

LoRA 组件

该实现提供了两个关键的 LoRA 组件:

LoRA 线性层

将标准线性层替换为 LoRA 增强版本。它将冻结的预训练权重与可训练的低秩矩阵相结合。

实现细节

  • 预训练权重矩阵 $W_0$ 被冻结
  • 矩阵 $A$ 和 $B$ 是可训练参数
  • $A$ 使用 Kaiming 初始化
  • $B$ 初始化为零,确保在初始化时 $\Delta W = BA$ 为零
  • 输出计算方式为 x @ W0.T + (x @ A.T @ B.T) * scaling

来源:labml_nn/lora/__init__.py28-98

LoRA 嵌入层

将相同的概念应用于嵌入层,用 LoRA 增强版本替换标准嵌入。

实现细节

  • 预训练嵌入矩阵被冻结
  • 矩阵 $A$ 和 $B$ 是可训练参数
  • $A$ 使用正态分布初始化
  • $B$ 初始化为零
  • 输出计算方式为 embedding(x, W0) + (embedding(x, A.T) @ B.T) * scaling

来源:labml_nn/lora/__init__.py101-151

GPT-2 架构中的 LoRA

该实现展示了在 GPT-2 模型中使用 LoRA 模块,将 LoRA 应用于模型中的所有线性变换。

使用 LoRA 的关键组件

  1. 词元和位置嵌入
  2. 注意力 QKV 投影
  3. 注意力输出投影
  4. 前馈网络线性层
  5. 最终语言模型头部

来源:labml_nn/lora/gpt2.py19-192

GPT-2 中的 LoRA 增强组件

组件描述实现
词元嵌入将词元 ID 转换为嵌入GPTModel.token_embedding 使用带 LoRA 的 Embedding
位置嵌入添加位置信息GPTModel.position_embedding 使用带 LoRA 的 Embedding
QKV 投影将输入投影到查询、键、值MultiHeadAttention.qkv_projection 使用带 LoRA 的 Linear
注意力输出将注意力输出投影到模型维度MultiHeadAttention.output_projection 使用带 LoRA 的 Linear
FFN 输入前馈网络中的第一个线性层FFN.linear_in 使用带 LoRA 的 Linear
FFN 输出前馈网络中的第二个线性层FFN.linear_out 使用带 LoRA 的 Linear
LM 头部投影到词汇表进行词元预测GPTModel.lm_head 使用带 LoRA 的 Linear

来源:labml_nn/lora/gpt2.py19-192

微调流程

该实现包含一个训练流程,用于在 Tiny Shakespeare 数据集上使用 LoRA 微调 GPT-2。

训练过程的关键特点

  1. Trainer 训练器类管理微调过程
  2. 预训练的 GPT-2 权重从 Hugging Face 加载
  3. 进行参数映射以匹配模型的架构
  4. 仅训练 LoRA 参数,原始权重保持冻结
  5. 使用学习率为 1e-4 的 Adam 优化器
  6. 训练中使用交叉熵损失函数

实现细节

  • _load_pretrained_weights 方法加载预训练权重并将其映射到模型
  • 卷积层的权重矩阵被转置以匹配线性层
  • 加载权重时,使用 strict=False,因为 LoRA 参数不在原始模型中
  • 训练器验证加载状态中仅缺少 LoRA 参数

来源:labml_nn/lora/experiment.py26-169

超参数

LoRA 微调的关键超参数

参数描述默认值
lora_r低秩矩阵的秩32
alpha缩放参数如果未指定,则等于 r
learning_rateAdam 优化器的学习率1e-4
batch_size训练批大小32
epochs训练轮数10
context_len训练的上下文长度512

来源:labml_nn/lora/experiment.py26-52

LoRA 的优势

  1. 参数效率:

    • LoRA 显著减少了可训练参数的数量。
    • 对于具有 dk 维度以及秩 r 的模型,参数减少的因子为 dk/(r(d+k))
  2. 内存效率:

    • 与完整微调相比,所需 GPU 内存大大减少
    • 无需为冻结参数存储优化器状态
  3. 适应性:

    • 可以为不同的任务训练多个 LoRA 适配器
    • 适配器可以在不改变基础模型的情况下进行切换
  4. 性能:

    • 尽管参数较少,LoRA 通常能达到与完整微调相当的性能。

评估与使用

评估微调模型或用于推理时:

  1. 前向传播包含冻结权重和 LoRA 适应
  2. 在推理期间,LoRA 矩阵可以与原始权重合并。
    • W_merged = W₀ + (α/r) * BA
    • 这消除了推理时的开销。

使用示例

该实现包含一个在 Tiny Shakespeare 数据集上使用 LoRA 微调 GPT-2 的实际示例,可作为 Colab 笔记本使用。该过程包括:

  1. 加载预训练的 GPT-2 模型
  2. 向模型添加 LoRA 模块
  3. 配置训练过程
  4. 在 Tiny Shakespeare 上进行微调
  5. 评估微调模型

来源:labml_nn/lora/experiment.py1-190 labml_nn/lora/experiment.ipynb1-97

与其他微调技术的关系

LoRA 属于参数高效微调方法家族。

LoRA 与其他方法的不同之处在于:

  1. 不增加额外的推理延迟(当矩阵合并时)
  2. 保持完全相同的模型架构
  3. 适用于任何权重矩阵,而不仅仅是注意力
  4. 通过秩 r 提供对参数数量的精确控制