On-Policy Self-Distillation in Diffusion Models

Wei Zhou, Xiongwei Zhu, Lingdong Kong, Bo Chen, Lei Zhang, Yongyuan Liang, Xiaoxia Hou, Ye Tian, Xian Sun, Yingshuo Wang, Linfeng Li, Shengqiong Wu, Leigang Qu, Feng Li, Wei Liu, Julian McAuley, Tat-Seng Chua

DiffusionOPSD aligns diffusion models by converting reward gradients into explicit, detached intermediate targets.

How can we convert image-level reward feedback into precise, intermediate supervision for diffusion model denoising steps?

Diffusion models are typically aligned using endpoint rewards, but these rewards provide no direct signal for how intermediate denoising steps should change to improve the final output. DiffusionOPSD treats alignment as an on-policy self-distillation loop: it uses reward gradients to construct explicit "positive" and "negative" targets for intermediate clean-output predictions, then fits the model to these targets as detached supervision. This approach achieves the best held-out scores in 19 of 20 reward-matched settings across two backbones, while reducing training compute by up to 63% compared to existing methods.

Paper Primer

The core mechanism is a three-stage loop: a frozen behavior policy collects trajectories, reward gradients construct bounded targets around the current prediction, and the trainable policy fits these targets as detached supervision. This is like a coach who doesn't just grade the final score, but provides specific, actionable corrections for each intermediate move in a sequence.

DiffusionOPSD consistently outperforms existing alignment methods in held-out quality.

Across SD3.5-M and Z-Image-Turbo backbones, the method achieved the best held-out score in 19 of 20 reward-matched settings. Up to 44.0% improvement over the strongest competing method.

The method significantly improves training efficiency.

Measured GPU-hours per 100 updates on SD3.5-M and Z-Image-Turbo. 40% and 63% reduction in training cost relative to DiffusionNFT.

Why is it necessary to construct intermediate targets instead of just backpropagating the reward?

Endpoint rewards create a structural mismatch because the reward is only observed after a full rollout, whereas the model acts at every intermediate denoising step. Constructing explicit targets allows the model to receive actionable supervision at every step of the generation process.

How does this method handle the risk of "over-fitting" to a specific reward?

The method uses a negative fitting branch that acts as a repulsive reference, and it continually rebuilds targets as the behavior policy evolves, preventing the model from collapsing into a single, potentially brittle, reward-maximizing point.

Paper Primer

Introducing DiffusionOPSD, an on‑policy self‑distillation method that turns reward gradients into intermediate targets.

Endpoint reward supervision supplies a single scalar score for the final clean image, leaving the intermediate denoising steps without guidance.

It treats the diffusion model as a behavior policy and, at each training step, turns the gradient of an image‑level reward into a concrete target for the clean‑output prediction at a sampled intermediate state.

How does DiffusionOPSD differ from standard on‑policy RL applied to diffusion models?

Standard on‑policy RL would treat the reward as a scalar loss applied after the full denoising trajectory, whereas DiffusionOPSD converts the reward gradient into explicit intermediate targets, providing concrete supervision for each sampled denoising step.

**Figure 2.** Samples generated by DiffusionOPSD from held-out text prompts.

DiffusionOPSD converts reward gradients into explicit intermediate denoising targets, enabling efficient and analyzable alignment of diffusion models.

The Reward Alignment Problem

We expose the mismatch between endpoint rewards and intermediate denoising steps and motivate explicit intermediate targets.

Diffusion and flow models generate high‑fidelity images, yet reinforcement‑learning objectives are typically supplied only at the rollout endpoint. This creates a structural mismatch: the reward is observed after the final decode, while the policy makes a sequence of interdependent denoising predictions.

Endpoint reward supervision provides a single scalar signal after the entire denoising trajectory, ignoring the intermediate queries that actually drive the model.

The paper separates the problem into two coupled stages. First, a target is constructed at a sampled query using the reward gradient; second, a finite fitting step updates the model to match that target. “Finite fitting” denotes the parameter‑update procedure, while “finite realization” denotes the actual reward change induced by that update. Because the same model is queried repeatedly across noise levels, fitting a target at one query can unintentionally affect predictions at many others, leading to under‑realization, rotation, or overshoot of the intended change.

