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.
Paper
Similar papers
Peer review
Summary
in this paper, the authors propose to improve the pretrained text-to-image diffusion model with human feedback in an online-reinforcement learning manner. The problem formulation, the differences between the online update and the supervised finetuning are clearly stated. Experiments show that the proposed method effectively improves the performance of the text-to-image generation ability.
Strengths
1. The paper is well-written and easy to understand. 2. The comparison and the potential advantages compared to related works are clearly discussed. 3. Experiments show that the proposed method improves the performance of the pretrained text-to-image model, also outperforms the supervised human-feedback finetuning baseline.
Weaknesses
1. Some claims and the results are connected weakly. The claims in Sec.4.3 should be better aligned with the experimental results. However, its current version makes me hard to find the direct mapping. 2. The title of this paper seems not to reflect the main contribution. In my opinion, the main contribution of this work lies in the online manner, since using reinforcement learning to fine-tune the text-to-image model has been explored before. It needs a better title. 3. There are some experimental problems. In my experience, the over-saturated problem in the SFT model can be alleviated by using small classifier-free guidance. Therefore, the necessity of using the online optimization is questionable. Also, why the authors do not directly compare with [17]?
Questions
Please refer to the weakness part.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
The authors discussed some limitations in the last part of the paper.
Summary
The paper introduces the idea of finetuning text-to-image diffusion models using reinforcement learning. The core idea is relatively simple: generate samples from a trained diffusion model, use the samples and a reward function to update the model's parameters, and iterate. The paper further introduces a KL divergence loss to prevent diverging too far from the original model's weights. evaluations are done on different prompts with respect to different capabilities/ performance metrics, such as colour or composition.
Strengths
- an interesting and novel idea - simple method and easy algorithm - qualitatively pleasing results
Weaknesses
- considering only 4 prompts for experimental evaluation is not sufficient - considering only a single prompt for bias evaluation is not sufficient - authors mention that "longer training time, hyper-parameter tuning and engineering efforts are required as the number of training text prompts increases". This indicates that their method does not scale to a larger set of training prompts, which questions its relevance - investigation of generalisation to prompts that are out of distribution should be conducted
Questions
see weaknesses
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
limitations are briefly discussed, especially limitations mentioned in weaknesses could be discussed in more detail
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.
Strengths
The motivation for using RL techniques to fine-tune T2I model is reasonable and compelling. The paper is well-written and presented. It is very easy to follow. The idea is simple and easy to implement. The 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]. I 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. In 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. The 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. [1] Training diffusion models with reinforcement learning
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. I 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?
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
Yes.
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).
Strengths
* The paper is clear, the theoretical section is sound, and the experiments are convincing. * This is arguably the first work applying online RL to fine-tune diffusion model, which is a key contribution. * I appreciated the MDP formulation of the diffusion process. * 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. - 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. - 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. - 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. * More generally more experiments are required to ensure that online RL for diffusion models consistently helps. - The experiments are made only for a very small number of prompts: 4! This limits applications to real-world applications. - The experiments are made with a single reward model. More experiments with diverse (less robust) reward models would help. - Therefore the empirical contribution may be seen as marginal. - 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.
Questions
* Could you elaborate on the difference between KL and reverse KL. * 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. * 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. * Have you explored refined online RL optimization strategies such as PPO? * 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.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
3 good
Presentation
3 good
Contribution
2 fair
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.
Summary
The paper proposes to adapt the popular RLHF framework used for fine-tuning LLMs to fine-tuning of diffusion generative models. Given a reward model, the goal is to finetune the parameters of the diffusion model such that the resulting images from the sampling 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.
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. - It is good that the authors added KL regularization to supervised finetuning to establish a better baseline for their RLHF framework. - 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. - 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. - 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. - 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. - 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). - 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.
Questions
- Is there a reason why eq (8) should not include a min, similar to (4)? - 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? - 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?
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
The main limitation of the paper is the restriction to a single prompt which is openly discussed by the authors.
Thanks for the clarifications and rebuttal. You state that "One possible approach is to conduct offline RL by collecting trajectories generated from the original model and comparing it to the online RL"; that would indeed be necessary to validate that online RL performs better than offline RL. Moreover, I still think that more ablation studies and additional experiments on diverse rewards is necessary. Overall, I read all responses and decided to keep my score.
Thank you for your additional comments! **Q1. Offline RL vs Online RL** Following your suggestion, we fine-tune the model using offline RL: for the prompt “A green colored rabbit”, we employ an advantage-weighted regression as an offline policy learning algorithm [2] with trajectories generated by the original model. Specifically, we use an exponential weighted advantage objective that learns a policy maximizing the Q-values subject to a distribution constraint (Eq.(7) in [2]) and try different $\beta$ parameters (i.e., $\beta \in 1,2,3,10$). In our experiment, **offline RL only improves the average reward from ~0.1 to ~0.3 (after 5000 gradient updates), while our online RL approach can quickly improve the reward to ~1.5 before 5000 gradient updates.** Because offline RL is usually unstable and its performance is upper-bounded by data coverage, we find that offline RL is worse than online RL. We hope that this additional result can clarify your question. [2] Offline Reinforcement Learning with Implicit Q-Learning. Ilya Kostrikov, Ashvin Nair, Sergey Levine **Q2. Ablation study and rewards** Would you be able to provide more details regarding the specific ablation studies and reward functions that we should consider? We are committed to addressing this concern to the best of our ability. If you have any other questions or suggestions, please do not hesitate to let us know.
Thank you for this additional and informative baseline. Many rewards could be considered (not only aesthetic ones, eg, nsfw-filtering ones or personalized ones), and all may not be robust. And as stated in the initial review: "the fact that your online RL requires a more robust reward model is actually a drawback/limitation, creating new challenges in robust reward design". Thus experiments with rewards weaker than ImageReward may be interesting to show when online learning actually succeeds, and in what situations the proposed method might fail.
Clarification on reward models by authors
Thank you for your reply! In our initial experiments, we actually used a weak reward model from [17], which is trained on small-scale human datasets compared to ImageReward. Our online RL fine-tuning worked well on tested text prompts, surpassing the performance of supervised fine-tuning. **This result shows the robustness of our method on weak reward function**. However, the reward model in [17] is only trained with a limited category of prompts, so we adopt ImageReward for a wider range of prompts. We agree that the robustness of the reward model is important in the RLHF framework. However, **learning a reward function from human feedback has shown to be effective and robust in many prior works**; Besides, there is some additional component like **KL regularization** to mitigate the issues from reward models, which is also adopted in our work. Thank you for your valuable comments on reward models! We leave the exploration of other reward choices as future work since our current version mainly focus on the RLHF setting, and we will add related discussions in the final version. Please let us know if you still have other questions! [17] Aligning Text-to-Image Models using Human Feedback. Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Shixiang Shane Gu
Thank you for the remarks and the additional experiments run. I have updated my score accordingly.
Thank you again for the review and we hope that our individual response (https://openreview.net/forum?id=8OTPepXzeh¬eId=usVNThWza4) and overall rebuttal (https://openreview.net/forum?id=8OTPepXzeh¬eId=UFj7ZsHm99) have addressed your main concerns. We would greatly appreciate your feedback and please feel free to let us know if you have any other questions in the discussion period!
Thank you again for the review and we hope that our individual response (https://openreview.net/forum?id=8OTPepXzeh¬eId=pZZwGatqVk) and overall rebuttal (https://openreview.net/forum?id=8OTPepXzeh¬eId=UFj7ZsHm99) have addressed your main concerns. We would greatly appreciate your feedback and please feel free to let us know if you have any other questions in the discussion period!
Thanks for the authors' efforts spending on the rebuttal! Similar to reviewer fxhn and YMPp, I still have concerns about the reward used in this paper for RL training. Moreover, even though the paper conducts LoRA fine-tuning for both the RL setting and SFT, there is no insight given on these methods.. In my view, using LoRA for the RL setting is completely different and may bring additional uncertainty. Overall, I read all responses and I would like to keep my score.
Thanks for your reply! 1) Could you elaborate on your concern about the reward used in the paper? About **Reviewer fxhn**'s concern on extra evaluations besides ImageReward and aesthetic score, in our rebuttal **we added new human evaluation**; We also show that **training with reward in [17] also works well for our RL method** since **Reviewer YMPp** asked about training with other reward choices, and we also include a discussion on our method with KL regularization will help in our RLHF setting, etc (https://openreview.net/forum?id=8OTPepXzeh¬eId=8z8kat5mBG). **We believe that we have addressed the above concerns and it would be extremely helpful if you could point out what is still unclear to you!** 2) As in the rebuttal to Reviewer fxhn (Comment 3. LoRA fine-tuning vs. full fine-tuning), we mentioned that "**In 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**". **In our experiments, fine-tuning with or without LoRA does not make a noticeable difference except for the GPU occupation**. [17] Aligning Text-to-Image Models using Human Feedback. Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Shixiang Shane Gu
Thank you again for the review and we hope that our individual response (https://openreview.net/forum?id=8OTPepXzeh¬eId=7cJZSAKzdA) and overall rebuttal (https://openreview.net/forum?id=8OTPepXzeh¬eId=UFj7ZsHm99) have addressed your main concerns. We would greatly appreciate your feedback and please feel free to let us know if you have any other questions in the discussion period!
Thank the author for the rebuttal content. I have read all the contents and the reviews from other reviewers. There are 3 suggestions, + One is about the claim and the experiments, although it is hard to conduct a clean ablation. It will be more useful to show some generated examples to support the claim. + It is clear to cite the paper [17] in the experimental part. + The online manner is the main contribution, and I think the title may include this keyword. After careful consideration, I keep borderline accept to this paper.
Thank you for the reply! >One is about the claim and the experiments, although it is hard to conduct a clean ablation. It will be more useful to show some generated examples to support the claim. As discussed with reviewer YMPp (https://openreview.net/forum?id=8OTPepXzeh¬eId=gDVthccGwl), **we already added some new results comparing offline and online RL methods, which shows that online RL is indeed better** as one ablation showing the differences between online and offline samples. We will also add such ablation results in our draft. >It is clear to cite the paper [17] in the experimental part. >The online manner is the main contribution, and I think the title may include this keyword. Thank you for the suggestions and we will modify them correspondingly in our draft. Please feel free to let us know if you still have any questions!
Dear AC and reviewers: Thank you so much for your time and the engagement in the discussion! We summarize the **extra problems raised by reviewers and our overall response** as follows: **1) Ablation study to verify the claim in Section 4.3 about online vs. offline samples**: **We conduct experiments to compare online and offline RL methods** (which is not the SFT method used in the original submission, see details here: https://openreview.net/forum?id=8OTPepXzeh¬eId=gDVthccGwl), which shows that **the online RL approach improves the average reward much faster than the offline approach**. It is not surprising since the performance of the offline approach is upper-bounded by data coverage. **2) Ablation study to use other reward models in RL training**: **Weaker rewards from human feedback**: In our initial experiments, we also used a weak reward model from [17], which is trained on a smaller-scale human dataset compared to ImageReward. Our online RL fine-tuning also worked well on tested text prompts, surpassing the performance of supervised fine-tuning. **This result shows the robustness of our method on weaker reward functions trained with human comparison.** Besides, we propose online KL regularization designed for diffusion models, showing that it is also useful to mitigate the issue of evaluating out-of-distribution samples using the reward and improve the image quality. [17] Aligning Text-to-Image Models using Human Feedback. Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Shixiang Shane Gu **For other reward functions**: Our work mainly focuses on the RLHF setting (where the reward function is trained with human comparison like ImageReward): The main reason is that this kind of reward is a good fit for policy gradient methods, since it is not a sparse reward and has shown great success in RLHF in the language domain. It is well-known that policy gradient would fail with a very sparse reward, so it is not our claim to train with arbitrary reward functions, but it could be interesting for future exploration. We believe that we have resolved the extra issues raised by the reviewers during the discussion period, and we will include all the results and necessary discussions in the revised draft. We also believe that our proposed method of online fine-tuning of diffusion models with KL regularization provides meaningful insight and important contribution to the community. Best, Authors
Decision
Accept (poster)