PaperlayerView exact revision on arXiv

Source-linked AI summary

DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models

Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, Kimin Lee

arXiv:2305.16381v3cs.LGcs.CV

TL;DR

Text-to-image models remain difficult to fine-tune directly with human-feedback rewards, especially without sacrificing image quality. DPOK formulates diffusion fine-tuning as online RL with policy gradients and KL regularization, and experiments report stronger image-text alignment and image quality than supervised fine-tuning. The paper also analyzes KL regularization for both approaches while noting practical limitations for multi-prompt training.

  • Problem

    Existing text-to-image models have weaknesses in compositional generation and specified colors or counts, while reward-weighted supervised fine-tuning can improve alignment but degrade image quality.

  • Method

    DPOK formulates diffusion-model fine-tuning as online reinforcement learning using policy gradients and KL regularization relative to the pre-trained model.

  • Results

    DPOK is generally superior to supervised fine-tuning in image-text alignment and image quality, with online RL maintaining high image fidelity while optimizing reward.

  • Takeaways & Limitations

    Online RL fine-tuning can improve text-to-image alignment while preserving image fidelity, and KL regularization provides a shared basis for comparing RL and supervised fine-tuning.

  • Takeaways & Limitations

    Fine-tuning on multiple prompts requires longer training time, hyperparameter tuning, and engineering effort.

Abstract

from arXiv · show

Learning from human feedback has been shown to improve text-to-image models. These techniques first learn a reward function that captures what humans care about in the task and then improve the models based on the learned reward function. Even though relatively simple approaches (e.g., rejection sampling based on reward scores) have been investigated, fine-tuning text-to-image models with the reward function remains challenging. In this work, we propose using online reinforcement learning (RL) to fine-tune text-to-image models. We focus on diffusion models, defining the fine-tuning task as an RL problem, and updating the pre-trained text-to-image diffusion models using policy gradient to maximize the feedback-trained reward. Our approach, coined DPOK, integrates policy optimization with KL regularization. We conduct an analysis of KL regularization for both RL fine-tuning and supervised fine-tuning. In our experiments, we show that DPOK is generally superior to supervised fine-tuning with respect to both image-text alignment and image quality. Our code is available at https://github.com/google-research/google-research/tree/master/dpok.

1 Introduction

Existing text-to-image models struggle with compositional and attribute-specific generation, while reward-weighted supervised fine-tuning can improve alignment but degrade image quality. DPOK instead uses online RL with KL regularization to fine-tune diffusion models, and experiments report stronger alignment while maintaining image fidelity.

  • Current text-to-image models have systematic weaknesses in composing multiple objects and generating specified colors and counts.
  • Reward-weighted supervised fine-tuning improves reward scores and image-text alignment but can produce over-saturated or non-photorealistic images.The degradation is attributed to training on a fixed dataset generated by the pre-trained model.
  • DPOK frames diffusion-model fine-tuning as online reinforcement learning and uses policy optimization with KL regularization relative to the pre-trained model.The expected reward is connected to policy gradient under regularity assumptions, while KL is treated as an implicit reward.
  • Experiments report that online RL achieves strong text-image alignment while maintaining high image fidelity, outperforming supervised fine-tuning in empirical comparisons.The experiments fine-tune Stable Diffusion with ImageReward and also add KL regularization to supervised fine-tuning for fair comparison.
  • The paper studies KL regularization in supervised fine-tuning and analyzes differences between supervised and online fine-tuning.

2 Related Work

The related work situates DPOK within diffusion-based generation, learning from human feedback, and earlier RL approaches for diffusion models. Prior studies establish reward modeling and RL fine-tuning as relevant routes for improving text-to-image systems.

  • Diffusion models use iterative denoising and, with large-scale language encoders, have achieved strong performance in text-to-image generation.
  • Learning from human feedback commonly trains reward models from annotator comparisons to produce scalar rewards aligned with human preferences.
  • Earlier RL work combined policy gradients with GAN training to improve pre-trained diffusion models on relatively simple image domains.
  • Concurrent work also framed text-to-image fine-tuning as multi-step decision-making and found RL could outperform reward-weighted supervised fine-tuning in reward optimization.

3 Problem Setting

The paper formulates conditional diffusion generation through forward noising and reverse denoising processes. Given a text prompt, the trained model produces conditional samples by iteratively denoising latent representations.

  • A DDPM approximates a data distribution with a parameterized reverse Markov chain for generation.
  • The forward diffusion process is a Markov chain that adds Gaussian noise according to a variance schedule.
  • Training optimizes a variational bound on the negative log-likelihood of the data.
  • The terminal noisy state is approximately Gaussian when the cumulative signal coefficient is near zero.
  • The model uses latent diffusion, so the intermediate x_t variables are latent representations rather than image-space variables.
  • At test time, classifier-free guidance combines conditional and unconditional denoising predictions to generate p_θ(x_0|z) for a text prompt.