**Figure 3.** Reward-to-policy paradigms. We contrast trajectory credit, late-state reward backpropagation, and endpoint supervision with our explicit intermediate targets constructed before finite fitting.

Reward-to-Target Paradigms

We contextualize prior diffusion supervision methods and introduce key interface concepts.

Reward-to-Target Interfaces describe how outcome-level feedback is turned into concrete supervision for each denoising step.

It maps a scalar reward signal onto a concrete target that each intermediate diffusion step should aim for.

It treats the diffusion trajectory as a latent space where a learned “NFT” (non‑fixed target) guides intermediate predictions.

ReFL backpropagates a reward through a differentiable clean‑output prediction at a sampled denoising state.

FlowGRPO assigns credit to groups of trajectories by comparing reverse‑process likelihood ratios.

On-Policy Distillation methods collect supervision from the learner’s own states to reduce distribution shift.

Reward Gradients approaches such as classifier guidance, DRaFT, and AlignProp directly translate rewards on generated samples into parameter updates.

Prior methods differ in how they convert outcome feedback into supervision.

On-Policy Self-Distillation

How DiffusionOPSD turns reward gradients into on‑policy self‑distillation targets.

Endpoint‑only rewards leave the intermediate denoising steps under‑specified, so the model has no guidance on how to improve the trajectory itself. DiffusionOPSD fills this gap by turning the reward gradient at a query into a concrete target that the model can chase. The following blocks spell out the mechanism.

Think of a teacher that watches the student’s current answer, asks “what would make this answer better?” (the reward gradient), and then writes a short‑term “homework” target for the student to hit before the next lesson; the teacher never changes while the student works on that homework.

How does OPSD differ from ordinary diffusion training that only uses an endpoint loss?

Ordinary diffusion training treats the final decoded image as the only supervision signal; OPSD additionally creates intermediate supervision by projecting the reward gradient onto the latent space and using the resulting positive/negative points as explicit targets for the denoising steps.

Freeze the current behavior policy and sample $K$ trajectories per prompt to obtain low‑noise queries.

Select the query whose noise level best matches the desired $\sigma^{\star}$ and compute its anchor $y_{0}$.

From $y_{0}$ run $M_{\text{tgt}}$ normalized reward‑gradient steps to obtain detached targets $\bar{y}_{+}$ and $\bar{y}_{-}$.

Fit the trainable velocity field $v_{\theta}$ to those targets using the branch loss $\mathcal{L}_{\text{OPSD}}$ for a fixed number $M_{\text{fit}}$ of optimizer updates.

Update the behavior policy by an exponential moving average (EMA) and repeat.

**Figure 4.** DiffusionOPSD Overview. The behavior policy collects low-noise queries, reward gradients construct bounded positive and negative targets, and the trainable policy fits the detached supervision. The updated behavior policy supplies new anchors for the next iteration. This procedure separates target construction from finite fitting.

Compute the clean output: $y_{\theta} = z_{\sigma} - \sigma v_{\theta} = (0.6,0.8) - 0.2\cdot(0.5,0.5) = (0.5,0.7)$.

Recover the velocity from the clean output: $v_{\theta} = (z_{\sigma} - y_{\theta})/\sigma = ((0.6,0.8)-(0.5,0.7))/0.2 = (0.5,0.5)$, confirming consistency.

This shows that the clean‑output prediction is simply the noisy latent shifted back along the velocity direction, and the relationship is invertible.

Compute the gradient at $y_{0}$: $\nabla_{y}R = -2\,(y_{0}-(2,2)) = (2,2)$.

Positive update: $y^{(1)}_{+}=y_{0}+0.1\cdot\frac{(2,2)}{\| (2,2)\|^{2}+10^{-6}} = (1.0,1.0)+0.1\cdot\frac{(2,2)}{8}= (1.025,1.025)$.

Negative update: $y^{(1)}_{-}=y_{0}-0.1\cdot\frac{(2,2)}{8}= (0.975,0.975)$.

After two steps the positive target moves toward $(2,2)$ while the negative target moves away, staying symmetric around $y_{0}$.

The normalization makes each step size inversely proportional to the gradient magnitude, preventing runaway updates when the reward surface is steep.

Empirical Evaluation

DiffusionOPSD beats baselines in quality, efficiency, and reward stability.

