使用 TRL 透過 DDPO 微調 Stable Diffusion

Hugging Face 透過 DDPOTrainer 將 Denoising Diffusion Policy Optimization (DDPO) 整合到 trl 函式庫中，使得 Stable Diffusion 模型能夠使用強化學習 (RL) 進行微調，以更好地符合人類偏好與審美標準。

DDPO 與獎勵加權回歸的比較

DDPO 透過降低運算開銷並消除近似誤差，改進了擴散模型的先前 RL 方法，例如 Reward-weighted regression (RWR)。雖然 RWR 重複使用去噪損失函數並根據最終樣本的獎勵進行加權——實際上忽略了中間步驟——DDPO 將整個去噪過程視為多步驟馬可夫決策過程 (MDP)。

透過將過程建模為 MDP 並使用固定採樣器，DDPO 使得代理政策可以是各向同性的高斯分布。這使得我們能夠計算每個去噪步驟的精確 likelihood，而非依賴最終樣本的近似 likelihood，從而提升性能並能處理更複雜的目標。

DDPO 演算法與 RLHF 工作流程

DDPO 使用政策梯度方法，具體來說是 Proximal Policy Optimization (PPO)，來解決優化問題。DDPO 實作中的主要自訂是 PPO 演算法的軌跡收集部分。

當整合到 Reinforcement Learning from Human Feedback (RLHF) 工作流程時，對擴散模型進行對齊的流程會被簡化為三個步驟：

  1. 預訓練模型：從一個預訓練的擴散模型開始（例如 Stable Diffusion）。
  2. 獎勵模型：收集偏好資料並訓練獎勵模型以作為信號。
  3. DDPO 微調：使用 DDPO 並以獎勵模型作為信號來微調模型。

在提升圖像審美性的情況下，Hugging Face 使用了一個凍結的 CLIP 模型，並搭配一個在 Aesthetic Visual Analysis (AVA) 資料集上訓練的可訓練回歸頭，作為獎勵指示器。

實作與訓練指南

使用 DDPO 訓練 Stable Diffusion 需要大量硬體資源；NVIDIA A100 GPU 是避免記憶體不足 (OOM) 錯誤的最低需求。實作是透過 trl 函式庫中的 DDPOTrainer 和 DDPOConfig 類別來處理的。

建議超參數

針對單 GPU 訓練，建議使用以下超參數：

參數 建議值
num_epochs 200
train_batch_size 3
sample_batch_size 6
gradient_accumulation_steps 1
sample_num_steps 50
sample_num_batches_per_epoch 4
per_prompt_stat_tracking True
per_prompt_stat_tracking_buffer_size 32
mixed_precision True
train_learning_rate 3e-4

主要發現與學到的教訓

來自 Hugging Face 的實驗結果顯示，經 DDPO 調整的模型在各種提示詞上具有良好的泛化能力，即使訓練提示詞集合非常小。

LoRA 與完整微調的比較

  • LoRA：建議且在多次測試中證實穩定。
  • 完整微調 (non-LoRA)：相較於 LoRA 能產生更複雜的圖像，但穩定起來顯著更具挑戰性。對於 non-LoRA 的運行，建議將學習率設定得更低（約 1e-5），並將 mixed_precision 設為 None。

目前限制

trl 函式庫中目前的 DDPOTrainer 實作僅限於微調原始 Stable Diffusion 模型。雖然 LoRA 是主要焦點且運作良好，但完整訓練也是可能的，只需要更精確的超參數調整。

Sources