Author response
Thans for the helpful comments and suggestions. We address your questions and concerns below:
**Q: Investigating backpropagation through sampling**
**A:** We provide an analysis of gradient norms, the impact of $K$, and the impact of deactivating LoRA parameters for different ranges of steps in Figures 8-10. We also note that DRaFT-$K$ is an approach to mitigate the issues of full backpropagation.
Additionally, we used gradient clipping throughout our experiments. We have added an ablation over the gradient clipping norm hyperparameter in Appendix B.7 (see Figure 28). We compare the reward values achieved over the course of fine-tuning with gradient clipping norms $c \in \{ 0.001, 0.01, 0.1, 1.0, 10.0, 100.0, 1000.0 \}$. We found that, when using DRaFT-50, smaller gradient clipping norms ($c = 0.001$) improved optimization substantially, while large clipping norms ($c = 1000$) impeded training. This further supports the observation that backpropagating through long diffusion sampling chains can lead to exploding gradients. We performed a similar ablation for DRaFT-1, and found that the fine-tuning performance was nearly identical for all gradient clipping coefficients.
Empirically, while vanilla DRaFT works better than RL, it is much less efficient than DRaFT-$K$ with small $K$, and thus we focused on DRaFT-$K$. However, we agree that further understanding the gradient flow through long sampling chains is an interesting direction for future work.
**Q: Dealing with over-optimization and reward hacking**
**A:** We agree that mitigating reward overfitting is an important issue, so we have added new experiments on this topic to Appendix B.8, Figure 29. We compared three strategies for mitigating reward hacking: early stopping, KL regularization towards the pre-trained model, and our LoRA scaling approach (discussed in Section 5.2 of the main text). We found LoRA scaling to be the most effective.
We view optimizing reward functions and mitigating overfitting as separate goals that can be studied independently, and primarily focus on the former in this paper. As we mention in Section 5.1, we believe that efficient methods like DRaFT could be a key stepping stone towards designing improved rewards (or regularizers), as DRaFT allows one to quickly evaluate the effect of “over-optimizing” a reward.
**Q: Suggested methods for mitigating reward overfitting.**
**A:** Interesting suggestions! We run have run several experiments on this:
**KL Regularization.** We agree that regularizing the model to stay close to the real image distribution is a good idea to reduce reward hacking. We experimented with one approach in this direction: adding KL regularization towards the pre-trained model. The pre-trained model implicitly captures the training image distribution, and this kind of regularization is standard practice for avoiding reward overfitting when performing RLHF with LLMs. We found that this approach helps mitigate reward overfitting (see Figure 29 in Appendix B.8), although it did not appear to work as well as the LoRA scaling approach we discuss in Section 5.2.
**$\mathbf{x}_0$ Idea.** We also experimented with using fixed $\mathbf{x}_0$s with DRaFT-LV rather than ones generated online by the current policy. Counterintuitively, this approach actually makes reward overfitting worse and leads to poor results (0.333 HPSv2, 0.216 PickScore in the same setting as Figure 29). The main issue is with distribution shift: the latents produced during sampling at test time will not match the ones seen in training. In particular, the model has no way to correct itself if it starts following the reward model’s score function too strongly, which results in artifacts in the samples. Previous work has also found offline policy learning to not be as effective as using policy-generated images; for example, see the comparison between DDPO and RWR in arxiv.org/abs/2305.13301.
**Q: Prompts for generated images**
**A:** Thank you for the suggestion! For Figure 7, we used prompts from the HPDv2 dataset, which are quite long and hard to fit due to space constraints; Figures 19 and 20 in Appendix B.5 provide expanded visualizations that include the prompt on the left-hand side. We will also add more prompt text to the main paper, where space allows.
**Q: Prompts for Figure 7.**
**A:** Figures 19 and 20 in Appendix B.5 show the samples generated for the same set of prompts when fine-tuning for HPSv2 and PickScore, respectively. We felt it was useful to show a diversity of samples in the main paper (Figure 7). Also, Figure 5 in the main paper compares the effects of different reward functions on the same prompts (“a photo of a deer” and “a painting of a deer”).
Thank you again for taking the time to read our paper and for providing thoughtful feedback. Please let us know if you have any other questions, and we would be happy to answer them.