Title: DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Diffusion models

RLHF
Abstract: 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.

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

Abstract

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 feedbacktrained 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/googleresearch/google-research/tree/master/dpok.

1 Introduction

Recent advances in diffusion models [10, 37, 38], together with pre-trained text encoders (e.g., CLIP [27], T5 [28]) have led to impressive results in text-to-image generation. Large-scale text-toimage models, such as Imagen [32], Dalle-2 [29], and Stable Diffusion [30], generate high-quality, creative images given novel text prompts. However, despite these advances, current models have systematic weaknesses. For example, current models have a limited ability to compose multiple objects [6, 7, 25]. They also frequently encounter difficulties when generating objects with specified colors and counts [12, 17].

Learning from human feedback (LHF) has proven to be an effective means to overcome these limitations [13, 17, 42, 43]. Lee et al. [17] demonstrate that certain properties, such as generating objects with specific colors, counts, and backgrounds, can be improved by learning a reward function from human feedback, followed by fine-tuning the text-to-image model using supervised learning. They show that simple supervised fine-tuning based on reward-weighted loss can improve the reward scores, leading to better image-text alignment. However, supervised fine-tuning often induces a deterioration in image quality (e.g., over-saturated or non-photorealistic images). This is likely due to the model being fine-tuned on a fixed dataset that is generated by a pre-trained model (Figure 1(a)).

In this work, we explore using online reinforcement learning (RL) for fine-tuning text-to-image diffusion models (Figure 1(b)). We show that optimizing the expected reward of a diffusion model's image output is equivalent to performing policy gradient on a multi-step diffusion model under certain

Reward function

Images from pre-trained model

Model update

(a) Supervised fine-tuning

Reward

function

Reward

function

Reward

function

Model

update

Model

update

(b) RL fine-tuning

D

Ppre

pθ

pe

pe

pe

pe

Ppre

regularity assumptions. We also incorporate Kullback-Leibler (KL) divergence with respect to the pre-trained model as regularization in an online manner, treating this as an implicit reward.

In our experiments, we fine-tune the Stable Diffusion model [30] using ImageReward [43], an opensource reward model trained on a large dataset comprised of human assessments of (text, image) pairs. We show that online RL fine-tuning achieves strong text-image alignment while maintaining high image fidelity by optimizing its objective in an online manner. Crucially, online training allows evaluation of the reward model and conditional KL divergence beyond the (supervised) training dataset . This offers distinct advantages over supervised fine-tuning, a point we demonstrate empirically. In our empirical comparisons, we also incorporate the KL regularizer in a supervised fine-tuning method for a fair comparison.

Our contributions are as follows:

We frame the optimization of the expected reward (w.r.t. an LHF-reward) of the images generated by a diffusion model given text prompts as an online RL problem. Moreover, we present DPOK: D iffusion P olicy O ptimization with K L regularization, which utilizes KL regularization w.r.t. the pre-trained text-to-image model as an implicit reward to stabilize RL fine-tuning.

We study incorporating KL regularization into supervised fine-tuning of diffusion models, which can mitigate some failure modes (e.g., generating over-saturated images) in [17]. This also allows a fairer comparison with our RL technique.

We discuss the key differences between supervised fine-tuning and online fine-tuning of text-toimage models (Section 4.3).

Empirically, we show that online fine-tuning is effective in optimizing rewards, which improves text-to-image alignment while maintaining high image fidelity.

2 Related Work

Text-to-image diffusion models. Diffusion models [10, 37, 39] are a class of generative models that use an iterative denoising process to transform Gaussian noise into samples that follow a learned data distribution. These models have proven to be highly effective in a range of domains, including image generation [4], audio generation [14], 3D synthesis [26], and robotics [3]. When combined with large-scale language encoders [27, 28], diffusion models have demonstrated impressive performance in text-to-image generation [29, 30, 32]. However, there are still many known weaknesses of existing text-to-image models, such as compositionality and attribute binding [6, 20] or text rendering [21].

Learning from human feedback. Human assessments of (or preferences over) learned model outcomes have been used to guide learning on a variety of tasks, ranging from learning behaviors [16] to language modeling [1, 23, 19, 41]. Recent work has also applied such methods to improve the alignment of text-to-image models. Human preferences are typically gathered at scale by asking annotators to compare generations, and a reward model is trained (e.g., by fine-tuning a visionlanguage model such as CLIP [27] or BLIP [18]) to produce scalar rewards well-aligned with the human feedback [43, 13]. The reward model is used to improve text-to-image model quality by fine-tuning a pre-trained generative model [17, 42]. Unlike prior approaches, which typically focus on reward-filtered or reward-weighted supervised learning, we develop an online fine-tuning framework with an RL-based objective.

RLfine-tuning of diffusion models. Fan & Lee [5] first introduced a method to improve pre-trained diffusion models by integrating policy gradient and GAN training [8]. They used policy gradient with reward signals from the discriminator to update the diffusion model and demonstrated that the fine-tuned model can generate realistic samples with few diffusion steps with DDPM sampling [10] on relatively simple domains (e.g., CIFAR [15] and CelebA [22]). In this work, we explore RL finetuning especially for large-scale text-to-image models using human rewards. We also consider several design choices like adding KL regularization as an implicit reward, and compare RL fine-tuning to supervised fine-tuning.

Concurrent and independent from our work, Black et al. [2] have also investigated RL fine-tuning to fine-tune text-to-image diffusion models. They similarly frame the fine-tuning problem as a multi-step decision-making problem, and demonstrate that RL fine-tuning can outperform supervised fine-tuning with reward-weighted loss [17] in optimizing the reward, which aligns with our own observations. Furthermore, our work analyzes KL regularization for both supervised fine-tuning and RL fine-tuning with theoretical justifications, and shows that adopting KL regularization is useful in addressing some failure modes (e.g., deterioration in image quality) of fine-tuned models. For a comprehensive discussion of prior work, we refer readers to Appendix D.

3 Problem Setting

In this section, we describe our basic problem setting for text-to-image generation of diffusion models.

Diffusion models. We consider the use of denoising diffusion probabilistic models (DDPMs) [10] for image generation and draw our notation and problem formulation from [10]. Let q 0 be the data distribution, i.e., x 0 ' q 0 p x 0 q , x 0 P R n . A DDPM approximates q 0 with a parameterized model of the form p θ p x 0 q ' ş p θ p x 0: T q dx 1: T , where p θ p x 0: T q ' p T p x T q ś T t ' 1 p θ p x t ´ 1 | x t q and the reverse process is a Markov chain with the following dynamics:

A unique characteristic of DDPMs is the exploitation of an approximate posterior q p x 1: T | x 0 q , known as the forward or diffusion process , which itself is a Markov chain that adds Gaussian noise to the data according to a variance schedule β 1 , . . . , β T :

Let α t ' 1 ´ β t , ¯ α t ' ś t s ' 1 α s , and ˜ β t ' 1 ´ ¯ α t ´ 1 1 ´ ¯ α t β t . Ho et al. [10] adopt the parameterization µ θ p x t , t q ' 1 ? α t ´ x t ´ β t ? 1 ´ ¯ α t ϵ θ p x t , t q ¯ .

