Paper: Training Diffusion Models with Reinforcement Learning (ICLR-2024)

Reviewer comment (Reproducibility):
As shown in Appendix B, there is a subtle bug in the implementation. This bug maybe affect the quantitative comparisons and lower the support for the effectiveness of the proposed method. It is one of my major concerns about this paper.

Author response:
Apologies for any confusion, but this bug only affected the leftmost chart in Figure 4 (primary comparison) and the leftmost chart in the RWR interleaving ablation in the appendix. No other results, qualitative or quantitative, were affected. While the affected results still represented a fair scientific comparison, we decided to mention the bug in the appendix in the interest of full transparency.

We have re-run the primary comparison in Figure 4 and updated the submission. Unfortunately, it is infeasible to re-run the extra interleaving baseline/ablation in the appendix due to compute limitations. However, the leftmost chart still provides a fair comparison, since all methods were tested on the same reward function, and the rightmost chart supports the same conclusion while being unaffected by the bug.

Reviewer comment (Writing):
In the related work, ‘Reward-Weighted Regression’ should be used before using its abbreviation (RWR). If not, it maybe a misleading phrase for the readers, especially ones without context.

Author response:
Thanks for the catch! We have fixed this in the updated submission.

Reviewer comment (Theory):
The proposed method does not consider the problem of overoptimisation, instead the authors argue that early stopping from visual inspection is sufficient. This makes the method applicable to problems where visual inspection is possible (which is likely the case for many image tasks). However, it renders the method inapplicable to problems where the needed visual inspection is not possible. (E.g. one might not be able to apply this method to a medical imaging task where visual inspection by a human supervisor does not allow to determine when to stop the optimisation process).

Author response:
Thank you for your thorough review. It seems that your primary concern is about overoptimization. It’s worth noting that we found it reasonably straightforward to run our method for a fixed number of iterations (e.g., 15k reward queries) and recover good generations; we also did not observe completely degenerated images in the VLM or aesthetic quality experiments, only in compressibility and incompressibility (when run for long enough). We have added an additional quantitative experiment that we hope addresses your concern, which we discuss more below.

