💻 computer science

On-Policy Self-Distillation in Diffusion Models

The paper introduces DiffusionOPSD, an on-policy self-distillation framework that converts image-level rewards into explicit intermediate supervision targets for diffusion models, achieving superior alignment performance and significant training efficiency gains compared to existing methods.

Original authors: 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-S

Published 2026-08-26

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: 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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of artificial intelligence, there is a growing class of systems capable of painting pictures from thin air, turning a simple sentence like "a cat wearing a hat" into a vivid image. These systems, known as diffusion models, work by starting with a chaotic cloud of static noise and gradually cleaning it up, step by step, until a clear picture emerges. For years, researchers have focused on making these images look sharper or more realistic. More recently, the goal has shifted to making them align with human preferences, ensuring the AI understands not just what an object looks like, but what humans find beautiful or useful. To teach a model these preferences, scientists often use a method called reinforcement learning. In this setup, the AI generates an image, a computer program scores it based on how well it matches a desired outcome, and the AI tries to improve its next attempt based on that score. However, a fundamental problem exists in this process: the score is only given at the very end, after the picture is finished. The AI is left guessing how to adjust its intermediate steps—the messy, blurry stages of creation—to reach that final good score. It is like trying to learn a complex skill by only being told if you won or lost at the end of the game, without any feedback on your specific moves during the match.

A team of researchers has introduced a new method called DiffusionOPSD to solve this specific problem. Instead of waiting until the image is complete to offer feedback, their approach breaks the process down into manageable moments. They treat the AI's generation process as a series of snapshots. At a specific, early stage of creating an image, they pause and ask the computer program to evaluate what the image would look like if it stopped right there. Using this evaluation, they calculate a precise direction for improvement, creating a clear target for the AI to aim for. This target is not just a vague suggestion; it is a concrete, bounded instruction that tells the AI exactly how to shift its current prediction to make the final result better. The researchers then train the AI to move toward this target. Crucially, they do this in a loop: the AI makes a move, the system calculates a new target based on the AI's own current behavior, and the AI learns from that specific instruction before the system updates its own understanding for the next round. This cycle allows the AI to learn from immediate, actionable feedback rather than distant, final outcomes.

The researchers tested this method on two different powerful image-generation systems. In almost every scenario they tried, involving ten different ways to judge image quality, their new method produced better results than the strongest existing techniques. In nineteen out of twenty comparisons, the AI trained with this new approach created images that were rated higher by both automated scoring systems and human evaluators. The improvement was significant; in some cases, the quality of the images increased by nearly forty-four percent compared to the previous best method. Beyond just making better pictures, the new approach is also much more efficient. It required forty percent less computing time on one system and sixty-three percent less on the other to achieve these results. This efficiency matters because training these models usually demands massive amounts of energy and expensive hardware. By reducing the time needed, the method makes it more feasible to refine these tools for specific tasks without an overwhelming cost.

A key discovery in this work was that simply having a better idea of what to improve does not guarantee the AI will actually get better immediately. The researchers found that they could measure two distinct things: the quality of the instruction given to the AI, and how well the AI actually followed that instruction in a single training step. Sometimes, a very strong instruction led to only a small improvement, while a weaker instruction led to a larger one. This separation allowed them to see that the success of their method came from both creating clear instructions and ensuring the AI could effectively learn from them in a single go. They also tested whether the AI needed to see the exact same messy image it had generated before, or if it could learn from a slightly different version of that image. They found that the source of the image mattered very little; what truly drove the improvement was the direction of the feedback itself. This suggests that the method is robust and relies on the quality of the guidance rather than the specific details of the image being analyzed.

When the researchers looked at the actual images produced, the differences were clear. The AI trained with this new method was better at following complex instructions, such as including specific text, maintaining the identity of objects, or capturing motion and lighting effects. In head-to-head comparisons, human judges preferred the images from this new method over those from the previous best techniques in a majority of cases. The images were more consistent with the prompts, preserving details that other methods often lost or distorted. The researchers also demonstrated that this method could handle multiple goals at once. They trained a single AI to satisfy three different criteria simultaneously, and it succeeded in improving all three without sacrificing performance in any of them. This shows that the approach is flexible enough to handle complex, real-world requirements where an image needs to be good in several ways at the same time.

The work represents a shift in how we teach these powerful systems. By converting a final score into a clear, intermediate instruction, the researchers have provided a way to guide the AI's thinking process step by step. This approach avoids the guesswork of traditional methods and allows for a more direct and efficient path to better results. The findings suggest that the future of training these models may lie in breaking down complex goals into small, explicit steps that the AI can understand and act upon immediately. The method does not rely on magic or complex tricks; it relies on a clear, repeatable process of setting a target, measuring the distance to it, and taking a step. The result is a system that learns faster, uses less energy, and creates images that are more faithful to human intent.

Technical Summary: On-Policy Self-Distillation in Diffusion Models (DiffusionOPSD)

1. Problem Statement

Reinforcement learning (RL) is increasingly used to align diffusion and flow models with human preferences and task-specific objectives. However, a fundamental structural mismatch exists between endpoint rewards and intermediate denoising predictions. In standard diffusion generation, a sample is produced through a sequence of interdependent denoising steps, yet the reward signal is only observed after the final rollout endpoint is decoded.

Existing methods attempt to bridge this gap but face specific limitations:

