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 减少了 40 %–63 % 的 GPU 小时数，同时在 19/20 个奖励匹配设置上实现了更高的保留分数。
丰富的脚本	提供了用于公开训练、基线对比、烟雾测试和混合奖励实验的即用型启动脚本；还包含用于检查奖励模型设置的工具。

🛠️ 安装与快速入门

先决条件 – 支持 CUDA 的 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 次优化器步骤 ≈ 数百次 rollout 轮次）。

# 使用 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 是一个文档齐全、开源的实现，采用一种新颖的在线策略自蒸馏技术，根据多个学习到的奖励函数改进扩散图像生成，同时相比先前方法降低了训练成本。它包含即用型脚本、支持两种主要扩散主干，并具备灵活的奖励混合系统。