Enhancing Diffusion Models with Text-Encoder Reinforcement Learning

TL;DR

This work tackles the mismatch between diffusion-model training objectives and downstream task requirements by finetuning the text encoder with reinforcement learning. TexForce uses low-rank adaptation (LoRA) and a PPO-based DDPO framework to optimize task-specific rewards, improving text-image alignment while preserving semantic quality. Importantly, it can be seamlessly combined with existing finetuned U-Net models without extra training, yielding substantial gains across backbones and prompts, including challenging face and hand generation scenarios. The approach is validated through diverse prompts, cross-backbone experiments, GPT-4V evaluations, and ablations, highlighting its practical impact for controllable, high-quality diffusion synthesis.

Abstract

Text-to-image diffusion models are typically trained to optimize the log-likelihood objective, which presents challenges in meeting specific requirements for downstream tasks, such as image aesthetics and image-text alignment. Recent research addresses this issue by refining the diffusion U-Net using human rewards through reinforcement learning or direct backpropagation. However, many of them overlook the importance of the text encoder, which is typically pretrained and fixed during training. In this paper, we demonstrate that by finetuning the text encoder through reinforcement learning, we can enhance the text-image alignment of the results, thereby improving the visual quality. Our primary motivation comes from the observation that the current text encoder is suboptimal, often requiring careful prompt adjustment. While fine-tuning the U-Net can partially improve performance, it remains suffering from the suboptimal text encoder. Therefore, we propose to use reinforcement learning with low-rank adaptation to finetune the text encoder based on task-specific rewards, referred as \textbf{TexForce}. We first show that finetuning the text encoder can improve the performance of diffusion models. Then, we illustrate that TexForce can be simply combined with existing U-Net finetuned models to get much better results without additional training. Finally, we showcase the adaptability of our method in diverse applications, including the generation of high-quality face and hand images.

Figures (12)

• Figure 1: By refining the text encoder through reinforcement learning, the proposed TexForce with Stable Diffusion v1.4 can generate images that align better with human quality preference. The compared images are generated with the same seed and prompts. (a)(b): "Impressionist painting of a cat, high quality"; (c)(d): "A photo of a hand" & "A complete face of a man".
• Figure 2: The qualities of outputs from pretrained diffusion models vary a lot with different prompts. Through the reinforcement learning, we can finetune text encoder to better align with images.
• Figure 3: Illustration of text encoder finetune with PPO algorithm.
• Figure 4: Comparison of training progress between finetuning text encoder and U-Net with LoRA. The image size after JPEG compression is marked on the top-left corner as "kb".
• Figure 5: Qualitative and quantitative comparisons with SDv1.4 and DPOK on individual scenarios. Images for comparison are generated with the same random seed. The results show that TexForce can generate more consistent images with better quality than SDv1.4 and DPOK, and simple combination of DPOK and TexForce gives even better performance without any additional training.