Policy-gradient methods (e.g., FlowGRPO): Rely on sampled advantage estimates and trajectory credit assignment, which depend heavily on sample budgets, likelihood estimation, and discretization choices. They do not construct explicit local improvement targets.
Differentiable-reward methods (e.g., ReFL): Backpropagate rewards through a single late-state prediction. While they couple reward evaluation with optimization, they often execute the prefix without gradients and do not decode the full rollout endpoint, limiting the scope of the update.
Endpoint-supervised methods (e.g., DiffusionNFT): Reweight rollout endpoints but fail to specify how an intermediate prediction should change to improve the final outcome.

The core challenge is translating outcome-level reward into actionable, explicit supervision for intermediate denoising queries, while distinguishing between the quality of the constructed target and the model's ability to realize that target through finite updates.

2. Methodology: DiffusionOPSD

The authors introduce DiffusionOPSD, an on-policy self-distillation framework that converts image-level reward gradients into explicit targets for clean-output predictions. The method operates in an iterative loop consisting of three main stages:

A. Query Collection and Anchor Generation

At each outer iteration $i$ , a frozen behavior policy ( $v_{\theta_{old}}$ ) generates trajectories for a set of prompts.

Anchors: For a sampled low-noise query state $s = (c, z_q, \sigma_q)$ , the behavior policy produces a "clean-output prediction" (anchor) $y_0 = z_q - \sigma_q v_{\theta_{old}}(s)$ . This $y_0$ serves as the reference point for target construction.
Rewards: The policy also generates rollout endpoints to compute group-normalized fitting weights ( $\omega$ ) based on endpoint rewards, which control the relative strength of positive vs. negative fitting.

B. Target Construction

Instead of using the endpoint directly, the method constructs bounded positive and negative targets around the anchor $y_0$ using reward gradients:

Positive Target ( $\bar{y}_+$ ): Constructed by taking stabilized, normalized steps in the direction of the reward gradient ascent ( $\nabla_y \tilde{R}$ ).
Negative Target ( $\bar{y}_-$ ): Constructed by taking steps in the opposite direction (reward descent), acting as a repulsive reference.
Boundedness: These targets are projected onto a trust-region ball of radius $\rho \|y_0\|$ to ensure stability.
Detachment: The constructed targets are detached (stop-gradient) from the reward computation graph. This separates the construction of the target from the fitting process.

C. Finite Fitting and Policy Update

A trainable policy ( $v_\theta$ ) is updated to fit these detached targets:

Fitting Branches: The model output is split into two branches: an attractive branch moving toward $\bar{y}_+$ and a repulsive branch moving away from $\bar{y}_-$ .
Loss Function: The objective minimizes the distance between the model's predictions and the detached targets, weighted by the group-normalized reward weights.
Behavior Policy Refresh: After a finite number of fitting updates (e.g., $M_{fit}=1$ ), the behavior policy is updated via an Exponential Moving Average (EMA) of the trainable policy. This ensures the next iteration's queries and anchors are generated by a policy that has incorporated the previous learning, maintaining the "on-policy" nature of the distillation.

3. Key Contributions

Methodological Innovation: The introduction of an on-policy self-distillation loop that constructs explicit, bounded positive and negative targets from image-level reward gradients. Unlike prior work, it does not rely on implicit advantage weights or endpoint-only supervision.
Decoupled Analysis: The framework explicitly separates target construction (the potential gain from the reward gradient direction) from finite realization (the actual gain achieved after a model update). This allows for diagnosing whether performance bottlenecks stem from poor target construction or the model's inability to fit the target.
Efficiency and Performance: The method achieves state-of-the-art results with significantly reduced computational costs compared to existing baselines.

4. Experimental Results

The authors evaluated DiffusionOPSD on two backbones: SD3.5-M and the step-distilled Z-Image-Turbo, across ten evaluators (including Pick-a-Pic, CLIPScore, HPSv2.1, and internal preference models).

Held-Out Performance: DiffusionOPSD achieved the best final held-out scores in 19 out of 20 reward-matched settings.
- On SD3.5-M, it outperformed the strongest competing method by up to 44.0% (on VLM-Pairwise) and 43.0% (on HPSv3).
- On Z-Image-Turbo, it led all ten reward-matched comparisons, with gains ranging from 3.9% to 30.7% over the best baselines.
Training Efficiency: DiffusionOPSD reduced training GPU-hours relative to DiffusionNFT by 40% on SD3.5-M and 63% on Z-Image-Turbo.
Multi-Reward Training: A single policy jointly trained on PickScore, CLIPScore, and HPSv2.1 outperformed DiffusionNFT on all three objectives simultaneously, retaining nearly all the gains of separately trained specialists.
Qualitative Improvements: Human evaluations showed a preference for DiffusionOPSD over the base model, FlowGRPO, DiffusionNFT, and ReFL on 64% to 90% of prompts, with significant improvements in prompt fidelity, text rendering, and fine structural details.

5. Significance and Claims

The paper claims that DiffusionOPSD establishes on-policy self-distillation as an efficient and analyzable approach to diffusion post-training. By converting image-level reward guidance into explicit intermediate supervision that is continually refreshed as the policy evolves, the method:

Resolves the Structural Mismatch: It provides a direct mechanism to translate endpoint rewards into actionable intermediate targets without relying on trajectory-level credit assignment or endpoint-only supervision.
Enables Diagnosis: The separation of target construction and finite fitting allows researchers to measure "construction gain" vs. "realized gain," revealing that larger target-level gains do not always translate to larger realized gains after a single update.
Offers Efficiency: It achieves superior alignment with human preferences while significantly reducing the computational cost compared to existing state-of-the-art methods like DiffusionNFT.

The authors conclude that this approach opens a path toward more efficient and diagnosable alignment for diffusion models, moving beyond the limitations of endpoint-reward-only supervision.