Paper: Aligning Text-to-Image Diffusion Models with Reward Backpropagation (ICLR-2024)

Reviewer comment (Writing):
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.

Author response:
**Q4.1) Typos**
Thanks for pointing this, please check our updated pdf, we believe we have 
fixed all of them.

Reviewer comment (Evaluation):
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.

Author response:
**Q4.2) 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?**
**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.

Reviewer comment (Experiments):
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.

Author response:
**Q3.3) Compute Efficiency. Fair comparison should be between AlignProp 
and DDPO with fully fine-tuning and no gradient checkpoint.**
    
We use the official pytorch repository of DDPO for all our experiments.  
DDPO (in their updated version) uses LoRA adapter weights. AlignProp 
however also uses gradient checkpointing while DDPO doesn’t, note that we 
only use gradient checkpointing to fit our model into GPU memory. Gradient 
checkpointing trades off compute with memory thus making AlignProp slower. 
Including gradient checkpointing in DDPO will only make its speed worse 
without any benefit, therefore we don't run this ablation. As per the 
perspective you mentioned, DDPO should indeed be 50x faster than AlignProp 
as it does a single step of backprop while we do 50 steps of 
backpropogation. However in DDPO's implementation they accumulate 
gradients over all denoising timesteps, to make the noisy gradients more 
stable. As can be seen in their code 
[here](https://github.com/kvablack/ddpo-pytorch/blob/1958463f020112c9a7bc85768d296daacc2e1b4b/scripts/train.py#L501), 
this makes learning from N samples 50x slower. Further they have an 
additional sampling loop which generates RGB images to obtain it's reward. 
We on the other hand don't have a seperate sampling and training loop. 
Overall to learn from 128 samples, DDPO takes 1 minute 20 seconds to 
sample them, then it takes 8 minute 10 seconds to train on them. We on the 
other hand take 6 minute 16 seconds to sample and learn from 128 samples.

Reviewer comment (Reproducibility):
About compute efficiency: 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?

Author response:
**Q3.3) Compute Efficiency. Fair comparison should be between AlignProp 
and DDPO with fully fine-tuning and no gradient checkpoint.**
    
We use the official pytorch repository of DDPO for all our experiments.  
DDPO (in their updated version) uses LoRA adapter weights. AlignProp 
however also uses gradient checkpointing while DDPO doesn’t, note that we 
only use gradient checkpointing to fit our model into GPU memory. Gradient 
checkpointing trades off compute with memory thus making AlignProp slower. 
Including gradient checkpointing in DDPO will only make its speed worse 
without any benefit, therefore we don't run this ablation. As per the 
perspective you mentioned, DDPO should indeed be 50x faster than AlignProp 
as it does a single step of backprop while we do 50 steps of 
backpropogation. However in DDPO's implementation they accumulate 
gradients over all denoising timesteps, to make the noisy gradients more 
stable. As can be seen in their code 
[here](https://github.com/kvablack/ddpo-pytorch/blob/1958463f020112c9a7bc85768d296daacc2e1b4b/scripts/train.py#L501), 
this makes learning from N samples 50x slower. Further they have an 
additional sampling loop which generates RGB images to obtain it's reward. 
We on the other hand don't have a seperate sampling and training loop. 
Overall to learn from 128 samples, DDPO takes 1 minute 20 seconds to 
sample them, then it takes 8 minute 10 seconds to train on them. We on the 
other hand take 6 minute 16 seconds to sample and learn from 128 samples.

Reviewer comment (Experiments):
About compute efficiency: Another perspective 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.

Author response:
**Q3.3) Compute Efficiency. Fair comparison should be between AlignProp 
and DDPO with fully fine-tuning and no gradient checkpoint.**
    
We use the official pytorch repository of DDPO for all our experiments.  
DDPO (in their updated version) uses LoRA adapter weights. AlignProp 
however also uses gradient checkpointing while DDPO doesn’t, note that we 
only use gradient checkpointing to fit our model into GPU memory. Gradient 
checkpointing trades off compute with memory thus making AlignProp slower. 
Including gradient checkpointing in DDPO will only make its speed worse 
without any benefit, therefore we don't run this ablation. As per the 
perspective you mentioned, DDPO should indeed be 50x faster than AlignProp 
as it does a single step of backprop while we do 50 steps of 
backpropogation. However in DDPO's implementation they accumulate 
gradients over all denoising timesteps, to make the noisy gradients more 
stable. As can be seen in their code 
[here](https://github.com/kvablack/ddpo-pytorch/blob/1958463f020112c9a7bc85768d296daacc2e1b4b/scripts/train.py#L501), 
this makes learning from N samples 50x slower. Further they have an 
additional sampling loop which generates RGB images to obtain it's reward. 
We on the other hand don't have a seperate sampling and training loop. 
Overall to learn from 128 samples, DDPO takes 1 minute 20 seconds to 
sample them, then it takes 8 minute 10 seconds to train on them. We on the 
other hand take 6 minute 16 seconds to sample and learn from 128 samples.

Reviewer comment (Experiments):
About compute efficiency: A more fair comparison should be between AlignProp and DDPO with fully fine-tuning and no gradient checkpoint.

Author response:
**Q3.3) Compute Efficiency. Fair comparison should be between AlignProp 
and DDPO with fully fine-tuning and no gradient checkpoint.**
    
