One-sentence Summary

Researchers from ByteDance Seed, National University of Singapore, UC San Diego, and other institutions introduce DiffusionOPSD, an on-policy self-distillation approach that translates image-level reward gradients into bounded positive and negative clean-output targets, fits them via finite fitting and EMA policy updates, and attains the best held-out scores in 19 of 20 reward-matched settings on SD3.5-M and Z-Image-Turbo while reducing training GPU-hours by up to $63\%$63%.

Key Contributions

• DiffusionOPSD is an on-policy self-distillation framework that converts image-level reward guidance into explicit, bounded positive and negative clean-output targets at sampled denoising queries, fits them as detached supervision, and refreshes the targets as the behavior policy evolves.
• The framework separates target construction from model fitting, and controlled same-query experiments show that larger target-construction gains do not necessarily translate into larger realized gains after a single fitting update.
• Across SD3.5-M and Z-Image-Turbo, DiffusionOPSD achieves the best final held-out scores in 19 of 20 reward-matched settings across ten evaluators, outperforms the strongest competing method by up to 44.0%, reduces training GPU-hours relative to DiffusionNFT by 40% on SD3.5-M and 63% on Z-Image-Turbo, and improves all three optimized rewards over DiffusionNFT in joint multi-reward training.

Introduction

Diffusion and flow models are widely used for high-fidelity image generation, and reinforcement learning is increasingly applied to align them with human preferences, aesthetics, and task-specific rewards. A core challenge is that reward is observed only after the final decoded image, while the diffusion policy acts through many intermediate denoising predictions, so outcome-level feedback must be converted into useful local supervision. Prior approaches rely on sampled trajectory credit, backpropagate through a single late clean-output prediction without decoding the full rollout, or reweight final endpoints while leaving the desired intermediate change implicit. The authors introduce DiffusionOPSD, an on-policy self-distillation method that constructs bounded positive and negative targets from image-level reward gradients at behavior-policy queries, fits those targets as detached intermediate supervision, and rebuilds them as the behavior policy evolves. This separates target construction from finite model fitting and enables both stages to be measured independently.

Method

The authors propose DiffusionOPSD, a framework that treats diffusion reward optimization as an on-policy self-distillation process. By converting image-level reward gradients into detached targets for intermediate predictions, the method continuously rebuilds these targets as the behavior policy evolves.

Refer to the framework diagram below for a comprehensive overview of the architecture and training pipeline.

The process begins by establishing a clean-output coordinate system to decode and evaluate predictions. Let $v_{\theta}$vθ​ be the trainable rectified-flow velocity field and $v_{\mathrm{old}}$vold​ denote the velocity field of the behavior policy. For a given prompt $c$c, noisy latent $z_{\sigma}$zσ​, and noise level $\sigma$σ, the query tuple is defined as $s = (c, z_{\sigma}, \sigma)$s=(c,zσ​,σ). Under the rectified-flow path, the clean-output prediction $y_{\theta}(s)$yθ​(s) corresponding to the velocity prediction is computed as:

This prediction is evaluated through the local reward $\widetilde{R}(y, c) = R(D(y), c)$R(y,c)=R(D(y),c), where $D$D is the latent decoder.

During each outer iteration, a frozen behavior policy collects trajectories to supply query states and anchors. A low-noise query state $z_q$zq​ is selected from the trajectory, and the behavior anchor $y_0$y0​ is computed as the clean-output prediction of the behavior policy at that query. Endpoint rewards from the collected trajectories determine a group-normalized fitting weight $\omega$ω. This weight centers each reward within its prompt group while using a global standard deviation over the complete rollout batch, ensuring that all prompt groups are placed on a common scale.

To guide the optimization, the authors construct bounded positive and negative targets starting from the behavior anchor $y_0$y0​. They apply stabilized normalized reward-gradient steps to iteratively update the positive and negative targets. Each update step is projected back onto a trust-region ball to ensure the displacement remains bounded by a radius $\rho \|y_0\|_2$ρ∥y0​∥2​. The final detached positive target $\bar{y}_+$yˉ​+​ is intended to improve the reward, while the negative target $\bar{y}_-$yˉ​−​ acts as a repulsive reference.

The trainable velocity field is then optimized using a finite fitting procedure. The model forms positive and negative fitting branches around the anchor $y_0$y0​ to encode both attraction toward the positive target and rejection of the negative target. Detached adaptive normalizers are utilized to scale the branch losses. The overall objective $\mathcal{L}_{\mathrm{OPSD}}$LOPSD​ is defined as a weighted sum of the squared fitting errors for both branches, where the weights are derived from the previously computed fitting weight $\omega$ω. Crucially, the query, anchor, weight, and targets remain fixed during this finite fitting phase, completely separating the target construction from the parameter updates.

Finally, the online training loop is completed by refreshing the behavior policy. After applying a finite number of optimizer updates to the trainable policy, the behavior policy is updated via an exponential moving average. This updated behavior policy then generates new trajectories for the subsequent outer iteration, supplying new query states and endpoint rewards. This continuous cycle of trajectory collection, target construction, finite fitting, and behavior policy refresh constitutes the core self-distillation loop of the system.

Experiment

The experiments evaluate DiffusionOPSD against FlowGRPO, ReFL, and DiffusionNFT on SD3.5-M and the step-distilled Z-Image-Turbo, measuring held-out reward-specific quality, training cost, and optimization stability. DiffusionOPSD achieves the best final scores in 19 of 20 reward-matched settings, reduces training cost by 40–63%, and sustains reward gains through the final update, with human annotators preferring its outputs over all baselines. Ablations confirm that the reward-gradient direction in target construction is essential, while most implementation choices and query-state provenance have minor effects, and the method remains effective with native few-step samplers and multi-reward training.

DiffusionOPSD achieves the best held-out evaluator scores in 19 of 20 reward-matched settings across a standard backbone and a native few-step distilled model, with gains reaching 44% on HPSv3 and VLM-Pairwise. It sustains improvement through the final training update, retaining 98% of within-run reward gains, while competing methods often plateau or degrade. A single jointly trained policy preserves near-specialist quality on three rewards, outperforming joint baselines by 8–14% and matching or exceeding individual specialists. DiffusionOPSD leads in 19 of 20 reward-matched comparisons across SD3.5-M and Z-Image-Turbo, with relative gains up to 44% on HPSv3 and VLM-Pairwise. On Z-Image-Turbo, DiffusionOPSD outperforms the strongest baseline by 30.7% on ImageReward and 14.6% on VLM-Pairwise, while DiffusionNFT falls below the unadapted model on eight of ten objectives. DiffusionOPSD retains 98% of its within-run reward improvement at the final update, compared to 83% for FlowGRPO and 90% for DiffusionNFT. A single DiffusionOPSD policy trained on three rewards achieves 113% of the PickScore specialist's gain, statistically matches the HPSv2.1 specialist, and surpasses the joint DiffusionNFT policy by 8–14% across all three metrics.

DiffusionOPSD is evaluated on reward-matched image generation tasks using both a standard backbone and a native few-step distilled model, comparing held-out evaluator scores and training stability. It achieves leading performance in nearly all settings, with sustained improvement through the final training update while competing methods often plateau or degrade. A single jointly trained policy preserves near-specialist quality across multiple rewards, outperforming joint baselines and matching or exceeding individual specialists.