Aligning Text-to-Image Diffusion Models with Reward Backpropagation

Summary

The paper introduces AlignProp, a novel method for aligning text-to-image diffusion models with specific reward functions. The authors argue that existing reinforcement learning methods are inefficient due to high variance gradients. AlignProp aims to overcome these limitations by using end-to-end backpropagation to fine-tune the model. The paper also discusses techniques to manage memory overhead, such as fine-tuning low-rank adapter (LoRA) modules and using gradient checkpointing.

Strengths

Comparative Analysis: The paper does a good job of positioning AlignProp against existing methods, particularly reinforcement learning techniques. This helps in understanding the unique advantages of AlignProp. Results: The paper claims that AlignProp achieves higher rewards in fewer training steps and is preferred by human users, although the empirical evidence supporting these claims could be strengthened.

Weaknesses

* The writing is not good. \cite and \citep are different. And there are many other typos and missing/wrong references, which cause some difficulties in understanding the work. * The idea is simple and straightforward. If authors want to convince me to increase the score and demonstrate the effectiveness of the approach, can authors provide an anonymous website to list un-cherry-picked images across different iterations and different methods? It is good to use open-eval prompts, such as the ones provided in part image generation eval benchmark.

Questions

See above

Rating

3: reject, not good enough

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

2 fair

Presentation

2 fair

Contribution

2 fair

Summary

This paper aims at aligning a pre-trained text-to-image diffusion model with downstream objectives using the most straightforward way. The proposed AlignProp fully reconstructs the input images during training, which are then taken as input to an off-the-shelf reward model for end-to-end reward optimization. The authors utilize several well-acknowledged optimization tricks for better GPU memory management and aligning efficiency. The extensive experiments demonstrate the superiority of AlignProp.

Strengths

- The authors construct this paper with clear architecture and detailed discussion. - The proposed method is fairly simple with extensive experiment results. - The motivation is relatively clear and consistent with the human intuition.

Weaknesses

- Novelty would be a controversial problem of this paper: - Methodology: The main technical components of this paper consist of two parts: 1) directly propagating the reward back to the diffusion models without RL, following DDPO, and 2) spanning the whole denoising process and perform back propagation through time, which has been utilized for diffusion model guidance / alignment early in DiffusionCLIP published in CVPR 2022. - Implementation: According to the authors, the main difficulty of applying the methodology above in the considered datasets is the high GPU memory usage, which, however, can be simply solved by several commonly used optimization tricks of diffusion models including LoRA, gradient checkpointing and gradient truncation. - Therefore, although nobody has conducted experiments by combining all these things together before, it is still hard to convince me that there exists a novel research problem and this work should be considered as a research paper for top-tier conference like ICLR instead of a solid technical report. - About compute efficiency - In order to propagate gradient back to the whole sample chain, AlignProp needs to perform the whole T (=50 according to the authors) for each data sample during training even with DDIM, which instead requires only 2 steps for DDPO, and this is a huge efficiency gap. - Therefore, I cannot understand why in the 2nd row of Fig. 4, AlignProp can demonstrate efficiency even with respect to Time. Can you give more implementation details about how these experiments are conducted, and why the re-implemented DDPO converges so quickly? - Another perceptive to view DDPO and AlignProp together is that both of them are working on a practical estimation towards fully fine-tuning UNet with respect to reward models using BPTT. DDPO chooses to back propagate only one step, while AlignProp chooses to use LoRA and gradient checkpointing. - Therefore, a more fair comparison should be between AlignProp and DDPO with fully fine-tuning and no gradient checkpoint. - About generalization of the proposed AlignProp: - I wonder if AlignProp can generalize to any circumstances as long as there exists a reward model. - In other words, if the alignment ability requires the pre-trained diffusion models have the downstream-desired generation ability at first (e.g., Aesthetics). For example, can AlignProp be applied to allow Stable Diffusion to generate medical images with a reward model trained on medical images? - In the 2nd paragraph, the authors claim that the motivation to utilize reward models instead of supervised fine-tuning is the requirement of high-quality data samples. However, in my understanding even aligning with reward models still require these high-quality samples with high rewards for alignment optimization.

