Value gradient lookahead improves reward guided diffusion models performance
VaLoP: Value Gradient Lookahead and Pullback for Reward-Guided Diffusion
Artificial Intelligence
Summary
Some AI models create images or actions step-by-step by removing noise, guided by how much reward they might get. The usual ways to guide these models make bigger mistakes when there's a lot of noise, which is when the main shape is decided. The authors found these errors come from uncertainty about the final clean image. Their new method, called VaLoP, predicts the reward gradient one step ahead where the noise is less and adjusts it back to the current step, reducing errors. This leads to better results in generating images and controlling robots, all without needing to retrain the whole model.
What this means in practice
- •For machine learning engineers: Improve image generation quality and control by integrating VaLoP guidance in pretrained diffusion models without retraining for each reward.
- •For robotics developers: Enhance robot policy success rates by applying VaLoP-Fast for critic-guided motion planning and decision making.
Authors
Kwanyoung Kim
Abstract
Reward guidance steers a pretrained diffusion or flow model toward images or actions that a reward model prefers, at sampling time and without retraining the model for each reward. It follows the gradient of a value function, the expected reward given the noisy state. The two common estimates of this gradient, gradients through a one-step denoised estimate and learned value models, are least accurate at high noise, where the coarse structure of the sample is decided. Interestingly, we find that both errors are posterior covariances, which grow with the uncertainty about the clean sample, and hence at higher noise. To reduce them, we evaluate the value at a less noisy state, one sampler step ahead, where this uncertainty is smaller, and pull its gradient back to the current state. This lookahead places both estimators at the two ends of one family and reveals that an intermediate lookahead requires one additional term, a correction through the generator. We therefore propose Value Gradient Lookahead and Pullback (VaLoP), which computes this correction with one vector-Jacobian product, and VaLoP-Fast, which amortizes it with a small head that, for text-to-image models, is trained with a single reward and reused for the others. Across various models, rewards, and benchmark datasets, both robustly outperform existing value-gradient guidance on the target reward, VaLoP-Fast at about the cost of learned-value guidance. Both extend to critic-guided robot policies, where VaLoP-Fast attains the highest success rates on LIBERO and RoboCasa.