使用 LoRA 高效微调 Stable Diffusion

Hugging Face 已将低秩适配（LoRA）集成到 diffusers 库中，实现了更快的 Stable Diffusion 微调，显著降低了显存需求，模型权重小至 3 MB。此实现使用户能够在不承担完整模型微调高昂计算成本的情况下，将大规模扩散模型适配到特定风格或概念。

LoRA 架构及其在 Stable Diffusion 中的应用

低秩适配（LoRA）是一种最初由微软研究员为大型语言模型（LLM）如 GPT-3 开发的技术。它通过冻结预训练模型权重，并向 transformer 块注入可训练的秩分解矩阵来实现。由于大部分模型权重不计算梯度，这种方法降低了可训练参数的数量和 GPU 内存需求。

在 Stable Diffusion 中，LoRA 应用于交叉注意力层。这些层负责建立图像表示与描述它们的文本提示之间的关系。通过针对这些特定层，LoRA 实现了与完整模型微调相当的微调质量，同时保持更快的速度和更高的计算效率。

LoRA 微调的主要优势

将 LoRA 集成到 diffusers 库中，相较于传统微调方法提供了三大优势：

  • 降低计算需求： 微调可以在消费级硬件上进行。例如，可使用配备 11 GB 显存的 NVIDIA 2080 Ti 创建完整的微调模型。
  • 更快的训练： 可训练参数的减少显著加快了训练周期。
  • 极大缩小的权重文件： 由于原始模型保持冻结，仅需保存注入层的权重。这导致权重文件约为 3 MB（某些情况下具体为 3.29 MB），大约比原始 UNet 模型小 1,000 倍。

实现与推理工作流

微调过程

Hugging Face 提供了专用的 LoRA 微调脚本（train_text_to_image_lora.py），可在仅 11 GB GPU 内存的情况下运行。一个值得注意的技术细节是 LoRA 通常需要比常规微调更高的学习率；例如，使用 1e-4 的学习率，而完整微调通常使用约 ~1e-6 的学习率。

推理与加载

推理过程旨在在未修改的 Stable Diffusion 基础模型之上加载 LoRA 权重。工作流包括：

  1. 基础模型识别： 使用 Hub API 确定训练时使用的基础模型（例如 CompVis/stable-diffusion-v1-4 或 runwayml/stable-diffusion-v1-5）。
  2. 管道初始化： 使用已识别的基础模型加载标准的 StableDiffusionPipeline。
  3. 权重注入： 使用 pipe.unet.load_attn_procs(model_path) 将 LoRA 权重直接从 Hub 加载到常规模型权重上。

与 Dreambooth 的兼容性

LoRA 与 Dreambooth 兼容，Dreambooth 是一种用于教会 Stable Diffusion 新特定概念的方法。将这两种技术结合可带来多项优势：

  • 比标准 Dreambooth 更快的训练速度。
  • 低数据需求，通常仅需 5 到 10 张目标图像。
  • 可选的文本编码器微调，以提升对目标的忠实度。

与其他微调方法的比较

虽然还有其他方法用于适配 Stable Diffusion，但 LoRA 根据使用场景提供了独特的优势：

  • 文本反演（Textual Inversion）： 生成类似 LoRA 的小型、易于共享的权重，但通常仅限于单一主体或少量概念。
  • LoRA： 适用于通用微调，使模型能够适配全新的领域或数据集。
  • 关键调优（Pivotal Tuning）： 一种混合方法，先使用文本反演获取 token 嵌入，再利用 LoRA 对该嵌入进行训练。

Sources