 |
 |
 |
 |
Recently, Direct Preference Optimization (DPO) has extended its success from aligning large language models (LLMs) to aligning text-to-image diffusion models with human preferences. Unlike most existing DPO methods that assume all diffusion steps share a consistent preference order with the final generated images, we argue that this assumption neglects step-specific denoising performance and that preference labels should be tailored to each step's contribution.
To address this limitation, we propose Step-aware Preference Optimization (SPO), a novel post-training approach that independently evaluates and adjusts the denoising performance at each step, using a step-aware preference model and a step-wise resampler to ensure accurate step-aware supervision. Specifically, at each denoising step, we sample a pool of images, find a suitable win-lose pair, and, most importantly, randomly select a single image from the pool to initialize the next denoising step. This step-wise resampler process ensures the next win-lose image pair comes from the same image, making the win-lose comparison independent of the previous step. To assess the preferences at each step, we train a separate step-aware preference model that can be applied to both noisy and clean images.
Our experiments with Stable Diffusion v1.5 and SDXL demonstrate that SPO significantly outperforms the latest Diffusion-DPO in aligning generated images with complex, detailed prompts and enhancing aesthetics, while also achieving more than 20× times faster in training efficiency. Code and model: https://rockeycoss.github.io/spo.github.io/
## Model Description This model is fine-tuned from [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5). It has been trained on 4,000 prompts for 10 epochs. This checkpoint is a LoRA checkpoint. We also provide a LoRA checkpoint compatible with [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui), which can be accessed [here](https://civitai.com/models/526379/spo-sd-v1-54k-p10eplorawebui) If you want to access the merged checkpoint that combines the LoRA checkpoint with the base model [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5), please visit [SPO-SD-v1-5_4k-p_10ep](https://huggingface.co/SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep). ## A quick example ```python from diffusers import StableDiffusionPipeline import torch # load pipeline inference_dtype = torch.float16 pipe = StableDiffusionPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", torch_dtype=inference_dtype, ) pipe.load_lora_weights("SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep_LoRA") pipe.to('cuda') generator=torch.Generator(device='cuda').manual_seed(42) image = pipe( prompt='an image of a beautiful lake', generator=generator, guidance_scale=7.5, output_type='pil', ).images[0] image.save('lake.png') ``` ## Citation If you find our work or codebase useful, please consider giving us a star and citing our work. ``` @article{liang2024step, title={Step-aware Preference Optimization: Aligning Preference with Denoising Performance at Each Step}, author={Liang, Zhanhao and Yuan, Yuhui and Gu, Shuyang and Chen, Bohan and Hang, Tiankai and Li, Ji and Zheng, Liang}, journal={arXiv preprint arXiv:2406.04314}, year={2024} } ```