Training a DDPM is performed by optimizing a variational bound on the negative log-likelihood E q r´ log p θ p x 0 qs , which is equivalent to optimizing:

Note that the variance sequence p β t q T t ' 1 P p 0 , 1 q T is chosen such that ¯ α T « 0 , and thus, q p x T | x 0 q « N p 0 , I q . The covariance matrix Σ t in (1) is often set to σ 2 t I , where σ 2 t is either β t or ˜ β t , which is not trainable. Unlike the original DDPM, we use a latent diffusion model [30], so x t 's are latent.

Text-to-image diffusion models. Diffusion models are especially well-suited to conditional data generation, as required by text-to-image models: one can plug in a classifier as guidance function [4], or can directly train the diffusion model's conditional distribution with classifier-free guidance [9].

Given text prompt z ' p p z q , let q p x 0 | z q be the data distribution conditioned on z . This induces a joint distribution p p x 0 , z q . During training, the same noising process q is used regardless of input z , and both the unconditional ϵ θ p x t , t q and conditional ϵ θ p x t , t, z q denoising models are learned. For data sampling, let ¯ ϵ θ ' wϵ θ p x t , t, z q ` p 1 ´ w q ϵ θ p x t , t q , where w ě 1 is the guidance scale. At test time, given a text prompt z , the model generates conditional data according to p θ p x 0 | z q .

4 Fine-tuning of Diffusion Models

In this section, we describe our approach for online RL fine-tuning of diffusion models. We first propose a Markov decision process (MDP) formulation for the denoising phase. We then use this MDP and present a policy gradient RL algorithm to update the original diffusion model. The RL algorithm optimizes an objective consisting of the reward and a KL term that ensures the updated model is not too far from the original one. We also present a modified supervised fine-tuning method with KL regularization and compares it with the RL approach.

4.1 RL Fine-tuning with KL Regularization

Let p θ p x 0: T | z q be a text-to-image diffusion model where z is some text prompt distributed according to p p z q , and r p x 0 , z q be a reward model (typically trained using human assessment of images).

MDPformulation: The denoising process of DDPMs can be modeled as a T -horizon MDP:

in which s t and a t are the state and action at time-step t , P 0 and P are the initial state distribution and the dynamics, R is the reward function, and π θ is the parameterized policy. As a result, optimizing policy π θ in (4) is equivalent to fine-tuning the underlying DDPM. 1 Finally, we denote by δ z the Dirac distribution at z .

It can be seen from the MDP formulation in (4) that the system starts by sampling its initial state s 0 from the Gaussian distribution N p 0 , I q , similar to the first state of the dinoising process x T . Given the MDP state s t , which corresponds to state x T ´ t of the denoising process, the policy takes the action at time-step t as the next denoising state, i.e., a t ' x T ´ t ´ 1 . As a result of this action, the system transitions deterministically to a state identified by the action (i.e., the next state of the denoising process). The reward is zero, except at the final step in which the quality of the image at the end of the denoising process is evaluated w.r.t. the prompt, i.e., r p x 0 , z q .

A common goal in re-training/fine-tuning the diffusion models is to maximize the expected reward of the generated images given the prompt distribution, i . e .,

The gradient of this objective function can be obtained as follows:

Lemma4.1 (A modification of Theorem 4.1 in [5]) . If p θ p x 0: T | z q r p x 0 , z q and ∇ θ p θ p x 0: T | z q r p x 0 , z q are continuous functions of θ , then we can write the gradient of the objective in (5) as

Proof. We present the proof in Appendix A.1.

Algorithm 1 DPOK: Diffusion policy optimization with KL regularization

Equation (6) is equivalent to the gradient used by the popular policy gradient algorithm, REINFORCE, to update a policy in the MDP (4). The gradient in (6) is estimated from trajectories p θ p x 0: T | z q generated by the current policy, and then used to update the policy p θ p x t ´ 1 | x t , z q in an online fashion.

Note that REINFORCE is not the only way to solve (5). Alternatively, one could compute the gradient through the trajectories to update the model; but the multi-step nature of diffusion models makes this approach memory inefficient and potentially prone to numerical instability. Consequently, scaling it to high-resolution images becomes challenging. For this reason, we adopt policy gradient to train large-scale diffusion models like Stable Diffusion [30].

Adding KL regularization. The risk of fine-tuning purely based on the reward model learned from human or AI feedback is that the model may overfit to the reward and discount the 'skill' of the initial diffusion model to a greater degree than warranted. To avoid this phenomenon, similar to [23, 41], we add the KL between the fine-tuned and pre-trained models as a regularizer to the objective function. Unlike the language models in which the KL regularizer is computed over the entire sequence/trajectory (of tokens), in text-to-image models, it makes sense to compute it only for the final image, i.e., KL ` p θ p x 0 | z q} p pre p x 0 | z q ˘ . Unfortunately, p θ p x 0 | z q is a marginal (see the integral in Section 3) and its closed-form is unknown. As a result, we propose to add an upper-bound of this KL-term to the objective function.

Lemma 4.2. Suppose p pre p x 0: T | z q and p θ p x 0: T | z q are Markov chains conditioned on the text prompt z that both start at x T ' N p 0 , I q . Then, we have

We report the proof of Lemma 4.2 in Appendix A.2. Intuitively, this lemma tells us that the divergence between the two distributions over the output image x 0 is upper-bounded by the sum of the divergences between the distributions over latent x t at each diffusion step.

Using the KL upper-bound in (7), we propose the following objective for regularized training:

ff

where α, β are the reward and KL weights, respectively. We use the following gradient to optimize the objective (8):

Note that (9) has one term missing from the exact gradient of (8) (see Appendix A.3). Removing this term is for efficient training. The pseudo-code of our algorithm, which we refer to as DPOK, is summarized in Algorithm 1. To reuse historical trajectories and be more sample efficient, we can also use importance sampling and clipped gradient, similar to [36]. We refer readers to Appendix A.6 for these details.

4.2 Supervised Learning with KL Regularization

We now introduce KL regularization into supervised fine-tuning (SFT), which allows for a more meaningful comparison with our KL-regularized RL algorithm (DPOK). We begin with a supervised

fine-tuning objective similar to that used in [17], i.e.,

To compare with RL fine-tuning, we augment the supervised objective with a similar KL regularization term. Under the supervised learning setting, we consider KL p p pre p x 0 | z q} p θ p x 0 | z qq , which is equivalent to minimizing E p pre p x 0 | z q r´ log p θ p x 0 | z qs , given any prompt z . Let q p x t q be the forward process used for training. Since the distribution of x 0 is generally not tractable, we use approximate upper-bounds in the lemma below (see derivation in Appendix A.4).

