PaperlayerView exact revision on arXiv

Source-linked AI summary

Enhancing Diffusion Models with Text-Encoder Reinforcement Learning

Chaofeng Chen, Annan Wang, Haoning Wu, Liang Liao, Wenxiu Sun, Qiong Yan, Weisi Lin

arXiv:2311.15657v2cs.CV

TL;DR

Diffusion models are commonly trained with a likelihood objective and often keep the pretrained text encoder fixed, while task-specific image rewards may be non-differentiable. The paper fine-tunes the text encoder with reinforcement learning and LoRA, reporting improved alignment and visual quality, including when combined with fine-tuned U-Nets without extra training.

  • Problem

    Diffusion models usually keep the pretrained text encoder fixed, and with limited fine-tuning text data its embedding distribution may be suboptimal, increasing the KL term.

    Source: §2.3 · p. 6

  • Method

    The method uses PPO-based reinforcement learning with LoRA to fine-tune the text encoder using task-specific rewards.

    Source: §2.2 · p. 5

  • Results

    Fine-tuning the text encoder improves diffusion-model performance, and combining it with fine-tuned U-Nets yields significant performance improvements without extra training.

    Source: §4 · pp. 14–15

  • Takeaways & Limitations

    The approach is shown across face and hand image generation, and task-specific LoRA weights can be combined to enhance image-generation quality.

    Source: §4 · p. 15

    Like other reinforcement-learning methods, the approach faces challenges in sample efficiency and reward-function engineering.

    Source: F Limitations · p. 26

Abstract

from arXiv · show

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.

1 Related Works

The related work surveys ways to improve text-to-image generation through model controls, feedback-based optimization, and image-quality evaluation. It notes that pretrained models and prompt engineering can still struggle with complex demands and text-prompt alignment.

1.1 Text-to-Image Diffusion Models

  • Prompt engineering and auxiliary image encoders offer control, but can be constrained by prompt expressiveness, paired-image requirements, and complex task demands.The cited examples include aesthetic quality and object composition.Source: §1.1 · p. 3

1.2 Learning from Feedback in Diffusion Models

  • Reward-weighted regression, reinforcement learning, and direct backpropagation optimize diffusion models from feedback, but U-Net-only tuning can limit prompt alignment.TextCraftor also fine-tunes the text encoder, but uses direct backpropagation and does not support non-differentiable rewards.Source: §1.2 · p. 4

1.3 Quality Metrics for Generative Models

  • ImageReward, PickScore, HPS, and the LAION Aesthetics Predictor are cited as commonly used benchmarks or image-aesthetic metrics.The passage also notes the use of aesthetic metrics in data filtering and results assessment.Source: §1.3 · p. 4

2 Method

The method treats text embeddings as the diffusion model’s policy and fine-tunes the text encoder with reward-based reinforcement learning and LoRA. It motivates this choice by the fixed pretrained encoder and the mismatch between likelihood training and task-specific objectives.

2.1 Preliminaries on Diffusion Models

  • Text-to-image diffusion conditions the denoising U-Net on text-encoder representations, while its usual training objective optimizes a log-likelihood bound and keeps the pretrained encoder fixed.The denoising network estimates cleaner samples from noisy ones during reverse diffusion.Source: §2.1 · p. 4

2.2 Reinforcement Learning with LoRA

  • The reinforcement-learning policy is the text encoder, which maps descriptions to embeddings that influence image generation and is optimized for expected reward.The reward can reflect image-text alignment, image quality, or desired attributes.Source: §2.2 · p. 5
  • DDPO applies PPO with clipped probability ratios to update the text-encoder policy, while LoRA adds trainable low-rank matrices to pretrained weights.The advantage is normalized reward over a buffer of generated samples; LoRA’s low-rank matrices are initialized to zero.Source: §2.2 · p. 5

2.3 Discussion of Finetuning for Diffusion Model

  • With limited fine-tuning data, the paper argues that a fixed text encoder may poorly estimate the embedding prior and increase the KL term, motivating encoder fine-tuning.The stated rationale concerns the fine-tuning stage, where the available text data are small relative to pretraining.Source: §2.3 · p. 6
  • The paper says direct reward backpropagation can overfit the reward and cause mode collapse, whereas reinforcement learning does not require differentiable rewards.DRaFT is given as an example of collapse to a single image for high aesthetic rewards.Source: §2.3 · p. 6

