使用 LoRA 進行高效的 Stable Diffusion 微調

Hugging Face 已將低秩適應（LoRA）整合到 diffusers 套件中，使 Stable Diffusion 的微調更快，所需的 VRAM 大幅降低，模型權重僅約 3 MB。此實作讓使用者能在不需全模型微調的高昂計算成本下，將大型擴散模型調整至特定風格或概念。

LoRA 架構與在 Stable Diffusion 中的應用

低秩適應（LoRA）是一項最初由微軟研究人員為大型語言模型（LLM）如 GPT-3 開發的技術。它透過凍結預訓練模型權重，並在 transformer 區塊中注入可訓練的秩分解矩陣來運作。此方法減少了可訓練參數的數量與 GPU 記憶體需求，因為大部分模型權重不會計算梯度。

在 Stable Diffusion 的情境下，LoRA 被應用於 cross‑attention 層。這些層負責建立影像表徵與描述它們的文字提示之間的關係。透過針對這些特定層，LoRA 能達到與全模型微調相當的微調品質，同時更快且計算更有效率。

LoRA 微調的主要優勢

  • 降低計算需求： 微調可以在消費級硬體上執行。例如，可使用具 11 GB VRAM 的 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) 從 Hub 直接將 LoRA 權重載入至常規模型權重上。

與 Dreambooth 的相容性

LoRA 與 Dreambooth 相容，Dreambooth 是用於教導 Stable Diffusion 新的特定概念的方法。結合這兩種技術可帶來多項好處：

  • 較快的訓練速度，相較於標準 Dreambooth。
  • 低資料需求，通常僅需 5 至 10 張主體圖像。
  • 可選的文字編碼器微調，以提升對主體的忠實度。

與其他微調方法的比較

雖然有其他方法可用於調整 Stable Diffusion，但 LoRA 依據使用情境提供了明顯的優勢：

  • Textual Inversion（文字反轉）： 產生類似 LoRA 的小型、易於分享的權重，但通常僅限於單一主體或少數概念。
  • LoRA： 適用於通用微調，允許模型適應全新領域或資料集。
  • Pivotal Tuning（關鍵微調）： 結合 Textual Inversion 取得 token 嵌入，然後使用 LoRA 訓練該嵌入的混合方法。

Sources