We evaluate DiffusionOPSD on two backbones (SD3.5‑M and 9‑step Z‑Image‑Turbo) across quality, efficiency, and reward‑gradient stability.

DiffusionOPSD attains the best final held‑out score in $19/20$ reward‑matched settings, beating every baseline.

Across both backbones and ten evaluators, DiffusionOPSD leads in 19 of 20 cases; only the Aesthetic metric on SD3.5‑M is marginally lower than ReFL.

**Figure 1.** Training and held-out quality curves. Normalized gains are averaged across matched 10 reward and 2 backbone settings. DiffusionOPSD improves fastest and reaches the highest final gains.

**Figure 5. Training dynamics.** Across 71 single-reward runs, DiffusionOPSD shows the strongest normalized progress and reaches a median terminal position of 98%. The two joint-reward runs are shown separately.

**Figure 8.** Held-out quality and compute. DiffusionOPSD reaches the highest held-out score on every Z-Image-Turbo frontier and on nine of ten SD3.5-M frontiers. The exception is SD3.5-M Aesthetic, where ReFL leads by 0.01. Markers show checkpoint means and bands show one standard error.

Table 1 lists absolute held‑out scores for all ten evaluators; DiffusionOPSD rows are bolded as the best in each block.

Table 2 breaks down training efficiency: DiffusionOPSD records the lowest step time (126.9 s) and highest images‑per‑second (9.08) on SD3.5‑M, and the smallest GPU‑hour budget (149.8) on Z‑Image‑Turbo.

Figures 6–7 illustrate joint three‑reward dynamics and native training rewards; in every panel DiffusionOPSD (dark blue) stays above DiffusionNFT (teal).

Figures 11–32 provide qualitative examples and robustness analyses, confirming that the gains persist across diverse prompts, noise levels, and hyper‑parameter settings.

DiffusionOPSD consistently outperforms baselines in training efficiency and final reward.

Mathematical Foundations

Mathematical derivations for the affine clean‑output map and endpoint‑weighted objectives.

This appendix supplies the formal derivations that underpin the affine clean‑output map, its low‑noise stability properties, and the endpoint‑weighted objectives used by DiffusionOPSD.

The model predicts a velocity $v_t$; solving the linear relation between $(z_t, v_t)$ and the latent clean output $y$ yields an explicit map $y_t(z_t, v_t)$ that can be evaluated at any diffusion time.

For the rectified‑flow schedule we set the coordinate $t=\sigma$, choose $\alpha(t)=1-t$ and $\sigma(t)=t$, yielding $\dot{\alpha}(t)=-1$ and $\dot{\sigma}(t)=1$, so $\Delta_\sigma=1$.

The map is a bijection between velocity and clean‑output coordinates but not an isometry; a perturbation $\delta v$ induces a scaled perturbation $\delta y = -\sigma\,\delta v$, and the squared norm scales as $\|\delta v\|_2^2 = \|\delta y\|_2^2 / \sigma$.

When the diffusion noise $\sigma$ is small, errors in the velocity prediction translate into proportionally smaller errors in the recovered clean output.

For rectified flow $\Delta_\sigma=1$, so the bound simplifies to $\|y_v-y_{v^{*}}\|_2^2 \le \sigma\,\|v-v^{*}\|_2^2$, confirming that low but non‑zero noise yields a stable clean‑output reconstruction.

When $\Delta_t$ approaches zero the bound blows up, so a schedule must keep $\Delta_t$ bounded away from zero in addition to using low noise.

DiffusionNFT weights samples by a scalar derived from the endpoint reward; DiffusionOPSD replaces that scalar weighting with reward‑gradient‑derived target maps while keeping the same weighting scheme.

To isolate the contribution of query‑state distribution versus target construction, the paper defines a factorial family $J_{a,b}$ with $a,b\in\{0,1\}$, where $a$ selects the source law ($P$ vs. $Q$) and $b$ selects the target map (endpoint vs. reward‑gradient).

The attribution formulas

Rewards are turned into a binary optimality variable $o\in\{0,1\}$; the distribution over endpoints is decomposed into a positive branch $\pi^{+}$ and a negative branch $\pi^{-}$, each re‑weighted by the conditional probability of $o$.

Within each prompt group the rewards are centered, then a global standard deviation rescales them; the result is clipped to $[0,1]$ to produce a stable per‑sample weight.

