On-Policy Self-Distillation in Diffusion Models
Paper • 2608.24646 • Published • 42
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{zhou2026onpolicy,
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 Li, Linfeng and Wu, Shengqiong and Qu, Leigang and Li, Feng and Liu, Wei and McAuley, Julian and Chua, Tat-Seng},
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