Deep Reward Supervisions for Tuning Text-to-Image Diffusion Models

TL;DR

This work tackles training text-to-image diffusion models under differentiable reward signals by introducing DRTune, a method that deep-supervises early denoising steps through a stop-gradient on the denoiser input and training a sparse, evenly spaced subset of steps. This design mitigates gradient explosion and reduces memory demands, enabling effective optimization of low-level rewards like symmetry. Across seven reward models, DRTune outperforms baselines and enables fine-tuning of SDXL 1.0 to produce FDXL 1.0, achieving image quality that is competitive with Midjourney in several categories. The paper also provides thorough ablations and discusses limitations such as reward bias and potential reward hacking, outlining a practical path toward reward-aligned diffusion models.

Abstract

Optimizing a text-to-image diffusion model with a given reward function is an important but underexplored research area. In this study, we propose Deep Reward Tuning (DRTune), an algorithm that directly supervises the final output image of a text-to-image diffusion model and back-propagates through the iterative sampling process to the input noise. We find that training earlier steps in the sampling process is crucial for low-level rewards, and deep supervision can be achieved efficiently and effectively by stopping the gradient of the denoising network input. DRTune is extensively evaluated on various reward models. It consistently outperforms other algorithms, particularly for low-level control signals, where all shallow supervision methods fail. Additionally, we fine-tune Stable Diffusion XL 1.0 (SDXL 1.0) model via DRTune to optimize Human Preference Score v2.1, resulting in the Favorable Diffusion XL 1.0 (FDXL 1.0) model. FDXL 1.0 significantly enhances image quality compared to SDXL 1.0 and reaches comparable quality compared with Midjourney v5.2.

Figures (11)

• Figure 1: Images generated by Favorable Diffusion XL 1.0 (FDXL 1.0). FDXL 1.0 is initialized from Stable Diffusion XL 1.0 (SDXL 1.0) and then trained to optimize Human Preference Score v2.1 via the proposed Deep Reward Tuning.
• Figure 2: Illustration of the iterative denoising sampling pipeline. When trained, gradients flow through the green branch and the red branch.
• Figure 3: Comparison between reward training algorithms.
• Figure 4: Comparison between images generated with DRTune or other baseline methods and symmetry supervision.
• Figure 5: Comparison between images generated with Stable Diffusion tuned by DRTune or other baseline methods. Supervised by PickScore.