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

TL;DR

<3-5 sentence high-level summary> The paper tackles the challenge of aligning diffusion models to human preferences at step level using pixel-space reward models, which struggle with noisy timesteps and transformation overhead. It proposes a latent-space reward model (LRM) that leverages pre-trained diffusion models to predict step-level preferences directly from noisy latent images, along with Multi-Preference Consistent Filtering (MPCF) and Latent Preference Optimization (LPO) to operate entirely in latent space. Empirical results show LPO yields superior general, aesthetic, and text-image alignment with substantial training speedups (2.5-28x) compared to existing methods. The work demonstrates a practical, efficient path for preference optimization in diffusion-based image generation by exploiting latent-space capabilities of diffusion models themselves.

Abstract

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. Our code and models are available at https://github.com/Kwai-Kolors/LPO.

Figures (15)

• Figure 1: Qualitative comparison among different preference optimization methods based on SDXL sdxl. LPO excels in both aesthetics and text-image alignment, resulting in improved overall quality. Larger versions with specific prompts are provided in Fig. \ref{['fig:vis_xl_1']} and Fig. \ref{['fig:vis_xl_2']}.
• Figure 2: The comparison between the pixel-level reward model (a) and the latent reward model (b). DMO denotes the diffusion model to be optimized.
• Figure 3: The architecture of LRM. The VAE encoder is only used during training.
• Figure 4: The training pipelines of PRM (a) and LRM (b), and sampling pipelines of SPO (c) and LPO (d). Compared to PRM, LRM can directly process latent images $x_t$ at various noise levels without requiring $\hat{x}_{0,t}$ prediction and VAE decoding. Therefore, LPO can perform sampling entirely within the noisy latent space throughout the full denoising process, covering $t\in[0,950]$.
• Figure 5: The distribution of preference consistency between human labels and reward scores in Pick-a-Pic.