On-Policy Self-Distillation in Diffusion Models
Paper • 2608.24646 • Published • 66
Text-to-Image
Diffusers
Safetensors
diffusion
image-generation
reinforcement-learning
self-distillation
lora
import torch
from diffusers import DiffusionPipeline
# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium,Tongyi-MAI/Z-Image-Turbo", dtype=torch.bfloat16, device_map="cuda")
pipe.load_lora_weights("WeiChow/DiffusionOPSD")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]DiffusionOPSD: On-Policy Self-Distillation in Diffusion Models
Reward-guided diffusion post-training through explicit, continually refreshed intermediate targets
Overview
DiffusionOPSD is an on-policy self-distillation framework for reward-guided diffusion post-training. A frozen behavior policy collects on-policy denoising states and clean-output anchors; differentiable reward gradients construct bounded positive and negative targets around each anchor; and the trainable policy fits these detached targets before an EMA update refreshes the behavior policy.
By turning image-level rewards into explicit, continually refreshed intermediate supervision, DiffusionOPSD makes target construction and finite realization separately observable. Across SD3.5-M and Z-Image-Turbo, it achieves the best final held-out score in 19 of 20 reward-matched settings and reduces training GPU-hours relative to DiffusionNFT by 40% and 63%, respectively.
Released Checkpoints
This repository provides three rank-32 LoRA adapters:
| Checkpoint | Backbone | Training objective |
|---|---|---|
sd35-m-hpsv3 |
Stable Diffusion 3.5 Medium | HPSv3 |
z-image-turbo-hpsv3 |
Z-Image-Turbo | HPSv3 |
z-image-turbo-pointwise |
Z-Image-Turbo | Pointwise reward |
Download all released adapters with:
hf download WeiChow/DiffusionOPSD --local-dir checkpoints/diffusionopsd
Resources
- Paper: On-Policy Self-Distillation in Diffusion Models
- Code: worldbench/DiffusionOPSD
- Project page: diffusionopsd.github.io
Please refer to the GitHub repository for installation, inference, evaluation, and training instructions.
Citation
@article{zhou2026policy,
title={On-Policy Self-Distillation in Diffusion Models},
author={Zhou, Wei and Zhu, Xiongwei and Kong, Lingdong and Chen, Bo and Zhang, Lei and Liang, Yongyuan and Hou, Xiaoxia and Tian, Ye and Sun, Xian and Wang, Yingshuo and others},
journal={arXiv preprint arXiv:2608.24646},
year={2026}
}
License
The released adapters are provided under the Apache License 2.0. Users must also comply with the licenses of the corresponding base models.
- Downloads last month
- -
Model tree for WeiChow/DiffusionOPSD
Base model
Tongyi-MAI/Z-Image-Turbo