Lemma 4.3. Let γ be the regularization weight and ˜ µ t p x t , x 0 q : ' ? ¯ α t ´ 1 β t 1 ´ ¯ α t x 0 ` ? α t p 1 ´ ¯ α t ´ 1 q 1 ´ ¯ α t x t . Assume r p x 0 , z q ` γ ą 0 , for all x 0 and z . Then, we have

Moreover, we also have another weaker upper-bound in which C 1 and C 2 are two constants:

In Lemma 4.3, we introduce KL regularization (Eq. (11)) for supervised fine-tuning, which can be incorporated by adjusting the original reward with a shift factor γ in the reward-weighted loss, smoothing the weighting of each sample towards the uniform distribution. 2 We refer to this regularization as KL-D since it is based only on data from the pre-trained model. The induced supervised training objective for KL-D is as follows:

We also consider another KL regularization presented in Eq. (12). This KL regularization can be implemented by introducing an additional term in the reward-weighted loss. This extra term penalizes the L 2 -distance between the denoising directions derived from the pre-trained and current models. We refer to it as KL-O because it also regularizes the output from the current model to be close to that from the pre-trained model. The induced supervised training objective for KL-O is as follows:

We summarize our supervised fine-tuning in Algorithm 2. Note that in comparison to online RL training, our supervised setting only requires a pre-trained diffusion model and no extra/new datasets.

4.3 Online RL vs. Supervised Fine-tuning

We outline key differences between online RL and supervised fine-tuning:

Online versus offline distribution. We first contrast their objectives. The online RL objective is to find a new image distribution that maximizes expected reward-this can have much different support than the pre-trained distribution. In supervised fine-tuning, the objective only encourages the model to 'imitate' good examples from the supervised dataset, which always lies in the support of the pre-trained distribution.

Different KL regularization. The methods also differ in their use of KL regularization. For online training, we evaluate conditional KL in each step using online samples, while for supervised training we evaluate conditional KL only using supervised samples. Moreover, online KL regularization can be seen as an extra reward function to encourage small divergence w.r.t. the pre-trained model for online optimization, while supervised KL induces an extra shift in the original reward for supervised training as shown in Lemma 4.3.

Different evaluation of the reward model. Online RL fine-tuning evaluates the reward model using the updated distribution, while supervised fine-tuning evaluates the reward on the fixed pre-training data distribution. As a consequence, our online RL optimization should derive greater benefit from the generalization ability of the reward model.

For these reasons, we expect online RL fine-tuning and supervised fine-tuning to generate rather different behaviors. Specifically, online fine-tuning should be better at maximizing the combined reward (human reward and implicit KL reward) than the supervised approach (similar to the difference between online learning and weighted behavior cloning in reinforcement learning).

5 Experimental Evaluation

We now describe a set of experiments designed to test the efficacy of different fine-tuning methods.

5.1 Experimental Setup

As our baseline generative model, we use Stable Diffusion v1.5 [30], which has been pre-trained on large image-text datasets [33, 34]. For compute-efficient fine-tuning, we use Low-Rank Adaption (LoRA) [11], which freezes the parameters of the pre-trained model and introduces low-rank trainable weights. We apply LoRA to the UNet [31] module and only update the added weights. For the reward model, we use ImageReward [43] which is trained on a large dataset comprised of human assessments of images. Compared to other scoring functions such as CLIP [27] or BLIP [18], ImageReward has a better correlation with human judgments, making it the preferred choice for fine-tuning our baseline diffusion model (see Appendix C for further justification). Further experimental details (e.g., model architectures, final hyper-parameters) are provided in Appendix B. We also provide more samples for qualitative comparison in Appendix E.6.

5.2 Comparison of Supervised and RL Fine-tuning

We first evaluate the performance of the original and fine-tuned text-to-image models w.r.t. specific capabilities, such as generating objects with specific colors, counts, or locations; and composing multiple objects (or composition ). For both RL and SFT training we include KL regularization (see hyperparameters in Appendix B). For SFT, we adopt KL-O since we found it is more effective than KL-D in improving the visual quality of the SFT model (see the comparison between KL-O and KL-D in Appendix E.1). To systematically analyze the effects of different fine-tuning methods, we adopt a straightforward setup that uses one text prompt during fine-tuning. As training text prompts,

Color

Composition

Count

Location

Original model

SFT model

RL model

Color

Count

Composition

Location

0.5

0.0

0.5

1.0

1.5

ImageReward score

Original model

SFT model

RL model

(a) ImageReward score

Color

Count

Composition

Location

5.0

5.2

5.4

5.6

5.8

6.0

6.2

Aesthetic score

Original model

SFT model

RL model

(b) Aesthetic score

Alignment

Image quality

0

20

40

60

80

100

Win (RL model)

Lose (SFT model)

Tie

(c) Human evalution

we use 'A green colored rabbit', 'A cat and a dog', 'Four wolves in the park', and 'A dog on the moon'. These test the models' ability to handle prompts involving color , composition , counting and location , respectively. For supervised fine-tuning, we use 20K images generated by the original model, which is the same number of (online) images used by RL fine-tuning.

Figure 3(a) compares ImageReward scores of images generated by the different models (with the same random seed). We see that both SFT and RL fine-tuning improve the ImageReward scores on the training text prompt. This implies the fine-tuned models can generate images that are better aligned with the input text prompts than the original model because ImageReward is trained on human feedback datasets to evaluate image-text alignment. Figure 2 indeed shows that fine-tuned models add objects to match the number (e.g., adding more wolves in 'Four wolves in the park'), and replace incorrect objects with target objects (e.g., replacing an astronaut with a dog in 'A dog on the moon') compared to images from the original model. They also avoid obvious mistakes like generating a rabbit with a green background given the prompt 'A green colored rabbit'. Of special note, we find that the fine-tuned models generate better images than the original model on several unseen text prompts consisting of unseen objects in terms of image-text alignment (see Figure 10 in Appendix E).

As we can observe in Figure 3(a), RL models enjoy higher ImageReward than SFT models when evaluated with the same training prompt in different categories respectively, due to the benefit of online training as we discuss in Section 4.3. We also evaluate the image quality of the three models using the aesthetic predictor [34], which is trained to predict the aesthetic aspects of generated images. 3 Figure 3(b) shows that supervised fine-tuning degrades image quality relative to the RL approach (and sometimes relative to the pre-trained model), even with KL regularization which is

RL model

SFT model

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

ImageReward score

RL model

SFT model

5.20

5.25

5.30

5.35

5.40

5.45

5.50

5.55

Aesthetic score

without KL

with KL

RL model

SFT model

Without KL

With KL

(a) ImageReward and Aesthetic scores

(b) Generated images

Original model

FOURs

COUES

ROU'S2

BOucO

TOURSES

Eouoss

FOUR'

FeNOLAAS

FesIAAs

FOsUE

FieocAss

FIOUR

H

Fine-tuned model

intended to blunt such effects. For example, the supervised model often generates over-saturated images, corroborating the observations of Lee et al. [17]. By contrast, the RL model generates more natural images that are at least as well-aligned with text prompts.