Questions

- Implementation details: - During fine-tuning, are the reward model parameters $\phi$ fixed or also tunable? - What is the specific T value (i.e., total denoising steps) utilized in your experiments? According to the "Baselines" paragraph in Sec. 5, the T value is set to be 50, which, however, is 1000 for Stable Diffusion trained with DDPM during pre-training. - Writing: - Fix the citation style by using the correct latex command. - Typo in the 2nd line of Page 5 ($\pi_{\theta}$ instead of $\pi_theta$)

Rating

5: marginally below the acceptance threshold

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

2 fair

Presentation

3 good

Contribution

1 poor

Summary

This paper introduces "AlignProp," a technique that refines text-to-image diffusion models using end-to-end backpropagation of the reward gradient during the denoising phase. Rather than consuming excessive memory, AlignProp fine-tunes low-rank weight modules and uses gradient checkpointing. When tested, AlignProp efficiently optimized diffusion models for various objectives like semantic alignment and aesthetic enhancement. It outperformed existing methods, achieving better results in fewer steps, making it a preferred choice for optimizing diffusion models.

Strengths

The experiments conducted are comprehensive and of high quality.

Weaknesses

1. **Originality & Novelty:** The paper seems to lack significant originality and novelty. Implementing the two memory-saving techniques - finetuning with LoRA and gradient checkpointing - does not appear challenging, especially since they are already available in the “diffusers” package. Further, randomizing the number of denoising steps appears to be a straightforward approach, and it's not guaranteed to address the collapsing issue. 2. **Previous Work Reference:** The concept of using a differentiable reward function and backpropagating the gradient directly to each timestep was earlier introduced in the Diffusion-QL[1] paper. It would be beneficial to cite the D-QL paper. There's also another very recent concurrent work on the subject [2]. 3. **Issue of Collapsing:** Both the present paper and [2] have highlighted the issue of collapsing when using a differentiable reward function for backpropagation. As demonstrated in the DDPO paper, using policy gradient through a non-differentiable reward signal doesn't present this issue. A more detailed exploration of the collapsing issue, along with effective mitigation strategies, would be a valuable addition. The currently proposed method of randomized denoising length seems somewhat simplistic. A more robust solution, accompanied by a thorough analysis, is anticipated. [1] Wang, Zhendong, Jonathan J. Hunt, and Mingyuan Zhou. "Diffusion policies as an expressive policy class for offline reinforcement learning." arXiv preprint arXiv:2208.06193 (2022). [2] Clark, Kevin, et al. "Directly Fine-Tuning Diffusion Models on Differentiable Rewards." arXiv preprint arXiv:2309.17400 (2023).

Questions

1. In Figure 4, the HPSv2_Score is noted as ranging from 2.4 to 3.6. However, to the left of Table 1, the paper mentions a score of 0.28. Is this a typographical error? 2. In Figure 5, the color map and object shapes from AlignProp appear quite similar. Yet, the results from Stable Diffusion and DDPO exhibit greater diversity. Could you clarify this?

Rating

5: marginally below the acceptance threshold

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

2 fair

Summary

In this paper proposes AlignProp, a method that aligns diffusion models to downstream reward functions using end-to-end backpropagation of the reward gradient through the inference chain. The main challenge the paper is trying to solve is the prohibitive memory cost required by naive implementation of such backpropagation.

Strengths

- The paper studies an important problem of end-to-end backpropagating a reward function through the denoising process. - The presented results look promising and the experiments are extensive and convincing.

Weaknesses

- Clarification: in eq 3, does the first term come from weight decay? - Typos: 1) eq 3 and 4, cdot notations are not consistent; 2) page 5 "policy \pi_{theta}"; 3) page 5 "k"m". - Figure 3 presents visual results on a single image, which seems not "comprehensive" enough to study the impact of value of K (as stated in the last paragraph in page 5).

