Paper: AdjointDPM: Adjoint Sensitivity Method for Gradient Backpropagation of Diffusion Probabilistic Models (ICLR-2024)

Reviewer comment (Evaluation):
NFE and wall-clock time for AdjointDPM and the baselines is missing, making it difficult to gauge the efficiency of AdjointDPM.

Author response:
> NFE and wall-clock time for AdjointDPM and the baselines are missing, so it is difficult to gauge the efficiency of AdjointDPM.
 
**Reply:** For the analysis of NFE, please refer to the global response. The actual training time highly depends on the hardware configuration.

When we finetune parameters of neural networks, for example, the stylization task, it takes around 18 hours to optimize in a A100 GPU. But once trained, this network can be used to generate any images with a certain style. The training time is acceptible in this case.

For guided sampling tasks (such as vocabulary expansion, security auditing, and text embedding inversion), as the key time consumption is in traversing sampling steps and doing gradient backpropagation, here we present the wall-clock time required for a single optimization step. In these tasks, we run the code in a V100 GPU. For optimization of initial noises and text embedding inversion in Stable diffusion, it takes around 30 seconds per optimization step. For optimization of initial noises in ImageNet 256x256 generation, it takes around 5.3 seconds per optimization step.

Reviewer comment (Novelty):
A background section explaining and comparing other related backpropagation methods, such as DOODLE, FlowGrad, and DEQ_DDIM, in detail would help readers understand the position of AdjointDPM with respect to previous work. Specifically, the paper lacks a discussion of theoretical and practical advantages and disadvantages of AdjointDPM compared to previous work on backpropagation through diffusion sampling. For instance, the authors state DEQ-DDIM requires the diffusion sampling process to have equilibrium points—whether this is a significant drawback is not discussed. It is also unclear whether there are certain tasks where this equilibrium assumption does not hold, making AdjointDPM applicable while DEQ-DDIM is not.

Author response:
> A background section explaining and comparing other related backpropagation methods, such as DOODLE, FlowGrad, DEQ_DDIM in detail would help readers understand the position of AdjointDPM w.r.t. previous work. Specifically, this paper lacks a discussion of theoretical and practical advantages/disadvantages of AdjointDPM w.r.t previous work on backpropagation through diffusion sampling. For instance, the authors state DEQ-DDIM require the diffusion sampling process to have equilibrium points -- is this a significant drawback? Are there certain tasks where this equilibrium assumption do not hold, so AdjointDPM is applicable while DEQ-DDIM is not?
 > 
**Reply:** We have presented part of comparsion with existing models, mainly including DOODL and FlowGrad in Appendix B. Here we have a brief comparison: 

1. **Comparison with [1] End-to-End Diffusion Latent Optimization Improves Classifier Guidance**
    DOODL in [1] optimizes the initial diffusion noise vectors w.r.t a model-based loss on images generated from the full-chain diffusion process. In their work, they obtain the gradients of loss w.r.t noise vectors by using invertible neural networks (INNs). There are two main differences between [1] and our work:
  - While DOODL optimizes the initial diffusion noise vectors, our work optimizes related variables, including network parameters, initial noises and textual embeddings w.r.t a model-based loss on images generated from the full-chain diffusion process. Thus, we consider the broader cases of DOODL. 
  - In the calculation of gradients w.r.t initial noises, [1] uses the invertibility of EDICT, i.e., $x_0$ and $x_T$ are invertible. This method does not apply to the calculation of gradients w.r.t. the network parameters and textual embeddings as they share across the full-chain diffusion process. Finally, with regard to the memory consumption when calculating gradients with respect to the initial noise, our experimental results are as follows: We utilized the stable diffusion v1.4 checkpoint to run both the AdjointDPM and DOODL models on a V100 GPU (32GB memory). For the AdjointDPM method, backpropagating the gradients with respect to the initial noise required 19.63GB of memory. In comparison, the DOODL method consumed 23.5GB for the same operation.  Thus, our method is more efficient in terms of memory consumption. In terms of time consumption, DOODL relies on the invertibility of EDICT, resulting in identical computation steps for both the backward gradient calculation and the forward sampling process. However, our AdjointDPM methods have the flexibility to design adaptive steps, allowing for faster backward gradient calculation.

