DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization

We would like to express our sincere gratitude for the detailed and professional attention you have given to our work during the review process.

Re: Weakness#1

    Yes, the vast search space actually introduces challenges to the learning process for obtaining an optimal SubNet. Therefore, we choose the gradient-based optimization method to tackle this challenge instead of the traditional search methods and design the pruner network with specific optimization losses. And the Table 6 of main submission shows that our method surpasses the search methods in both accuracy and efficiency. Moreover, Table 1, 2, 3 and 4 in main submission show that our method can obtain optimal SubNets in the search space, which is benefited to the gradient-based optimization method. We can reduce the search space by reducing the number of pruning steps and pruning blocks, which may hinder the achievement of an optimal SubNet than before and sacrifices accuracy of pruning result.

Re: Weakness #2

    Yes, although our method can maintain the main object content of images under very high compression ratios, subtle changes in the background still occur. This is a common issue in model acceleration tasks, and it is essential to consider the accuracy-speed trade-off for the deployment scenario when using our method. We also show this phenomenon in Figure 3 of the main submission, and our method can maintain the generation quality with a significant acceleration on SD1.5 with 50-step sampler.

Re: Question #1

    The additional backup connections are predefined in the dense SuperNet before training the pruner network. As shown in Figure 1b of main submission, a block would build backup connections for its dependent blocks and the backup connections are from the corresponding dependent blocks in previous timestep to it. The pruner network would select either the original connection or the backup connection, but not both simultaneously, during training and testing phases based on the importance score predicted by it.

Re: Question #2

   The training cost of our method is not high. For example, using SD-1.5 with a 50-step sampler, our method requires approximately 2.3 hours of training on a single MI250 GPU as shown in Table 6, which is significantly less than the time cost of traditional pruning methods (10% to 20% of the cost compared to pre-training) [1] and the SD-1.5 training takes about 150K GPU hours for pretraining [2]. And our method requires about 4 hours of training a pruner network for DiT with 250-step sampler, which is also efficient.  

Re: Question #3

    For different datasets, the threshold τ is set to 0.2, it's related to the pruned ratio requirement.  The τ should be sat smaller at higher pruning ratios and bigger at lower pruning ratios.

[1] Structural Pruning for Diffusion Models. NeurIPS 2023.

[2] Stable-diffusion-v1-5 in hugging face.

Dear Reviewer JBdU,

Thank you for reviewing our paper and providing thoughtful feedback. We have addressed your concerns in our rebuttal, which are summarized as follows:

1. Vast Search Space Challenge: We detailed how our gradient-based optimization method overcomes this challenge and outperforms existing methods, such as GA Search.

2. Image Quality under Higher Pruning Ratios: We provided a more detailed analysis of our method under high pruning ratios and clarified that considering the accuracy-speed trade-off is crucial for deployment scenarios when using our method.

3. Clarification: We explained the concept of backup connections in the SuperNet and provided more details about the threshold $\tau$.

4. Computational Resources: We detailed the computational resources required to train the pruner network, demonstrating that our method is significantly more efficient than traditional pruning methods.

Given these clarifications and the positive aspects of our work that you have acknowledged, we kindly request that you reconsider your evaluation. We believe our research contributes valuable insights to the field and addresses key challenges in diffusion model acceleration.

With approximately 15 hours remaining in the discussion phase, we hope our rebuttal has resolved your concerns. If so, we would appreciate if you could consider raising your score.

If you have any further questions or need additional information, please let us know. We are committed to providing any further clarification needed within the remaining time.

Thank you again for your time and feedback.

Best regards,

Authors

We would like to express our sincere gratitude for the detailed and professional attention you have given to our work during the review process.

Re: Weakness #1

    $D$ is the embedding dimension of a learnable query in the pruner network, and the dimension of the prune queries is $T \times N \times D$ due to the $N \times T$ queries mentioned in the caption of Figure 2.  We clarify its meaning and value in the experimental section on Line 267-268 of the paper.  But there is a typo, which changes $D$ to $d$. We will fix it in the revised version.

Re: Weakness #2

    We tested our method on SD-1.5 with 50 steps to align with the configuration of DeepCache [1], the main SOTA method for comparison. Additionally, we tested our method on DiT with 250 steps, as specified in the official DiT paper [2]. Our method achieved significant acceleration while minimizing loss on DiT with the 250-step sampler, demonstrating the effectiveness of our gradient-based optimization method. For models with fewer steps, it is easier to obtain the optimal SubNet to reduce redundant computations due to a smaller search space. We also tested our method on SD-2.1 with a 25-step DPM-Solver sampler and PixArt-$\alpha$ with a 20-step DPM-Solver sampler, achieved excellent pruning results, as shown in Table 1 and Table 2 below.

