worldbench/DiffusionOPSD

📚 什麼是 DiffusionOPSD？

DiffusionOPSD（在策略自我蒸餾）是一個研究級別的程式碼庫，實現了論文 「DiffusionOPSD: 在擴散模型中的在策略自我蒸餾」（arXiv 2608.24646）所描述的演算法。它提供了一種方法，可對大型擴散圖像生成器（例如 Stable Diffusion 3.5‑M 和 Z‑Image‑Turbo）進行微調，使生成的圖像在多種學習到的圖像品質或偏好指標（PickScore、CLIPScore、HPS 等）上獲得更高分數。

其核心思想是彌補通常在整個去噪軌跡結束後才觀察到獎勵時所存在的監督差距。DiffusionOPSD 重複執行以下步驟：

從當前（凍結）的「行為」策略中收集低雜訊狀態。
透過在每個狀態的乾淨輸出錨點附近進行幾步歸一化的獎勵梯度步驟，建立明確的正負目標。
訓練一個新的「可訓練」策略，使其在有限次優化器更新內匹配這些分離的目標。
透過可訓練策略的指數移動平均（EMA）來更新行為策略，然後重複。

由於目標構建與模型擬合是分離的，此方法可被仔細檢視、公平地與先前方法（如 DiffusionNFT 和 FlowGRPO）進行基準測試與比較。

🚀 主要功能

功能	重要性
在策略查詢收集	目標是從當前模型實際訪問的狀態中建立，避免使用離線或人工加噪狀態時產生的不匹配問題。
明確的獎勵引導目標	正負錨點透過沿獎勵梯度進行有界步驟生成，讓模型有明確的改進方向。
分離的有限擬合	在策略更新前，獎勵/解碼器圖形被分離，因此目標品質可獨立於模型學習效果進行衡量。
基於 EMA 的行為刷新	每次外層迴圈後，行為策略會透過指數移動平均更新，確保訓練全程都有新鮮的監督。
支援兩種主幹模型	可直接使用 Stable Diffusion 3.5‑M（512²）和 Z‑Image‑Turbo（1024²），涵蓋標準與少步驟擴散架構。
單一與混合獎勵訓練	提供論文中使用的七個公開評估器的預設設定，並支援任意正權重組合（例如 PickScore/26 + CLIPScore + HPSv2.1）。
效率提升	相較於 DiffusionNFT，報告顯示 GPU 小時數減少 40 %–63 %，同時在 19/20 的獎勵匹配設定中獲得更高保留分數。
豐富的腳本工具	提供公開訓練、基線對照、煙霧測試與混合獎勵實驗的預設啟動腳本；也包含檢查獎勵模型設定的實用工具。

🛠️ 安裝與快速入門

先決條件 – 具備 CUDA 支援 GPU 的 Linux 系統、Python 3.10‑3.11，以及對應的 PyTorch 建置版本。

# 克隆程式碼庫
git clone https://github.com/worldbench/DiffusionOPSD.git
cd DiffusionOPSD

# 建立新的 conda 環境（可選但建議）
conda create -n diffusionopsd python=3.11 -y
conda activate diffusionopsd

# 安裝核心套件（包含 reward-stack 額外功能）
pip install -e "[rewards]"
# ImageReward 會鎖定舊版 timm；請不帶依賴安裝
pip install --no-deps 'image-reward==1.5'

# 下載公開預設使用的獎勵模型檢查點
export REWARD_CKPT_PATH="$PWD/reward_ckpts"
bash scripts/download_reward_weights.sh

# 準備 Pick‑a‑Pic 提示清單（論文中實驗使用）
python scripts/prepare_pickapic_prompts.py

若想在 Z‑Image‑Turbo 上訓練，需要一個包含 ZImagePipeline 的 Diffusers 建置版本：

git clone https://github.com/huggingface/diffusers.git
pip install -e "./diffusers[torch]"

可選環境調整

export HF_HOME=/path/to/huggingface-cache   # 模型權重快取位置
export WANDB_MODE=offline                    # 預設關閉線上記錄

▶️ 執行小型煙霧測試

