DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models
Authors: Ying Fan
Olivia Watkins
Yuqing Du
Hao Liu
Moonkyung Ryu
Craig Boutilier
Pieter Abbeel
Mohammad Ghavamzadeh
Kangwook Lee
Kimin Lee
Venue: NeurIPS 2023 poster (2023)
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-lmage Diffusion Models
Ying Fan Olivia Watkins', Yuging Duu, Hao Liu Moonkyung Ryu', Craig Boutilier',
Pieter Abbeel?, Mohammad Ghavamzadeh Kangwook Lee?. Kimin Lee
*Equal technical contribution "Wrr was done at Google Research
'Google Research Uniiversity of Wisconsin-Madison *UC Berkeley *Amazon SKAIST
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, tne-ttnnng exxt-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 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 http:://gituh..comggoooo wwwpeh//ggooge--eeeaccceeeeeeee
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 text-to-image generation. Large-scale text-to- image 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 be an effective means 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, from human feedback, followed by fine-tuning the text-to-image model using supervised learning. They show that simple supervised min-ttunig based on wwerh-we loss can improve the reward scores, leading to better image-text alignment. However, supervised fine-tuning often induces deterioration image quality (e.g., eeturratr or non- ponopoooorealistic images). This is likely due to the model being fine-tuned on fixed dataset that is generated by pre-trained model (Figure (())..
be improved by learning reward function
In this work, 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 diffusion model's image output is equivalent performing policy gradient multi-step diffusion model under certain
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Imagee from pre-trained model
Pore
Reward function
Model update
P6
Reward function
Reward function
Model update
Model update
P6
Reward function
Pore
Po
po
(a) Superviscd fine-tuning
(b) RL fine-tuning
Figure 1: Illustration of (a) reward-weite weight- supervised fine-tuning and (b) RL fine-tuning. Both start with the same pre-trained model (the blue rectangle). In supervised fine-tuning, the model is updated on fixed dataset generated by the pre-trained model. In contrast, the model is updated using new samples from the previously trained model during online RL fine-tuning.
regularity assumptions. We also incorporate Kullback-Leibler (KL) divergence with respect to the pre-trained model regularization an online manner, treating this as implicit reward.
In our experiments, we fine-tune the Stable Diffusion model [30] using ImageReward [43]. an open- source reward model trained 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 online manner. Crucially, online training allows evaluation of the reward model and conditional KL divergence beyond the (supervised) training dataset. This offers distinet advantages supervised fine-tuning, point demonstrate empirically. 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.rtt an LHF-reward) of the images generated by diffusion model given text prompts as an online RL problem. Moreover, we present DPOK: Diffusion Policy Optimization with KL 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 fairer comparison with RL technique.
We discuss the key differences between supervised fine-tuning and online fine-tuning of text-to- image 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.
Related Work
Text-to-image diffusion models. Diffusion models [10, 37, 39] are class of generative models that use iterative denoising process transform Gaussian noise into samples that follow learned data distribution. These models have proven to be highly effective in 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 eext-to-image generation [29, 30, 32]. However, there are still many known weaknesses of existing text-to-image models, such compositionality and attribute binding [6, 20] 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 vision- language 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-fittrred eward-weighted supervised learning, we develop an online fine-tuning framework with an RL- ba-dd objective.
RL fine-tunme tume-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 fine- tuning especially for large-scale text-to-image models using human rewards. We also consider several design choices like adding KL regularization as implicit reward, and compare RL fine-tuning to supervised tne-ttnnng.
Concurrent and independent from our work, Black et al. [2] have also investigated RL fine-tuning to fine-tune gext-to-image diffusion models. They similarly frame the fine-tuning problem 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 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 probubilistic models (DDPMs) [10] for image generation and draw our notation and problem formulation from [10]. Let go be the data
distribution, i.c., zo go (zo), zo E R A DDPM approximates qu with parameterized model of the form po (zo) J po (2orr) dzi.7. where (=(:r) = pr (xr) IT- 10(1))1 and the reverse
process is a Markov chain with the following dynamics:
p(xr)= = =N I), Po( = N (HO(TE ()(((,,))))) E,)
unique characteristic of DDPMs is the exploitation of an approximate posterior q( (mn), known as the the forward diffusion process, which itself is Markov chain that adds Gaussian noise to the data according to a variance schedule $1....,Brr Br:
((F1TT100) = [
a(27(==-1) = WWVIIIIII B, 1.-1,
(2)
Let 04 1 - Br, ā. I1:m Css. and 3r (x(()))) Vac
1-a--1 I- B. Ho et al. [10] adopt the parameterization
Ft. VI-a, tn (x,,,))
Training a DDPM is performed by optimizing a variational bound on the negative log- likelioeddoood E.[- log po (xo)], which is equivalent to optimizing:
T
E
KL Σ
Note that the variance sequence (B:)T_ E (0, 1)T is chosen such that ār & O, and thus, (( x (O, I). The covariance matrix E, in (1) often set to of where of either Ba or 3r, which not trainable. Unlike the original DDPM, use a latent diffusion model [30], so 's are latent.
3

---

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 s conditional distribution with classifier- guidance [9].
Given text prompt p(z), let g<20|2) be the data distribution conditioned on z. This induces joint distribution p(zo. z). During training, the same noising process is used regardless of input and both the unconditional eo (Ia, t) and conditional c0(2,, t, 2) denoising models are learned. For data sampling, let Ég = weal (Ze, z) + (1 - w)ea( (Zr, t), where wo > 1 is the guidance scale. At test time, given text prompt 2, the model generates conditional data according pa (zol2).
Fine-tuning of Diffusion Models
In this section, we describe our approach for online RL fine-tuning diffusion models. We first propose Markov decision process (MDP) formulation for the denoising phase. We then this MDP and present a policy gradient RL algorithm to update the original diffusion model. The RL algorithm optimizes objective consisting of the reward and 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 po(zurr(2) (2) be text-to-image diffusion model where is some text prompt distributed according to p(z), and r(ze, 2) be a reward model (typically trained using human assessment of images).
MDP formulation: The denoising process of DDPMs can be modeled as T T-horizon MDP:
(2)
Po(so) ((0, P(s++1 St. ax) (6., da.);
IT-t-1,
(r(s(+1) 2(),,) if t T {1(
R(s,, 0z)
Ta(a 8t) = po(xr---1 XT oT=
otherwise.
in which 8 and ag are the state and action at time-step 4, Po and Pae the initial state distribution and the dynamics, R is the reward function, and मe is the parameterized policy. As a result, optimizing policy xe in (4) is equivalent to fine-tuning the underlying DDPM. Finally, we denote by 6₂ the Dirac distribution at 2.
It can be seen from the MDP formulation in (4) that the system starts by sampling its initial state So from the Gaussian distribution /V0 I), similar to the first state of the dinoising process zy- Given the MDP state st, which corresponds to state zT of the denoising process, the policy takes action time-step as the next denoising state, i.e., a, = #T-t-1- As a result of this action, the system transitions deterministically to state identified by the action (i.e., the next state of the denoising process). The reward zero, except at the final step which the quality of the image at the end of the denoising process is evaluated w.r.t. the prompt, i.c., r(zo, 2).
common goal in re-training/line-tumee the diffusion models is maximize the expected reward of the generated images given the prompt distribution, i.e.,
min Ep(z) Ppo(eols)(- =)].
(5)
The gradient of this objective function can be obtained as follows:
Lemma 4.1 (A modification of Theorem 4.1 in [5]). If pe ((((ooTT())((())) (2)r(x), and V Vepo(zo:r>>) continuous functions of 0, then can write the gradient of the objective in (5) as
T
VeLggt))Ennttssss-r(aaa 2)] Epttorppltoorlll)
-r(z0,2) 2)
Vo log pe (2
(6)
Proof. We present the proof in Appendix A.
'ee keep the covariance in DDPM constant, Σg E, and only train (see Eq. 1). This would naturally
provide stochastic policy for online exploration.

---

Algorithm DPOK: Diffusion policy optimization with KL regularization
Input:
reward model r, pre-trained model Ppre. current model pe, batch size m, text distribution p(z)
Initialize po = Pyee
while / not converged do
Obtain m i.i.d. samples by first sampling 2 ~ p(z) and then #0.7 ~ pe (0)
Compute the gradient using Eq. (9) and update 0 end while
Output: Fine-tuned diffusion model pa
Equation (6) is equivalent to the gradient used by the popular policy gradient algorithm, REINFORCE, to update policy in the MDP (4). The gradient (6) is estimated from trajectories pe (Toor (2) generated by the current policy, and then used to update the policy po(x;-1/=z. 2) in an online fashion.
Note that REINFORCE is not the only way 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 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 Al feedback is that the model may overfit 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), text-to-image models, makes sense to compute it only for the final image, i.e., KLL(po zo 2²] Pore (Io) Unfortunately, Po(zul 2) is a marginal (see the integral in Section 3) and its closed-form form is unknown. As result, we propose to add an upper- bound of this KL-term to the objective function.
Lemma 4.2. Suppose Ppre P (NO.T Z) and pl( are Markov chains conditioned on the text prompt that both start at xT (O, I). Then, we have
(L 2) (pre (ft-1 Ft, =)]]
(7)
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 zn is upper-bounded by the sum of the divergences between the distributions over latent x, at each diffusion step.
Using the KL upper-bound in (7), propose the following objective for regularized training:
T
ppts) (eo =)] +
Σ E
(x (po 20)
(8)
where at the reward and KL weights, respectively, We use the following gradient to optimize the objective (8):
ar arr (ro, 2) V log 2) +B Σ V (poL 2) (2x (2r 1
Ir, 2))
(9)
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 pscudo-code 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 KI. regularization into supervised fine-tuning (SFT), which allows for a more meaningful comparison with our KL-regularized RL algorithm (DPOK). We begin with supervised
5

---

fine-tuning objective similar to that used in [17], i.e.
pp(s) BB (x00=)
2) log (=)o=)))
(10)
To compare with RL fine-tuning, augment the supervised objective with similar KL regular- ization term. Under the supervised learning setting, we consider (xoo (x))))))))) which is equivalent to minimizing [x)[(=)))) log Pe [2o]2)], given any prompt Let ((xz) be the forward process used for training. Since the distribution of Zh is generally not tractable, we use approximate upper-bounds in the lemma below (see derivation in Appendix A.4).
Lemma 4.3. Let y be the regularization weight and jıa (11, : va 1-a,
ro yar(1-ān-1) St.
1-a;
Assume r r 2)
for all zo and Then, have
Bp (xo())) (r(z0: 2) 7) log pe (o(oolz))
(1t, - 2o(xx,,,2)]111 + C2. (11)
Ep(s) (on (= 2) + 7) Σ Bo
2
Moreover, we also have another weaker upper- in which C1 and C2 constants:
Eppx E Pa (wo a) (r(oo,,
log log
r 3) (2 zo) uo tx,,,
Ep(x) E 2)
E (y(#r(uop.)))
20
(12)
All (I t 2) ( 2)1
]] C2.
207
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 swward-weighted loss, smoothing the weighting of each sample towards the uniform distribution." We refer this regular- ization as KL-D since it is based only on data from the pre-traine trained model. The induced supervised training objective for KL.-D is as follows:
||:00-
fol 1, 212 20
Ep(x) ppec(rol a) (r(20: 2) + +)) Σ (x(21(001e)
(13)
We also consider another KI. regularization presented Eq. (12). This KI. regularization can be implemented by introducing additional term in the reward-weightn loss. This extra term penalizes the Le-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:
r(zo, 2| fo t, 2)1 + yllSpre (Ie, 2) fo fo (Z,, 2)
Σ E (a (no.=)
Ep(=) Bpe (xol=))
20
We summarize our supervised fine-tuning in Algorithm 2. Note that in comparison to online RL training, our supervised setting only requires pre-trained diffusion model and 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.
Intuitivelyy adjusting is similar adjusting the temperature parameter in reward-weel weighted regression [24].

---

Algorithm 2 Supervised fine-tuning with KL regularization
Reward model pre-trained diffusion model Pure. diffusion model po, regularization weight y, batch size n, the number of training samples M
Input:
Initialize Pe = Ppre
Collect M samples from pre-trained model: ( Pon(=oo/2)p(2)]We MM)
while & not converged do
Obtain i.i.d. samples (zo: 2) from training dataset D Sample 24 given To, ~ [1,
For KL-D, compute the gradient Vp (r(ao.2)+1)]lxg fa 1.l update
20
For KL-O, compute the gradient V r(mo,)) fo fo(2)...)) ( P (el f0)))
update 0
20
end while
Output: Fine-tuned diffusion model pe
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 evaluate conditional KL only using supervised samples. Moreover, online KL regularization can be seen as 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 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-ter datasets [33, 34]. For compute- compute-efficient fine-tuning,
Low-Rank Adaption (LoRA) [11], which freezes the parameters 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 IeageReward [43] which is trained on a large dataset comprised of human assessments of images. Compared to other scoring functions such CLIP [27] or BLIP [18], ImageReward has better correlation with human judgments, making 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.
5.2 Comparison of Supervised and RL Fine-tuning
We first evaluate the performance of the original and fine-tuned text-to- rext-to-image models w.r.t. specific capabilities, such 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 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 1). To systematically analyze the effects of different fine-tuning methods, we adopt straightforward setup that uses one text prompt during fine-tuning, As training text prompts,
7