Table 1: Comparison with a 25-step DPM-Solver sampler for SD model. We evaluate the effectiveness of our methods on COCO2017 validation set.

Table 2: Comparison with a 20-step DPM-Solver sampler for diffusion transformer model. We evaluate the effectiveness of our methods on COCO2017 validation set.

Re: Weakness #3

   The speedup mentioned in the paper is based on actual measurements, not theoretical measurements. Detailed descriptions of the environment and platform used for these measurements can be found on Line 272 of the main submission. The testing method adheres to the Deepcache method for fair comparison, and the actual inference latency is presented in Table 3 of the main submission. We will include further details about the measurement method in the revised version.

Re: Weakness #4

We apologize for missing this data. We have now included two additional few-step results in Table 3, and will include these updates in the revised version of the paper.

Table 3: Comparison of fast sampler methods under pruned-0.6.

Re: Weakness #5

    We first choose 50-step DPM sampler due to its higher quality with more sampling steps, and the DPM sampler using 50 steps is also found in the Table 2 of the reference paper [3].

Further, we provide the pruning result on SD-2.1 with 25-step sampler in above Table 1. As shown in Table 1, our method can prune 50% computation nearly without loss.

Re: Weakness #6

   Yes, you are right. Line 298 in the paper explains this.  Our method benefits from information redundancy in multi-step optimization processes, the LCM model has fewer redundancy in features across adjacent timesteps due to its efficiency.  For better pruning results, we think it's necessary to train the original diffusion model.

[1] Deepcache: Accelerating Diffusion Models for Free, CVPR 2024.

[2] Scalable Diffusion Models with Transformers, ICCV 2023.

[3] Cross-Attention Makes Inference Cumbersome in Text-to-Image Diffusion Models, arXiv.

[4] PixArt-$\alpha$: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis, ICLR24.

We suspect that there might have been a system issue that caused our previous response to not be sent successfully. Therefore, we are resending our response and hope it hasn’t caused you any inconvenience. Below is the main content:

We would like to express our sincere gratitude for your detailed review.

Re: Weakness #1

$D$ is the embedding dimension of a learnable query in the pruner network, and the dimension of the prune queries is $T \times N \times D$ due to the $N \times T$ queries mentioned in the caption of Figure 2. We clarify its meaning and value in the experimental section on Line 267-268 of the paper. But there is a typo, which changes $D$ to $d$. We will fix it in the revised version.

Re: Weakness #2

We tested our method on SD-1.5 with 50 steps to align with the configuration of DeepCache [1], the main SOTA method for comparison. Additionally, we tested our method on DiT with 250 steps, as specified in the official DiT paper [2]. Our method achieved significant acceleration while minimizing loss on DiT with the 250-step sampler, demonstrating the effectiveness of our gradient-based optimization method. For models with fewer steps, it is easier to obtain the optimal SubNet to reduce redundant computations due to a smaller search space. We also tested our method on SD-2.1 with a 25-step DPM-Solver sampler and PixArt-$\alpha$ with a 20-step DPM-Solver sampler, achieved excellent pruning results, as shown in Table 1 and Table 2 below.

Table 1: Comparison with a 25-step DPM-Solver sampler for SD model. We evaluate the effectiveness of our methods on COCO2017 validation set.

Table 2: Comparison with a 20-step DPM-Solver sampler for diffusion transformer model. We evaluate the effectiveness of our methods on COCO2017 validation set.

Re: Weakness #3

The speedup mentioned in the paper is based on actual measurements, not theoretical measurements. Detailed descriptions of the environment and platform used for these measurements can be found on Line 272 of the main submission. The testing method adheres to the Deepcache method for fair comparison, and the actual inference latency is presented in Table 3 of the main submission. We will include further details about the measurement method in the revised version.

Re: Weakness #4

We apologize for missing this data. We have now included two additional few-step results in Table 3, and will include these updates in the revised version of the paper.

Table 3: Comparison of fast sampler methods under pruned-0.6.

Re: Weakness #5

We first choose 50-step DPM sampler due to its higher quality with more sampling steps, and the DPM sampler using 50 steps is also found in the Table 2 of the reference paper [3].

Further, we provide the pruning result on SD-2.1 with 25-step sampler in above Table 1. As shown in Table 1, our method can prune 50% computation nearly without loss.

Re: Weakness #6

Yes, you are right. Line 298 in the paper explains this. Our method benefits from information redundancy in multi-step optimization processes, the LCM model has fewer redundancy in features across adjacent timesteps due to its efficiency. For better pruning results, we think it's necessary to train the original diffusion model.

[1] Deepcache: Accelerating Diffusion Models for Free, CVPR 2024.

