SparseDM: Toward Sparse Efficient Diffusion Models

Authors: Kafeng Wang, Jianfei Chen, He Li, Zhenpeng Mi, Jun Zhu
Primary area: generative models
Keywords: Diffusion models, sparse pruning, 2:4 sparsity

Abstract:
Diffusion models have been extensively used in data generation tasks and are recognized as one of the best generative models. However, their time-consuming deployment, long inference time, and requirements on large memory limit their application. In this paper, we propose a method based on the improved Straight-Through Estimator to improve the deployment efficiency of diffusion models. Specifically, we add sparse masks to the Convolution and Linear layers in a pre-trained diffusion model, then transfer learn the sparse model during the fine-tuning stage and turn on the sparse masks during inference. Experimental results on a Transformer and UNet-based diffusion models demonstrate that our method reduces MACs by 50% while increasing FID by only 0.44 on average.  Sparse models are accelerated by approximately 1.2x on the GPU. Under other MACs conditions, the FID is also lower than 1 compared to other methods.

Decision: reject
Reviewer scores: 3, 5, 3, 5

Reviews and discussion:
<thread>
creation date:
1729410496976

thread_type:
official_review

<thread_object>
creator: reviewer_x5e7

<content>
summary:
This paper introduces SparseDM, which converts existing diffusion models into sparse models that fit in a 2:4 sparse operator on the GPU. Specifically, the authors propose a Straight-Through Estimator (STE)-based fine-tuning framework that learns sparse masks. These sparse masks accelerate GPU inference speed up to 1.2. Comprehensive experiments validate the effectiveness of the proposed method.

strengths:
* The paper introduces a simple fine-tuning method that converts existing diffusion models into sparse models, enabling them to be used in scenarios with limited computing power, such as on mobile devices.

* The observations about fixed sparse training are interesting.

* Experiments on various generation scenarios verify the effectiveness of SparseDM compared to baselines.

weaknesses:
**Weakness 1: More clarifications on Section 2.3.**

In Section 2.3, the authors claim that diffusion models only consider the distribution shift of the noisy data while sparse pruning methods only consider the model's weight change. Then, referring to RFR, the authors convert the model's weight changes resulting from sparse pruning methods into data changes for the diffusion model's training process. However, typical diffusion models have indicators for perturbed data (such as the noise schedule and timestep embedding), and it is unclear how these relate to perturbations caused by sparse training.

**Weakness 2: Lack of analysis of fixed sparse training**

I am not sure why fixed sparse training would be more effective than traditional progressive sparse training. Based on the experimental results, it seems that fixed sparsity applies a consistent distribution shift across all noise levels in diffusion training, whereas progressive sparse training gradually shifts the predefined noise levels, which may hinder the diffusion training process. However, this claim has not been theoretically verified, so the authors should provide theoretical proof to demonstrate the relationship between diffusion training and sparse training.

questions:
* In Table 3, some variants (e.g., patch size = 2 and mlp_ratio = 2) are slower than the dense model, why do you think this is?
* I think it would strengthen the effectiveness of SparseDM if the author show that it can also be applied to models like Stable Diffusion.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730280326718

thread_type:
official_review

<thread_object>
creator: reviewer_q9z9

<content>
summary:
This work aims to reduce the computation of Diffusion Models during inference. The authors suggest a method of straight-through estimation, which applies sparse masks to layers of a pretrained diffusion model and then employs transfer learning for training. Then, they use the same sparse mask during inference to improve compute efficiency.

strengths:
- The 2:4 sparse model calculation offers practical values for practitioners using NVIDIA Ampere architecture GPUs.

weaknesses:
- While it may have some practical value for practitioners using NVIDIA Ampere architecture, the same technique may not benefit other practitioners or general researchers without access to Ampere architecture. 

- Besides, the straightforward idea of using masked training is neither interesting nor technically new. 

- More disappointingly, the speed acceleration due to this customized training for a particular architecture increases by x1.2 only. Studies related to reducing time steps for Diffusion inference or diffusion quantization/pruning methods may be more effective in achieving the same purpose.

questions:
Please address the weakness stated above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730359724788

thread_type:
official_review

<thread_object>
creator: reviewer_f6dv

<content>
summary:
This paper proposes to improve the efficiency of DM by sparse matrix for 2:4 sparse acceleration GPU. The authors improve the STE method and propose to gradually transfer knowledge from dense models to sparse models.

strengths:
1.	This paper is well-written.
2.	The motivation is clear enough.
3.	The organization of this paper is great.

