worldbench/DiffusionOPSD

📚 What is DiffusionOPSD?

DiffusionOPSD (On‑Policy Self‑Distillation) is a research‑grade codebase that implements the algorithm described in the paper “DiffusionOPSD: On‑Policy Self‑Distillation in Diffusion Models” (arXiv 2608.24646). It provides a way to fine‑tune large diffusion image generators (e.g., Stable Diffusion 3.5‑M and Z‑Image‑Turbo) so that the generated images score higher on a variety of learned image‑quality or preference metrics (PickScore, CLIPScore, HPS, etc.).

The key idea is to close the supervision gap that normally exists when a reward is only observed after an entire denoising trajectory. DiffusionOPSD repeatedly:

Collects low‑noise states from the current (frozen) “behavior” policy.
Creates explicit positive and negative targets by taking a few normalized reward‑gradient steps around the clean‑output anchor of each state.
Trains a new “trainable” policy to match those detached targets within a limited number of optimizer updates.
Refreshes the behavior policy with an exponential‑moving‑average (EMA) of the trainable policy, then repeats.

Because target construction and model fitting are separated, the method can be inspected, benchmarked, and compared fairly against prior approaches such as DiffusionNFT and FlowGRPO.

🚀 Main Features

Feature	Why it matters
On‑policy query collection	Targets are built from states actually visited by the current model, avoiding the mismatch that occurs when using offline or artificially noised states.
Explicit reward‑guided targets	Positive/negative anchors are generated by taking bounded steps along the reward gradient, giving the model a clear direction to improve.
Detached finite fitting	The reward/decoder graph is detached before the policy update, so the quality of the target can be measured independently of how well the model learns it.
EMA‑based behavior refresh	After each outer loop the behavior policy is updated with an exponential moving average, ensuring fresh supervision throughout training.
Supports two backbones	Works out‑of‑the‑box with Stable Diffusion 3.5‑M (512²) and Z‑Image‑Turbo (1024²), covering both standard and few‑step diffusion regimes.
Single‑ and mixed‑reward training	Pre‑defined presets for each of the seven public evaluators used in the paper, plus the ability to combine arbitrary positive‑weighted sums (e.g., PickScore/26 + CLIPScore + HPSv2.1).
Efficiency gains	Reported 40 %–63 % reduction in GPU‑hours compared to DiffusionNFT while achieving higher held‑out scores on 19/20 reward‑matched settings.
Extensive scripts	Ready‑made launch scripts for public training, baseline baselines, smoke‑tests, and mixed‑reward experiments; also utilities for checking reward‑model setup.

🛠️ Installation & Quick‑Start

Prerequisites – Linux with CUDA‑capable GPUs, Python 3.10‑3.11, and a matching PyTorch build.

# Clone the repo
git clone https://github.com/worldbench/DiffusionOPSD.git
cd DiffusionOPSD

# Create a fresh conda env (optional but recommended)
conda create -n diffusionopsd python=3.11 -y
conda activate diffusionopsd

# Install the core package (includes the reward‑stack extra)
pip install -e "[rewards]"
# ImageReward pins an old timm version; install it without deps.
pip install --no-deps 'image-reward==1.5'

# Download reward‑model checkpoints used by the public presets
export REWARD_CKPT_PATH="$PWD/reward_ckpts"
bash scripts/download_reward_weights.sh

# Prepare the Pick‑a‑Pic prompt manifest (used by the paper’s experiments)
python scripts/prepare_pickapic_prompts.py

If you want to train on Z‑Image‑Turbo, you need a Diffusers build that contains ZImagePipeline:

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

Optional environment tweaks

export HF_HOME=/path/to/huggingface-cache   # cache location for model weights
export WANDB_MODE=offline                    # disables online logging by default

▶️ Running a Small Smoke Test

The repository ships tiny “smoke” scripts that run a single optimizer update on a single GPU‑node, useful for verifying that the reward‑gradient pipeline works on your machine.

# SD3.5‑M backbone, HPSv2.1 reward (fast check)
SMOKE_TEST=1 NPROC=8 UPDATES=1 bash scripts/train_public.sh sd35 hpsv2

# Z‑Image‑Turbo backbone, CLIPScore reward
SMOKE_TEST=1 NPROC=8 UPDATES=1 bash scripts/train_public.sh zimage clipscore

You can also test the heavy rewards (HPSv3, DeQA) which require a separate environment because they depend on large 7‑/8‑B models:

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

📊 Full‑scale Training Example

Below is a typical command to train a public‑reward‑specific model for the full 100‑update budget used in the paper (≈ 100 optimizer steps ≈ several hundred rollout rounds).

# Train SD3.5‑M with the PickScore/26 + CLIPScore + HPSv2.1 mixed objective
NPROC=8 UPDATES=100 OUTPUT_DIR=outputs/sd35_mixed_opsd \
  bash scripts/train_public.sh sd35 mixed

# Train Z‑Image‑Turbo with the HPSv3 evaluator (6 policy ranks + 1 reward‑server rank)
NPROC=7 UPDATES=100 OUTPUT_DIR=outputs/zimage_hpsv3 \
  bash scripts/train_public.sh zimage hpsv3

Checkpoints are saved every 10 updates (configurable) and the final model is written automatically.

📦 What’s Inside the Repo?

Directory / File	Purpose
`scripts/`	Helper scripts for downloading rewards, preparing prompts, running smoke tests, and launching distributed training.
`config/`	YAML‑style configuration files for each backbone, each reward, and mixed‑reward setups.
`scripts/train_*.py`	Entry points that set up the distributed run, load the behavior and trainable policies, collect trajectories, build targets, and perform the optimizer step.
`scripts/check_reward_setup.py`	Verifies that the required reward‑model checkpoints and library versions are present before a long job.
`scripts/smoke_reward_gradient.py`	Loads a single reward model and confirms that a non‑zero image‑space gradient can be computed.
`assets/`	Figures used in the README (qualitative gallery, training curves, ablations).
`LICENSE`	Apache 2.0 – permissive open‑source license.

📜 License

The code is released under the Apache 2.0 License, allowing free use, modification, and distribution (with attribution). Model weights themselves are subject to the original licenses of the upstream diffusion and reward models and are not redistributed in the repo.

🎯 Who Should Use This?

Researchers exploring reward‑guided diffusion fine‑tuning, self‑distillation, or on‑policy reinforcement‑learning style methods for generative models.
Practitioners who need a high‑quality, controllable diffusion model tuned to a specific aesthetic or preference metric and want a reproducible baseline.
Students looking for a concrete, end‑to‑end example of how to integrate differentiable reward models (including large language‑vision evaluators) into diffusion training pipelines.

📚 Further Reading

The full paper (arXiv 2608.24646) explains the mathematical formulation, theoretical motivations, and extensive ablations.
The project page (https://diffusionopsd.github.io/) hosts additional visual results and a FAQ.
For background on competing methods, see the cited works DiffusionNFT and FlowGRPO.

TL;DR: DiffusionOPSD is a well‑documented, open‑source implementation of a novel on‑policy self‑distillation technique that improves diffusion image generation according to multiple learned reward functions, while cutting training cost compared to prior methods. It includes ready‑to‑run scripts, support for two major diffusion backbones, and a flexible reward‑mixing system.