We also conduct human evaluation as follows: We collect 40 images randomly generated from each prompt, resulting in a total of 160 images for each model. Given two (anonymized) sets of four images from the same random seeds, one from RL fine-tuned model and one from the SFT model, we ask human raters to assess which one is better w.r.t. image-text alignment and image quality. Each query is evaluated by 8 independent raters and we report the average win/lose rate in Figure 3(c). The RL model consistently outperforms the SFT model on both alignment and image quality.

5.3 The Effect of KL Regularization

To demonstrate the impact of KL regularization on both supervised fine-tuning (SFT) and online RL fine-tuning, we conduct an ablation study specifically fine-tuning the pre-trained model with and without KL regularization on 'A green colored rabbit'.

Figure 4(a) shows that KL regularization in online RL is effective in attaining both high reward and aesthetic scores. We observe that the RL model without KL regularization can generate lower-quality images (e.g., over-saturated colors and unnatural shapes) as shown in Figure 4(b). In the case of SFT with KL-O (where we use the same configuration as Section 5.2), we find that the KL regularization can mitigate some failure modes of SFT without KL and improve aesthetic scores, but generally suffers from lower ImageReward. We expect that this difference in the impact of KL regularization is due to the different nature of online and offline training-KL regularization is only applied to fixed samples in the case of SFT while KL regularization is applied to new samples per each update in the case of online RL (see Section 4.3 for related discussions).

5.4 Reducing Bias in the Pre-trained Model

To see the benefits of optimizing for ImageReward, we explore the effect of RL fine-tuning in reducing bias in the pre-trained model. Because the original Stable Diffusion model is trained on

large-scale datasets extracted from the web [33, 34], it can encode some biases in the training data. As shown in Figure 5, we see that the original model tends to produce whiskey-related images given the prompt 'Four roses' (which happens to be the brand name of a whiskey), which is not aligned with users' intention in general. To verify whether maximizing a reward model derived from human feedback can mitigate this issue, we fine-tune the original model on the 'Four roses' prompt using RL fine-tuning. Figure 5 shows that the fine-tuned model generates images with roses (the flower) because the ImageReward score is low on the biased images (ImageReward score is increased to 1.12 from -0.52 after fine-tuning). This shows the clear benefits of learning from human feedback to better align existing text-to-image models with human intentions.

5.5 Fine-tuning on Multiple Prompts

We further verify the effectiveness of the proposed techniques for fine-tuning text-to-image models on multiple prompts simultaneously. We conduct online RL training with 104 MS-CoCo prompts and 183 Drawbench prompts, respectively (the prompts are randomly sampled during training). Detailed configurations are provided in Appendix E. Specifically, we also learn a value function for variance reduction in policy gradient which shows benefit in further improving the final reward (see Appendix A.5 for details.) We report both ImageReward and the aesthetic score of the original and the RL fine-tuned models. For evaluation, we generate 30 images from each prompt and report the average scores of all images. The evaluation result is reported in Table 1 with sample images in Figure 11 in Appendix E, showing that RL training can also significantly improve the ImageReward score while maintaining a high aesthetic score with much larger sets of training prompts.

6 Discussions

In this work, we propose DPOK, an algorithm to fine-tune a text-to-image diffusion model using policy gradient with KL regularization. We show that online RL fine-tuning outperforms simple supervised fine-tuning in improving the model's performance. Also, we conduct an analysis of KL regularization for both methods and discuss the key differences between RL fine-tuning and supervised fine-tuning. We believe our work demonstrates the potential of reinforcement learning from human feedback in improving text-to-image diffusion models.

Limitations, future directions. Several limitations of our work suggest interesting future directions: (a) As discussed in Section 5.5, fine-tuning on multiple prompts requires longer training time, hyperparameter tuning, and engineering efforts. More efficient training with a broader range of diverse and complex prompts would be an interesting future direction to explore. (b) Exploring advanced policy gradient methods can be useful for further performance improvement. Investigating the applicability and benefits of these methods could lead to improvements in the fine-tuning process.

Broader Impacts Text-to-image models can offer societal benefits across fields such as art and entertainment, but they also carry the potential for misuse. Our research allows users to finetune these models towards arbitrary reward functions. This process can be socially beneficial or harmful depending on the reward function used. Users could train models to produce less biased or offensive imagery, but they could also train them to produce misinformation or deep fakes. Since many users may follow our approach of using open-sourced reward functions for this finetuning, it is critical that the biases and failure modes of publicly available reward models are thoroughly documented.

References

Bai, Yuntao, Jones, Andy, Ndousse, Kamal, Askell, Amanda, Chen, Anna, DasSarma, Nova, Drain, Dawn, Fort, Stanislav, Ganguli, Deep, Henighan, Tom, et al. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862 , 2022. 3

Black, Kevin, Janner, Michael, Du, Yilun, Kostrikov, Ilya, and Levine, Sergey. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13301 , 2023. 3, 18

Chi, Cheng, Feng, Siyuan, Du, Yilun, Xu, Zhenjia, Cousineau, Eric, Burchfiel, Benjamin, and Song, Shuran. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137 , 2023. 2

Dhariwal, Prafulla and Nichol, Alexander. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems , 2021. 2, 4

Fan, Ying and Lee, Kangwook. Optimizing ddpm sampling with shortcut fine-tuning. arXiv preprint arXiv:2301.13362 , 2023. 3, 4

Feng, Weixi, He, Xuehai, Fu, Tsu-Jui, Jampani, Varun, Akula, Arjun, Narayana, Pradyumna, Basu, Sugato, Wang, Xin Eric, and Wang, William Yang. Training-free structured diffusion guidance for compositional text-to-image synthesis. arXiv preprint arXiv:2212.05032 , 2022. 1, 2

Gokhale, Tejas, Palangi, Hamid, Nushi, Besmira, Vineet, Vibhav, Horvitz, Eric, Kamar, Ece, Baral, Chitta, and Yang, Yezhou. Benchmarking spatial relationships in text-to-image generation. arXiv preprint arXiv:2212.10015 , 2022. 1

Goodfellow, Ian, Pouget-Abadie, Jean, Mirza, Mehdi, Xu, Bing, Warde-Farley, David, Ozair, Sherjil, Courville, Aaron, and Bengio, Yoshua. Generative adversarial networks. Communications of the ACM , 63(11):139-144, 2020. 3

Ho, Jonathan and Salimans, Tim. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022. 4

Ho, Jonathan, Jain, Ajay, and Abbeel, Pieter. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems , 2020. 1, 2, 3, 15, 16

Hu, Edward J, Shen, Yelong, Wallis, Phillip, Allen-Zhu, Zeyuan, Li, Yuanzhi, Wang, Shean, Wang, Lu, and Chen, Weizhu. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021. 7

Hu, Yushi, Liu, Benlin, Kasai, Jungo, Wang, Yizhong, Ostendorf, Mari, Krishna, Ranjay, and Smith, Noah A. Tifa: Accurate and interpretable text-to-image faithfulness evaluation with question answering. arXiv preprint arXiv:2303.11897 , 2023. 1