If all rewards in a batch are transformed by a positive affine map $r\mapsto c\,r+d$, the normalized advantages and clipped weights remain unchanged (provided $\epsilon_{Z}=0$).

The affine clean‑output map provides a closed‑form bridge from velocity predictions to clean outputs, and its low‑noise stability guarantees that small velocity errors produce proportionally small output errors.

Target Construction Details

This appendix derives the bounded target construction and its trust‑region guarantees.

We now formalize how the intermediate denoising targets are constructed in a way that respects a trust‑region around the current clean‑output prediction.

Instead of moving arbitrarily far toward higher reward, the method takes a small, norm‑bounded step in the direction of the reward gradient, then repeats this safely many times.

Lemma B.4 quantifies the reward change achieved by a single normalized step.

Proposition B.5 extends the single‑step bound to $M_{\text{tgt}}$ steps, showing a cumulative reward gain that scales with the sum of normalized gradient magnitudes.

Section B.4 connects the local target moves to the final endpoint reward by introducing a fixed‑suffix reward.

Finally, Section B.5 analyzes the geometry of the two‑branch fitting objective, showing how the adaptive normalizers shape the ideal output improvement.

Branch Residual Loss

Derives the branch‑residual loss optimum and its implications for reward‑gradient‑guided denoising.

The branch residual loss measures how far a clean‑output displacement $\delta$ deviates from two target residuals $d^{+}$ (positive) and $d^{-}$ (negative), weighting each with coefficients $a$ and $b$ and scaling by $\beta$.

If the two targets are exact opposites along the reward‑gradient direction—$d^{+}= \bar{h}\,\hat{u}_{\text{grad}}$ and $d^{-}= -\bar{h}\,\hat{u}_{\text{grad}}$—the optimal displacement simplifies dramatically.

At the anchor ($\delta=0$) the gradient points opposite to the combined force $F$, so a gradient‑descent step moves the model output toward $a d^{+}-b d^{-}$.

When the targets are not perfectly symmetric we write them as an ideal component plus an error:

Using the smoothness of the reward function $R$, the improvement from applying $\delta^{*}$ is lower‑bounded by a first‑order term minus a curvature penalty.

The branch optimum improves reward whenever its alignment with the gradient dominates the curvature term:

Switching to the velocity‑space formulation, the clean‑output gradient propagates through the affine map $y_{t}= (\dot\sigma_{t}z_{t}-\sigma_{t}v_{t})/\Delta_{t}$.

A gradient step in velocity space, $\Delta v_{t}= -\eta_{v}\nabla_{v_{t}}L$, translates to a clean‑output step $\Delta y_{t}= -\eta_{v}\,\sigma_{t}/\Delta_{t}\,\nabla_{y_{t}}L$, showing that the effective learning rate varies with the schedule noise.

Derivation of Target Dynamics

Appendix B Part 4 details adaptive scaling, finite‑fitting analysis, EMA updates, and computational cost of DiffusionOPSD.

Adaptive normalizers $\gamma$⁺ and $\gamma$⁻ appear in the branch loss to make the gradients insensitive to the absolute scale of the residuals.

The stabilizer $\epsilon_{\gamma}$ caps the denominator for very small $m(r)$, which in turn bounds the gradient norm.

Because the two branches (positive and negative) have independent normalizers, each can retain influence even when the other’s residual is much smaller.

At the symmetric anchor where $r^{+}= -\bar{h}_{\text{grad}}$ and $r^{-}= \bar{h}_{\text{grad}}$, the expression collapses to a pair of opposite vectors scaled by $\gamma^{\pm}$.

Proposition B.8 formalises when the ordering of candidate displacements is preserved after a single optimizer step.

The frozen behavior policy is refreshed by an exponential moving average after each outer iteration.

The computational cost of DiffusionOPSD can be expressed in closed form, separating rollout, target‑construction, and policy‑update components.

Conclusion

Appendix A records authors, affiliations, and supplemental details.

We introduced DiffusionOPSD, an on‑policy self‑distillation method that converts image‑level reward gradients into intermediate denoising targets, achieving the best held‑out scores in 19 of 20 reward‑matched settings and reducing GPU‑hours by up to 63 % versus DiffusionNFT.