We use the official pytorch repository of DDPO for all our experiments.  
DDPO (in their updated version) uses LoRA adapter weights. AlignProp 
however also uses gradient checkpointing while DDPO doesn’t, note that we 
only use gradient checkpointing to fit our model into GPU memory. Gradient 
checkpointing trades off compute with memory thus making AlignProp slower. 
Including gradient checkpointing in DDPO will only make its speed worse 
without any benefit, therefore we don't run this ablation. As per the 
perspective you mentioned, DDPO should indeed be 50x faster than AlignProp 
as it does a single step of backprop while we do 50 steps of 
backpropogation. However in DDPO's implementation they accumulate 
gradients over all denoising timesteps, to make the noisy gradients more 
stable. As can be seen in their code 
[here](https://github.com/kvablack/ddpo-pytorch/blob/1958463f020112c9a7bc85768d296daacc2e1b4b/scripts/train.py#L501), 
this makes learning from N samples 50x slower. Further they have an 
additional sampling loop which generates RGB images to obtain it's reward. 
We on the other hand don't have a seperate sampling and training loop. 
Overall to learn from 128 samples, DDPO takes 1 minute 20 seconds to 
sample them, then it takes 8 minute 10 seconds to train on them. We on the 
other hand take 6 minute 16 seconds to sample and learn from 128 samples.

Reviewer comment (Experiments):
About generalization of the proposed AlignProp: 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?

Author response:
**Q3.4)  Can AlignProp generalize to any circumstances as long as there 
exists a reward model:**
We think farther the reward model from the training distribution of Stable 
Diffusion the more sample inefficient it becomes to tune it. To study this 
hypothesis, we ran a toy experiment where we took 10 random images from 
[CT Medical 
Images](https://www.kaggle.com/datasets/kmader/siim-medical-images) 
dataset, and 10 random images from ImageNet. Our reward model is simply 
the similarity with the best matching image from the 10 examples. 
Emperically, we find that AlignProp achieves similar reward on both 
datasets, however it ends up being 60x more sample efficient on ImageNet 
than the medical dataset.

Reviewer comment (Theory):
About generalization of the proposed AlignProp: 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.

Author response:
**Q3.5) Why not supervised finetuning? Aligning with reward models still 
require these high-quality samples.**


Reward models are trained by humans ranking **images generated by the 
generative model, therefore the process does not require access to any 
real-world high-quality  samples**. Supervised fine-tuning on the other 
hand requires access to real-world high-quality samples. This is a  
central benefit of  reward-guided alignment.

Reviewer comment (Reproducibility):
Implementation details question: During fine-tuning, are the reward model parameters $\phi$ fixed or also tunable?

Author response:
**Q3.6) During fine-tuning, are the reward model parameters  fixed or also 
tunable?**
The reward models are  fixed.

Reviewer comment (Reproducibility):
Implementation details question: 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.

Author response:
**Q3.7)  T value is set to be 50, which, however, is 1000 for Stable 
Diffusion**
We use DDIM sampling which reduces the sampling steps from 1000 to 50 by 
making the generation process deterministic and non-markovian.

Reviewer comment (Writing):
Writing: Typo in the 2nd line of Page 5 ($\pi_{\theta}$ instead of $\pi_theta$)

Author response:
**Q3.7) Mentioned Typos**
Thanks you, we will fix these typos.

Reviewer comment (Novelty):
**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].

Author response:
**Q2.3) Comments on Diffusion-QL[1] and concurrent work of DRAFT[2]**
    
Thanks for the reference! We will cite DQL, we were not aware of the work. 
As per the differences, DQL focuses on offline reinforcement learning for 
robotics tasks, we on the other hand focus on aligning large-scale 
foundational text-to-image models. 

DRAFT[2] indeed is a concurrent work , it in fact came out a week after 
the ICLR paper submission deadline. Having said that, they still don’t do 
full backpropagation through time, which might prevent them from easily 
tuning at a semantic level, as they don’t update the earlier steps of 
denoising.

Reviewer comment (Experiments):
**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.

Author response:
**Q2.4) Issue of Collapsing, DDPO doesn't collapse.**

Yes indeed DDPO doesn’t collapse with full back propagation. We think 
policy gradient methods are inherently regularized due to noise from the 
credit assignment issue thus don’t end up collapsing. However the noise 
also results in sample inefficiency as we show in our results. We instead 
do direct reward backpropagation while keeping the full backprop of DDPO 
by randomizing the truncation steps.

Reviewer comment (Presentation):
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?

Author response:
**Q2.5) Typographical error with HPS score in text**
Yes this is an error, thanks for pointing. We meant 2.8.

Reviewer comment (Theory):
Clarification: in eq 3, does the first term come from weight decay?

Author response:
**Q1.1) Does the first term in Eqn 3 come from weight decay?**
Yes, the first term is a result of the weight decay loss.

Reviewer comment (Experiments):
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).

Author response:
**Q1.4) More examples for Figure 3 to study the impact of value of K**

Thanks for the suggestion, you can find more examples in the following 
[link](https://alignprop-iclr.github.io/index.html)

Reviewer comment (Experiments):
It might be interesting to study other PEFT methods besides LoRA.

Author response:
**Q1.2) Might be interesting to study other PEFT methods besides LoRA.**
Thanks for the suggestion, we tried using 
[IA3](https://arxiv.org/abs/2205.05638), we indeed get some improvement in 
performance. For instance with Aesthetic reward function, AlingProp is 
about 2% more sample efficient, we will include this result in our final 
submission.