---

Color
Composition
Comnt
Location
Original modeel
SPT medel
RL middel
Figure 2: Comparison of images generated by the original Stable Diffusion model, supervised fine- tuned (SFT) model, and RL fine-tuned model. Images in the same column are generated with the same random seed. Images from seen text prompts: "A green colored rabbit" (color). "A cat and dog" (composition), "Four wolves in the park" (count), and "A dog on the moon" (location).
15
Lock CSF modelt
The
(a) IewgeReward score
(b) Aesthetic score
(c) Human evalution
Figure 3: (a) ImageReward scores and (b) Aesthetic scores of three models: the original model, supervised fine-tuned (SFT) model, and RI. fine-tuned model. ImageReward and Aesthetic scores are averaged over 50 samples from each model. (c) Human preference rates between RL. model and SFT model terms for image-text alignment and image quality. The results show the mean and standard deviation averaged over eight independent human raters.
use "A green colored rabbit", "A cat and 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, 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 (c.g., adding more wolves in "Four wolves in the park"), and replace incorrect objects with target objects (c.g., replacing astronaut with dog in "A dog on the moon") compared to images from the original model. They also avoid obvious mistakes like generating rabbit with green background given the prompt "A green colored rabbit".. Of special note, we find that the mentuned models generate better images than the original model on several unseen prompts consisting of unseen objects in terms of image-text alignment (see Figure 10 in Appendix E).
As 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 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 the RL approach (and sometimes relative to the pre-trate pre-trained model), even with KL regularization which is
discuss in Section 4.3. We also evaluate the image quality of the three models
'The aesthetic predictor had been utilized to filter out the low-quality images in training data for Stable Diffusion models.

