Value-as-Return: A Two-Stage Framework to Align on the Optimal Score Function
Shikun Sun ⋅ Shuo Huang ⋅ Yiding Chen ⋅ Wen Sun ⋅ Jia Jia
Abstract
Reinforcement learning with diffusion models has shown strong potential, but existing approaches such as variants of Direct Preference Optimization (DPO) often rely on an inaccurate simplification: they equate trajectory likelihoods with final-state probabilities. This mismatch leads to suboptimal alignment. We address this limitation with a principled framework that leverages the optimal value function as the return for short trajectory segments. Our approach follows a two-stage procedure: (i) learning a value-distribution function to estimate segment-level returns, and (ii) applying our VRPO to refine the score function. We prove that, under sufficient model capacity, the resulting model is equivalent to training a diffusion process on the tilted distribution proportional to $p(x)\exp(\eta r(x))$. Experiments on large-scale diffusion models validate our analysis and show stable and consistent improvements over prior methods.
Lay Summary
Teaching AI Image Generators to Better Understand Human Preferences Artificial Intelligence (AI) has become incredibly good at generating stunning images from simple text descriptions. However, making these models truly understand and follow human preferences—like visual quality, realistic compositions, or precise details—remains a major challenge. Currently, popular methods try to improve these AI models by looking only at the final image and making broad guesses about the entire step-by-step creation process. This is like grading a student's entire semester based solely on their final exam score, without looking at how they solved problems along the way. Because diffusion models generate images gradually over multiple steps, this "all-or-nothing" approach often confuses the AI, leading to awkward image layouts, missing details, or training instability. To fix this, we introduce Value-as-Return, a new two-stage framework that acts like a supportive tutor guiding the AI at every single step of image generation. * Stage 1 (The Foresight Stage): We first train a lightweight evaluator model to look at an incomplete, noisy image at any point in the generation process and predict how good the final picture will be. * Stage 2 (The Step-by-Step Training Stage): Using this evaluator, we break down the long image-creation process into very short segments. The AI creates alternative options for these brief segments, and our evaluator instantly tells it which path leads to a better final image. The AI then updates itself based on these quick, precise feedback loops. By utilizing rewards at a micro-level rather than guessing from the final image alone, our method completely bypasses the mathematical contradictions of prior techniques. In our extensive tests, images generated using our approach achieved significantly higher human and automated preference scores. Best of all, because our framework is highly focused and efficient, it achieves state-of-the-art image quality while consuming the least training time among all competing baselines.
Successful Page Load