SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {\mathrm {p r e}} (x _ {0} | z)} \left[ \sum_ {t > 1} \mathbb {E} _ {q (x _ {t} | x _ {0}, z)} \left[ \frac {r (x _ {0} , z) \| x _ {0} - f _ {\theta} (x _ {t} , t , z) \| ^ {2} + \gamma \| f _ {\mathrm {p r e}} (x _ {t} , t , z) - f _ {\theta} (x _ {t} , t , z) \| ^ {2}}{2 \sigma_ {t} ^ {2}} \right] \right].
$$
We summarize our supervised fine-tuning in Algorithm 2. Note that in comparison to online RL training, our supervised setting only requires a pre-trained diffusion model and no extra/new datasets.
4.3 Online RL vs. Supervised Fine-tuning
We outline key differences between online RL and supervised fine-tuning:
- Online versus offline distribution. We first contrast their objectives. The online RL objective is to find a new image distribution that maximizes expected reward—this can have much different support than the pre-trained distribution. In supervised fine-tuning, the objective only encourages the model to "imitate" good examples from the supervised dataset, which always lies in the support of the pre-trained distribution.
${}^{2}$  Intuitively,adjusting  $\gamma$  is similar to adjusting the temperature parameter in reward-weighted regression [24].
Algorithm 2 Supervised fine-tuning with KL regularization
Input: Reward model  $r$ , pre-trained diffusion model  $p_{\mathrm{pre}}$ , diffusion model  $p_{\theta}$ , regularization weight  $\gamma$ , batch size  $n$ , the number of training samples  $M$   
Initialize  $p_{\theta} = p_{\mathrm{pre}}$   
Collect  $M$  samples from pre-trained model:  $\mathcal{D} = \{(x_0, z) \sim p_{\mathrm{pre}}(x_0 | z) p(z) | \forall i \in \{1, \dots, M\}\}$   
while  $\theta$  not converged  
Obtain  $n$  i.i.d. samples  $\{x_0, z\}$  from training dataset  $\mathcal{D}$   
Sample  $x_t$  given  $x_0, t \sim [1, T]$   
For KL-D, compute the gradient  $\nabla_{\theta} \frac{(r(x_0, z) + \gamma) \| x_0 - f_{\theta}(x_t, t, z) \|^2}{2\sigma_t^2}$ , update  $\theta$   
For KL-O, compute the gradient  $\nabla_{\theta} \frac{r(x_0, z) \| x_0 - f_{\theta}(x_t, t, z) \|^2 + \gamma \| f_{\mathrm{pre}}(x_t, t, z) - f_{\theta}(x_t, t, z) \|^2}{2\sigma_t^2}$ , update  $\theta$   
end while  
Output: Fine-tuned diffusion model  $p_{\theta}$
- Different KL regularization. The methods also differ in their use of KL regularization. For online training, we evaluate conditional KL in each step using online samples, while for supervised training we evaluate conditional KL only using supervised samples. Moreover, online KL regularization can be seen as an extra reward function to encourage small divergence w.r.t. the pre-trained model for online optimization, while supervised KL induces an extra shift in the original reward for supervised training as shown in Lemma 4.3.
- Different evaluation of the reward model. Online RL fine-tuning evaluates the reward model using the updated distribution, while supervised fine-tuning evaluates the reward on the fixed pre-training data distribution. As a consequence, our online RL optimization should derive greater benefit from the generalization ability of the reward model.
For these reasons, we expect online RL fine-tuning and supervised fine-tuning to generate rather different behaviors. Specifically, online fine-tuning should be better at maximizing the combined reward (human reward and implicit KL reward) than the supervised approach (similar to the difference between online learning and weighted behavior cloning in reinforcement learning).
We now describe a set of experiments designed to test the efficacy of different fine-tuning methods.
As our baseline generative model, we use Stable Diffusion v1.5 [30], which has been pre-trained on large image-text datasets [33, 34]. For compute-efficient fine-tuning, we use Low-Rank Adaption (LoRA) [11], which freezes the parameters of the pre-trained model and introduces low-rank trainable weights. We apply LoRA to the UNet [31] module and only update the added weights. For the reward model, we use ImageReward [43] which is trained on a large dataset comprised of human assessments of images. Compared to other scoring functions such as CLIP [27] or BLIP [18], ImageReward has a better correlation with human judgments, making it the preferred choice for fine-tuning our baseline diffusion model (see Appendix C for further justification). Further experimental details (e.g., model architectures, final hyper-parameters) are provided in Appendix B. We also provide more samples for qualitative comparison in Appendix E.6.
5.2 Comparison of Supervised and RL Fine-tuning
We first evaluate the performance of the original and fine-tuned text-to-image models w.r.t. specific capabilities, such as generating objects with specific colors, counts, or locations; and composing multiple objects (or composition). For both RL and SFT training we include KL regularization (see hyperparameters in Appendix B). For SFT, we adopt KL-O since we found it is more effective than KL-D in improving the visual quality of the SFT model (see the comparison between KL-O and KL-D in Appendix E.1). To systematically analyze the effects of different fine-tuning methods, we adopt a straightforward setup that uses one text prompt during fine-tuning. As training text prompts,
Figure 3: (a) ImageReward scores and (b) Aesthetic scores of three models: the original model, supervised fine-tuned (SFT) model, and RL 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 in terms for image-text alignment and image quality. The results show the mean and standard deviation averaged over eight independent human raters.
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 a dog" (composition), "Four wolves in the park" (count), and "A dog on the moon" (location).
we use "A green colored rabbit", "A cat and a dog", "Four wolves in the park", and "A dog on the moon". These test the models' ability to handle prompts involving color, composition, counting and location, respectively. For supervised fine-tuning, we use 20K images generated by the original model, which is the same number of (online) images used by RL fine-tuning.
Figure 3(a) compares ImageReward scores of images generated by the different models (with the same random seed). We see that both SFT and RL fine-tuning improve the ImageReward scores on the training text prompt. This implies the fine-tuned models can generate images that are better aligned with the input text prompts than the original model because ImageReward is trained on human feedback datasets to evaluate image-text alignment. Figure 2 indeed shows that fine-tuned models add objects to match the number (e.g., adding more wolves in "Four wolves in the park"), and replace incorrect objects with target objects (e.g., replacing an astronaut with a dog in "A dog on the moon") compared to images from the original model. They also avoid obvious mistakes like generating a rabbit with a green background given the prompt "A green colored rabbit". Of special note, we find that the fine-tuned models generate better images than the original model on several unseen text prompts consisting of unseen objects in terms of image-text alignment (see Figure 10 in Appendix E).
As we can observe in Figure 3(a), RL models enjoy higher ImageReward than SFT models when evaluated with the same training prompt in different categories respectively, due to the benefit of online training as we discuss in Section 4.3. We also evaluate the image quality of the three models using the aesthetic predictor [34], which is trained to predict the aesthetic aspects of generated images. Figure 3(b) shows that supervised fine-tuning degrades image quality relative to the RL approach (and sometimes relative to the pre-trained model), even with KL regularization which is
3The aesthetic predictor had been utilized to filter out the low-quality images in training data for Stable Diffusion models.
(a) ImageReward and Aesthetic scores
(b) Generated images
Figure 4: Ablation study of KL regularization in both SFT and RL training, trained on a single prompt "A green colored rabbit". (a) ImageReward and Aesthetic scores are 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. In 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 at least as well-aligned with text prompts.
We also conduct human evaluation as follows: We collect 40 images randomly generated from each prompt, resulting in a total of 160 images for each model. Given two (anonymized) sets of four images from the same random seeds, one from RL fine-tuned model and one from the SFT model, we ask human raters to assess which one is better w.r.t. image-text alignment and image quality. Each query is evaluated by 8 independent raters and we report the average win/lose rate in Figure 3(c). The RL model consistently outperforms the SFT model on both alignment and image quality.
5.3 The Effect of KL Regularization
To demonstrate the impact of KL regularization on both supervised fine-tuning (SFT) and online RL fine-tuning, we conduct an ablation study specifically fine-tuning the pre-trained model with and without KL regularization on "A green colored rabbit".
Figure 4(a) shows that KL regularization in online RL is effective in attaining both high reward and aesthetic scores. We observe that the RL model without KL regularization can generate lower-quality images (e.g., over-saturated colors and unnatural shapes) as shown in Figure 4(b). In the case of SFT with KL-O (where we use the same configuration as Section 5.2), we find that the KL regularization can mitigate some failure modes of SFT without KL and improve aesthetic scores, but generally suffers from lower ImageReward. We expect that this difference in the impact of KL regularization is due to the different nature of online and offline training—KL regularization is only applied to fixed samples in the case of SFT while KL regularization is applied to new samples per each update in the case of online RL (see Section 4.3 for related discussions).
5.4 Reducing Bias in the Pre-trained Model
To see the benefits of optimizing for ImageReward, we explore the effect of RL fine-tuning in reducing bias in the pre-trained model. Because the original Stable Diffusion model is trained on
large-scale datasets extracted from the web [33, 34], it can encode some biases in the training data. As shown in Figure 5, we see that the original model tends to produce whiskey-related images given the prompt "Four roses" (which happens to be the brand name of a whiskey), which is not aligned with users' intention in general. To verify whether maximizing a reward model derived from human feedback can mitigate this issue, we fine-tune the original model on the "Four roses" prompt using RL fine-tuning. Figure 5 shows that the fine-tuned model generates images with roses (the flower) because the ImageReward score is low on the biased images (ImageReward score is increased to 1.12 from -0.52 after fine-tuning). This shows the clear benefits of learning from human feedback to better align existing text-to-image models with human intentions.
5.5 Fine-tuning on Multiple Prompts
We further verify the effectiveness of the proposed techniques for fine-tuning text-to-image models on multiple prompts simultaneously. We conduct online RL training with 104 MS-CoCo prompts and 183 Drawbench prompts, respectively (the prompts are randomly sampled during training). Detailed configurations are provided in Appendix E. Specifically, we also learn a value function for variance reduction in policy gradient which shows benefit in further improving the final reward (see Appendix A.5 for details.) We report both ImageReward and the aesthetic score of the original and the RL fine-tuned models. For evaluation, we generate 30 images from each prompt and report the average scores of all images. The evaluation result is reported in Table 1 with sample images in Figure 11 in Appendix E, showing that RL training can also significantly improve the ImageReward score while maintaining a high aesthetic score with much larger sets of training prompts.
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 1: 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).
In this work, we propose DPOK, an algorithm to fine-tune a text-to-image diffusion model using policy gradient with KL regularization. We show that online RL fine-tuning outperforms simple supervised fine-tuning in improving the model's performance. Also, we conduct an analysis of KL regularization for both methods and discuss the key differences between RL fine-tuning and supervised fine-tuning. We believe our work demonstrates the potential of reinforcement learning from human feedback in improving text-to-image diffusion models.
Limitations, future directions. Several limitations of our work suggest interesting future directions: (a) As discussed in Section 5.5, fine-tuning on multiple prompts requires longer training time, hyperparameter tuning, and engineering efforts. More efficient training with a broader range of diverse and complex prompts would be an interesting future direction to explore. (b) Exploring advanced policy gradient methods can be useful for further performance improvement. Investigating the applicability and benefits of these methods could lead to improvements in the fine-tuning process.
Broader Impacts Text-to-image models can offer societal benefits across fields such as art and entertainment, but they also carry the potential for misuse. Our research allows users to finetune these models towards arbitrary reward functions. This process can be socially beneficial or harmful depending on the reward function used. Users could train models to produce less biased or offensive imagery, but they could also train them to produce misinformation or deep fakes. Since many users may follow our approach of using open-sourced reward functions for this finetuning, it is critical that the biases and failure modes of publicly available reward models are thoroughly documented.
We thank Deepak Ramachandran and Jason Baldridge for providing helpful comments and suggestions. Support for this research was also provided by the University of Wisconsin-Madison Office of the Vice Chancellor for Research and Graduate Education, NSF Award DMS-2023239, NSF/Intel Partnership on Machine Learning for Wireless Networking Program under Grant No. CNS-2003129, and FuriosaAI. OW and YD are funded by the Center for Human Compatible Artificial Intelligence.
[1] Bai, Yuntao, Jones, Andy, Ndousse, Kamal, Askell, Amanda, Chen, Anna, DasSarma, Nova, Drain, Dawn, Fort, Stanislav, Ganguli, Deep, Henighan, Tom, et al. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862, 2022. 3
[2] Black, Kevin, Janner, Michael, Du, Yilun, Kostrikov, Ilya, and Levine, Sergey. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13301, 2023. 3, 18
[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. 2
[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 arXiv:2301.13362, 2023. 3, 4
[6] Feng, Weixi, He, Xuehai, Fu, Tsu-Jui, Jampani, Varun, Akula, Arjun, Narayana, Pradyumna, Basu, Sugato, Wang, Xin Eric, and Wang, William Yang. Training-free structured diffusion guidance for compositional text-to-image synthesis. arXiv preprint arXiv:2212.05032, 2022. 1, 2
[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 arXiv:2212.10015, 2022. 1
[8] Goodfellow, Ian, Pouget-Abadie, Jean, Mirza, Mehdi, Xu, Bing, Warde-Farley, David, Ozair, Sherjil, Courville, Aaron, and Bengio, Yoshua. Generative adversarial networks. Communications of the ACM, 63(11):139–144, 2020. 3
[9] Ho, Jonathan and Salimans, Tim. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022. 4
[10] Ho, Jonathan, Jain, Ajay, and Abbeel, Pieter. Denoising diffusion probabilistic models. In Advances in 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.09685, 2021. 7
[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.11897, 2023. 1
[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. In Advances in Neural Information Processing Systems, 2023. 1, 3
[14] Kong, Zhifeng, Ping, Wei, Huang, Jiaji, Zhao, Kexin, and Catanzaro, Bryan. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint arXiv:2009.09761, 2020. 2
[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. 7
[32] Sahara, Chitwan, Chan, William, Saxena, Saurabh, Li, Lala, Whang, Jay, Denton, Emily, Ghasemipour, Seyed Kamyar Seyed, Ayan, Burcu Karagol, Mahdavi, S Sara, Lopes, Rapha Gontijo, et al. Photorealistic text-to-image diffusion models with deep language understanding. In Advances in Neural Information Processing Systems, 2022. 1, 2
[33] Schuhmann, Christoph, Vencu, Richard, Beaumont, Romain, Kaczmarczyk, Robert, Mullis, Clayton, Katta, Aarush, Coombes, Theo, Jitsev, Jenia, and Komatsuzaki, Aran. Laion-400m: Open dataset of clip-filtered 400 million image-text pairs. arXiv preprint arXiv:2111.02114, 2021. 7, 9, 10
[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:2210.08402, 2022. 7, 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.06347, 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.02502, 2020. 1
[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, Iain, 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:2303.14420, 2023. 1, 3
[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, 2, 3, 7, 18
Appendix:
A.1 Lemma 4.1
$$
\begin{array}{l} \nabla_ {\theta} \mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {\theta} (x _ {0} | z)} \left[ - r (x _ {0}, z) \right] = - \mathbb {E} _ {p (z)} \left[ \nabla_ {\theta} \int p _ {\theta} (x _ {0} | z) r (x _ {0}, z) d x _ {0} \right] \\ = - \mathbb {E} _ {p (z)} \left[ \nabla_ {\theta} \int \left(\int p _ {\theta} \left(x _ {0: T} | z\right) d x _ {1: T}\right) r \left(x _ {0}, z\right) d x _ {0} \right] \\ = - \mathbb {E} _ {p (z)} \left[ \int \nabla_ {\theta} \log p _ {\theta} \left(x _ {0: T} | z\right) \times r \left(x _ {0}, z\right) \times p _ {\theta} \left(x _ {0: T} | z\right) d x _ {0: T} \right] \\ = - \mathbb {E} _ {p (z)} \left[ \int \nabla_ {\theta} \log \left(p _ {T} \left(x _ {T} | z\right) \prod_ {t = 1} ^ {T} p _ {\theta} \left(x _ {t - 1} | x _ {t}, z\right)\right) \times r \left(x _ {0}, z\right) \times p _ {\theta} \left(x _ {0: T} | z\right) d x _ {0: T} \right] \\ = - \mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {\theta} (x _ {0: T} | z)} \left[ r (x _ {0}, z) \sum_ {t = 1} ^ {T} \nabla_ {\theta} \log p _ {\theta} (x _ {t - 1} | x _ {t}, z) \right], \\ \end{array}
$$
where the second to last equality is from the continuous assumptions of  $p_{\theta}(x_{0:T}|z)r(x_0,z)$  and  $\nabla_{\theta}p_{\theta}(x_{0:T}|z)r(x_0,z)$ .
A.2 Lemma 4.2
Similar to the proof in Theorem 1 in [40], from data processing inequality with the Markov kernel, given any  $z$  we have
$$
\mathrm {K L} \left(p _ {\theta} \left(x _ {0} | z\right)\right) \| p _ {\text {p r e}} \left(x _ {0} | z\right)) \leqslant \mathrm {K L} \left(p _ {\theta} \left(x _ {0: T} | z\right) \| p _ {\text {p r e}} \left(x _ {0: T} | z\right)\right). \tag {14}
$$
Using the Markov property of  $p_{\theta}$  and  $p_{\mathrm{pre}}$  we have