Aligning Few-Step Diffusion Models With Dense Reward Difference Learning

Area of Science:

• Artificial Intelligence
• Computer Vision
• Machine Learning

Background:

• Few-step diffusion models offer efficient high-resolution image synthesis.
• Existing reinforcement learning (RL) methods struggle with alignment in low-step diffusion models due to limited states and sample quality.

Purpose of the Study:

• Introduce Stepwise Diffusion Policy Optimization (SDPO), a novel RL framework for few-step diffusion models.
• Address limitations in aligning diffusion models with downstream objectives in low-step regimes.

Main Methods:

• SDPO employs a dual-state trajectory sampling mechanism (noisy and clean states) for dense reward feedback.
• A latent similarity-based dense reward prediction strategy minimizes costly reward queries.
• Utilizes dense reward difference learning, stepwise advantage estimates, temporal importance weighting, and step-shuffled gradient updates.

Main Results:

• SDPO enables low-variance, mixed-step optimization with more frequent and granular policy updates.
• Experimental results show consistent superior reward-aligned outcomes across various few-step tasks.
• Demonstrates enhanced long-term dependency, low-step priority, and gradient stability.

Conclusions:

• SDPO effectively overcomes the limitations of traditional RL methods in few-step diffusion model optimization.
• The proposed framework significantly improves the alignment of synthesized images with specific downstream objectives.
• SDPO represents a substantial advancement in efficient and effective high-resolution image synthesis.