weaknesses:
1.	There is a typo in Eq5. Please also check all equations. Moreover, not all symbols have been explained.
2.	The experiments are relatively limited. Specifically, only two U-ViT and DDPM are tested on the proposed pruning, which are proposed in 2022 and 2020 respectively. More recently proposed DiT or other methods should also be included.
3.	The limitation and discussion are missing in this paper.

questions:
1.	The authors mentioned that “it does not mean that the greater the sparsity, the better the FID”. Please discuss the reason and why you choose 2:4 sparse.
2.	Please discuss the reason ASP performs so worse in all experiments.
3.	Please also clarify why your method and STE-based pruning fulfill the same MACs.
4.	Please explain the reason that the FID of the proposed method in Fig. 3a obtain a lower FID in the first several steps.
5.	Why the initial FID of 2:4 sparse in Fig.3b and Fig.3d is different?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700332292

thread_type:
official_review

<thread_object>
creator: reviewer_qpma

<content>
summary:
This paper proposes a pruning strategy for Diffusion models, using mask pruning to achieve progressive multi-step pruning. Ultimately, it realizes 1:2 pruning according to the Ampere architecture. During training, knowledge distillation is used to transfer knowledge from the full model to the pruned model.

strengths:
The writing is very clear, and the main idea is highlighted effectively.

weaknesses:
1. The pruning strategy is based on existing structures, with a relatively simple motivation. There are already other methods that achieve similar results, such as using linear attention or directly training a smaller model with distillation. 
2. Compared to directly using STE-based pruning, it does not further reduce the computational load.
3. In Section 3.2, "Transfer learn sparse diffusion models" strategy is mentioned, but it does not explain the significant differences between this strategy and the progressive sparse training strategy discussed in Section 2.2. If the focus is solely on testing with perturbed datasets, it may not constitute a significant contribution.
4. A generalized pruning strategy suitable for Transformer networks has not been proposed; simply relying on data perturbations is insufficient to demonstrate applicability to other datasets. Further testing on additional datasets, such as CelebA-HQ, LSUN Church, would be beneficial.
5. Many of the latest comparative algorithms from 2024 are not mentioned, such as "Pruning for Robust Concept Erasing in Diffusion Models" and "LD-Pruner: Efficient Pruning of Latent Diffusion Models using Task-Agnostic Insights."
6. There is no comparison of the parameter counts for each layer of the SD model before and after sparse pruning. It is recommended to include a chart in the appendix to illustrate this.
7. While Section 2.3 mentions applying perturbations to the dataset, it does not provide specific details on how the perturbations were implemented.
8. The experiments only validate the FID score as a single metric; it is advisable to explore additional metrics, such as SSIM.

questions:
Could the authors provide the parameter counts for each layer of the SD model before and after sparse pruning?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732637524235

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
JOINT REBUTTAL

comment:
Dear Reviewers:

We want to thank you all for the time spent reviewing our paper and for the constructive comments and feedback provided. 
We are pleased that our paper was found to be a good presentation (Reviewers QPmA, F6DV, and Q9Z9), and a good contribution (Reviewer X5E7). We also appreciate that our research has been recognized as well-motivated (Reviewer F6DV) and highlighted effectively (Reviewer QPmA). 

We have noticed that the most common concern is about the novelty of our work and the need to compare our approach with existing Stable Diffusion (SD) pruning methods.

Our transfer learning diffusion model strategy is significantly different from the progressive sparsity strategy. There are two main differences. 1. Our method must perform progressive sparsity across diffusion models. 2. Our method must transfer knowledge across models. 
These two improvements are aimed at solving the problem that the progressive sparsity strategy fails in the diffusion model.

Our approach achieves a general-purpose diffusion model acceleration. It is not targeted at a specific task.
We are performing sparse acceleration on the general basic diffusion model, so it has been verified to be effective on the diffusion models based on Transformer and U-Net.
The basic diffusion model of SD 1.5 is also based on U-Net, and the basic diffusion model of SD 3.0 is also based on Transformer.
Therefore, our method can be generalized to SD. 
To prevent the influence of multi-modal models, such as CLIP, in SD, we currently only perform experiments on the diffusion models, U-ViT and DDPM.
Existing SD pruning methods, such as "Pruning for Robust Concept Erasing in Diffusion Models" and "LD-Pruner: Efficient Pruning of Latent Diffusion Models using Task-Agnostic Insights" are a speedup after fine-tuning for specific tasks.

Overall, we hope that the reviewers will reconsider raising their scores, as we truly believe that our approach provides novel and valuable information to the community.

Below, we discuss each reviewer's concerns and indicate how we have addressed them in the revised version of the paper.

The Authors
</content>
</thread_object>

</thread>


Metareview: