Diffusion Model as a Noise-Aware Latent Reward Model for Step-Level Preference Optimization

Yan Li (Mohamed bin Zayed University of Artificial Intelligence) · Di ZHANG (Kuaishou Technology) · Kun Ding (Institute of Automation, Chinese Academy of Sciences) · SHIMING XIANG (Chinese Academy of Sciences, China) · Tingting Gao (Kuaishou- 快手科技) · Tao Zhang (Institute of Automation, Chinese Academy of Sciences) · Cheng Da (Kuaishou) · Huan Yang (Kuaishou- 快手科技) · kun jin (Tsinghua University, Tsinghua University) · Chunhong Pan (Institute of Automation, Chinese Academy of Sciences)
aesthetic preferencesdiffusion modelslatent imageslatent preference optimizationlatent reward modellatent spacemodel alignmentnoise levelsnoisy imagespixel-level reward modelspreference optimizationstep-level preference optimizationtext-image alignmenttraining speedupvision-language models

Preference optimization for diffusion models aims to align them with human preferences for images. Previous methods typically use Vision-Language Models (VLMs) as pixel-level reward models to approximate human preferences. However, when used for step-level preference optimization, these models face challenges in handling noisy images of different timesteps and require complex transformations into pixel space. In this work, we show that pre-trained diffusion models are naturally suited for step-level reward modeling in the noisy latent space, as they are explicitly designed to process latent images at various noise levels. Accordingly, we propose the **Latent Reward Model (LRM)**, which repurposes components of the diffusion model to predict preferences of latent images at arbitrary timesteps. Building on LRM, we introduce **Latent Preference Optimization (LPO)**, a step-level preference optimization method conducted directly in the noisy latent space. Experimental results indicate that LPO significantly improves the model's alignment with general, aesthetic, and text-image alignment preferences, while achieving a 2.5-28x training speedup over existing preference optimization methods.