On-Policy Self-Distillation in Diffusion Models

TL;DR

DiffusionOPSD is an on-policy self-distillation framework for RL-style alignment of diffusion models. Instead of trying to translate endpoint rewards into diffuse per-step gradients directly, it converts image-level reward signal into explicit targets for the clean-output prediction at sampled query states: a frozen behavior policy generates trajectories, reward gradients construct bounded positive/negative anchors around each denoised prediction, and the trainable policy fits those targets under detached supervision before an EMA refreshes the behavior policy. The framing separates target construction from finite realization, which lets the paper measure them independently. On SD 3.5-M and the step-distilled Z-Image-Turbo it wins 19 of 20 reward-matched settings across two backbones and ten evaluators, beats the strongest competitor by up to 44%, and cuts GPU-hours vs DiffusionNFT by 40–63%.

Key claims

Endpoint rewards do not specify how intermediate denoising predictions should change; DiffusionOPSD closes this by turning reward guidance at the image level into explicit clean-output targets at sampled query states [Abstract].
The outer training loop is fully on-policy self-distillation: a frozen behavior policy supplies query states + anchors each iteration, and an EMA update refreshes the behavior policy after the trainable policy fits [Abstract].
Reward gradients construct bounded positive and negative targets around each anchor, giving the fitting stage a stable, discretized supervision signal rather than an unbounded gradient [Abstract].
The target-construction step and the finite-fitting step are cleanly separable in this framework, enabling controlled same-query experiments that measure them independently [Abstract].
Controlled same-query experiments show larger target-construction gains do not necessarily translate into larger realized gains after a single fitting update — the finite-realization gap is a distinct axis, not a proxy for target quality [Abstract].
On SD 3.5-M and the step-distilled Z-Image-Turbo, DiffusionOPSD achieves the best final held-out scores in 19 of 20 reward-matched settings across two backbones and ten evaluators [Abstract].
Outperforms the strongest competing method by up to 44.0% on the same reward-matched protocol [Abstract].
Reduces training GPU-hours relative to DiffusionNFT by 40% on SD 3.5-M and 63% on Z-Image-Turbo [Abstract].

Method

DiffusionOPSD structures diffusion RL as a two-stage outer loop rather than a single policy-gradient objective. At each outer iteration a frozen behavior policy generates rollouts and produces two things: (a) a set of sampled query states (intermediate noisy latents along the trajectory), and (b) an anchor at each query — the behavior policy’s own current clean-output prediction. Reward gradients (computed at the image level, via the same reward as any GRPO/DiffusionNFT baseline) are used to construct bounded positive and negative targets around each anchor, one on each side. The trainable policy then fits these targets as detached supervision — the targets are treated as fixed labels during the inner fitting phase, so the trainable policy’s clean-output prediction is pulled toward the positive target and away from the negative target at the sampled queries, but does not propagate gradients back into the anchor construction. After the inner fitting completes, an EMA update refreshes the behavior policy toward the trainable policy for the next outer iteration.

The decomposition is the paper’s central design: target construction (how reward gradients turn into positive/negative anchors) is one axis, finite realization (how well one fitting update actually moves the student toward the constructed targets) is another. Controlled same-query experiments (identical query states across variants) let the paper measure the two axes independently — and the finding is that they don’t move together, i.e. constructing better targets does not automatically buy larger realized gains after one fitting update.

Experiments cover two backbones — SD 3.5-M (a standard flow-matching image generator) and Z-Image-Turbo (a step-distilled few-step student) — evaluated across ten reward evaluators under reward-matched settings against DiffusionNFT and other RL post-training baselines.

Results

Headline: DiffusionOPSD achieves the best final held-out scores in 19 of 20 reward-matched settings (2 backbones × 10 evaluators), outperforming the strongest competing method by up to 44.0% [Abstract]. Compute-side: 40% GPU-hour reduction relative to DiffusionNFT on SD 3.5-M and 63% reduction on Z-Image-Turbo — the wins are larger on the already step-distilled backbone [Abstract]. Per-benchmark breakdowns, ablations isolating target-construction from finite-realization, and reward-evaluator identities were not surfaced in the abstract-level fetch at filing time.

Why it’s interesting

DiffusionOPSD sits at the intersection of two threads the wiki has been tracking hard. On the Reasoning RL side, it’s the diffusion-side counterpart to On-Policy Distillation: same recipe shape — sample on-policy from a behavior policy, convert reward signal into dense per-step supervision via a distillation loss, EMA-refresh — but with the “teacher distribution” replaced by reward-gradient-constructed anchors around the student’s own clean-output prediction. It’s the diffusion analog of self-distillation where the earlier checkpoint is the reward-perturbed self. It also gives an explicit answer to the “why on-policy distillation instead of on-policy RL” question in the diffusion setting: the endpoint reward doesn’t specify how intermediate predictions should change, and reward-anchored bounded targets are a way to make that specification.

On the Diffusion Distillation side, this is the reverse composition to TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward. TDM-R1 first distills a many-step teacher into a few-step student, then RL-post-trains the student against non-differentiable rewards. DiffusionOPSD instead uses the self-distillation objective as the RL signal itself, on top of both an un-distilled generator (SD 3.5-M) and an already step-distilled generator (Z-Image-Turbo). The finding that gains are larger (63% GPU-hour reduction vs 40%) on the already-distilled backbone suggests the two recipes compose additively rather than substituting — a natural next ablation is TDM-R1’s Surrogate Reward machinery dropped inside DiffusionOPSD’s outer loop.

Third, the target-construction vs finite-realization decomposition is a diagnostic other diffusion RL work doesn’t offer. GRPO / DiffusionNFT / Pref-GRPO / DGPO all measure the objective end-to-end; DiffusionOPSD’s same-query protocol is the first filed setup where the two stages can be independently probed. The empirical finding — better targets don’t automatically yield better realized gains — is a warning shot for any future work that tries to improve reward-guidance quality without also improving fitting efficiency.