2. **Comparison with [2] FlowGrad: Controlling the Output of Generative ODEs With Gradients**
        FlowGrad efficiently backpropagates the output to any intermediate time steps on the ODE trajectory, by decomposing the backpropagation and computing vector Jacobian products. FlowGrad focuses on refining the ODE generation paths to the desired direction. This is different from our work, which focuses on the finetuning of related variables, including network parameters, textual embedding and initial noises of diffusion models for customization. Besides, FlowGrad methods also can not obtain the gradients of loss w.r.t. textual embeddings and neural variables as these variables share across the whole generation path. Then for the gradients w.r.t the latent variables, we could show the memory consumption of our methods is constant while they need to store the intermediate results. Since their method needs to store the intermediate results, in the paper, they present examples on Latent Diffusion Model (LDM) on FFHQ and LSUN Church, and Rectified Flow (RF) on CelebA-HQ, which are not large-scale diffusion models. Additionally, their officially released code for FlowGrad is only on pre-trained Rectified Flow [4], where the diffusion ODE path is a straight flow. Thus, when dealing with stable diffusion models, the use of their method may lead to inefficient memory utilization.
        
        
    [4] Liu, X., Gong, C., & Liu, Q. (2022). Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003.
        
3. **Comparison with the paper [3] Deep Equilibrium Approaches to Diffusion Models (DEQ-DDIM)**
    Deep equilibrium approaches model the entire DDIM sampling chain as a joint, multivariate fixed point system. It solves all the equilibria simultaneously, which leads to faster convergence to final outputs. When doing model inversion (i.e., optimizing the initial noise), it directly calculates the gradients on the obtained fixed point. This method needs to store intermediate results, resulting in high memory consumption.  

    But for Adjoint method, it is general for all ODE models and suitable for different numerical solvers, including first-order (DDIM) and higher-order solvers. Besides, to obtain the gradients, adjoint sensitivity methods in our paper solve a backward ODE with constant memory. Thus, DEQ-DDIM is more time-efficient than AdjointDPM, but AdjointDPM is more memory-efficient than DEQ-DDIM.


    We conducted a comparison of the memory consumption between our method and DEQ when performing model inversion. Utilizing the official code from DEQ and the CIFAR10 checkpoint, we ran the model on a V100 GPU (32GB memory). When sampling an image using DEQ, the memory consumption reached 19GB, and the system exhausted its memory during the model inversion process. 

     For the AdjointDPM methods, we applied the stable diffusion v1.5 checkpoint for model inversion. The memory consumption for sampling an image was 10.3GB, and it required 19.63GB to perform the model inversion. As the stable diffusion v1.5 checkpoint is much larger than the CIFAR10 checkpoint, our method is shown to be much more memory efficient. Regarding computation time, we did not evaluate this aspect due to limitations in our GPU resources as it ran out of memory when we run the model inversion of DEQ.

Reviewer comment (Theory):
The proposed method cannot guarantee the property of diffusion models, i.e., the noise-pred network corresponds to the score functions, because it directly train the neural ODE. Thus, it may be hard to leverage the other properties of diffusion models, such as classifier / classifier-free guidance, and it may be hard to further use diffusion SDEs for better sample quality. Instead, other guided sampling methods introduce another guidance model at each time step, which is based on the score functions and thus can maintain the diffusion property.

Author response:
Thanks for raising this issue. 

**For finetuning of initial noises**, as long as we control the scale of perturbation on initial noises, the trained noises have little influence on the sampling process of diffusion models. Then for finetuning of text embedding, it also has little influence on score estimation as we just want to learn a proper text embedding for specific visual effects. Thus, CG/CFG and SDEs still apply in these two cases.

In related experiments, besides the guidance from the loss function defined on final outputs, we still can use CG/CFG and SDEs in sampling. For example, in security auditing, we use CG to generate images with prescribed classes and find an initial noise which can mislead pretrained ImageNet classifier while visually similar to the original one. Besides, after finding this initial noise and text embedding, we could still use SDEs (such as DDPM sampler) to generate better-quality images. The results are added to Appendix.