Kirstain, Yuval, Polyak, Adam, Singer, Uriel, Matiana, Shahbuland, Penna, Joe, and Levy, Omer. Pick-a-pic: An open dataset of user preferences for text-to-image generation. In Advances in Neural Information Processing Systems , 2023. 1, 3

Kong, Zhifeng, Ping, Wei, Huang, Jiaji, Zhao, Kexin, and Catanzaro, Bryan. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint arXiv:2009.09761 , 2020. 2

Krizhevsky, Alex, Hinton, Geoffrey, et al. Learning multiple layers of features from tiny images. 2009. 3

Lee, Kimin, Smith, Laura, and Abbeel, Pieter. Pebble: Feedback-efficient interactive reinforcement learning via relabeling experience and unsupervised pre-training. arXiv preprint arXiv:2106.05091 , 2021. 3

Lee, Kimin, Liu, Hao, Ryu, Moonkyung, Watkins, Olivia, Du, Yuqing, Boutilier, Craig, Abbeel, Pieter, Ghavamzadeh, Mohammad, and Gu, Shixiang Shane. Aligning text-to-image models using human feedback. arXiv preprint arXiv:2302.12192 , 2023. 1, 2, 3, 6, 9, 18

Li, Junnan, Li, Dongxu, Xiong, Caiming, and Hoi, Steven. Blip: Bootstrapping languageimage pre-training for unified vision-language understanding and generation. In International Conference on Machine Learning , 2022. 3, 7, 18

Liu, Hao, Sferrazza, Carmelo, and Abbeel, Pieter. Chain of hindsight aligns language models with feedback. arXiv preprint arXiv:2302.02676 , 2023. 3

Liu, Nan, Li, Shuang, Du, Yilun, Torralba, Antonio, and Tenenbaum, Joshua B. Compositional visual generation with composable diffusion models. In European Conference on Computer Vision , 2022. 2

Liu, Rosanne, Garrette, Dan, Saharia, Chitwan, Chan, William, Roberts, Adam, Narang, Sharan, Blok, Irina, Mical, RJ, Norouzi, Mohammad, and Constant, Noah. Character-aware models improve visual text rendering. arXiv preprint arXiv:2212.10562 , 2022. 2

Liu, Ziwei, Luo, Ping, Wang, Xiaogang, and Tang, Xiaoou. Deep learning face attributes in the wild. In International Conference on Computer Vision , 2015. 3

Ouyang, Long, Wu, Jeffrey, Jiang, Xu, Almeida, Diogo, Wainwright, Carroll, Mishkin, Pamela, Zhang, Chong, Agarwal, Sandhini, Slama, Katarina, Ray, Alex, et al. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems , 2022. 3, 5

Peters, Jan and Schaal, Stefan. Reinforcement learning by reward-weighted regression for operational space control. In Proceedings of the 24th international conference on Machine learning , pp. 745-750, 2007. 6

Petsiuk, Vitali, Siemenn, Alexander E, Surbehera, Saisamrit, Chin, Zad, Tyser, Keith, Hunter, Gregory, Raghavan, Arvind, Hicke, Yann, Plummer, Bryan A, Kerret, Ori, et al. Human evaluation of text-to-image models on a multi-task benchmark. arXiv preprint arXiv:2211.12112 , 2022. 1

Poole, Ben, Jain, Ajay, Barron, Jonathan T, and Mildenhall, Ben. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209.14988 , 2022. 2

Radford, Alec, Kim, Jong Wook, Hallacy, Chris, Ramesh, Aditya, Goh, Gabriel, Agarwal, Sandhini, Sastry, Girish, Askell, Amanda, Mishkin, Pamela, Clark, Jack, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning , 2021. 1, 2, 3, 7, 18

Raffel, Colin, Shazeer, Noam, Roberts, Adam, Lee, Katherine, Narang, Sharan, Matena, Michael, Zhou, Yanqi, Li, Wei, and Liu, Peter J. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research , 21(1):5485-5551, 2020. 1, 2

Ramesh, Aditya, Dhariwal, Prafulla, Nichol, Alex, Chu, Casey, and Chen, Mark. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022. 1, 2

Rombach, Robin, Blattmann, Andreas, Lorenz, Dominik, Esser, Patrick, and Ommer, Björn. High-resolution image synthesis with latent diffusion models. In Conference on Computer Vision and Pattern Recognition , 2022. 1, 2, 3, 5, 7

Ronneberger, Olaf, Fischer, Philipp, and Brox, Thomas. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical Image Computing and Computer Assisted Intervention , 2015. 7

Saharia, Chitwan, Chan, William, Saxena, Saurabh, Li, Lala, Whang, Jay, Denton, Emily, Ghasemipour, Seyed Kamyar Seyed, Ayan, Burcu Karagol, Mahdavi, S Sara, Lopes, Rapha Gontijo, et al. Photorealistic text-to-image diffusion models with deep language understanding. In Advances in Neural Information Processing Systems , 2022. 1, 2

Schuhmann, Christoph, Vencu, Richard, Beaumont, Romain, Kaczmarczyk, Robert, Mullis, Clayton, Katta, Aarush, Coombes, Theo, Jitsev, Jenia, and Komatsuzaki, Aran. Laion-400m: Open dataset of clip-filtered 400 million image-text pairs. arXiv preprint arXiv:2111.02114 , 2021. 7, 9, 10

Schuhmann, Christoph, Beaumont, Romain, Vencu, Richard, Gordon, Cade, Wightman, Ross, Cherti, Mehdi, Coombes, Theo, Katta, Aarush, Mullis, Clayton, Wortsman, Mitchell, et al. Laion-5b: An open large-scale dataset for training next generation image-text models. arXiv preprint arXiv:2210.08402 , 2022. 7, 8, 9, 10

Schulman, John, Moritz, Philipp, Levine, Sergey, Jordan, Michael, and Abbeel, Pieter. Highdimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438 , 2015. 16

Schulman, John, Wolski, Filip, Dhariwal, Prafulla, Radford, Alec, and Klimov, Oleg. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017. 5, 17

Sohl-Dickstein, Jascha, Weiss, Eric, Maheswaranathan, Niru, and Ganguli, Surya. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , 2015. 1, 2

Song, Jiaming, Meng, Chenlin, and Ermon, Stefano. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020. 1

Song, Yang and Ermon, Stefano. Improved techniques for training score-based generative models. In Advances in neural information processing systems , 2020. 2

Song, Yang, Durkan, Conor, Murray, Iain, and Ermon, Stefano. Maximum likelihood training of score-based diffusion models. In Advances in Neural Information Processing Systems , 2021. 14

Stiennon, Nisan, Ouyang, Long, Wu, Jeffrey, Ziegler, Daniel, Lowe, Ryan, Voss, Chelsea, Radford, Alec, Amodei, Dario, and Christiano, Paul F. Learning to summarize with human feedback. In Advances in Neural Information Processing Systems , 2020. 3, 5