4 Fine-tuning of Diffusion Models

The paper models diffusion denoising as an MDP and optimizes the resulting policy with reward and KL regularization. It develops DPOK for online RL and KL-regularized supervised alternatives, arguing that online training better optimizes the combined objective.

  • RL Fine-tuning with KL Regularization: The terminal reward evaluates the generated image, so optimizing the MDP policy is equivalent to fine-tuning the underlying diffusion model.
  • RL Fine-tuning with KL Regularization: Diffusion denoising is formulated as a finite-horizon MDP whose policy transitions between successive denoising states.
  • RL Fine-tuning with KL Regularization: Policy-gradient estimates from trajectories sampled by the current model update the diffusion policy online, making REINFORCE suitable for large-scale models.Direct differentiation through multi-step trajectories is described as memory-inefficient and potentially numerically unstable.
  • RL Fine-tuning with KL Regularization: DPOK combines the reward with an upper-bounded KL penalty relative to the pre-trained model to limit excessive deviation during fine-tuning.The output-image KL is intractable, so the method uses an upper bound based on divergences across diffusion steps.
  • RL Fine-tuning with KL Regularization: The implemented DPOK gradient omits one exact-gradient term for efficiency, with importance sampling and clipped gradients available for reusing historical trajectories.
  • Supervised Learning with KL Regularization: KL-D modifies reward-weighted supervised learning by shifting rewards with γ, smoothing sample weights toward a uniform distribution.
  • Supervised Learning with KL Regularization: KL-O adds a penalty on the L2 distance between current and pre-trained denoising directions, regularizing the supervised model's output toward the pre-trained model.
  • Comparison of Online and Supervised Fine-tuning: Online fine-tuning is expected to optimize the combined human and implicit KL reward more effectively than supervised fine-tuning.

5 Experimental Evaluation

Experiments compare Stable Diffusion, supervised fine-tuning, and online RL fine-tuning across alignment, image quality, KL regularization, bias reduction, and multi-prompt training. Online RL generally improves alignment while preserving image quality, with KL regularization helping stabilize training.

  • Comparison of Supervised and RL Fine-tuning: Online RL and SFT improve ImageReward on prompts testing color, composition, counting, and location.The evaluation uses four prompts and 20K images for each fine-tuning approach.
  • Comparison of Supervised and RL Fine-tuning: RL models achieve higher ImageReward than SFT models while maintaining more natural image quality.SFT can produce over-saturated images, whereas RL generates images at least as well-aligned with prompts.
  • Comparison of Supervised and RL Fine-tuning: Human raters consistently prefer RL over SFT for both image-text alignment and image quality.The comparison uses images from the same random seeds and eight independent raters.
  • The Effect of KL Regularization: KL regularization gives online RL both high reward and aesthetic scores, while mitigating some SFT failure modes but generally lowering SFT ImageReward.Without KL, RL can produce over-saturated colors and unnatural shapes; the differing effects are attributed to online versus fixed-sample training.
  • Reducing Bias in the Pre-trained Model: RL fine-tuning changes the “Four roses” interpretation from whiskey-related imagery toward flower roses.The original model’s association reflects a whiskey brand present in web-scale training data.
  • Fine-tuning on Multiple Prompts: With 104 MS-CoCo or 183 Drawbench prompts, RL significantly improves ImageReward while maintaining a high aesthetic score.The evaluation averages scores across 3120 MS-CoCo images and 5490 Drawbench images.

6 Discussions

The discussion presents DPOK as policy-gradient fine-tuning with KL regularization and reports online RL as stronger than simple supervised fine-tuning. It also identifies training complexity and reward-model risks as important boundaries.

  • Contributions: DPOK fine-tunes text-to-image diffusion models with policy gradient and KL regularization.The KL term is defined relative to the pre-trained model and acts as an implicit reward.
  • Conclusions: Online RL outperforms simple supervised fine-tuning in improving text-to-image model performance.The paper analyzes KL regularization for both methods and highlights RL from human feedback as promising for diffusion models.
  • Limitations and Future Directions: Multiple-prompt fine-tuning requires longer training, hyperparameter tuning, and engineering effort.The paper suggests more efficient training over broader and more diverse prompts as future work.
  • Broader Impacts: Arbitrary reward functions can support beneficial or harmful fine-tuning, making documentation of reward-model biases and failure modes critical.Potential misuse includes generating misinformation or deep fakes.

A Derivations

