Reproduce “Aligning Few-Step Diffusion Models with Dense Reward Difference Learning” on OpenResearch

About this paper

Ziyi Zhang, Li Shen, Sen Zhang, Deheng Ye, Yong Luo, Miaojing Shi, Dongjing Shan, Bo Du and 1 more

Researchers from Wuhan University and collaborators developed Stepwise Diffusion Policy Optimization (SDPO), a reinforcement learning framework to robustly align few-step diffusion models with reward objectives. SDPO introduces a dual-state trajectory sampling mechanism and latent similarity-based dense reward prediction, outperforming existing methods in sample efficiency and generalization across various few-step text-to-image and multiview generation tasks.

Key insights

• Leveraging predicted clean states (`_x_t_0`) as reliable intermediate representations allows for dense, low-variance reward feedback without full image decoding, crucial for few-step regimes.
• Interpolating rewards based on latent space similarity, justified by a smoothness assumption, provides an efficient and implicitly regularized approach to dense reward estimation.
• Combining stepwise advantage difference learning, temporal importance weighting, and step-shuffled gradient updates significantly stabilizes and improves policy optimization for few-step diffusion models by addressing variance and temporal dependencies.

Summary by alphaXiv.