AdjointDPM: Adjoint Sensitivity Method for Gradient Backpropagation of Diffusion Probabilistic Models

Thank you for carefully reviewing this work and asking us very insightful questions! We have tried our best to answer your questions below.

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.

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.

Thank you for carefully reviewing this work. We have tried our best to answer your questions below.

Weakness Major 1. The proposed method seems to be quite inefficient ... the training cost (i.e., total training time) seems to be quite expensive.

Reply: AdjointDPM can be used for both finetuning network weights and guided sampling.

Regarding fine-tuning network weights, we only need to train the model once for a specific task. After training, we can use any type of sampling method to generate new results. The efficiency of finetuing is not an issue.

On the other hand, in guided sampling (i.e., fine-tuning noisy samples during the sampling process and text embedding inversion), AdjointDPM is much more flexible than CG or CFG. We need pairs of data of image and guidance to train a classifier or CFG. In contrast, we do not need paired data in AdjointDPM since it supports any type of computable metrics defined on the generated results. Besides, it is not possible for CG or CFG to perform some tasks like security auditing or single image guided generation.

Weakness Major 2. The proposed method cannot guarantee the property of diffusion models..., which is based on the score functions and thus can maintain the diffusion property.

Question 2. After training, can the model be used for classifier / classifier-free guidance and diffusion SDEs?

Reply: 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.

Weakness Minor 1. Equation(9) is exactly the EDM sampler so it is not a new method. It should be compared and discussed in detail.

Reply: 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.

Question 1. What is the training time for each experiment?

Reply: For the analysis of NFE, please refer to the global response. The actual training time highly depends on 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), 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.

Thank you for carefully reviewing this work and asking us very insightful questions! We have tried our best to answer your questions below.

Weakness 1. While the proposal is quite novel, one might still argue that it is not really necessary to use Adjoint Backprop....While I understand the memory advantage, sadly, the paper barely talks anything about the computational requirements of the method. BTW, Neural ODEs (and Adjoint Backprop) are known to be not very scalable. Weakness 3. No comparison or mention of methods that DO backprop through the ODE solver.

Reply:

• 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.

Weakness 2. Experiments are okay-ish, but not really extensive. Qualitative samples are lacking in some experiments (e.g. vocabulary expansion). Also Vocabulary Expansion is shown for only two classes.

Reply: Thanks for raising this point. We extend our results comparison to 90 classes in Standford dogs Dataset. We compute the change of FID relative to the original Stable Diffusion generations for comparison. The results are shown in the following table. As there is limited time to run DOODL's result for us, we use the value reported in [1]. We also did not extend the comparison to more classes on the Birds dataset. Complete quantitative results will be added to the final version of the manuscript in Table 2. More qualitative results are added in Appendix.

[1] Bram Wallace, et al. End-to-end diffusion latent optimization improves classifier guidance.

Q1. 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?

Reply: 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.

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

Reply: 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.