此程式碼庫附帶微型「煙霧」腳本，可在單一 GPU 節點上執行單一優化器更新，非常適合驗證你的機器上獎勵梯度管道是否正常運作。

# SD3.5‑M 主幹，HPSv2.1 獎勵（快速檢查）
SMOKE_TEST=1 NPROC=8 UPDATES=1 bash scripts/train_public.sh sd35 hpsv2

# Z‑Image‑Turbo 主幹，CLIPScore 獎勵
SMOKE_TEST=1 NPROC=8 UPDATES=1 bash scripts/train_public.sh zimage clipscore

你也可以測試耗資源的獎勵（HPSv3、DeQA），它們需要獨立環境，因為依賴大型 7‑/8‑B 模型：

SMOKE_TEST=1 NPROC=7 UPDATES=1 bash scripts/train_public.sh zimage hpsv3
SMOKE_TEST=1 NPROC=7 UPDATES=1 bash scripts/train_public.sh zimage deqa

📊 全規模訓練範例

以下是訓練公開獎勵特定模型的典型指令，使用論文中使用的完整 100 次更新預算（≈ 100 次優化器步驟 ≈ 數百次滾動回合）。

# 使用 PickScore/26 + CLIPScore + HPSv2.1 混合目標訓練 SD3.5‑M
NPROC=8 UPDATES=100 OUTPUT_DIR=outputs/sd35_mixed_opsd \
  bash scripts/train_public.sh sd35 mixed

# 使用 HPSv3 評估器訓練 Z‑Image‑Turbo（6 個策略等級 + 1 個獎勵伺服器等級）
NPROC=7 UPDATES=100 OUTPUT_DIR=outputs/zimage_hpsv3 \
  bash scripts/train_public.sh zimage hpsv3

檢查點每 10 次更新儲存一次（可設定），最終模型會自動寫入。

📦 程式碼庫內含內容？

目錄 / 檔案	目的
`scripts/`	用於下載獎勵、準備提示、執行煙霧測試與啟動分散式訓練的輔助腳本。
`config/`	每個主幹、每種獎勵與混合獎勵設定的 YAML 格式設定檔。
`scripts/train_*.py`	入口點，負責設定分散式執行、載入行為與可訓練策略、收集軌跡、建立目標並執行優化器步驟。
`scripts/check_reward_setup.py`	在長時間作業前，驗證所需獎勵模型檢查點與程式庫版本是否存在。
`scripts/smoke_reward_gradient.py`	載入單一獎勵模型，確認可計算非零的影像空間梯度。
`assets/`	README 中使用的圖示（質性圖庫、訓練曲線、消融實驗）。
`LICENSE`	Apache 2.0 – 容許度高的開源授權。

📜 授權條款

程式碼以 Apache 2.0 授權釋出，允許自由使用、修改與分發（需註明來源）。模型權重本身受上游擴散模型與獎勵模型原始授權約束，不會在程式碼庫中重新分發。

🎯 誰應該使用此工具？

研究人員：探索獎勵引導的擴散模型微調、自我蒸餾，或針對生成模型的在策略強化學習方法。
實務工作者：需要一個高品質、可控制的擴散模型，針對特定美學或偏好指標進行調校，並希望獲得可重現的基線。
學生：尋找一個具體、端到端的範例，了解如何將可微分獎勵模型（包括大型語言-視覺評估器）整合至擴散訓練流程中。

📚 更多閱讀

完整論文（arXiv 2608.24646）詳細說明數學公式、理論動機與廣泛的消融實驗。
專案頁面（https://diffusionopsd.github.io/）提供額外視覺結果與常見問題解答。
關於競爭方法的背景，請參閱引用的 DiffusionNFT 與 FlowGRPO。

TL;DR： DiffusionOPSD 是一個文件完整、開源的新型在策略自我蒸餾技術實現，能根據多個學習到的獎勵函數提升擴散圖像生成品質，同時相比先前方法大幅降低訓練成本。它包含可立即執行的腳本、支援兩大主要擴散主幹模型，以及靈活的獎勵混合系統。