Papers
arxiv:2608.24646

On-Policy Self-Distillation in Diffusion Models

Published on Aug 25
ยท Submitted by
Wei Chow
on Aug 26
Authors:
,
,
,
,
,
,
,
,
,
,
,
,
,
,
,
,

Abstract

DiffusionOPSD uses on-policy self-distillation to turn image-level rewards into explicit intermediate targets for diffusion models, improving alignment efficiency and enabling separate analysis of target construction and policy fitting.

Reinforcement learning can align diffusion models with human preferences and task-specific objectives, but endpoint rewards do not specify how an intermediate denoising prediction should change. We introduce DiffusionOPSD as an on-policy self-distillation framework that converts image-level reward guidance into explicit targets for clean-output predictions at sampled queries. At each outer iteration, a frozen behavior policy generates trajectories and supplies query states and anchors. Reward gradients construct bounded positive and negative targets around each anchor. The trainable policy fits these targets as detached supervision through finite fitting before an exponential moving average update refreshes the behavior policy. This setup lets us measure target construction and finite realization separately. Controlled same-query experiments show that larger target-construction gains do not necessarily translate into larger realized gains after a single fitting update. Across SD 3.5-M and the step-distilled Z-Image-Turbo, our approach achieves the best final held-out scores in 19 of 20 reward-matched settings across two backbones and ten evaluators. It outperforms the strongest competing method by up to 44.0% and reduces training GPU-hours relative to DiffusionNFT by 40% on SD 3.5-M and 63% on Z-Image-Turbo. These results support on-policy self-distillation as an efficient and analyzable approach to diffusion post-training by converting image-level reward guidance into explicit and continually refreshed intermediate supervision, thereby opening a path toward more efficient and diagnosable alignment.

Community

Paper submitter

On-Policy Self-Distillation in Diffusion Models

Iโ€™m excited to introduce our latest work at ByteDance Seed, DiffusionOPSD, a follow-up to DanceOPD.
Diffusion and flow models generate images through multi-step denoising trajectories, while mainstream reward models usually provide only an image-level score at the rollout endpoint. Such a score tells us whether the final result is good, but not how an intermediate denoising prediction should change.

The idea grew naturally out of DanceOPD. In DanceOPD, we found that querying a teacher only once at a low-noise state along the student rollout already provides effective intermediate supervision. This led us to ask: Can we remove the external teacher, use reward gradients to improve the modelโ€™s own prediction at a single query, and then distill that improvement back into the model?

This is the starting point of DiffusionOPSD.
We introduce an On-Policy Self-Distillation framework:

  • A frozen behavior policy generates trajectories and provides low-noise query states.
  • Its clean-output prediction at each query serves as an anchor.
  • Reward gradients construct bounded positive and negative targets around the anchor.
  • The trainable policy fits these detached targets under a finite update budget.
  • An EMA update refreshes the behavior policy before the next iteration regenerates trajectories, anchors, and targets.

Intuitively, DiffusionOPSD first turns image-level reward guidance into an explicit description of how the current prediction should improve, and then distills that target back into the model. This design also makes target construction and finite realization separately observable and diagnosable.
Our experiments cover SD3.5-M and Z-Image-Turbo across 10 evaluators. DiffusionOPSD achieves the best final held-out score in 19 of 20 reward-matched settings, outperforming the strongest competing method by up to 44.0%. Compared with DiffusionNFT, it reduces training GPU-hours by 40% on SD3.5-M and 63% on Z-Image-Turbo ๐Ÿš€

For joint training on PickScore, CLIPScore, and HPSv2.1, a single DiffusionOPSD policy outperforms the corresponding DiffusionNFT policy on all three reward objectives while retaining most of the gains achieved by separate specialists. In our OPD experiments, direct joint multi-reward training with DiffusionOPSD also outperforms the two-stage DanceOPD, DiffusionOPD, and FlowOPD students on all three optimized rewards.

We also observe an interesting phenomenon: a larger reward gain during target construction does not necessarily lead to a larger realized gain after model updates. Constructing a better target and successfully realizing it through finite fitting are different problemsโ€”and should be evaluated separately.
The paper and code are now publicly available. The code supports single-reward and mixed-reward training with seven open-weight evaluators. Feedback and discussions are very welcome ๐Ÿ™Œ

๐ŸŒ Project: https://diffusionopsd.github.io/
๐Ÿ“„ Paper: https://arxiv.org/pdf/2608.24646
๐Ÿ’ป Code: https://github.com/worldbench/DiffusionOPSD

Sign up or log in to comment

Models citing this paper 1

Datasets citing this paper 0

No dataset linking this paper

Cite arxiv.org/abs/2608.24646 in a dataset README.md to link it from this page.

Spaces citing this paper 0

No Space linking this paper

Cite arxiv.org/abs/2608.24646 in a Space README.md to link it from this page.

Collections including this paper 0

No Collection including this paper

Add this paper to a collection to link it from this page.