---

RL model
SFT model
Imagellewand score
Aesthetic score
1.
with
Widbous
1.4
With KL
3:2
RL model
SPT model
RL model
SPT model
(a) ImageReward and Aesthetic scores
(b) Generated images
Figure 4: Ablation study KL. regularization in both SFT and RL training, trained on a single prompt "A green colored rabbit". (a) ImageReward and Aesthetic scores averaged over 50 samples from each model. (b) Images generated by RL models and SFT models optimized with and without KL regularization. Images in the same column are generated with the same random seed.
Original model
Fine-tuned model
Figure 5: Comparison of images generated by the original model and RL fine-tuned model on text prompt "Four roses". The original model, which is trained on large-scale datasets from the web [34, 33], tends to produce whiskey-related images from "Four roses" due to the existence of a whiskey brand bearing the same name as the prompt. contrast, RL. fine-tuned model with ImageReward generates images associated with the flower "rose".
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 least as well-aligned with text prompts.
We also conduct human evaluation as follows: We collect 40 images randomly generated from each prompt, resulting
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 assess which one is better w.r.t. image-fext 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 both alignment and image quality.
5.3 The Effect of KL Regularization
To demonstrate the impact of KL regularization 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 KI. regularization "A green colored rabbit".
Figure 4(a) shows that KI. 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). 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 RI. (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 RI. 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 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 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 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.
MS-CoCo
Drawbench Original model RL model
Original model RL model
ImageReward score
0.22
0.55
0.13
0.58
Aesthetic score
5.39
5.43
5.31
5.35
Table I: ImageReward scores and Aesthetic scores from the original model, and RL fine-tuned model on multiple prompts from MS-CoCo (104 prompts) and Drawbench (183 prompts). We report the average ImageReward and Aesthetic scores across 3120 and 5490 images on MS-CoCo and Drawbench. respectively (30 images per each prompt).
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 multiple prompts requires longer training time, hyper- parameter tuning, and engineering efforts. More efficient training with a broader range of diverse and complex prompts would be interesting future direction 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-ime 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 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, is critical that the biases and failure modes of publicly available reward models are thoroughly documented.
10

---

Acknowledgements
We thank Deepak Ramachandran and Jason Baldridge for providing helpful comments and sugges- tions. Support for this research was also provided by the University of Wisconsin-Madisoo Office of the Vice Chancellor for Research and Graduate Education, NSF Award DMS-2023239, NSElInnte Partnership on Machine Learning for Wireless Networking Program under Grant No. CNS-2003 12, and FuriosaAl. OW and YD are funded by the Center for Human Compatible Artificial Intelligence.
References
[1] Bai. Yuntao, Jones, Andy, Ndousse, Kamal, Askell, Amanda, Chen, Anna. DasSarma, Nova, Drain, Dawn, Fort, Stanislav, Ganguli, Deep, Henighan, Tom, al. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint (xXiv:2204.05862 2022.
[2] Black, Kevin, Janner, Michael, Du, Yilun, Kostrikov, lla, and Levine, Sergey. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305. 13301, 2023. 3, 18
[3] 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.
[4] Dhariwal, Prafulla and Nichol, Alexander. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems, 2021. 2, 4
[5] Fan, Ying and Lee, Kangwook. Optimizing ddpm sampling with shortcut fine-tuning, arXiv preprint 2rXiv:2301.13362 2023, 3,4
[6] Feng. Weixi, He, Xuchai, Fu, Tsu-Jui, Jampani, Varun, Akula, Arjun, Narayana, Pradyumna, Basu, Sugato, Wang, Xin Eric, and Wang, William Yang. Training-frre structured diffusion guidance for compositional text-to-image synthesis. arXiv preprint arXiv:2212.05032, 2022. 1, 2
[7] 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 1rXiv:2212.10011 2022.
[8] Goodfellow, Ian, Pouget-Abadie, Jean, Mirza, Mehdi, Xu, Bing, Warde-Farles David, Ozair, Sherjil, Courville, Aaron, and Bengio, Yoshua. Generative adversarial networks. Communica- lions of the ACM, 63(11))13--144, 2020.
[9] Ho, Jonathan and Salimans, Tim. Classifier-free diffusion guidance. wwXiv:2207.12598 2022. 4
arXiv preprint
[10] Ho, Jonathan, Jain, Ajay, and Abbeel, Pieter. Denoising diffusion probabilistic models. In Advances Neural Information Processing Systems, 2020. 1, 2, 3, 15, 16
[11] 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.000 09685, 2021.
[12] 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. 1897, 2023.
[13] 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. Advances in Neural Information Processing Systems, 2023. 1,
[14] Kong, Zhifeng, Ping, Wei, Huang, Jiaji, Zhao, Kexin, and Catanzaro, Bryan. Diffwave: versatile diffusion model for audio synthesis. arXiv preprint wwXiv:2009.09761. 2020. 2
11

---

[15] Krizhevsky, Alex, Hinton, Geoffrey, et al. Learning multiple layers of features from tiny images. 2009.
[16] Lee, Kimin, Smith, Laura, and Abbeel, Pieter. Pebble: Feedback- interactive rein- forcement learning via relabeling experience and unsupervised pre-training. arXiv preprint arXiv:2106.05091, 2021. 3
[17] 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: 12192, 2023. 1, 2, 6, 18
[18] Li, Jumnan, Li, Dongxu, Xiong, Caiming, and Hoi, Steven. Blip: Bootstrapping language- image pre-training for unified vision-language understanding and generation. In International Conference on Machine Learning, 2022. 3, 7,
[19] Liu, Hao, Sferrazza, Carmelo, and Abbeel, Pieter. Chain of hindsight aligns language models with feedback. arXiv preprini arXiv:2302.00266 2023. 3
[20] 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
[21] 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.105622 2022.
[22] Liu, Ziwei, Luo, Ping, Wang, Xiaogang, and Tang, Xiaoou. Deep learning face attributes in the wild. In International Conference on Computer Vision, 2015.
[23] 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. 5.5
[24] Peters, Jan and Schaal, Stefan. Reinforcement learning by reward-weighter regression for operational space control. In Proceedings of the 24th international conference on Machine learning, pp. 745-750, 2007. 6
[25] 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 text-to-image models on a multi-task benchmark. arXiv preprint 2rx11:221..1211. 2022.
[26] Poole, Ben. Jain, Ajay. Barron, Jonathan and Mildenhall, Ben. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209. 14988, 2022. 2
[27] 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, 7,
[28] 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()):5485-5511, 2020. 1.2
[29] Ramesh, Aditya, Dhariwal, Prafulla, Nichol, Alex, Chu, Casey, and Chen, Mark. Hierarchical text-conditional image generation with clip latents. arXiv preprint 0066:22204.06//5, 2022. 1,
2
[30] Rombach, Robin, Blattmann, Andreas, Lorenz, Dominik, Esser, Patrick, and Ommer, Bjöm. High-resolution image synthesis with latent diffusion models. In Conference on Computer Vision and Pattern Recognition, 2022. 1, 2,
12

