Reward Fine-Tuning Two-Step Diffusion Models via Learning Differentiable Latent-Space Surrogate Reward

TL;DR

This paper introduces LaSRO, a novel method for fine-tuning diffusion models using learned differentiable surrogate rewards in latent space, enabling ultra-fast image generation with improved efficiency and stability.

Contribution

LaSRO is the first approach to learn surrogate rewards in latent space for fine-tuning diffusion models, outperforming existing RL methods in ultra-fast image generation.

Findings

LaSRO outperforms PPO, DPO, DDPO, and Diffusion-DPO in ultra-fast image generation tasks.

LaSRO effectively converts arbitrary rewards into differentiable forms for gradient-based optimization.

Theoretical analysis links LaSRO to value-based reinforcement learning.

Abstract

Recent research has shown that fine-tuning diffusion models (DMs) with arbitrary rewards, including non-differentiable ones, is feasible with reinforcement learning (RL) techniques, enabling flexible model alignment. However, applying existing RL methods to step-distilled DMs is challenging for ultra-fast ($\le2$-step) image generation. Our analysis suggests several limitations of policy-based RL methods such as PPO or DPO toward this goal. Based on the insights, we propose fine-tuning DMs with learned differentiable surrogate rewards. Our method, named LaSRO, learns surrogate reward models in the latent space of SDXL to convert arbitrary rewards into differentiable ones for effective reward gradient guidance. LaSRO leverages pre-trained latent DMs for reward modeling and tailors reward optimization for $\le2$-step image generation with efficient off-policy exploration. LaSRO is effective and stable for improving ultra-fast image generation with different reward objectives, outperforming popular RL methods including DDPO and Diffusion-DPO. We further show LaSRO's connection to value-based RL, providing theoretical insights. See our webpage \href{https://sites.google.com/view/lasro}{here}.