Questions

- It might be interesting to study other PEFT methods besides LoRA.

Rating

6: marginally above the acceptance threshold

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

General Response

Dear Reviewers, Thank you for your comments. Many Reviewers brought up lack of novelty as the main weakness of the paper. **Contribution:** AlignProp addresses the task of aligning large-scale foundational text-to-image models using differentiable reward functions, by end-to-end backpropagation of reward gradients to diffusion model parameters. The state-of-the-art method for this task till our work was DDPO that uses policy gradients, **and thus does not use the gradients of the reward function**. Our method directly backpropagates gradients from the reward model to update the diffusion model. **Reviewer e5M5 says: "directly propagating the reward back to the diffusion models without RL, following DDPO" We believe this is incorrect:** there is a significant difference between the two methods, DDPO does not backpropogate gradients from the reward model (it does not use the gradient of the reward function, just its values in specific image samples), whereas we do. **Closest related works:** The closest related work to AlignProp is DiffusionQL and DiffusionCLIP, as also pointed out by reviewers JeN1 and e5M5?. While all the above mentioned methods use reward backpropogation to update the diffusion model, the use-cases are very different which results in very different set of challenges for the three works. DiffusionQL[1] addresses the problem of offline reinforcement learning where a small-scale policy neural network is trained via Q-learning to generate high rewarding robot actions. AlignProp instead focuses on aligning billion parameter large-scale pretrained text-to-image models, a very different setup and computational challenges. DiffusionCLIP[2] is a category-specific image translation model, which maps a set of source images (male face images) to a target domain (female). It does so by doing DDIM inversion on the source images to obtain latents z, it then updates the denoising process such that the generated image matches with the text by backpropagating gradients from a text-image matching CLIP objective. It uses unconditional diffusion models trained on category-specific datasets such as Dogs or faces. We on the other hand update large-scale text-conditioned diffusion models such as Stable Diffusion. Instead of doing DDIM inversion using a dataset of source images, we directly generate the images from the foundation diffusion model given a set of prompts. As one can see, these setups are very different, they address very different problems with different methodologies, and pose unique challenges for each, as well as very complementary benefits to the community. Reviewer JeN1 brought up DRAFT[3], which is a concurrent work, and in fact it came out a week after the ICLR paper submission deadline. **If AlignProp was indeed obvious, we do not see why a work such as DRAFT would exist.** Further, unlike AlignProp, DRAFT does not adapt earlier denoising steps which might prevent them from easily adapting at a semantic level, as we show in Figure 3 of our paper. Reviewer JeN1 says: "finetuning with LoRA and gradient checkpointing - does not appear challenging, especially since they are already available in the “diffusers” package." You are correct. However, **we do not think a method should be evaluated based on the number of line changes in the code, but based on how valuable of a contribution it makes in an important problem.** There are many useful works in the past such as beta-VAE, that have mainly changed a single hyperparameter value and still been very useful to the community. **Additional Dense Qualitative results:** **Aesthetics reward model:** In the following webpage [<link1>](https://wandb.ai/alignprop-iclr/alignprop/reports/Aesthetics-Results-different-Epochs---Vmlldzo2MDU5MzI4?accessToken=vph90cg4qa44h0dl15nbxfm4tu9b3efv6c55r8ssf39yv3e9t3mzon82q9eq221j), we show results for different epochs of training for our model and the baselines. As can be seen in the visuals, our model achieves better qualitative results even in the intermediate epoch of training. **HPS reward model:** In the following webpage [<link2>](https://alignprop-iclr.github.io/hps.html), we compare AlignProp with baselines over unseen prompts from the evaluation prompt set of HPS reward model. **Disabling LoRA weights:** In the following webpage [<link3>](https://alignprop-iclr.github.io/index.html), we expand Figure 3 of our paper and show more example cases for dropping lora weights over different timesteps of denoising.