---

[31] 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.
[32] Saharia, Chitwan, Chan, William, Saxena, Saurabh, Li, Lala, Whang, Jay, Denton, Emily, Ghasemipour, Seyed Kamyar Seyed, Ayan, Burcu Karagol, Mahdavi, S Sara, Lopes, Rapha Gon- tijo, al. Photorealistic text-to-image diffusion models with deep language understanding. In Advances in Neural Information Processing Systems, 2022. 1,
[33] Schuhmann, Christoph, Vencu, Richard, Beaumont, Romain, Kaczmarezyk, 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.02144 2021. 7, 9, 10
[34] 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:22 N..v:22200004002. 2022. 8, 9, 10
[35] Schulman, John, Moritz, Philipp, Levine, Sergey, Jordan, Michael, and Abbeel, Pieter. High- dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438 2015. 16
[36] Schulman, John, Wolski, Filip, Dhariwal, Prafulla, Radford, Alec, and Klimov, Oleg. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06341 2017. 5, 17
[37] 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
[38] Song, Jiaming, Meng. Chenlin, and Ermon, Stefano. Denoising diffusion implicit models. arXiv preprint arXiv:2010.002502, 2020.
[39] Song. Yang and Ermon, Stefano. Improved techniques for training score-based generative models. In Advances in neural information processing systems, 2020. 2
[40] Song, Yang. Durkan, Conor, Murray, lain, and Ermon, Stefano. Maximum likelihood training of score-based diffusion models. In Advances in Neural Information Processing Systems, 2021.
14
[41] 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
[42] Wu, Xiaoshi, Sun, Keqiang, Zhu, Feng, Zhao, Rui, and Li, Hongsheng. Better aligning text-to-image models with human preference. arXiv preprint arXiv:2305 230.20000.44 2023. 1,
[43] 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, 3, 18
13

---

