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

TL;DR

This paper introduces Latent Reward Model and Latent Preference Optimization, leveraging diffusion models for step-level preference alignment in noisy latent space, resulting in faster training and improved preference alignment.

Contribution

It proposes a novel latent space reward model and preference optimization method that utilize diffusion models directly, enhancing efficiency and effectiveness over pixel-based approaches.

Findings

Significant improvement in preference alignment across multiple criteria.

Achieved 2.5-28x faster training speeds.

Effective in handling noisy images at various timesteps.

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.