[2] Scalable Diffusion Models with Transformers, ICCV 2023.

[3] Cross-Attention Makes Inference Cumbersome in Text-to-Image Diffusion Models, arXiv.

[4] PixArt-$\alpha$: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis, ICLR24.

We believe that the additional explanations and data substantiate a higher impact in these areas, reflecting the rigor and potential impact of our work. If our explanations have resolved your concerns, we would be grateful if you could reconsider your rating.

Once again, thank you for your valuable time and expert opinion.

Best regards, Authors

We sincerely appreciate the time and effort you invested in reviewing our paper. Your acknowledgment of the significant improvements in efficiency and effectiveness across different architectures is highly valued. We are also grateful for your recognition that our method avoids retraining diffusion models, thereby saving substantial computational resources and time.

We address the raised concerns as follows:

Re Weaknesses #1:

Gradient checkpointing and FP16 mixed precision are standard practices for training diffusion models, as recommended by Hugging Face. These techniques are crucial for long-timestep diffusion models (e.g., DiT, which requires 250 timesteps for inference). However, for models with fewer timesteps (e.g., SD-1.5 with 50 timesteps), these techniques are not necessary.

Additionally, we apply FP16 precision only to the diffusion model during the forward pass to reduce memory overhead, while keeping the pruner network and backward gradients in FP32 to maintain performance. Furthermore, additional acceleration techniques such as DeepSpeed [1] and ZeRO [2] can be employed to further reduce memory usage.

Re Questions #1:

Feature similarity across adjacent time steps in fast samplers has been confirmed in recent works. We have cited their conclusions in Line 156-157 of the manuscript.  We also analyzed the feature similarity between adjacent steps in the fast sampler. Specifically, we sampled 200 samples from COCO2017 validation set and calculate the average cosine similarity between the features of the penultimate upsample block across all $T$ steps for two typical fast samplers, creating a $T \times T$ similarity matrix shown in Figure 2 (refer to the rebuttal pdf). The heat map in Figure 2 illustrates the high degree of similarity between features of consecutive time steps.

Re Questions #2:

We conducted an ablation study on α and present the results below: our method achieves the best performance when α=1.0. We will include this additional comparison Table 1 in the next version of the manuscript.

Table 1: Comparison of different α values. Pruning experiments with 80% pruning ratio were conducted on COCO2017 validation using SD-1.5.

[1] DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters, KDD 2020, Tutorial.

[2] ZeRO: memory optimizations toward training trillion parameter models, SC 2020.

We are deeply thankful for the thorough and expert review of our work. Your acknowledgment of the timely and practically-relevant problem our research tackles is greatly appreciated. Your feedback highlights the superior performance and the extensive experiments that were conducted.

Below are our detailed responses to the weaknesses and questions raised:

Re: Weakness #1

We sincerely appreciate your valuable feedback. We believe your concerns primarily stem from two aspects:

1. Our pruner network uses learnable queries to predict the importance scores of $T \times N$ blocks, rather than directly predicting the similarity between adjacent steps. Although this importance score is influenced by the similarity between time steps during training.

2. The idea of using time embedding as an alternative to the T dimension is a good design choice for the queries. We incorporated time embedding as a condition, added it to the N queries, and then conducted experiments with a pruning network. On SD1.5, we achieved a 29.49 CLIP score with a pruning rate of 80%, which is lower than our DiP-GO with a 30.29 CLIP score. This design helps reduce the number of parameters when T is very large, requiring $N \times D + D_t \times D$ parameters (where $D_t$ is the time embedding dimension), while queries in our method require $T \times N \times D$ parameters. However, since $T$ is usually a constant, this design does not result in a significantly smaller network size or faster training.

Re: Weakness #2

Thank you for the suggestion. We provide qualitative comparisons with baseline, DeepCache, since Diff-pruning lacks pruning results for SD and DiT, as shown in Figure 3 (refer to the rebuttal PDF). And we will add them to the revised version.

Re: Weakness #3

1. Our method exhibits a specific pattern of pruning ratios with respect to the timesteps. As shown in Figure 1 (refer to the rebuttal pdf), fewer blocks are pruned during the middle denoising stage (approximately between steps 65 and 150), as this is when the image content is rapidly being generated. Conversely, the pruning ratio in the latter stage is higher since the content has already taken shape.

2. Yes, we conducted an ablation study on γ. Without gamma, pruning 80% on SD1.5 resulted in a CLIP score of 29.50 (w/ γ: 30.29).

Re: Weakness #4