However, **when using AdjointDPM to finetune or train the parameters of diffusion networks**, the question of whether this operation will change the property of the underlying diffusion models is still largely open. However, we surmise that finetuning a diffusion model on a small amount of data for a few steps will still preserve the applicability of CG/CFG techniques. In practice, after finetuning the parameters of the Unet only for 10 epochs in stylization, we still use the CFG to guide the sampling with given text prompts under a guidance scale of 7.5. We also test the generation in SDE solvers and it works. The results are added to Appendix.

Besides, we observe, in some other works that distill diffusion models based on Neural ODEs (e.g., consistency model, rectified flow), the resultant/distilled model lacks sampling diversity. This is a commonly-encountered problem concerning finetuning of diffusion models based on deterministic ODEs. In the future, we will explore techniques for effectively finetuing diffusion models and preserving their original properties.

Reviewer comment (Reproducibility):
After training, can the model be used for classifier / classifier-free guidance and diffusion SDEs?

Author response:
Thanks for raising this issue. 

**For finetuning of initial noises**, as long as we control the scale of perturbation on initial noises, the trained noises have little influence on the sampling process of diffusion models. Then for finetuning of text embedding, it also has little influence on score estimation as we just want to learn a proper text embedding for specific visual effects. Thus, CG/CFG and SDEs still apply in these two cases.

In related experiments, besides the guidance from the loss function defined on final outputs, we still can use CG/CFG and SDEs in sampling. For example, in security auditing, we use CG to generate images with prescribed classes and find an initial noise which can mislead pretrained ImageNet classifier while visually similar to the original one. Besides, after finding this initial noise and text embedding, we could still use SDEs (such as DDPM sampler) to generate better-quality images. The results are added to Appendix.

However, **when using AdjointDPM to finetune or train the parameters of diffusion networks**, the question of whether this operation will change the property of the underlying diffusion models is still largely open. However, we surmise that finetuning a diffusion model on a small amount of data for a few steps will still preserve the applicability of CG/CFG techniques. In practice, after finetuning the parameters of the Unet only for 10 epochs in stylization, we still use the CFG to guide the sampling with given text prompts under a guidance scale of 7.5. We also test the generation in SDE solvers and it works. The results are added to Appendix.

Besides, we observe, in some other works that distill diffusion models based on Neural ODEs (e.g., consistency model, rectified flow), the resultant/distilled model lacks sampling diversity. This is a commonly-encountered problem concerning finetuning of diffusion models based on deterministic ODEs. In the future, we will explore techniques for effectively finetuing diffusion models and preserving their original properties.

Reviewer comment (Evaluation):
While the proposal is quite novel, one might still argue that it is not really necessary to use Adjoint Backprop. One can very well accomplish the same task by backprop-ing through the solver machinery (maybe by deceasing sampling steps and using better sampler), which of course, won’t be very efficient. So, at the end, it all boils down to compute/memory efficiency. While I understand the memory advantage, sadly, the paper barely talks anything about computational requirements of the method. BTW, Neural ODEs (and Adjoint Backprop) are known to be not very scalable.

Author response:
- When we try to finetune the state-of-the-art text-to-image diffusion model (e.g., Stable Diffusion), it is nearly infeasible to do direct backpropagation. Stable Diffusion usually requires more than 25 steps to produce a high-quality image, however even V100-GPU cannot back-propagate gradients throughout 5 times of iterations (31913MiB / 32768MiB), let alone 25 iterations. 
- For the memory consumption, NFE analysis, and scalability, please refer to Global Response.

Reviewer comment (Experiments):
No comparison or mention of methods that DO backprop through the ODE solver.

Author response:
- When we try to finetune the state-of-the-art text-to-image diffusion model (e.g., Stable Diffusion), it is nearly infeasible to do direct backpropagation. Stable Diffusion usually requires more than 25 steps to produce a high-quality image, however even V100-GPU cannot back-propagate gradients throughout 5 times of iterations (31913MiB / 32768MiB), let alone 25 iterations. 
- For the memory consumption, NFE analysis, and scalability, please refer to Global Response.