The derivation relates diffusion-model reward optimization to KL-regularized objectives, using bounds and assumptions about the diffusion transition distributions.

  • The derivation invokes continuous assumptions for the diffusion-model and reference-distribution transition terms.
  • Regularizing the sum of future conditional KL-divergences is less efficient than regularizing only the first term, which the authors find works well empirically.
  • The derivation uses the ELBO to obtain an upper-bound relation for the KL-regularized diffusion objective.The first inequality is explicitly attributed to the ELBO.
  • The diffusion model uses a discrete decoder for pθ(x0|z, x1) and a fixed Gaussian for p(xT|z), neither of which is trainable.

A.5 Value function learning

Value-function learning reduces policy-gradient variance, while importance sampling and ratio clipping support reuse and stabilization during online fine-tuning.

  • Subtracting a learned value function minimizes gradient-estimation variance and is expected to improve policy-gradient training.
  • ImageReward increased from 0.86 to 1.51 for “A dog on the moon” when variance reduction was added.The aesthetic score also slightly increased from 5.57 to 5.60.
  • Similar improvements from value learning were observed in multi-prompt training.The corresponding learning curves are reported for the Drawbench prompt set.
  • Importance sampling reuses old trajectory samples, while clipped gradients constrain the updated policy relative to the old policy.The clipping parameter is denoted by ϵ.

B Experimental Details

The experiments specify separate online-RL and supervised-training configurations, including matched sample counts and reward filtering for supervised fine-tuning.

  • Online RL training: Online RL uses α = 10, β = 0.01, learning rate 10^-5, sampling batch size m = 10, and five gradient steps per sampling step.Gradient norms are clipped below 0.1, with importance sampling for off-policy samples and batch size n = 32 per gradient step.
  • Online RL training: The online-RL learning curves report that adding value learning could achieve higher reward using less time.
  • Supervised training: Supervised training uses γ = 2.0, learning rate 2 × 10^-5, batch size n = 128, M = 20000, and 8K gradient steps.
  • Supervised training: Reward filtering sets r + γ to zero when it is negative because supervised fine-tuning could otherwise fail during training.This filtering is required by Lemma 4.3’s non-negativity condition.

C Investigation on ImageReward

The study evaluates ImageReward against human labeler assessments across four text-prompt categories using CLIP, BLIP, and ImageReward scores.

  • ImageReward quality is evaluated by testing whether its scores agree with human preferences over generated images.
  • The evaluation covers color, count, location, and composition prompt categories generated by combining category phrases with various objects.
  • Table 2 reports the accuracy of CLIP score, BLIP score, and ImageReward in predicting human labeler assessments.

D A Comprehensive Discussion of Related Work

The paper distinguishes DPOK from concurrent online-RL work by adding theoretical policy-gradient analysis and a systematic KL-regularization study. It targets supervised fine-tuning failure modes while seeking high rewards and preserved image quality.

  • Unlike concurrent work, this paper provides theoretical conditions under which policy-gradient optimization of diffusion-model rewards is equivalent to the intended objective.
  • DPOK extends online RL fine-tuning with KL regularization relative to the pre-trained model, treating KL as an implicit reward.
  • The method is motivated by supervised fine-tuning failure cases including over-saturated and non-photorealistic images.
  • The paper studies KL regularization in both supervised and online fine-tuning, with theoretical justifications for the comparison.
  • The proposed online KL regularization is reported to achieve high rewards while maintaining image quality without over-optimization.

E.1 KL-D vs KL-O: Ablation Study on Supervised KL Regularization

The ablation compares KL-D and KL-O for supervised fine-tuning and finds a tradeoff between alignment and visual quality. In contrast, RL with KL regularization retains both properties more effectively, including on complex prompts.

  • KL-D vs KL-O: KL-D can slightly increase aesthetic scores but does not significantly improve visual quality, even with a large γ, and tends to lower ImageReward.
  • KL-D vs KL-O: KL-O addresses supervised fine-tuning failure modes more noticeably than KL-D at relatively large γ, but significantly reduces ImageReward.
  • KL-D vs KL-O: Supervised fine-tuning generally struggles to achieve high ImageReward and aesthetic scores simultaneously, unlike RL fine-tuning.
  • KL-D vs KL-O: Figure 9 compares an RL-with-KL sample against an SFT-without-KL sample and shows SFT samples across regularization choices and γ values.
  • Qualitative Comparisons: The appendix includes qualitative comparisons across original, SFT, RL, KL-regularized, and non-KL models for multiple prompts and random samples.
  • Long and Complex Prompts: The long-prompt example reports that online training produces a different painting style and more fine-grained details than supervised fine-tuning with KL-O and γ = 2.0.
Loading 2305.16381v3…