We also tried simply skipping time steps for every N steps. When N = 2, the CLIP score on COCO was 19.74 (Ours: 30.29), which resulted in much lower speedup and accuracy compared to our method. This indicates that simply skipping steps leads to significant performance loss, while our method of searching for the pruner gate provides a more effective pruning strategy.

Re: Question #1

Thank you for pointing this out. We will refine the definition of the covariance constant β in the final version.

Re: Question #2

Yes. The flops ratio $\gamma$ is in the range [0, 1].

Re: Question #3

The sparsity threshold 0.2 is a hyperparameter used to prevent the network from overly optimizing the sparsity loss, which could lead to all-zero gates predictions during training.

We suspect that there might have been a system issue that caused our previous response to not be sent successfully. Therefore, we are resending our response and hope it hasn’t caused you any inconvenience. Below is the main content:

We are deeply thankful for the thorough and expert review of our work. Your acknowledgment of the timely and practically-relevant problem our research tackles is greatly appreciated. Your feedback highlights the superior performance and the extensive experiments that were conducted.

Below are our detailed responses to the weaknesses and questions raised:

Re: Weakness #1

We sincerely appreciate your valuable feedback. We believe your concerns primarily stem from two aspects:

1. Our pruner network uses learnable queries to predict the importance scores of $T\times N$ blocks, rather than directly predicting the similarity between adjacent steps. Although this importance score is influenced by the similarity between time steps during training.

2. The idea of using time embedding as an alternative to the $T$ dimension is a good design choice for the queries. We incorporated time embedding as a condition, added it to the $N$ queries, and then conducted experiments with a pruning network. On SD-v1.5, we achieved a $29.49$ CLIP score with a pruning rate of 80%, which is lower than our DiP-GO with a $30.29$ CLIP score. This design helps reduce the number of parameters when $T$ is very large, requiring $N\times D + D_t \times D$ parameters (where $D_t$ is the time embedding dimension), while queries in our method require $T\times N \times D$ parameters. However, since $T$ is usually a constant, this design does not result in a significantly smaller network size or faster training.

Re: Weakness #2

Thank you for the suggestion. We provide qualitative comparisons with baseline, DeepCache, since Diff-pruning lacks pruning results for SD and DiT, as shown in Figure 3 (refer to the rebuttal PDF). And we will add them to the revised version.

Re: Weakness #3

1. Our method exhibits a specific pattern of pruning ratios with respect to the timesteps. As shown in Figure 1 (refer to the rebuttal pdf), fewer blocks are pruned during the middle denoising stage (approximately between steps 65 and 150), as this is when the image content is rapidly being generated. Conversely, the pruning ratio in the latter stage is higher since the content has already taken shape.

2. Yes, we conducted an ablation study on γ. Without gamma, pruning 80% on SD1.5 resulted in a CLIP score of 29.50 (w/ γ: 30.29).

Re: Weakness #4

We also tried simply skipping time steps for every N steps. When $N = 2$, the CLIP score on COCO was 19.74 (Ours: 30.29), which resulted in much lower speedup and accuracy compared to our method. This indicates that simply skipping steps leads to significant performance loss, while our method of searching for the pruner gate provides a more effective pruning strategy.

Re: Question #1

Thank you for pointing this out. We will refine the definition of the covariance constant β in the final version.

Re: Question #2

Yes. The flops ratio $gamma$ is in the range [0, 1].

Re: Question #3

The sparsity threshold 0.2 is a hyperparameter used to prevent the network from overly optimizing the sparsity loss, which could lead to all-zero gates predictions during training.

We appreciate the time and effort you have invested in reviewing our paper. We have provided a comprehensive rebuttal addressing your concerns, which can be summarized as follows:

1. Alternative Design: We explained the motivation behind our learnable queries design and incorporated your advice by providing an additional comparison with time embedding conditioned queries. This supports the rationale of our design both theoretically and experimentally.

2. More Qualitative Comparison: We included a more detailed qualitative comparison with respect to baselines (see the rebuttal PDF), demonstrating that our method generates higher fidelity samples with better content consistency compared to the baselines.

3. Pruning Patterns: We presented the pruning gates of our learned pruner network and analyzed the pruning ratios with respect to the timesteps.

4. Hyperparameters Ablation: We conducted an additional ablation study on the parameter $\gamma$, demonstrating its effectiveness.

5. Direct Skip Steps Comparison: We evaluated a direct skip step design and found that our method achieves better performance.

6. Clarity: We provided additional details on the covariance constant $\beta$, FLOPs ratio $\gamma$, and the sparsity threshold used in our paper.

With remaining time in the discussion phase, we sincerely hope our rebuttal has addressed your concerns. If our responses have clarified the issues you raised, we kindly request that you consider raising your score.

Once again, thank you for your valuable time and expert opinion.

Best regards, Authors