Directly Fine-Tuning Diffusion Models on Differentiable Rewards

Thank you for the helpful comments and suggestions. We address your questions and concerns below:

Q: DRaFT-1 works the best but the improvement over ReFL is a bit marginal (Figs 3 and 4).

A: Please see the general response: DRaFT-LV is 2x more efficient than ReFL, leading to higher rewards. Fast learning may be particularly important for future reward fine-tuning using large reward models or on-the-fly learning to a specific user’s preferences.

Q: Furthermore, DRaFT-1 is the same as ReFL with K fixed to 1. Increasing K does not improve the performance.

A: One of our contributions is unifying gradient-based approaches for diffusion fine-tuning in a common framework, as shown in Algorithm 1. The fact that DRaFT is equivalent to ReFL when $K=1$ provides a useful connection between the approaches, which we hope helps to understand the space of possible approaches. In addition, we show that using $K=1$ works surprisingly well and yields a conceptually simpler and easier-to-implement approach than ReFL, which randomly samples $K$ in a range specified by two hyperparameters (the min and max step number from which to predict a one-step denoised image). Please see the general response for further discussion about ReFL.

Q: As DRaFT-1 (and LV version) already works well, in this case, we might not need gradient checkpointing and LoRA, right? How does finetuning the full parameter influence the result?

A: Yes, you are correct that for small $K$, we do not need gradient checkpointing; we have clarified this in the paper (on page 4). Fine-tuning the full model parameters performs comparably to using LoRA, but LoRA substantially reduces memory cost. Additionally, because all of the information incorporated during fine-tuning is contained within the LoRA parameters, we can more easily store and combine a set of LoRA weights adapted for different rewards and use them to obtain semantic interpolations as shown in Figure 7.

Q: Why do you choose Stable-Diffusion v1.4 instead of SD v1.5 or a higher version?

A: We used Stable Diffusion 1.4 to have direct comparisons with DDPO.

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.

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.

I thank the authors for the detailed reponse and for adding experiments related to gradient clipping showing the promise of gradient surgery to improve the optimizer performance. I also thank the authors for adding results related to mitigating reward hacking. Given the promise shown by a relatively simple fix on mitigating exploding gradients, I would have appreciated a more thorough treatment of the mechanisms involved and possible solutions to fix the exploding gradients issue. However, I understand the paper offers computational tools for future work to explore and understand those issues better. Thus, I'm bumping up my score to 7.

Edit : Oops, just realized there is no 7. So making it 8 instead with some hesitation.

Thank you for the helpful comments and suggestions. We address your questions and concerns below:

Q: Paper organization

A: Thank you for the suggestions for improving the organization of the paper! We used paragraphs rather than subsections due to space constraints. We have moved Figure 2 to Section 4 in the updated paper.

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.

Thank you for your helpful comments and suggestions. We address your questions and concerns below:

Q: The paper lacks a thorough discussion of related work.

A: Our related work section (Section 2) covers: 1) learning human preferences, 2) guidance, 3) backpropagation through diffusion sampling, 4) reward fine-tuning with supervised learning, 5) fine-tuning with RL, and 6) reward feedback learning. In our original submission, we also included an extended discussion of related work in Appendix D. We have updated the main text to cite and discuss the paper you pointed out.

Q: Could you elaborate on the relationship between your work and that of [1], especially regarding the idea of directly fine-tuning diffusion models on differentiable rewards?

A: We could not find a paper with the title you referenced, “Reinforcement learning for faster DDPM sampling.” The closest paper we found was Fan & Lee, “Optimizing DDPM Sampling with Shortcut Fine-Tuning,” ICML 2023. Is this the paper you are referring to? If so, then we note that its focus is different from our work. Fan & Lee (2023) propose an approach to speed up sampling: they consider the possibility of doing this using gradients by backpropagating through diffusion sampling, but due to memory and exploding/vanishing gradient concerns, they focus on an RL-based approach.

We propose a thorough framework for gradient-based diffusion fine-tuning that addresses the memory concern by leveraging LoRA parameters and gradient checkpointing, and addresses exploding gradients by using truncated backpropagation through sampling. Then, we show empirically that the method works well over a diverse set of challenging tasks that could not previously be optimized efficiently with RL-based approaches. Additionally, we propose several modifications (DRaFT-LV, LoRA mixing) that further improve performance over the base approach.

The idea of backpropagating through sample quality was studied empirically by Watson et al. (2022) (https://arxiv.org/abs/2202.05830), which we discuss in the paper. Importantly, the goals of Fan & Lee (2023) and Watson et al. (2022) are very different from ours: they focus on improving sampling speed, whereas we focus on optimizing arbitrary differentiable rewards. That being said, we agree that Fan & Lee (2023) is a very relevant prior work and we have added a citation and discussion of it in our related work (Section 2, in both the "backpropagation through diffusion sampling" and "reward fine-tuning with reinforcement learning" paragraphs). Thank you for pointing it out.

Q: Diffusion as a multi-step decision-making task.

A: Thank you for pointing out the Fan & Lee (2023) paper that was the first to interpret denoising as a multi-step decision-making task. We have updated the related work section to cite this paper (Section 2, in the paragraph on reinforcement learning).

Q: Can DRaFT work with other PEFT methods such as Adapters, Normalization training, and IA3?

A: It certainly could and this is a great direction to pursue, although we leave exploring that direction for future work.

Q: The paper does not report the performance of DPOK. Could you clarify how your approach compares to this method?

A: DPOK is an RL approach similar to DDPO (and done concurrently). While there are some differences (e.g. adding KL regularization but not using an importance sampling estimator to perform multiple updates), we expect that it will perform similarly to DDPO. In contrast, DRaFT learns 200x faster than DDPO because it leverages reward gradient information (see Figure 3).

Q: The use of the "stop_grad" function in the pseudocode is not properly explained

A: We explain the use of stop_grad in the “General Reward Fine-Tuning Algorithm” paragraph in Section 4. We have further clarified the role of stop_grad in the updated text. Thank you for the suggestion.

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.