3 Experiments

Experiments compare text-encoder and U-Net fine-tuning, evaluate TexForce against prior methods and across backbones, and test joint fine-tuning and task-specific applications. Across these settings, the reported results favor combining TexForce with existing U-Net fine-tunes, while text-encoder fine-tuning improves prompt alignment.

3.1 Implementation Details

  • Experiments use animal prompts, DPOK phrases, long prompts from ImageReward and HPSv2, and task-specific face and hand prompts.Rewards include ImageReward and HPSv2 as well as specific task rewards, including face and hand rewards.Source: §3.1 · p. 6

3.2 Finetuning Text Encoder v.s. U-Net

  • At matched incompression scores, text-encoder tuning produces more coherent novel concepts than U-Net tuning, and fusing their LoRAs achieves the highest score without extra training.Text-encoder optimization is reported as more challenging and time-consuming.Source: §3.2 · pp. 7–8

3.3 Comparison with Existing Works

  • TexForce combined with DPOK delivers the best ImageReward scores and visual quality on seen and unseen prompts, outperforming SDv1.4 and DPOK alone.The comparisons cover color consistency, object composition, count, and location; models use the same random seeds.Source: Figure 5 · p. 8, §3.3 · p. 9
  • On ImageReward prompts, ReFL mainly improves appearance while TexForce aligns images better with text; on HPSv2, TexForce better preserves prompt semantics than AlignProb.Combining TexForce with either method is reported to improve quantitative results and visual appearance or preserve both HPSv2-preferred styles and prompt meaning.Source: §3.3 · pp. 10–11

3.4 Experiments with Different Backbones

  • Across SDv1.5 and SDv2.1, TexForce improves text-image alignment, and combining it with ReFL yields the best performance for both backbones.ReFL scores higher on ImageReward but is described as less effective at aligning images with prompts.Source: §3.4 · p. 11

3.5 GPT-4V Evaluation

  • GPT-4V evaluations find TexForce improves text-image alignment, ReFL improves image appearance, and their combination yields the best results.The reported averages come from three evaluation rounds on the ImageReward test set.Source: §3.5 · p. 11

3.6 Ablation Study about Joint Finetune

  • Joint fine-tuning surpasses TexForce and ReFL individually but falls short of their simple combination; the authors hypothesize a fixed U-Net provides a pixel-generation prior.They suggest joint updates complicate text-encoder optimization.Source: §3.6 · pp. 12–13, Table 3 · p. 13

3.7 Applications

  • TexForce improves face and hand image quality using a face-quality metric and hand-detection confidence, and fusing ImageReward and face-quality LoRAs produces high-quality faces.The paper presents progressive training examples for both face and hand generation.Source: §3.7 · p. 13

4 Conclusion

TexForce uses reinforcement learning to refine diffusion-model text encoders, improving text-image alignment and visual quality. It also integrates with fine-tuned U-Nets without extra training and extends to face and hand generation.

4 Conclusion

  • TexForce refines diffusion-model text encoders with reinforcement learning to improve text-image alignment and visual quality.The conclusion also reports improved overall model performance.Source: §4 · p. 14
  • TexForce integrates with fine-tuned U-Net models without extra training and yields significant performance improvements.LoRA weights from different tasks can also be combined to enhance specific image-generation qualities.Source: §4 · p. 15
  • The method is shown across applications including high-quality face and hand image generation.The conclusion also describes the approach as versatile across applications.Source: §4 · p. 15

A Method Details

The appendix gives PPO-based diffusion formulations and describes task-specific face and hand reward construction.

A.1 Reinforcement Learning Formulations

  • The denoising step is conditioned on text embeddings, while the U-Net parameter is held fixed in the following formulations.The denoising transition is given in the displayed formulation.Source: §A.1 · p. 19
  • The diffusion formulation uses a clipped PPO objective based on the probability ratio between current and old policies.The reward is defined from the generated image and prompt as A = R(x0, s).Source: §A.1 · p. 19

A.2 Face and Hand Rewards

  • Face rewards use a TOPIQ model fine-tuned on GFIQA-20k, with generated faces aligned before scoring.The figure caption identifies the face-quality and hand-detection reward calculations.Source: §A.2 · p. 19, Figure 13 · p. 20
  • Hand rewards use detection confidence directly because the authors report no existing hand-quality model.The authors state that this confidence score can provide reliable generation-quality rewards.Source: §A.2 · p. 19