The appendix lists all contributors—Wei Zhou, Xiongwei Zhu, Lingdong Kong, Bo Chen, Lei Zhang, Yongyuan Liang, Xiaoxia Hou, Ye Tian, Xian Sun, Yingshuo Wang, Linfeng Li, Shengqiong Wu, Leigang Qu, Feng Li, Wei Liu (corresponding), Julian McAuley, and Tat‑Seng Chua—and their affiliations across ByteDance Seed, several universities, and research institutes.

An additional table (GradRandomNo‑opResidual …) reports ablation sensitivities such as CLIPScore changes, CFG dependence, and human‑preference percentages, illustrating the robustness of DiffusionOPSD under varied hyper‑parameters.

Section B, “Mathematical Foundations”, is reserved for formal derivations that underpin the target‑construction loss; the details are omitted here but referenced elsewhere in the manuscript.

Experiment Details

Appendix C records the experimental protocols, hyper‑parameters, and extra analyses that underpin the main results.

Table 1 reports model‑based text‑to‑image evaluator scores for every method; rows marked ✓ use a reward‑specific checkpoint, while rows marked × evaluate a single checkpoint across all columns. All locally trained SD3.5‑M rows fine‑tune LoRA adapters (rank 32, $\alpha$ 64) on 512 × 512 images drawn from Pick‑a‑Pic, using deterministic 10‑step DPM‑Solver++ 2M rollouts with guidance scale 1.0.

**Figure 9. Ablation dynamics.** Reward-gradient training improves throughout the run, while the rollout-residual target collapses and the forward-noised control remains close to the canonical query state.

**Figure 10. Target construction and finite fitting.** The reward-gradient target gains 0.03511, while the DiffusionNFT endpoint loses 0.03551. After one fitting update, HPSv2.1 target ordering reverses on 62.3% of prompts, showing that construction gain does not determine one-update realization.

C.2 describes a two‑stage OPD pipeline: three reward‑specific DiffusionOPSD specialists are trained for 100 updates, frozen as teachers, and then distilled into a shared student for 300 updates with equal teacher weights.

C.3 measures end‑to‑end training cost on eight‑GPU nodes; DiffusionOPSD runs 0.59× the step time of ReFL on SD3.5‑M (28.2 GPU‑hours vs 47.7) and 0.68× on Z‑Image‑Turbo (102.1 GPU‑hours vs 149.8), while preserving higher held‑out quality.

C.4 reports three component‑ablation variants: rollout query state, forward‑noised control, and a large branch‑coefficient; the latter harms prompt semantics under finite fitting, whereas the first two stay close to the original behavior.

C.5 presents normalized‑gain curves (Figures 20 & 21) that aggregate training‑reward trajectories across 73 runs; DiffusionOPSD achieves the strongest average gain and stable final rewards for both SD3.5‑M and Z‑Image‑Turbo backbones.

**Figure.** Iridescent origami crane on a dark surface, with crisp paper folds highlighted by soft macro lighting.

**Figure.** Comparison of image generation results across five methods: Base, ReFL, DiffusionNFT, FlowGRPO, and DiffusionOPSD, based on the prompt: "1950s American diner at night with chrome, neon, a roller-skating waitress, red Cadillac, warm Kodachrome nostalgia, and the restaurant name 'RRACE' visible."

C.6 details three audit protocols: endpoint attribution (fixed‑suffix reward differences), construction gain (cosine alignment of candidate displacement with reward gradient), and realized‑gain (single‑step AdamW update on a positive target).

C.7 runs a few‑step stress test on Z‑Image‑Turbo (9‑step FlowMatchEuler, guidance 0.0) for ten rewards, training each checkpoint for 100 updates and a joint policy for 300 updates, to verify multi‑reward retention under limited compute.

C.8 enumerates component‑ablation settings (e.g., disabling endpoint checking, varying target radius $\rho$, branch coefficient $\beta$) and four target variants (reward‑gradient, random‑direction, no‑op, rollout‑residual) evaluated on fixed‑query protocols.

C.9 derives Proposition C.1, showing that a mismatch between training guidance scale $\zeta_{\text{tr}}$ and evaluation scale $\zeta_{\text{ev}}$ incurs a quadratic error term proportional to the hidden branch difference $\|\Delta v_{\text{gap}}\|_2^2$.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers