Diffusionrvge speeds up training of image-based reinforcement learning models
DiffusionRVGE: Reinforcement Learning via Regression-Based Value Gradient Estimation
Artificial Intelligence
Summary
Improving image generation using feedback (or rewards) can be very slow because it involves repeating complex calculations many times. The authors introduce DiffusionRVGE, a method that learns better update directions more efficiently by linking rewards directly to the final images made. This method skips costly steps used before and adapts its approach depending on how smooth or jumpy the rewards are. It matches or beats existing methods on multiple tests while using less computing power, and it learns much faster when focusing on a single reward.
What this means in practice
- •For machine learning engineers: Train image generation models faster by efficiently estimating reward-guided updates without backpropagation through many sampling steps.
- •For computer vision teams: Improve image generation quality on several evaluation metrics using joint training with multiple rewards while reducing computational cost.
Authors
Jiayuan Wang, Tongyao Pang
Abstract
Reinforcement learning improves diffusion models using rewards on generated images, but repeated denoising and gradient estimation make training expensive. We introduce \method, an efficient method that constructs useful update directions by regressing centered rewards on centered final denoised images. It requires neither reward derivatives nor adjoint backpropagation through the sampling trajectory. A unified first-variation formulation identifies the value-gradient target with or without a trajectory-KL penalty. A population analysis under Gaussian noising explains the regression geometry, while a subspace analysis motivates clean-image coordinates. The implementation adapts rollout bandwidth and query weights to reward continuity: local rollouts suit smooth rewards, whereas broader groups provide informative contrasts for discrete rewards. Trained jointly on five rewards, one SD3.5-Medium model matches or exceeds DiffusionNFT's joint model on seven of eight metrics, reaching GenEval $0.94$, OCR $0.92$, and PickScore $23.82$ at about one fifth of its estimated training compute. In single-reward training, \method reaches matched reward thresholds $2.1$--$6.7\times$ faster than reproduced DiffusionNFT on identical hardware. Ablations support the clean-image estimator, direct direction transfer, and reward-dependent rollout choices. Code is available at https://github.com/jy-wang-tsinghua/DIFFUSIONRVGE.