B Experiment Settings

The experiments share hyperparameters except for batch size and samples per epoch, which differ between single-prompt and multi-prompt datasets.

B.1 Training Details

  • All experiments use the same hyperparameters except for batch size and samples per epoch.The passage identifies those two settings as the only differences.Source: §B.1 · p. 21
  • Single-prompt experiments use batch size 8 and 256 samples per epoch.Source: §B.1 · p. 21
  • Multi-prompt experiments use batch size 64 and 2048 samples per epoch.Source: §B.1 · p. 21

B.2 Evaluation Details

The evaluation details specify a common random-seed setup and different numbers of generated examples for single- and multi-prompt datasets.

B.2 Evaluation Details

  • PyTorch Lightning’s seed_everything() sets the starting random seed to 234 for all prompts.Source: §B.2 · p. 21
  • The single-prompt dataset generates 50 examples for each prompt.Source: §B.2 · p. 21
  • The multi-prompt ImageReward and HPSv2 datasets generate one example per prompt to save time.Source: §B.2 · p. 21

C GPT4V Evaluation

This section describes a GPT-4V evaluation of generated images for aesthetic quality and prompt coherence. It explains the ranking and scoring procedure and identifies the evaluation results reported for SDv1.4 on ImageReward.

C GPT4V Evaluation

  • GPT-4V ranks image names for a prompt by aesthetic quality and coherence with that prompt.The example asks it to rank four images from good to bad.Source: C GPT4V Evaluation · pp. 21–22
  • Figure 14 reports the complete GPT-4V evaluation results for SDv1.4 on the ImageReward dataset.The passage specifies the evaluation scope but not the results of individual comparisons.Source: Figure 14 · p. 22, C GPT4V Evaluation · p. 23
  • The evaluation assigns scores from 3 for the best image to 0 for the worst, then normalizes the final score to [0, 1].Each evaluation is run three times and the average is reported.Source: C GPT4V Evaluation · p. 23

D.1 More Results of Incompression Task

These supplementary experiments compare text-encoder and U-Net tuning on incompression and unseen-animal prompts, test aesthetic-reward tuning, and evaluate ReFL-LoRA with TexForce. They report stronger unseen-animal robustness for text-encoder tuning, lifelike outputs under aesthetic rewards, and better ReFL-LoRA performance when combined with TexForce.

D.1 More Results of Incompression Task

  • Text-encoder tuning is more robust on unseen animals and scores higher than U-Net tuning despite similar training rewards; combining them improves reward scores.The unseen-animal results are reported in Table 5.Source: §D.1 · p. 23, Table 5 · p. 23
  • On unseen-animal prompts, text-encoder tuning adds reasonable visual concepts, while U-Net tuning mainly changes appearance and can disrupt structure.Figure 15 shows examples at different checkpoints, including a disrupted hippo head.Source: §D.1 · p. 23, Figure 15 · p. 24

D.2 Results of Aesthetic Reward with Animal Prompts

  • TexForce produces more lifelike, prompt-aligned animal images than DDPO and AlignProb, which are described as over-optimizing aesthetic rewards; that reward does not measure prompt coherence and can be hacked.DDPO and TexForce use 10K samples, while AlignProb uses early stopping to avoid model collapse; Figure 16 compares jellyfish, penguin, and swan prompts.Source: §D.2 · p. 24, Figure 16 · p. 25

D.3 Training ReFL with LoRA

  • ReFL-LoRA performs similarly to ReFL, but performs much better when combined with TexForce.The quantitative results are presented in Table 6.Source: §D.3 · p. 25, Table 6 · p. 25

E More Qualitative Results

This section adds qualitative examples from different backbones, including Stable Diffusion v1.4 results on the ImageReward test dataset.

E More Qualitative Results

F Limitations

The paper identifies sample efficiency and reward-function engineering as challenges for its RL method. It also notes potential societal concerns related to misinformation, intellectual property rights, and illegal use.

F Limitations

  • Like other RL-based methods, TexForce faces challenges in sample efficiency and reward-function engineering complexity.Source: F Limitations · p. 26

G Broader Impacts

  • TexForce's ability to tune text-to-image models for specific rewards raises potential concerns about misinformation, intellectual property rights, and illegal use.Source: G Broader Impacts · p. 27

Related papers

Loading 2311.15657v2…