Reviewer comment (Novelty):
I am confused about section 3.4. That is not AdjointDPM — that is just unconditional generation with an already known sampler (DEIS, which used Adam-Bashforth) that exploits the semi-linear nature. What part of this is your contribution ? Am I missing something here ?

Author response:
Sorry for the confusion. The exponential integration technique has been used in many diffusion ODE-type solvers, including EDM [1], DPM-solver [2], DEIS [3]. Here our point is not to compare the difference on the exponential integration technique in different ODE solvers. We emphasize that AdjointDPM can retain comparable image generation quality with existing models (see FID results in Table 1), and at the same time, it can support obtaining the gradients by solving a backward ODE. Exploiting semi-linear accelerates both the forward and backward ODEs as there also exists a semi-linear structure in the backward ODE.

[1] Karras, T., et al. (2022). Elucidating the design space of diffusion-based generative models. 

[2] Lu, C., et al. (2022). Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. 

[3] Zhang, Q., et al. (2022). Fast Sampling of Diffusion Models with Exponential Integrator.

Reviewer comment (Reproducibility):
In the “security auditing” application, what exactly is the guidance $L$ ? What is the meaning of “distance between a harmful prompt and a prediction score” ? Also, what exactly is the NSFW filter $f(\cdot)$ ? Can you provide more details please ?

Author response:
Security auditing aims to search for adversarial initial noises $\mathbf{x}_T$ that will evolve into harmful images and bypass the filter. In the paper, we present two examples: the first one sets the filter to be a pretrained ImageNet classifier (in the experiment, we use the pretrained Resnet50) and the second one is the Not Safe For Work (NSFW) Filter [2] in Stable Diffusion. 

For the ImageNet classifier, the guidance is set to be the cross entropy loss between the classification results of the generated images and any class that is not equal to the classification results (we choose this class randomly). For the NSFW filter, we set the guidance to be a weighted combination of cosine similarities between CLIP image features of the generated unsafe image and CLIP text features of some harmful concepts [2]. Usually in an NSFW filter, when the similarity of the generated image to harmful features is greater than a threshold, the generated images will be filtered. Thus, we minimize the similarity to help the generated unsafe images skip the filter.

[2] Schramowski, et al. (2023). Safe latent diffusion: Mitigating inappropriate degeneration in diffusion models.

Reviewer comment (Writing):
The unnumbered eq b/w Eq. 5 & 6 — what is the meaning of that line (”Similarly, for $\theta$, we can regard ..”) ?

Author response:
Eq. (5) introduces the adjoint state corresponding to $\mathbf{x}$. The unnumbered eq b/w Eqs. (5) and (6) introduce the augmented states $\theta$ and $t$ and their corresponding adjoint states.

Reviewer comment (Writing):
“adjoint state $\mathbf{a}(t) = ..$, which represents how the loss depends on the state ..” → “adjoint state $\mathbf{a}(t)$ = .., which represents how the loss **changes w.r.t the** state ..”.

Author response:
Thanks for pointing this out. We will correct it in the paper.

Reviewer comment (Theory):
The AdjointDPM eq. 8 shows gradient w.r.t $t$ — is it acutally used anywhere ?

Author response:
No. Gradient w.r.t $t$ is not used. Here we just introduce it for completeness in the introduction of the adjont method.

Reviewer comment (Writing):
Why is it called “vocabulary expansion” ? I still don’t get it.

Author response:
“Vocabulary expansion” is a task where we use fine-grained visual classification (FGVC) models to guide the diffusion generation process. This can help refine the vocabulary of a pretrained diffusion models such that it can generate a specific class learned by an FGVC classifier. In detail, the majority of concepts present in FGVC models are rare or non-existent in the training data of diffusion models and as such, the pretrained diffusion models cannot generate accurate images for them. Then by using AdjointDPM to guide the sampling procedure, the text-to-image model can generate the target object. This is equivalent to expanding the vocabulary (knowledge base) of the text-to-image model.