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
cs.CV
2026-08-25
DiffusionOPSD builds bounded positive and negative clean-output targets from reward gradients on on-policy queries, then fits them as detached distillation. It wins 19 of 20 held-out settings, with up to 44.0% over the next-best method and 40–63% fewer GPU-hours than DiffusionNFT.
Diffusion and flow models usually see a reward only after the full image is decoded. The policy acts at every intermediate denoising step. Trajectory RL scores the whole path from the endpoint. ReFL backpropagates through a clean-output prediction. DiffusionNFT regresses toward a reward-selected endpoint. None of these states how the current intermediate prediction should move.
Building a better local target and actually fitting that target under a small update budget are different jobs. One network is queried at many noise levels, so a fit at one query shifts others. If those stages share one loss, failures are hard to locate.
DiffusionOPSD casts post-training as on-policy self-distillation. In each outer loop a frozen behavior policy rolls out trajectories and supplies query states plus clean-output anchors. Reward gradients around each anchor build a bounded positive target (reward ascent) and a negative target (repulsion). The trainable policy fits these detached targets for a finite number of steps. An EMA then refreshes the behavior policy, and the next round rebuilds queries and targets.
The clean-output map comes from a rectified-flow velocity field: at a noisy latent and noise level, velocity converts into a decoder-ready clean image, so an image-level reward can sit on an intermediate prediction. A group-normalized endpoint reward weights the two branches. Stop-gradient between construction and fitting lets the same query and the same suffix report construction gain and realized gain separately.
On SD3.5-M and step-distilled Z-Image-Turbo, across ten evaluators, DiffusionOPSD has the best held-out score in 19 of 20 reward-matched settings. The largest relative gain over the next-best method is 44.0%, on SD3.5-M Pair after reward-specific training: 0.465 versus DiffusionNFT at 0.323.
With a shared checkpoint and 300 updates on SD3.5-M, PickScore is 25.51 against DiffusionNFT at 23.62. After 100 reward-specific updates on Z-Image-Turbo, PickScore moves from 22.86 to 25.15 and HPSv2.1 from 0.296 to 0.390. One policy trained jointly on PickScore, CLIPScore, and HPSv2.1 beats the matching DiffusionNFT policy on all three.
| Backbone | Method | GPU-h per 100 updates | vs NFT |
| SD3.5-M | DiffusionNFT | 47.2 | 1.00× |
| SD3.5-M | DiffusionOPSD | 28.2 | 0.60× |
| Z-Image-Turbo | DiffusionNFT | 405.8 | 1.00× |
| Z-Image-Turbo | DiffusionOPSD | 149.8 | 0.37× |
The saving is backward count: about 16 diffusion backwards per update on SD3.5-M versus 144 for NFT. Same-query tests show a larger construction gain can yield a smaller gain after one fitting step. In a CLIPScore screen, replacing the reward-gradient direction with random, no-op, or rollout-residual targets drops 0.3117 to 0.2311, 0.2280, and 0.1456. Swapping only the query source to a forward-noised state barely moves (0.3103). Dropping the negative branch does not hurt (0.3137). Query noise 0.90 and target radius 0.02 are the clear failure settings.
Alignment for diffusion still has to turn a late reward into a local training signal. This method makes that signal explicit, rebuilds it as the policy moves, and reports construction and fitting as separate numbers. If a stack already runs DiffusionNFT or ReFL, the GPU-hour gap is the practical reason to look.
It does not claim to replace all trajectory RL. It claims intermediate supervision can be written down and audited.
44.0% is a peak relative gain, not an average across evaluators. Many main-table scores are single runs. The fitting gap is demonstrated, not closed by an extra algorithm. Reward models and the decoder can be overfit. Peak VRAM on Z-Image-Turbo is 61.5 GB versus 49.9 GB for NFT, so the saving is time, not memory. The component screen uses 20 DrawBench prompts and 100 images.