On-Policy Self-Distillation in Diffusion Models

The core problem with applying reinforcement learning to diffusion models is the mismatch between where rewards live and where training happens. A reward scores a finished image; the denoising process is a long chain of intermediate predictions. DiffusionOPSD addresses this by converting image-level reward signals into explicit supervision targets for the clean-output predictions made at intermediate timesteps — giving the model something concrete to fit at each step, not just nudging the final output.

The mechanism works in an outer loop. A frozen behavior policy generates trajectories and provides both query states and anchor predictions. Reward gradients then construct bounded positive and negative targets around each anchor. The trainable policy fits those targets as detached supervision through a finite number of update steps, after which an exponential moving average refreshes the behavior policy. The separation between target construction and target realization turns out to matter: controlled same-query experiments show that larger gains at the construction stage do not reliably produce larger gains after a single fitting update. That gap is worth knowing about — it means you cannot simply assume that better reward gradients translate into better training.

The empirical results are hard to dismiss. Across two backbones — SD 3.5-M and the step-distilled Z-Image-Turbo — DiffusionOPSD achieves the best held-out scores in 19 of 20 reward-matched settings spanning ten evaluators. The margin over the next-best competing method reaches up to 44.0%. On compute, the gains are substantial: training GPU-hours drop by 40% relative to DiffusionNFT on SD 3.5-M and by 63% on Z-Image-Turbo.

What makes this approach analytically interesting is the explicit decomposition of the training loop into two measurable phases. Most RL-for-diffusion methods are difficult to diagnose because the reward signal and the training update are entangled. Here, you can ask separately whether target construction is working and whether the policy is actually realizing those targets — and the paper's own controlled experiments show those questions have different answers. That is a useful property for anyone trying to understand why a run succeeded or failed.

The framing as on-policy self-distillation is apt: the behavior policy is always the model's own recent self, the targets are always freshly constructed from current trajectories, and the update is always bounded and detached. For practitioners working on post-training alignment of generative models, the combination of strong empirical performance and a diagnosable training loop is the real offering.