Wu, Xiaoshi, Sun, Keqiang, Zhu, Feng, Zhao, Rui, and Li, Hongsheng. Better aligning text-to-image models with human preference. arXiv preprint arXiv:2303.14420 , 2023. 1, 3

Xu, Jiazheng, Liu, Xiao, Wu, Yuchen, Tong, Yuxuan, Li, Qinkai, Ding, Ming, Tang, Jie, and Dong, Yuxiao. Imagereward: Learning and evaluating human preferences for text-to-image generation. In Advances in Neural Information Processing Systems , 2023. 1, 2, 3, 7, 18

Appendix:

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

A Derivations

A.1 Lemma 4.1

Proof.

where the second to last equality is from the continuous assumptions of p θ p x 0: T | z q r p x 0 , z q and ∇ θ p θ p x 0: T | z q r p x 0 , z q .

A.2 Lemma 4.2

Similar to the proof in Theorem 1 in [40], from data processing inequality with the Markov kernel, given any z we have

Using the Markov property of p θ and p pre, we have

Taking the expectation of p p z q on both sides, we have

A.3 The gradient of objective Eq. (8)

From Lemma 4.2, for online fine-tuning, we need to regularize ř T t ' 1 E p θ p x t q ' KL ` p θ p x t ´ 1 | x t , z q} p pre p x t ´ 1 | x t , z q ˘‰ . By the product rule, we have

which treats the sum of conditional KL-divergences along the future trajectory as a scalar reward at each step. However, computing these sums is more inefficient than just the first term in Eq. (16). Empirically, we find that regularizing only the first term in Eq. (16) already works well, so we adopt this approach for simplicity.

A.4 Lemma 4.3

Similar to derivation in [10], and notice that q is not dependent on z (i.e., q p¨ ¨ ¨ | z q ' q p¨ ¨ ¨ q ), we have

where the first inequality comes from ELBO.

where C is a constant.

Notice that p θ p x 0 | z, x 1 q is modeled as a discrete decoder as in [10] and p p x T | z q is a fixed Gaussian; neither is trainable.

As a result, we have

Furthermore, from triangle inequality we have

So another weaker upper bound is given by

Notice that p θ p x 0 | z, x 1 q is modeled as a discrete decoder as in [10] and p p x T | z q is a fixed Gaussian; neither is trainable. Then the proof is complete.

A.5 Value function learning

Similar to [35], we can also apply the variance reduction trick by subtracting a baseline function V p x t , z, θ q :

where

So we can learn a value function estimator by minimizing the objective function below for each t :

Since subtracting V p x t , θ, z q will minimize the variance of the gradient estimation, it is expected that such a trick can improve policy gradient training. In our experiments, we find that it can also improve the final reward: For example, for the prompt "A dog on the moon", adding variance reduction can improve ImageReward from 0.86 to 1.51, and also slightly improve the aesthetic score from 5.57 to 5.60. The generated images with and without value learning are shown in Figure 6. We also find similar improvements in multi-prompt training (see learning curves with and without value learning in Figure 7).

With value learning

Without value learning

A.6 Importance sampling and ratio clipping

In order to reuse old trajectory samples, we can apply the important sampling trick:

In order to constrain p θ to be close to p θ old , we can use the clipped gradient for policy gradient update similar to PPO [36]:

where ϵ is the clip hyperparameter.

B Experimental Details

Online RL training. For hyper-parameters of online RL training used in Section 5.2., we use α ' 10 , β ' 0 . 01 , learning rate = 10 ´ 5 and keep other default hyper-parameters in AdamW, sampling batch size m ' 10 . For policy gradient training, to perform more gradient steps without the need of re-sampling, we perform 5 gradient steps per sampling step, with gradient norm clipped to be smaller than 0 . 1 , and use importance sampling to handle off-policy samples, and use batch size n ' 32 in each gradient step. We train the models till generating 20000 online samples, which matches the number of samples in supervised finetuning and results in 10K gradient steps in total. For stable training, we freeze the weights in batch norm to be exactly the same as the pretrained diffusion model. For importance sampling, we apply the clip hyperparameter ϵ ' 10 ´ 4 .

Supervised training. For hyper-parameters of supervised training, we use γ ' 2 . 0 as the default option in Section 5.2, which is chosen from γ P

0.8

0.7

w/o value learning

0.6

0.5

Reward

0.4

0.3

0.2

0.1

0.0

2000

4000

6000

8000

10000

Step

t 0 . 1 , 1 . 0 , 2 . 0 , 5 . 0 u . We use learning rate ' 2 ˆ 10 ´ 5 and keep other default hyper-parameters in AdamW, which was chosen from t 5 ˆ 10 ´ 6 , 1 ˆ 10 ´ 5 , 2 ˆ 10 ´ 5 , 5 ˆ 10 ´ 5 u . We use batch size

n ' 128 and M ' 20000 such that both algorithms use the same number of samples, and train the SFT model for 8K gradient steps. Notice that Lemma 4.3 requires r ` γ to be non-negative, we apply a filter of reward such that if r ` γ ă 0 , we just set it to 0. In practice, we also observe that without this filtering, supervised fine-tuning could fail during training, which indicates that our assumption is indeed necessary.

C Investigation on ImageReward

We investigate the quality of ImageReward [43] by evaluating its prediction of the assessments of human labelers. Similar to Lee et al. [17], we check whether ImageReward generates a higher score for images preferred by humans. Our evaluation encompasses four text prompt categories: color, count, location, and composition. To generate prompts, we combine words or phrases from each category with various objects. These prompts are then used to generate corresponding images using a pre-trained text-to-image model. We collect binary feedback from human labelers on the image-text dataset and construct comparison pairs based on this feedback. Specifically, we utilize 804, 691, 1154 and 13228 comparison pairs obtained from 38, 29, 41 and 295 prompts for color, location, count, and composition, respectively. 4 Additionally, for evaluation on complex text prompts from human users, we also utilize the test set from ImageReward, which consists of 6399 comparison pairs from 466 prompts. 5 Table 2 compares the accuracy of ImageReward and baselines like CLIP [27] and BLIP [18] scores, which justifies the choice of using ImageReward for fine-tuning text-to-image models.

D A Comprehensive Discussion of Related Work

We summarize several similarities and differences between our work and a concurrent work [2] as follows:

Both Black et al. [2] and our work explored online RL fine-tuning for improving text-to-image diffusion models. However, in our work, we provide theoretical insights for optimizing the reward with policy gradient methods and provide conditions for the equivalence to hold.

Black et al. [2] demonstrated that RL fine-tuning can outperform supervised fine-tuning with reward-weighted loss [17] in optimizing the reward, which aligns with our own observations.

In our work, we do not only focus on reward optimization: inspired by failure cases (e.g., over-saturated or non-photorealistic images) in supervised fine-tuning [17], we aim at finding an RL solution with KL regularization to solve the problem.

Unique to our work, we systematically analyze KL regularization for both supervised and online fine-tuning with theoretical justifications. We show that KL regularization would be more effective in the online RL fine-tuning rather than the supervised fine-tuning. By adopting online KL regularization, our algorithm successfully achieves high rewards and maintains image quality without over-optimization issue.

E More Experimental Results

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

We also present the effects of both KL-D and KL-O with coefficient γ P t 0 . 1 , 1 . 0 , 2 . 0 , 5 . 0 u in Figure 8 and Figure 9. We observe the followings:

KL-D (eq. (11)) can slightly increase the aesthetic score. However, even with a large value of γ , it does not yield a significant improvement in visual quality. Moreover, it tends to result in lower ImageReward scores overall.

In Figure 9, it is evident that KL-O is more noticeably effective in addressing the failure modes in supervised fine-tuning (SFT) compared to KL-D, especially when γ is relatively large. However, as a tradeoff, KL-O significantly reduces the ImageReward scores in that case.

In general, achieving both high ImageReward and aesthetic scores simultaneously is challenging for supervised fine-tuning (SFT), unlike RL fine-tuning (blue bar in Figure 8).

We also present more qualitative examples from different choices of KL regularization in Appendix E.5.

RL

=0.0

=0.1

=0.1

=1.0

=1.0

=2.0

=2.0

=5.0

=5.0

0.6

0.8

1.0

1.2

1.4

1.6

ImageReward score

RL with KL

SFT without KL

SFT with KL-O

SFT with KL-D

(a) ImageReward score

RL

=0.0

=0.1

=0.1

=1.0

=1.0

=2.0

=2.0

=5.0

=5.0

5.20

5.25

5.30

5.35

5.40

5.45

5.50

5.55

Aesthetic score

(b) Aesthetic score

=0.1

=1.0

=2.0

=5.0

SFT with KL-O SFT with KL-D

RL with KL

E.2 Images from unseen text prompts

Figure 10 shows image samples from the original model, SFT model and RL fine-tuned models on unseen text prompts.

Color

Composition

Count

Location

Original model

SFT model

RL model

E.3 Multi-prompt Training

Here we report more details for multi-prompt training. We still use the same setting in Appendix B, but 1) increase the batch size per sampling step to n ' 45 for stable training; 2) use a smaller KL weight β ' 0 . 001 ; 3) use extra value learning for variance reduction. For value learning, we use a larger learning rate 10 ´ 4 and batch size 256. Also, we train for a longer time for multi-prompt such that it utilizes 50000 online samples. See sample images in Figure 11.

E.4 Long and Complex Prompts

We show that our method does not just work for short and simple prompts like "A green colored rabbit", but can also work with long and complex prompts. For example, in Fig. 12 we adopt a complex prompt and compare the images generated by the original model and the supervised finetuned model. We can observe that online training encourages the models to generate images with a different style of painting and more fine-grained details, which is generally hard to achieve by supervised fine-tuning only (with KL-O regularization, γ ' 2 . 0 ).

E.5 More Qualitative Results from Ablation Study for KL Regularization in SFT

Here we provide samples from more prompts in KL ablation for SFT: see Figure 13.

E.6 Qualitative Comparison

Here we provide more samples from the experiments in Section 5.2 and Section 5.3: see Figure 14, Figure 16, Figure 15, Figure 17, Figure 18 and Figure 19, with the same configuration as in Section 5.2 for both RL and SFT.

nal

(a)

(q)

nal

Original model

SFT model

RL model

=0.1

=1.0

=2.0

=5.0

SFT with KL-O SFT with KL-D

=0.1

=1.0

=2.0

=5.0

SFT with KL-D SFT with KL-O

=0.1

=1.0

=2.0

=5.0

）SFT with KL-D

SFT with KL-O

(a) Original model

(b) SFT model

(c) RL model

O

L

(a) RL model without KL regularization

(b) RL model with KL regularization

(b) SFT model with KL regularization

Peer reviews:
comments: SUMMARY:
in this paper, the authors propose to improve the pretrained text-to-image diffusion model with human feedback in an online-reinforcement learning manner. The problem formulation, the differences between the online update and the supervised finetuning are clearly stated. Experiments show that the proposed method effectively improves the performance of the text-to-image generation ability.

STRENGTHS:
1. The paper is well-written and easy to understand.

2. The comparison and the potential advantages compared to related works are clearly discussed.

3. Experiments show that the proposed method improves the performance of the pretrained text-to-image model, also outperforms the supervised human-feedback finetuning baseline.

WEAKNESSES:
1. Some claims and the results are connected weakly. The claims in Sec.4.3 should be better aligned with the experimental results. However, its current version makes me hard to find the direct mapping.

2. The title of this paper seems not to reflect the main contribution. In my opinion, the main contribution of this work lies in the online manner, since using reinforcement learning to fine-tune the text-to-image model has been explored before. It needs a better title.

3. There are some experimental problems. In my experience, the over-saturated problem in the SFT model can be alleviated by using small classifier-free guidance. Therefore, the necessity of using the online optimization is questionable. Also, why the authors do not directly compare with [17]?

QUESTIONS:
Please refer to the weakness part.

LIMITATIONS:
The authors discussed some limitations in the last part of the paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-15

comments: SUMMARY:
The paper introduces the idea of finetuning text-to-image diffusion models using reinforcement learning. The core idea is relatively simple: generate samples from a trained diffusion model, use the samples and a reward function to update the model's parameters, and iterate. The paper further introduces a KL divergence loss to prevent diverging too far from the original model's weights. evaluations are done on different prompts with respect to different capabilities/ performance metrics, such as colour or composition.

STRENGTHS:
- an interesting and novel idea
- simple method and easy algorithm
- qualitatively pleasing results

WEAKNESSES:
- considering only 4 prompts for experimental evaluation is not sufficient
- considering only a single prompt for bias evaluation is not sufficient
- authors mention that "longer training time, hyper-parameter tuning and engineering efforts are required as the number of training text prompts increases". This indicates that their method does not scale to a larger set of training prompts, which questions its relevance
- investigation of generalisation to prompts that are out of distribution should be conducted

QUESTIONS:
see weaknesses

LIMITATIONS:
limitations are briefly discussed, especially limitations mentioned in weaknesses could be discussed in more detail
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-28

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to use RL to finetune t2i diffusion models and use KL regularization into supervised fine-tuning of diffusion models.  The paper shows that RL fine-tuning can avoid the overfitting that arises in SFT, and is generally superior to SFT with respect to both image-text alignment and image quality.\n\nSTRENGTHS:\nThe motivation for using RL techniques to fine-tune T2I model is reasonable and compelling. \n\nThe paper is well-written and presented. It is very easy to follow.\n\nThe idea is simple and easy to implement. \n\nThe proposed KL regularization technique is effective. The evidence provided in the paper supports the efficacy of this technique\n\nWEAKNESSES:\nThe novelty of this paper is kind of limited. Similar ideas have already been proposed in [1].\n\nI feel the experiment parts are a bit lacking. Only one policy gradient approach is tried in the paper and only 20K images are used for fine-tuning.\n\nIn the contribution part, the author claimed \"online fine-tuning is critical to improving text-to-image alignment while maintaining high image fidelity\", but fidelity scores or IS scores are not reported in the paper. Also no human evaluation are conducted in the main paper.\n\nThe cost for online RL fine-tuning is not reported. I believe it will be more cost-expensive compared with supervised fine-tuning, especially for scenarios such as fine-tuning on multiple prompts.\n\n\n[1] Training diffusion models with reinforcement learning\n\nQUESTIONS:\nHow do you define the aesthetic scores and the ImageReward scores? Do you use the same metric as in the Imagereward paper? I am also curious about the CLIP/BLIP score with the proposed method which are often used for the evaluation.\n\nI think LoRA is proposed to supervised fine-tuning. Do you apply LoRA to both Supervised and RL fine-tuning? Also the UNet in the Stable Diffusion contains different layers including conv layers and cross-attention layers etc. Which part of the UNet module do you apply LoRA to?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies RLHF fine-tuning of diffusion models to learn from and align with human feedback. Specifically, they introduce (1) an online RL strategy and (2) a KL regularization (inspired by a similar regularization for the RLHF via PPO of LLMs). Empirically, they show better alignment with the optimized reward and increased aesthetic when compared to the more trivial offline imitation based RL (named SFT in the paper).\n\nSTRENGTHS:\n* The paper is clear, the theoretical section is sound, and the experiments are convincing.\n* This is arguably the first work applying online RL to fine-tune diffusion model, which is a key contribution.\n* I appreciated the MDP formulation of the diffusion process.\n* The inclusion of the KL divergence is interesting, though straightforward given the recent success of RLHF in NLP.\n\nWEAKNESSES:\n* Section 4.3 proposes 3 theoretical/intuitive reasons why online RL would perform better than SFT. I believe those reasons could be further validated in the experiments.\n    - The first argument is that online learning favors exploration away from the pre-trained distribution. This could be ablated by saving the generated images along the online RL, and then applying SFT on these images.\n    - The second argument is the difference in KL regularization, which actually theoretically comes to the difference between KL and reverse KL. Applying online RL with reverse KL (though more costly) could help ablate the importance of this component. Actually, I am sceptikal by the KL-D for SFT, as (1) the more standard KL-0 in App.E actually performs better and (2) offline RL usually requires baseline rewards for normalization, thus adjusting the original reward with a shift factor would not change anything.\n    - The third argument is the robustness of the reward model. Therefore, ablations could better analyze the quality difference of the reward model on the pre-train distribution and then on the updated distribution. More generally, the fact that your online RL requires a more robust reward model is actually a drawback/limitation, creating new challenges in robust reward design.\n* More generally more experiments are required to ensure that online RL for diffusion models consistently helps.\n    - The experiments are made only for a very small number of prompts: 4! This limits applications to real-world applications.\n    - The experiments are made with a single reward model. More experiments with diverse (less robust) reward models would help.\n    - Therefore the empirical contribution may be seen as marginal.\n    - The paper lacks human evaluation, even though this was done for the reward model. I believe human evaluation after online RL vs SFT would have been more interesting.\n\nQUESTIONS:\n* Could you elaborate on the difference between KL and reverse KL.\n* Do you think it may be possible to add \"online\" in the title \"Reinforcement Learning for Online Fine-tuning of Text-to-Image Diffusion Models\", to specify that your contribution is actually about online RL.\n* I believe your SFT baselines could be improved with some reward normalization, value function and with KL-D. Such improved SFT baseline would strengthen your experimental sections.\n* Have you explored refined online RL optimization strategies such as PPO?\n* Finally, I believe more prompts (for example those used in evaluation) and more reward models (for example those used for aesthetic evaluation) are needed in training.\n\nLIMITATIONS:\nThe limitations could be extended to the societal risks of generative AI, and also mention the fact that online RL requires more robust reward models.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to adapt the popular RLHF framework used for fine-tuning LLMs to fine-tuning of diffusion generative models.\nGiven a reward model, the goal is to finetune the parameters of the diffusion model such that the resulting images from the \nsampling process achieve high reward. Importantly, the reward is a function of the final output of the entire diffusion process (online finetuning) and can thus not be optimized independently across timesteps like the denoising loss used to train the LDM. The authors therefore adapt a previous result to compute the gradient without the need of backpropagating through the entire diffusion graph. Additionally, they regularize the optimisation via an upper bound on the KL divergence between the fine-tuned and the original model.\n\nSTRENGTHS:\n- I find the idea to adapt RLHF from language tasks to diffusion generative models to be very relevant as diffusion models clearly often misinterpret details like counts and colors about the prompt. \n - It is good that the authors added KL regularization to supervised finetuning to establish a better baseline for their RLHF framework.\n - Section 4.3 clearly distinguishes RLHF from supervised finetuning and gives several good reasons why RLHF is superior to supervised finetuning from a more theoretical perspective.\n - I find the quality of the results in the paper to be mostly convincing and I get the impression that RLHF significantly outperforms both baselines.\n\nWEAKNESSES:\n- The paper focuses mostly on finetuning on a single or few prompts at a time. Being able to train on a large variety of prompts to train a model that is overall better than SD instead of only focusing on a few specific aspects would greatly increase the usefulness of the paper. \n - The evaluation is somewhat limited and it is not clear how cherry-picked the prompts for the results are (for the given prompts the appendix contains non-cherry-picked results). It would be much more convincing if the authors could show results on challenging prompts that are randomly generated. Also, the quantitative results in Figure 3 would be more impressive if they were averaged over a larger selection of prompts. \n - Since the method is trained to improve ImageReward, I would prefer a different evaluation metric in Figure 3 to demonstrate generalization. If there is no automatic pipeline available, the authors could conduct a user study or simply generate prompts that are easy to evaluate by a human. For example, prompts of the form \"{N} {CLASS} {BACKGROUND}\" , where N is the number of objects and CLASS the object to generate. Then they could simply generate a certain number of images and manually count the ones where the model correctly generates the required number of objects. \n - In general, the evaluation is in my opinion the weakest part of the paper and I would gladly increase my score if the authors could present more convincing quantitative evaluations or a broader selection of qualitative results (or both). \n - While RLHF is significantly better than the baselines, to me some of the images, especially the \"green rabbit\" and \"green cat\" still look somewhat oversaturated.\n\nQUESTIONS:\n- Is there a reason why eq (8) should not include a min, similar to (4)?\n - Do you believe that LoRA helps achieve better results or would standard finetuning on the original dimensional weights yield better results if compute would not be a concern? \n - From my experience, diffusion models have issues with multiple colors in a single prompt. For example a \"green rabbit and a yellow bird  in front of a blue background\". You show that RLHF can improve colors but can it also handle multiple colors. Also diffusion models often fail with text, for example \"a tshirt with NeurIPS printed on it\", does RLHF help with this or does the reward function not yield useful feedback for this kind of prompt?\n\nLIMITATIONS:\nThe main limitation of the paper is the restriction to a single prompt which is openly discussed by the authors.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}