Appendix:
Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models
A Derivations
A.1 Lemma 4.1
Proof.
V [-((x, = -Ep()) Vo po po(xol=))(ao dxo
Vo ( po (p((urll)ddg r(xu, 2)dro
--p(za)
|
= --B((=)
Vo log pa (2)011)))) x r(zo, 2) po po(0rr1)) dro:r
po po 2) x r(z0,))) 2) po(worrlz) drper 11
T
= -Pp(s)
Vo Volog log
pr(xt|2)
T
= -Bp(=z) (xo
r(zo,2) 2)
V log Po (0
where the second to last equality is from the continuous assumptions of po (rrrr)))r((00 2) and Vopo (2)r)))))((
A.2
Lemma 4.2
Similar to the proof in Theorem in [40], from data processing inequality with the Markov kernel, given any we have
(L(pp((oll=))llppc(cc & (Z(p((zorll)) Ppoc (2o.T
(14)
Using the Markov property of pa and Ppre. we have
(Lpp(Zoo:1=))lppr (No.r (z)) (o( log (o(zo01(2) dro. Ppre (20rT(z)
- -f (o(oo)))2)) log
Po(xT| 2)
[t- Pa (11,,
dror
Ppre [rT 2) IT Ppre (are |1t, 2)
(Po(TO:T12)
log po(xr|2) z)
Σ log
Po (2(e--1 [2e 2)
dro.r
Ppre 2)
Ppre (I De, 2)
[po(zor11=) log po(x 122, 2)
log po p ITe, 2) (p 1x,, 2)
B pa (x
Bo (xoee
Ppre Fr
Ee, 2)
Σ (po(te)s)
log
Po Po Ir, 2)
E (xo(et sl log
po po ITe, 2) Ppre (3t-1 12t, 2)
Ppre (00 Fe,
(x.())))) [po [KL 2) 200 2))]
Taking the expectation of p(z) on both sides, we have
14

---

(LL(p)(()())) (pre ( Bp(s)
2
E p2)())))) ( 2)
P Ppre (zt (2ty 2)
(15)
A.3
The gradient of objective Eq. (8)
From
Lemma
4.2, for online fine-tuning,
need to regularize By the product rule, have
Epo(z.) (KL(po(0t [rze =) (x 2))]
Va
(xs) (mx) [L = 3
z])]
E
Epo(zi)
Vo
2]]]
(16)
T
r
+ Σ B po)(
Σ
Vo log (o (L(poxxx-1laa P
Ppre (T [21,
which treats the sum of conditional KL-divergences along the future trajectory as 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 derivation in [10], and notice that q is not dependent on (i.e., q(...-1.) g(...) )), we have
E Ep(s) E Pno (wo 2) (r(xo, 2) + log log
(r(30: +7) log po(FoTT(2) (
Ep(e) pno By(
p(xr(2)
pe) ppo E Belerr 1- 2) 7) log
g (0;
(r(zo, 2) + 7) log Po( [x,
(17)
81 12 To: z)
t>1
((((z0, 2) + 7) log Pog (1, [), =)]
Pp 1) ()(((((00 + /
+(r(zo. 2) 7) Σ E
ga (no 200 2))]
1>1
-(r(zo: 2) + (x pa ( 2)]]
where the first inequality comes from ELBO.
Let ie Tel 20) = vān- 1-ar +
Var(1-ār-1)
I-ai Dr and 3, 1-ā.-1 Br, we have:
T-as
E Bp (e [rr(oo 2) +7) Σ E Ba(ee /o..
(18)
= (eot 2) + + Eglas 120 [ 4, C,
where C is constant.
Notice that po (2oo2 mn) modeled as a discrete decoder in [10] and p(æ7(22) is fixed Gaussian; neither is trainable.
15

---

As result, we have
E By(s) (emt)) =) 7)
log (o(ol
7) Σ R (0
1 (are; (a)e no(a,, =)]]]] + C. (19)
201
Furthermore, from triangle inequality have
Eo (mo. 1> (2 (Ii, mo) 140 ( 1, 2)113
20
(20)
gl 110 mo)
Mpre ( L, =)11 (1t, - po(s1; [2)][²]
207
So another weaker upper bound is given by
Bp(s) E Pno(uoo =) (r(((( 7) log
r(zg 2) 2) Ø1, zo) no (o(
Ep(x) E pho 2)
E =x(
20² 11 C2.
(21)
(ll (Tre,
polar, e t, 2)1
207
Notice that (0
2, 2 is modeled as a discrete decoder as in [10] and p(ærl2) is 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( 2,0):
Ep(z) (.oo) (xro.r))))
r(xo.2 2)
V log (10
(22)
T
Ep(() Fooa Papa
(r(zo, 2) (i 2, 0)) Vo log 1o
where
Vi 2, := E (x [r[ <0r
(23)
So we can learn value function estimator by minimizing the objective function below for each t:
(r(so).)) - VV 0,2)) ((s,,,, l
E Epo(zox)
(24)
Since subtracting ((z,, f,, 2) will minimize the variance of the gradient estimation, it is expected that such trick can improve policy gradient training. our experiments, we find that 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 1.51, and also slightly improve 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 multi-prompt training (see learning curves with and without value learning in Figure 7).
16

---

With value learning
Without value learning
Figure 6: Text prompt: "A dog on the moon", generated from the same random seeds from the model trained with and without value learning, respectively.
A.6 Importance sampling and ratio clipping
In order to reuse old trajectory samples, we can apply the important sampling trick:
Ep(s) Blo
(r(20, 2) - V( 2, 0) Vo log log (o (r,, 2)
(r(wr: 2) 2( Po ( [e, 2)
= -Ep()) Bea (wo
Vo logpo(<t-11or
ao (3+- (ee,
(r(#0, (zz,,, V Po - [e; Paa (zt-1 [2t, z)
= -pp(s) Eeo (20
(25)
In order to constrain pa to close to Paas we can use the clipped gradient for policy gradient update similar to PPO [36]:
Po 2) Peay ( F 2)
E pp(=) pass (meer =)
Σ (r(10,2) - (z,,,, 0)) Voclip
1 1 1
(26)
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, W, sam- pling batch size m 10. For policy gradient train- ing, to perform more gradient steps without the need of re-sampling, we perform 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 in each gradient step. We train the mod- els till generating 20000 online samples, which matches the number of samples in supervised fine- tuning and results in 10K gradient steps in total. For stable training, we freeze the weights in batch norm to be exactly the same the pretrained dif- fusion model. For importance sampling, we apply the clip hyperparameter 10-4
UNTU Saaming VONKC 1C0
07-
0.4
02
0.0
Figure 7: Learning curves with and without value learning, trained on the Drawbench prompt set: Adding value learning could result in higher re- ward using less time.
Supervised training. For hyper- parameters of supervised training, we use > = 2.0 as the default option in Section 5.2, which is chosen from 7 (0.1, 1.0, 2.0, 5.0). We use learning rate 2 x 10-5 and keep other default hyper-parameters in
AdamW, which was chosen from (5 10 10-5 2 10-5,5 10-53. We use batch size
17

---

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 r+ y be non-negative, we apply filter of reward such that if + ? < O, we just set it to O. In practice, we also observe that without this filtering, supervised fine-tuning could fail during training, which indicates that our assumption 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 then used generate corresponding images using pre-trained text-to-image model. We collect binary feedback from human labelers on the image-text dataset and construct comparison pairs based this feedback. Specifically,
utilize 804, 691, 1 1154 and 13228 comparison pairs obtained from 38, 29, 41 and 295 prompts for color, location, count, and composition, respectively." 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 compares the accuracy ImageReward and baselines like CLIP [27] and BLIP [18] scores, which justifies the choice of using ImageReward for fine-tuning text-to-iman to-iiiaeee models.
Color
Count
Location
Composition
Complex
CLIP BLIP
82.46
66.11
80.60
70.98
54.83
77.36
61.61
78.14
70.98
57.79
ImageReward
86.06
73.65
79.73
79.65
65.13
Table 2: Accuracy (%) of CLIP score, BLIP score and ImageReward when predicting the assessments of human labelers.
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 al. [2] and our work explored online RL fine-tuning for improving text-to-image diffusion models. However, in our work, provide theoretical insights for optimizing the reward with policy gradient methods and provide conditions for the equivalence 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 own observations.
In our work, we do not only focus on reward optimization: inspired by failure cases (c.g.. over-saturated or pontphotoraalistic images) in supervised fine-tuning [17], we aim at finding an RL solution with KL regularization 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 tie-ttnnng rather than the supervised fine-tuning. By adopting online KL regularization, our algorithm successfully achieves high rewards and maintains image quality without ver-optimizztiin issue.
Full list text prompts available [link].
18

---

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 y 2.00,,00000 in Figure and Figure 9. We observe the followings:
KL-D (eq. (11)) can slightly increase the aesthetic score. However, with large value of it does not yield significant improvement in visual quality. Moreover, it tends to result in lower ImageReward scores overall.
In Figure 9, is evident that KL-O is more noticeably effective in addressing the failure modes in supervised fine-tuning (SFT) compared KL-D, especially when 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 Ap- pendix E.5.
5.55
RL with
without SFT RLO SPT wtb KLD
5.40
1.25
06
ImageReward score
(b) Aesthetic
Figure 8: (a) ImageReward scores and (b) Aesthetic scores of supervised fine-tuned (SFT) models with different choices regularization term and coefficient (?) on text prompt "A green colored rabbit". Each score is averaged over 50 samples from each model. KL-O and KL-D refer eq. (12) and eq. (11), respectively.
7= 0.1
1.0
2.0
5.0
SPT with KL-O SPT with KI-D
RL with KL
SFT w.o KL
Figure 9: (Left) Sample from the RL model with KL regularization and sample from the SFT model without KL regularization. (Right) Samples from supervised fine-tuned models with different choices of regularization term and coefficient (7) on "A green colored rabbit". As we increase 7 in the SFT case, there is tradeoff between alignment with the prompt and image quality (i.e. oversaturation). the other hand, the RL with KL sample able retain both alignment and image quality.
E.2 Images from unseen text prompts
Figure shows image samples from the original model, SFT model and RI. fine-tuned models on unseen text prompts.
19

---

Color
Composition
Comnt
Location
Original modell
SPT medel
RL middel
Figure 10: Comparison of images generated by the original Stable Diffusion model, supervised fine-tuned (SFT) model, and RL fine-tuned model. Images in the same column are generated with the same random seed. Images from unseen text prompts: "A green colored cat" (color), "A cat and cup" (composition), "Four birds in the park" (count), and "A lion the moon" (location).
E.3 Multi-prompt Training
Here report more details for multi-promptt 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 3 0.001; 3) use extra value learning for variance reduction. For value learning, we use larger learning rate 10 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 complex prompt and compare the images generated by the original model and the supervised fine- tuned model. We can observe that online training encourages the models generate images with a different style of painting and more fine mme-grained details, which is generally hard to achieve by supervised fine-tuning only (with KL-O regularization, 7 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 and Figure 19, with the same configuration as Section 5.2 for both RI. and SFT.
20

---

Original model
RL model
WOR
(a)
(b)
Original model
model
(c)
(d)
Figure 11: Sample images generated from prompts: (a) "A chair in the corner on boat"; (b) "A dog is laying with remote controller": (c) "A cube made of brick"; (d) "A red book and a yellow vase". from the original model and RL model respectively. Images in the same column are generated with the same random seed.
Original model
SFT model
RL model
Figure 12: Text prompt: "oil portrait of archie andrews holding picture of among intricate, elegant, highly detailed, lighting, painting, artstation, smooth, illustration, art by greg rutowski and alphonse mucha".
21

---

0.1
1.0
= 2.0
5.0
SFT with KL-O SPT with KI-D
RL with KL
SFT KL
(a) "Four wolves in the park"
0.1
1=
2.0
5.0
SFFT with OLT/ SET with KI.D
RL with KL
SFT w.o KL
(b) "A cat a dog"
0.1
7 1.0
7 2.0
7= 5.0
SET with KL-O SET with KI-D
RL with KL
SFT KL
(c) "A dog on the moon"
Figure 13: Samples from supervised fine-tuned models with different KL regularization and KL coefficients.
22

---

(a) Original model
(b) SFT model
(c) RL model
Figure 14: Randomly generated samples from (a) the original Stable Diffusion model, (b) supervised fine-tuned (SFT) model and (c) RL fine-tuned model.
23

---

(a) Original model
(b) SFT model
(c) RL model
Figure 15: Randomly generated samples from (a) the original Stable Diffusion model, (b) supervised fine-tuned (SFT) model and (c) RL fine-tuned model.
24

---

(a) Original model
(b) SFT model
(c) RL model
Figure 16: Randomly generated samples from (a) the original Stable Diffusion model, (b) supervised fine-tuned (SFT) model and (c) RL fine-tuned model.
25

---

Onl
(a) Original model
(b) SFT model
(c) RL model
Figure 17: Randomly generated samples from (a) the original Stable Diffusion model, (b) supervised fine-tuned (SFT) model and (c) RL fine-tuned model.
26

---

(a) RL model without KL regularization
(b) RL model with KL regularization
Figure 18: Randomly generated samples from RL fine-tuned models (a) without and (b) with KL
regularization.
27

---

(a) SFT model without KL regularization
(b) SFT model with KL regularization
Figure 19: Randomly generated samples from supervised fine-tuned (SFT) models (a) without and (b) with KL regularization.
28


Reviews:
[{"review_id": "NcwueVOh9h", "reviewer": "Reviewer_fxhn", "summary": "The 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.", "questions": " - 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? ", "limitations": "The main limitation of the paper is the restriction to a single prompt which is openly discussed by the authors.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": " - 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. ", "weaknesses": " - 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable comments. We found them extremely helpful in improving our draft. We address each comment in detail below.\n\n**Comment 1. Fine-tuning on a single or few prompts at a time**\n\nThanks for your suggestion! We added experiments that contain training with a large variety of prompts in global response pdf. The results show that **our method can improve the average reward on 104 MS-CoCo prompts and 183 Drawbench prompts**. For details, please check the overall author rebuttal and global response pdf (**Table 3 and Fig 1 in the response pdf**).\n\n**Comment 2. Human evaluation**\n\nThanks for your suggestion! We conducted **human evaluation as a supplement of Fig 3**  in global response pdf. In our human evaluation, **the RL model consistently outperforms the SFT model in terms of both alignment and image quality and also outperforms the original model in terms of alignment with comparable image quality**. For details, please check the overall author rebuttal and global response pdf (**Table 1&2 in the response pdf**).\n\n**Comment 3. LoRA fine-tuning vs. full fine-tuning**\n\nIn fact, we first tried fine-tuning with all parameters which also works well, but we finally switched to LoRA fine-tuning for more efficient training.\n\n**Comment 4. Multiple colors & text generation**\n\nThis is an interesting question! We believe the answer to the question “whether RLHF helps with such tasks” depends on the capability of the reward function. In the new experiments with multiple prompts, we find that **RL fine-tuning leads to enhanced reward scores (from 0.76 to 1.57) on text prompts with multiple colors (e.g., yellow vase and red book) and produces better-aligned images compared to the original model (see Fig 3(d) in the response PDF).** Regarding text rendering, although the reward scores show improvement (from 0.45 to 1.11), there is no substantial difference in the generated images. However, we expect that with a more sophisticated reward function that is capable of capturing finer details, RLHF can improve text rendering.\n\nThese results support our claim that RLHF remains effective even in highly challenging scenarios, encompassing multiple colors and text rendering.\n\n**Comment 5. Minor issues like the adding min in eq(8)**\n\nThanks for pointing them out! We will fix them in the draft.\n"}, {"review_id": "34Ef7kRrFq", "reviewer": "Reviewer_YMPp", "summary": "This 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).", "questions": "* 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.", "limitations": "The limitations could be extended to the societal risks of generative AI, and also mention the fact that online RL requires more robust reward models.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* 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.", "weaknesses": "* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable comments. We found them extremely helpful in improving our draft. We address each comment in detail, one by one below.\n\n**Comment 1. Ablation study to verify the effect of exploring online samples** \n\nResponse 1: This is an interesting question! However, we notice that **online/offline is not the only difference between RL & SFT training**: noticing that weighted ELBO (the RHS of Eq. (9)) is equivalent to a weighted score matching loss in [35], SFT is learning the score function with high rewards from a fixed denoising process, where RL model explores the trajectories with high rewards. As a result, even if we use SFT to train on images generated by the RL model, there are still other differences between such SFT and the RL training. \n\nIn fact, **RL & SFT are two very different methods with many different aspects, and the first difference discussed in Section 4.3 is only one aspect**, so it could be hard to just select one aspect and conduct a clean ablation study due to other differences. One possible approach is to conduct offline RL by collecting trajectories generated from the original model and comparing it to the online RL, but it would not be a comparison between SFT and RL. We also expect that such an offline method would be outperformed by the online method. This is because the performance of the offline method is upper-bounded by the data coverage, and thus, online exploration would improve the model performance when the offline data coverage is not good. Finally, we want to mention that even if we intuitively treat SFT as an offline method, it is not an offline RL method, which is also discussed in Response 2.\n\n**Comment 2. Question about KL-D for SFT:**\n\nResponse 2: Thank you for the question, but we think that your claim “the shift in the reward would not change anything in SFT” is not accurate. We would like to highlight that **there is no policy gradient in SFT**, and **subtracting a baseline without changing the expectation of the gradient estimation only applies to policy gradient methods**. Actually, the solution of the weighted ELBO (the RHS of Eq. (9), equivalent to the score matching loss in [35], which is an L2-loss between the actual score and the score function model) could change when $\\gamma$ changes: under the assumption that reward is non-negative, the lowest possible reward is 0. If there are images with 0 reward in the SFT dataset, we will give 0 weight to the score matching loss of such images. If we add a positive shift, the reward will be non-zero and it will give a positive weight on learning the score function of such images. \n\nEmpirically, as we can observe in Fig 7(a), increasing $\\gamma$ will result in a lower reward for SFT training. Moreover, since SFT is not a policy gradient method, it will not benefit from baseline/value function tricks. For more results in KL-D, we present our full ablations result in Section E in the appendix.\n\n**Comment 3. The robustness of the reward model:**\n\nResponse 3: To clarify, the third difference mentioned in Section 4.3 is under the assumption that we have a good reward model. Our point is that if we have a reward model which generalizes well or is even perfect, **SFT will not fully utilize that model by only evaluating the reward on a fixed dataset, whereas RL has the capability to leverage it more comprehensively**. \n\nAlso, we consider the case when the reward model could be imperfect, which is **why we add KL-regularization to avoid over-optimization and out-of-distribution reward evaluation**.\n\n**Comment 4. More experiments with multiple prompts and human evaluations:** \n\nResponse 4: Thanks for your feedback. Please check the overall author rebuttal and global response pdf. In these documents, we have included experiments that contain training with a large variety of prompts, showcasing that **our method can also improve the average reward on 104 MS-CoCo prompts and 183 Drawbench prompts (see Table 3 and Fig 1 in the response pdf).** Furthermore, we have included **human evaluation results on a single prompt (see Table 1&2 in the response pdf)**, which demonstrate that our RL model outperforms the SFT model in terms of both alignment and image quality. Moreover, our RL model outperforms the original model in terms of alignment while maintaining comparable image quality.\n\n**Comment 5. Elaborate on the difference between KL & reverse KL:** \n\nResponse 5: The key difference is which distribution is used to evaluate the difference in the log probability. In KL for SFT, we evaluate using the offline distribution, which is the training data; in KL for RL, we evaluate using online generated trajectories.\n\n**Comment 6. Minor issue about the title**\n\nResponse 6: Thanks for pointing it out! We will address this in the final draft.\n\n**Comment 7. Refined online RL optimization like PPO**\n\nResponse 7:  As mentioned in Appendix B, we tried a small learning rate and clipping the gradient norm to stabilize the training, which is similar to PPO. We also tried importance sampling and clipping the ratio, where we did not find much improvement in the learning process. We will clarify this in the final draft.\n"}, {"review_id": "dLs36uUHk0", "reviewer": "Reviewer_VuBv", "summary": "This 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.", "questions": "How 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? ", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable comments. We found them extremely helpful in improving our draft. We address each comment in detail, one by one below.\n\n**Comment 1. Novelty is limited, similar ideas have already been proposed in [1]**\n\nWe respectfully disagree with the reviewer’s assessment. In fact, it is unfair to list the similarity to DDPO [1] as a weakness of our work, because 1) our work and DDPO [1] were developed independently, and 2) DDPO is an unpublished work, which appeared online only after the NeurIPS-2023 submission deadline. Moreover, we study the effect of KL-regularization (whose importance in fine-tuning foundation models has been well-established in the literature) in both online and supervised learning approaches, supported by theoretical justifications and experimental results, which has not been explored in the DDPO paper [1].  \n\n\n[1] Training diffusion models with reinforcement learning\n\n**Comment 2. Lacking experiments and fidelity score** \n\nThanks for your feedback. Please check the overall author rebuttal and global response pdf. In these documents, we have included experiments that contain training with a large variety of prompts, showcasing that **our method can also improve the average reward on 104 MS-CoCo prompts and 183 Drawbench prompts (see Table 3 and Fig 1 in the response pdf).** Furthermore, we have included **human evaluation results on a single prompt (see Tables 1&2 in the response pdf)**, which show that our RL model outperforms the SFT model in terms of both alignment and image quality. Moreover, our RL model outperforms the original model in terms of alignment, while maintaining comparable image quality.\n\nAs mentioned in the paper, we use the aesthetic predictor from LAION [2] to measure the image quality (fidelity), since we select some synthetic and challenging prompts for which we do not have ground truth images to compute FID. \n\n[2] Schuhmann et al. Laion-5b: An open large-scale dataset for training next generation image-text models. arXiv preprint arXiv:2210.08402, 2022.\n\n**Comment 3. The cost for training:** \n\n**We reported the cost for RL training in Appendix B**. It is true that RL training requires a relatively smaller learning rate, making it computationally more expensive. However, the increase in computational requirements is not significantly large as we also need a smaller learning rate for supervised fine-tuning to achieve higher image quality.\n\n**Comment 4. ImageReward**\n\nYes, as mentioned in our paper, we use the same way to evaluate ImageReward and aesthetic scores as in the ImageReward paper. Also, it has been demonstrated that ImageReward is more correlated with real human decisions than CLIP/BLIP. As a result, we believe that the ImageReward score offers a more meaningful and relevant signal for evaluation purposes. \n\n**Comment 5. LoRA fine-tuning**\n\nWe use LoRA fine-tuning for both RL and SFT. Also, we applied LoRA to the cross-attention layers.\n"}, {"review_id": "PhMaso1kE6", "reviewer": "Reviewer_gRSW", "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.", "questions": "see weaknesses", "limitations": "limitations are briefly discussed, especially limitations mentioned in weaknesses could be discussed in more detail", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- an interesting and novel idea\n- simple method and easy algorithm\n- qualitatively pleasing results", "weaknesses": "- considering only 4 prompts for experimental evaluation is not sufficient\n- considering only a single prompt for bias evaluation is not sufficient\n- 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\n- investigation of generalisation to prompts that are out of distribution should be conducted\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable comments. We found them extremely helpful in improving our draft. We address each comment in detail, one by one below.\n\n**Comment 1. Lacking experiments, training with more prompts, not scale to a larger set of training prompts** \n\nThanks for your feedback. Please check our new results in the overall author rebuttal and global response pdf. In these documents, we have included experiments that contain training with a large variety of prompts, showcasing that **our method can also improve the average reward on much larger training sets with 104 MS-CoCo prompts and 183 Drawbench prompts (see Table 3 and Fig 1 in the response pdf).**  \n\n\n**Comment 2. Generalization of out-of-distribution prompts:**\n\n**We conducted tests on unseen prompts, which consist of unseen objects, in Fig. 2 (b).** Notably, the fine-tuned models produced better images than the original model for these unseen text prompts. However, for better generalization, we need to train on a larger and more diverse set of text prompts.\n\n**We hope our response resolves all the concerns in your review, and please feel free to let us know if you have any other questions in the discussion period.**\n"}, {"review_id": "Usvn6NckF9", "reviewer": "Reviewer_gCGU", "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.", "questions": "Please refer to the weakness part.", "limitations": "The authors discussed some limitations in the last part of the paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to understand.\n\n2. The comparison and the potential advantages compared to related works are clearly discussed.\n\n3. 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.\n\n2. 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.\n\n3. 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]?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable comments. We found them extremely helpful in improving our draft. We address each comment in detail, one by one below.\n\n**Comment 1. The connection between Section 4.3 and experiments:**\n\n\nIt is hard to make a one-to-one mapping between each claim in Section 4.3 and the experimental results since Section 4.3 provides the overall combination of these differences as a comprehensive explanation of the benefits of using RL over SFT. Since RL and SFT are two very different methods with several differences, it could be hard to just select one difference and conduct a clean ablation study, due to the existence of other differences. Also, we do not claim that the superiority of RL over SFT is due to a single factor (or difference between them). \n\n**Comment 2. The necessity for online optimization:**\n\nFirst, **we propose online fine-tuning to better optimize ground truth rewards, and not just for solving the image quality degradation in fine-tuning**. \n\nSecond, the over-saturated problem is just one aspect of downgraded image quality. There are other aspects that cannot be solved by adjusting the guidance weight (also note that only reducing the guidance weight could suffer from lower text-image alignment). In fact, fine-tuning with a limited distribution of images itself could suffer from worse image quality, so tricks like reconstruction loss have been proposed to solve such issues. Our KL regularization is performing a similar role.\n\nFinally, adjusting the classifier guidance scale did not make much difference for the SFT model in our experiments.\n\n**Comment 3. Comparison with [17]**\n\n**We did a comparison with [17] in our paper**, which corresponds to the SFT method without KL regularization. We added some theoretical justification and also tried some extra KL regularization methods on top of it. \n"}]

Author rebuttal:
Overall author rebuttal:

We thank all reviewers for their thoughtful comments. We greatly appreciate all the reviewers' acknowledgment that our method is **empirically effective with solid theories**. To address common concerns about scaling up the training prompts and human evaluation, we have add new experiments and evaluations:

**1. Training with a large variety of prompts:**

We conduct online RL training with 104 MS-CoCo prompts and 183 Drawbench prompts, respectively (the prompts are randomly sampled during training, and the full prompt dataset will be made public). We use the same configuration in Appendix B but with a longer training time till it generates 10K online samples (50K gradient steps). 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 3 with sample images in Fig 1 in the author rebuttal pdf, showing that RL training can also significantly improve the ImageReward score while maintaining a high aesthetic score with much larger sets of training prompts**.

**2. Human evaluation:**

As a supplementary evaluation to Fig. 3, we conduct extra human evaluation. We gather 40 images that are randomly generated from each prompt (“green rabbit”, “cat and dog”, “dog on the moon” and “four wolves”), resulting in a total of 160 images from each model (RL, SFT and the original model). Given two (anonymized) sets of four images from the same random seeds, one from RL fine-tuned model and one from the original model (or SFT model), we ask human raters to assess which one is better w.r.t. image-text alignment and fidelity (i.e., image quality). Each query is evaluated by 8 independent human raters and we report the average win/lose rate from RL&SFT comparison and RL&original-model comparison. **The results are presented in Table 1 & 2 in the author rebuttal pdf, showing that the RL model consistently outperforms the SFT model on both alignment and image quality and also outperforms the original model in the sense of alignment with comparable image quality**.

We will also include all the results in the final draft.

Meta-review: This paper proposes a reinforcement learning from human feedback (RLHF) method for diffusion models on text-to-image generation. 
From the method innovation perspective, RLHF for diffusion models is a novel idea; it is an "obvious" but meaningful next step for text-to-image generation given the success of RLHF for large language models (LLMs). The reviewers acknowledged the significance of this work and praised the use of KL divergence. They unanimously lean toward acceptance of this work (most are borderline accept though). The AC carefully read the rebuttal and author-reviewer discussions, and feels that most concerns are addressed in the rebuttal. One reviewer mentioned another concurrent work that was released after the NeurIPS deadline, which, however, should not invalidate this paper's novelty and contributions at all. Overall, this is a meaningful paper with novel method and promising results, despite some limitations. The AC tend to accept this work. The authors are encouraged to include the additional results and human evaluation into the revision.
Decision: Accept (poster)