We would also like to point out that overoptimization is an open research problem that is well outside the scope of this paper. While we agree that it is certainly a major issue that hinders many real-life applications (such as the provided medical imaging example), the issue of overoptimization also exists in language modeling to an equivalent degree ([Gao et. al., 2022](https://arxiv.org/abs/2210.10760)), yet it has certainly not prevented RLHF from having a huge impact in that domain ([Ouyang et. al., 2022](https://arxiv.org/abs/2203.02155); [Bai et. al., 2022](https://arxiv.org/abs/2212.08073)). Current techniques for combating overoptimization, such as KL-regularization or automated evaluation, do not completely eliminate the need for human evaluation in real-life systems.

Reviewer comment (Experiments):
The evaluation is missing a simple classifier guidance baseline, where samples are generated from the diffusion model with classifier guidance according to the given reward function. (Classifier guidance is mentioned in the related work but not evaluated as a baseline)

Author response:
> **The evaluation is missing a simple classifier guidance baseline**

We have added a comparison to classifier guidance in Appendix B of the updated submission ([screenshot](https://ibb.co/dfJbVTw)). We note that only the aesthetic predictor reward function is differentiable, meaning it is the only one where classifier guidance is applicable. See the top-level rebuttal for more information.

Reviewer comment (Experiments):
I am wondering why the lacking quantitative experiments have not been conducted (or why they cannot be conducted).

Author response:
Sorry, yes we regret that we couldn't get to your question sooner. It's a good suggestion that would definitely provide some good insights. Unfortunately we had to prioritize and were busy running quite a few other experiments for this rebuttal. We actually did start the aforementioned run yesterday, and it seems that the load on our machine is a bit less today, so it may finish in time after all. We'll post an update tonight if it has progressed enough to get some interesting results.

Reviewer comment (Experiments):
I am wondering how the authors think about the early stopping mechanism.

Author response:
Thank you for your thorough review. It seems that your primary concern is about overoptimization. It’s worth noting that we found it reasonably straightforward to run our method for a fixed number of iterations (e.g., 15k reward queries) and recover good generations; we also did not observe completely degenerated images in the VLM or aesthetic quality experiments, only in compressibility and incompressibility (when run for long enough). We have added an additional quantitative experiment that we hope addresses your concern, which we discuss more below.

We would also like to point out that overoptimization is an open research problem that is well outside the scope of this paper. While we agree that it is certainly a major issue that hinders many real-life applications (such as the provided medical imaging example), the issue of overoptimization also exists in language modeling to an equivalent degree ([Gao et. al., 2022](https://arxiv.org/abs/2210.10760)), yet it has certainly not prevented RLHF from having a huge impact in that domain ([Ouyang et. al., 2022](https://arxiv.org/abs/2203.02155); [Bai et. al., 2022](https://arxiv.org/abs/2212.08073)). Current techniques for combating overoptimization, such as KL-regularization or automated evaluation, do not completely eliminate the need for human evaluation in real-life systems.

Reviewer comment (Experiments):
I think one of the main contributions of the paper is on utilizing VLMs for optimizing text-to-image diffusion models. In this context, the discussion on the choice of reward function should be discussed more in detail. For example, I think different instruction prompts for the LLaVA model would make different results, but the paper lacks details in the choice of reward function.

Author response:
We briefly discuss the choice of LLaVA prompt in the third paragraph of Section 5.3, where we mention that the design space of VLM reward functions is virtually unlimited and that we look forward to future work that explores this direction. In the camera-ready, we will elaborate on why we chose LLaVA (it was the best open-source VLM at the time) and the evaluation scheme that we did (we chose to stick with the simplest, most general prompt, but found that BERTScore was necessary to increase reward granularity).

Reviewer comment (Theory):
It seems like the reward function for different tasks have different ranges, and thus DDPO uses running mean and variance for normalization. I guess the difficulty of solving MDP varies across different prompts. Could the author elaborate on the optimization analysis with respect to the difficulty of prompts?

Author response:
Some prompts may be naturally more difficult for the pretrained model than others. For example, the  model may easily generate “a man riding a horse,” but struggle to generate well-aligned images for “a horse riding a man.” Figure 5 provides a quantitative example of this: the “riding a bike” prompts achieve the highest average reward, followed by “playing chess,” and then “washing dishes.”

Reviewer comment (Experiments):
In implementation, the author chooses T=50 throughout experiments. What if the timesteps become smaller or larger? Definitely the denser timestep sampling would require more cost, but I believe it can reduce the variance of gradient. Could the author elaborate on this point?

Author response:
Using more generation timesteps would certainly improve policy gradient estimates, as well as improve the base generation quality of the model. 50 timesteps is a standard choice in the literature, and we found it to work well in our experiments, so we did not ablate this hyperparameter. We suspect that one would find diminishing returns if the number of timesteps was increased much beyond 50.

Reviewer comment (Experiments):
Do the authors try different reward models for image-text alignment? How robust is DDPO algorithm with respect to different reward functions?

Author response:
We have included an additional experiment with ImageReward in Appendix C ([screenshot](https://ibb.co/Jc0LmS0)). This and LLaVA are the only two models for image-text alignment that we have personally tried.

Reviewer comment (Experiments):
In terms of image generation quality, the paper lacks a quantitative and qualitative comparison with recent works. It fails to provide experimental support for its effectiveness. Specifically, in the absence of comparisons in image quality with all methods related to "Optimizing diffusion models using policy gradients," it is challenging to discern the improvements this paper offers over baseline approaches. This makes it difficult to evaluate the paper's contribution to the community.

Author response:
Thank you for your thorough review. It seems the novelty of the contribution with respect to DPOK is your primary concern; as we explain in the top-level rebuttal, our paper was posted to arXiv 3 days before DPOK, and as such the work should be considered concurrent. The overall algorithm, narrative, and writing structure of our paper have not changed since that initial posting, and as such, our imitation of DPOK is impossible.

That said, we have added an experiment directly comparing our method to DPOK in Appendix C of the updated PDF ([screenshot of quantitative results;](https://ibb.co/Jc0LmS0) [screenshot of qualitative results](https://ibb.co/z4shhBd)). Our method significantly outperforms the numbers reported by DPOK on the same task, even without KL regularization.

Reviewer comment (Evaluation):
The experimental evaluation criteria are unreasonable. Most pretrained reward evaluation models are not designed for robust data manifold, especially for treating it as a "fine-tuning teacher." Consequently, out-of-domain evaluation is necessarily needed. For example, the authors use Aesthetic Quality as a reward function to fine-tune the Diffusion model and employ the same Aesthetic Quality for scoring during evaluation. This approach does not allow for assessing whether an improvement in the Aesthetic Quality Score correlates with an enhancement in image generation quality. A more reasonable evaluation, as exemplified by DPOK, would involve fine-tuning with one Reward model and then evaluating it against both this model and a new, out-of-domain Reward model. DPOK's evaluating approach might provide a more substantial basis for assessment.

Author response:
> **out-of-domain evaluation is necessarily needed**

Our comparison to DPOK includes the out-of-domain evaluation you recommended; namely, training using ImageReward and evaluating using Aesthetic Score. We find that both metrics increase across all prompts ([screenshot](https://ibb.co/Jc0LmS0)).

Reviewer comment (Reproducibility):
Lack of code. The lack of open-source code, compounded by the absence of comparisons with recent works, makes it difficult to assess the practical feasibility of the approach.

Author response:
> **Lack of code**

Apologies for the oversight! Open-source code has been available for the past 6 months, but we forgot to provide an anonymized link in our submission. You can find anonymized code [here](https://anonymous.4open.science/r/ddpo-pytorch-1D5C/).