SparseDM: Toward Sparse Efficient Diffusion Models

## Abstract

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 use design progressive sparse model training in the fine-tuning stage, and switch the inference mask on, which supports sparse inference.
Experiments on a Transformer-based and UNet-based diffusion model demonstrate that our method reduces MACs by 50 % percent 50 50\% 50 % while increasing FID by only 1.22 on average. Sparse models are accelerated by approximately 1.2x on the GPU.
Under other MACs conditions, the FID is also lower than 1 ∼ similar-to \sim ∼ 137 compared to other methods.

## 1 Introduction

Diffusion Models (DM) Song et al. (); Karras et al. () have been one of the core generative modules in various computer vision tasks with their advantages on both sample quality and mode coverage. Generally, they are composed of two stages
: a forward/diffusion process that perturbs the data distribution to learn the time-dependent score functions, and a reverse/sampling process that generates data samples from a prior distribution in an iterative manner. However, their slow inference speed and heavy computational load during the inference process inevitably restrict their applications on most mobile devices.

To reduce the computational load in the inference process of diffusion models, various methods have been proposed to minimize the number of inference steps, such as the training-free samplers Bao et al. (); Zhao et al. (); Lu et al. (, ); Zheng et al. (); Chen et al. () and the distillation methods Salimans and Ho (); Luo et al. (, ). However, their sample quality is still unsatisfactory as a few sampling steps cannot faithfully reconstruct the high-dimensional data space, e.g., the image or video samples.
This problem is even more pronounced on mobile devices with limited computing capabilities.
Simultaneously, a few works explore reducing the Multiple-Accumulate operations (MACs) at each inference step Bolya and Hoffman (); Fang et al. (). However, their work cannot be accelerated on GPUs.
Since DM is an intensive model parameter calculation, NVIDIA Ampere architecture GPU supports a 2:4 sparse (4-weights contain 2 non-zero values) model calculation, which can achieve nearly 2 times calculation acceleration Pool and Yu (); Mishra et al. ().
Although structural pruning Fang et al. () has been used in DM, 2:4 structured sparsity inference has not been implemented.
Our method aims to reduce MACs at each step through 2:4 and other scale sparsity.

Recently, several popular computational architectures, such as NVIDIA Ampere architecture and Hopper GPUs, have developed acceleration methods for model inference and have been equipped with fine-grained structured sparse capabilities. A common requirement of these acceleration techniques is the 2:4 sparse mode, which only preserves 2 of the 4 adjacent weights of a pre-trained model, i.e., requires a sparsity rate of $50\%$. Given this sparsity, the acceleration techniques only process the non-zero values in matrix multiplications, theoretically achieving a 2x speedup. In essence, a diffusion model supporting the 2:4 sparsity mode can reduce $50\%$ computational load at each inference step, which is valuable when considering the iterative refinement process in a generation. To the best of our knowledge, previous works have not explored maintaining the sample quality of diffusion models with a 2:4 or other scale sparsity mode, which motivates us to present the techniques in this work.

We use a state-of-the-art Transformer-based DM, named U-ViT Bao et al. (), to analyze the shortcomings of existing 2:4 structured sparse tools.
Also, we manually design a half-size DM network to test FID and MACs.
From Table 1, we can see that simply reducing the U-ViT model parameters by nearly $50\%$ may cause FID to collapse catastrophically.
Automatic Sparsity (ASP) Pool and Yu (); Mishra et al. ()) for model sparse training consumes a lot of GPU time, but FID is poor.
Therefore, for DM, we need to redesign the pruning method to reduce the amount of calculation and maintain FID as much as possible.

Table: Table 1: The results of manually redesigning the UViT model and sparse pruning the UViT model using ASP. UViT Small is from the U-ViT method. Half UViT Small is set to half the depth and number of heads of UViT Small. ASP is the method of sparse pruning from Nvidia.

This paper proposes a new method for implementing 2:4 structured sparsity and other scale sparse inference.
(1) Sparse masks are applied to the Convolution and Linear layers.
(2) The STE method improved by sparse regularization is used to train the sparse model.
(3) Progressive sparse models are trained by transferring knowledge from a dense model to a sparse model.

Our contributions can be listed as follows:

- •
We propose a progressive sparse finetuning method that achieves 2:4 structured sparse inference and other scale sparse pruning for DMs with a Transformer or UNet backbones.
- •
We conduct experiments on four datasets. The average FID of 2:4 sparse DM is only increased by 1.22 compared with the dense model. The inference acceleration on the GPU is approximately 1.2x.
- •
Testing nine scales, with similar MACs for each scale, our sparse inference FID is $1\sim 137$ lower than other methods.

## 2 Problem Formulation

### 2.1 Diffusion Models

The diffusion model is divided into a forward process and a backward process.
The forward process is a step-by-step process of adding noise to the original image to generate a noisy image,
generally formalized as a Markov chain process.
The reverse process is to remove noise from a noisy image and restore the original image as much as possible.
Gaussian mode was adopted to approximate the ground truth reverse transition of the Markov chain.
The training process of diffusion models is the process of establishing noise prediction models.
The loss function of DM is minimizing a noise prediction objective.
The forward process is formalized as:

$$ $q({\bm{x}}_{1:T}|{\bm{x}}_{0})=\prod_{t=1}^{T}q({\bm{x}}_{t}|{\bm{x}}_{t-1}),$ (1) $$

where ${\bm{x}}_{t}$ is input data at $t$.
The backward process is formalized as:

$$ ${\bm{\mu}}_{t}^{*}({\bm{x}}_{t})=\frac{1}{\sqrt{\alpha_{t}}}\left({\bm{x}}_{t} -\frac{\beta_{t}}{\sqrt{1-\overline{\alpha}_{t}}}\mathbb{E}[{\bm{\epsilon}}|{ \bm{x}}_{t}]\right),$ (2) $$

where $\alpha_{t}$ and $\beta_{t}$ are the noise schedule at $t$, $\alpha_{t}+\beta_{t}=1$. $\overline{\alpha}_{t}=\prod_{i=1}^{t}\alpha_{i}$, and ${\bm{\epsilon}}$ is the standard Gaussian noises added to ${\bm{x}}_{t}$.
The DM training is a task of minimizing the noisy prediction errors, expressed as:

$$ $\min_{\mathcal{W}}\mathcal{L}(\mathcal{W};\mathcal{D}),$ (3) $$

where $\mathcal{D}$ is dataset, $\mathcal{L}$ is loss function, $\mathcal{W}$ is dense weights.
The DM inference, data generation, is expressed as:

$$ ${\bm{y}}_{t}=\mathcal{F}(\mathcal{W};{\bm{x}}_{t}).$ (4) $$

where $\mathcal{F}$ is trained DM, ${\bm{y}}_{t}$ is the output of DM inference.

### 2.2 Sparse Pruning

Sparse network computation is an effective method to reduce MACs in deep networks and accelerate computation.
Currently, one-shot sparse training and progressive sparse training are commonly used.

The one-shot training method is easy to use, and the steps are as follows Mishra et al. (); Pool and Yu ():
Train a regular dense model.
Prune the weights on the fully connected and convolutional layers in a 2:4 sparse mode.
Retrain the pruned model.
The one-shot pruning is expressed as:

$$ $\begin{split}\widetilde{\mathcal{W}}&=\mathcal{W}\odot\mathcal{M},\end{split}$ (5) $$

where, $\widetilde{\mathcal{W}}$ is the sparse weight after the dense weight $\mathcal{W}$ is pruned, recorded as
$\widetilde{\mathcal{W}}\leftarrow\text{Pruning}(\mathcal{W})$, $\mathcal{M}$ is a 0-1 mask, $\mathcal{M}$ is the ($N$:$M$) sparse mask ($M$-weights contain $N$ non-zero values), $\odot$ represents element-wise multiplication.

One-shot sparse pruning can cover most tasks and achieve speedup without losing accuracy. However, for some difficult tasks that are sensitive to changes in weight values, doing sparse training for all weights at once will result in a large amount of information loss Han et al. ().
With the same number of tuning iterations, progressive sparse training can achieve higher model accuracy than one-shot sparse training.
Suppose there are $k$ masks $[\mathcal{M}_{1},\mathcal{M}_{2},...,\mathcal{M}_{k}]$,
the progressive sparse process is formalized as:

$$ $\widetilde{\mathcal{W}}=\mathcal{W}\odot\mathcal{M}_{i},~{}~{}\mathcal{M}_{i} \in[\mathcal{M}_{1},\mathcal{M}_{2},...,\mathcal{M}_{k}].$ (6) $$

The reason why traditional progressive sparse training methods work well is to reuse the knowledge of dense models as much as possible.
However, this progressive sparse training method has only proven effective when the training data distribution is stable, such as when training classification models.

### 2.3 Distribution Shift on DM Training and Sparse Pruning

Robust Fairness Regularization (RFR) Jiang et al. () improved that distribution shift can be transformed as data perturbation, and data perturbation and model weight perturbation are equivalent for classifier models. Assuming that the perturbation of the label is not considered, the equivalent of data perturbation and model weight perturbation is formalized as:

$$ $\mathbb{E}_{\delta_{\mathcal{D}}(\mathcal{D})}\mathbb{E}_{(\mathcal{D})\sim \mathcal{P}}[\mathcal{L}(\mathcal{F}_{\mathcal{W}}(\mathcal{D}+\delta_{ \mathcal{D}}(\mathcal{D}))]=\mathbb{E}_{(\mathcal{D})\sim\mathcal{P}}[\mathcal {L}(\mathcal{F}_{\mathcal{W}+\Delta\mathcal{W}}(\mathcal{D})).$ (7) $$

where the training dataset $\mathcal{D}$ with distribution $\mathcal{P}$, suppose the training dataset is perturbed with data perturbation $\delta$, and the neural network is given by $\mathcal{F}_{\mathcal{W}}(\cdot)$, for the general case, there exists model weight pruned as perturbation $\Delta\mathcal{W}$, so that the training loss on perturbed training dataset is the same with that for model weight perturbation $\Delta\mathcal{W}$ on training distribution.

The training process of diffusion models generally involves only the distribution shift of the noisy data.
Existing sparse pruning methods generally do not consider the distribution shift of noisy data but only the model’s weight changes.
Inspired by RFR’s conclusions, we convert the distribution shifts caused by changes in model weights into distribution shifts caused by data changes for the diffusion model’s training process.

## 3 Sparse Finetuning DM

In this section, we introduce our proposed framework in detail. We start with the overall framework of our proposal. Then, we present the finetuning DM with sparse masks to reduce MACs and prepare for 2:4 sparse GPU acceleration. Moreover, we introduce the training and inference with sparse masks to enhance the DM.
Fig. 1 shows an overview of the proposed framework.

Figure: Figure 1: Framework Overview. Training is the process of training with a progress sparse mask.
Refer to caption: x1.png

### 3.1 Straight-Through Estimator for Sparse Training

A straightforward solution for training an $N$:$M$ sparsity network is to extend the Straight-through Estimator (STE) Bengio et al. () to perform online magnitude-based pruning and sparse parameter updating. The concrete formula is described as:

$$ $\mathcal{W}_{t+1}\leftarrow\mathcal{W}_{t}-\gamma_{t}g(\widetilde{\mathcal{W}^ {(N:M)}_{t}})$ (8) $$

where $\gamma$ is learning rate, $g$ is gradient.
In STE, a dense network is maintained during the training process.
In the forward pass, the dense weights $\mathcal{W}$ are projected into sparse weights
$\widetilde{\mathcal{W}^{(N:M)}}=S(\mathcal{W},N,M)$
satisfying $N$:$M$ sparsity. $S(\cdot)$ is projection function. The sparse DM training task is expressed as:

$$ $\min_{S(\mathcal{W},N,M)}\mathcal{L}(\mathcal{W};\mathcal{D})$ (9) $$

The sparse DM inference is expressed as:

$$ ${\bm{y}}_{t}^{(N:M)}=\mathcal{F}(\widetilde{\mathcal{W}^{(N:M)}};{\bm{x}}_{t})$ (10) $$

To mitigate the negative impact of the approximate gradient calculated by vanilla STE,
we incorporate sparse mask information into the backward propagation process.
The updated formula is modified as follows:

$$ $\mathcal{W}_{t+1}\leftarrow\mathcal{W}_{t}-\gamma_{t}\left(g(\widetilde{ \mathcal{W}^{(N:M)}_{t}})+\lambda_{W}(\mathcal{W}_{t}-\widetilde{\mathcal{W}^{ (N:M)}_{t}})\right)$ (11) $$

where $\lambda$ is coefficient.
To obtain extremely sparse DM, jumping directly from the dense model to the extremely sparse model, such as 1:32 ($\frac{1}{32}=0.03125$) sparsity, for training and using the extremely sparse weight gradient for reverse propagation can easily lead to the knowledge loss of dense model or even collapse of the trained model.

### 3.2 Progressive Sparse Finetuning Diffusion Model

Traditional progressive sparse training works well for extremely sparse model training. However, its training data distribution does not change.
Based on the analysis of two distribution shifts in subsection 3.3, only one distribution shift is generally optimized during training.
If two distribution shifts are optimized at the same time, such as switching the sparsity rate when training a DM, the optimization will fail.
If a highly sparse DM is trained from a dense DM, the traditional progressive sparse training method needs to be changed.
The sparsity rate can only be fixed during the DM training. Similar to the traditional progressive sparse training method, we transfer knowledge to the sparse model through the samples generated by the dense model to accelerate the sparse training of DM and reuse the knowledge of the dense model as much as possible.

The prediction of a dense model transfers the original knowledge to a sparsely structured model and does not involve sampling compression; that is, the number of sampling steps is consistent by default.
During sparse training, a neural network $\mathcal{F}_{\mathcal{W}}({\bm{x}}_{t},t)$ is trained to predict the noise in ${\bm{x}}_{t}$ by minimizing the L2 loss between them. The loss is formulated as follows:

$$ $\mathcal{L}_{\mathrm{diff}}:=||\bm{\epsilon}_{t}-\mathcal{F}_{\mathcal{W}}({ \bm{x}}_{t},t)||_{2}^{2}.$ (12) $$

The default pre-trained model is a dense model.
The overall loss function of progressive sparse training is composed of the original dense task loss, a diffusion loss that optimizes the diffusion model. It can be formulated as:

$$ $\mathcal{L}_{\mathrm{sparse}}^{i}=\mathcal{L}_{\mathrm{dense}}+\lambda\mathcal {L}_{\mathrm{diff}},~{}~{}i\in[1,k]$ (13) $$

where $\lambda$ is the loss weight to balance the losses. We simply set $\lambda=1$ in all experiments. $k$ is the sparse mask number.

To achieve 2:4 sparse GPU acceleration, the current sparse training method is to directly train from a dense model with a sparsity rate of 0 to a sparse model with a sparsity rate of $50\%$, which can easily cause information loss. Due to the lengthy training process aimed at obtaining $50\%$ of the sparse model, we adopt a progressive sparse training process, ensuring we can obtain $50\%$ of the sparse model with minimal information loss. We add progressive sparse masks to existing STE-based methods. The progressive sparse DM training task is expressed as follows:

$$ $\min_{S(\mathcal{W}\odot\mathcal{M})}\mathcal{L}(\mathcal{W};\mathcal{D};[ \mathcal{M}_{1},\mathcal{M}_{2},...,\mathcal{M}_{k}]),$ (14) $$

The projection function $S(\cdot)$, which is non-differentiable during back-propagation, generates the $N$:$M$ sparse sub-network on the fly. To get gradients during back-propagation, STE computes the gradients of the sub-network $g(\widetilde{\mathcal{W}})=\bigtriangledown_{\widetilde{\mathcal{W}}}\mathcal{
L}(\widetilde{\mathcal{W}};\mathcal{D})$
based on the sparse sub-network $\widetilde{\mathcal{W}}$,
which can be directly back-projected to the dense network as the approximated gradients of the dense parameters. The approximated parameter update rule for the dense network can be formulated as:

$$ $\begin{split}\mathcal{W}_{t+1}&\leftarrow\mathcal{W}_{t}-\gamma_{t}g( \widetilde{\mathcal{W}^{(N_{i}:M_{i})}_{t}}),~{}~{}i\in[1,k]\end{split}$ (15) $$

This STE-based method could be easily improved by sparse mask regularization as follows:

$$ $\begin{split}\mathcal{W}_{t+1}&\leftarrow\mathcal{W}_{t}-\gamma_{t}\left(g( \widetilde{\mathcal{W}^{(N_{i}:M_{i})}_{t}})+\lambda_{W}(\mathcal{W}_{t}- \widetilde{\mathcal{W}^{(N_{i}:M_{i})}_{t}}))\right),~{}~{}i\in[1,k]\end{split}$ (16) $$

### 3.3 2:4 Sparse Mask for Sampling on GPU

The structured sparse function provides fully connected layers and convolutional layers with 2:4 sparse weights to achieve 2:4 sparse acceleration on the GPU. If their weights are pruned ahead of time, these layers can be accelerated using structured sparse functions on the GPU.
In order to test real inference acceleration on an NVIDIA GPU, we used the 2:4 sparse operator in the acceleration library provided by NVIDIA. NVIDIA Ampere architecture GPUs have implemented CUDA operators to accelerate the multiplication of such matrices.
2:4 sparsity inference is expressed as:

$$ ${\bm{y}}^{(2:4)}_{t}=\mathcal{F}(\widetilde{\mathcal{W}^{(2:4)}};{\bm{x}}_{t}).$ (17) $$

Transposable masks Hubara et al. () suggest that a weight matrix and its transpose can be simply pruned by multiplying binary masks, so the backward pass shares the same sparse weight matrix with the forward pass. Transposable masks are one of the important technologies for NVIDIA Ampere architecture GPUs to accelerate sparse matrices.

## 4 Results and Discussion

In this section, we will compare and discuss the different sparsity rate optimization results and some ablation study results.

### 4.1 Experimental Settings

Evaluation Metrics.
In this paper, we concentrate primarily on two types of metrics:
The efficiency metric is MACs;
The quality metric is FID, which Equ. (18) calculates.

$$ $FID_{(x,g)}=\|\mu_{x}-\mu_{g}\|^{2}_{2}+T_{r}(\sigma_{x}+\sigma_{g}-2(\sigma_{ x}\sigma_{g})^{0.5}).$ (18) $$

where $T_{r}$ represents the trace of the matrix. $x$ and $g$ represent real pictures and generated pictures, $\mu$ represents the mean, and $\sigma$ is the covariance matrix.
We generated 50,000 images and calculated FID together with the original images.

Baseline Methods.
We compare the performance of our method with the following baseline algorithm at 2:4 or other sparsity.
Diff-Pruning Fang et al. () is a method of structural pruning of diffusion models.
ASP (Automatic SParsity) Pool and Yu (); Mishra et al. () is a tool that enables 2:4 sparse training and inference for PyTorch models provided by Nvidia.
Nvidia developed a simple training workflow that can easily generate a 2:4 structured sparse network that matches the accuracy of the dense network.
To compare the performance of other sparse masks, we give the results of inference with different sparsity rates.
Finally, we provide the ablation study results of STE-based for the training process and traditional progressive sparse training.
We give the ablation study results of untrained and STE-trained masks for the inference process.

Implementation Details.
All of the experiments are implemented by Pytorch and conducted on servers with 24 GPUs, eight NVIDIA 4090, eight A40, and eight 3090 GPUs.
We evaluate the FID score every 50K training iterations on 10K generated samples (instead of 50K samples for efficiency).

### 4.2 Result Comparison

2:4 Sparsity Results

**Table 2: The comparison of 2: 4 ($50\%$ ) sparsity results. No Pruning is the method of U-ViT. U-ViT with LDM is U-ViT using Latent Diffusion Models (LDM) Rombach et al. () for data process. Diff-Pruning is a structural pruning method. ASP is the method of sparse pruning from NVIDIA. DDPM Ho et al. () is a diffusion model based on U-Net.**
| Start DM | Datasets | Methods | FID | MACs (G) |
| --- | --- | --- | --- | --- |
| U-ViT | CIFAR10<br>32x32 | No Pruning | 3.20 | 11.34 |
| Diff-Pruning | 12.63 | 5.32 |  |  |
| ASP | 319.87 | 5.76 |  |  |
| Ours | 4.23 | 5.67 |  |  |
| CelebA<br>64x64 | No Pruning | 2.87 | 11.34 |  |
| Diff-Pruning | 11.35 | 5.32 |  |  |
| ASP | 438.31 | 5.76 |  |  |
| Ours | 3.52 | 5.67 |  |  |
| U-ViT<br>with<br>LDM | MS-COCO<br>256x256 | No Pruning | 5.95 | 11.34 |
| Diff-Pruning | 15.20 | 5.43 |  |  |
| ASP | 350.87 | 5.79 |  |  |
| Ours | 8.14 | 5.68 |  |  |
| MS-COCO<br>256x256<br>(Deep) | No Pruning | 5.48 | 14.72 |  |
| Diff-Pruning | 16.92 | 7.12 |  |  |
| ASP | 340.69 | 7.43 |  |  |
| Ours | 6.82 | 7.37 |  |  |
| ImageNet<br>256x256 | No Pruning | 3.40 | 76.66 |  |
| Diff-Pruning | 14.28 | 34.06 |  |  |
| ASP | 367.41 | 37.93 |  |  |
| Ours | 5.43 | 36.84 |  |  |
| DDPM | CIFAR10<br>32x32 | No Pruning | 3.23 | 9.44 |
| Diff-Pruning | 12.31 | 4.75 |  |  |
| ASP | 328.47 | 4.91 |  |  |
| Ours | 3.83 | 4.87 |  |  |
| CelebA<br>64x64 | No Pruning | 2.94 | 9.44 |  |
| Diff-Pruning | 11.19 | 4.75 |  |  |
| ASP | 441.52 | 4.91 |  |  |
| Ours | 3.66 | 4.87 |  |  |

We experimented with four datasets (CIFAR10, CelebA, MS COCO 2014, and ImageNet) and three resolutions (32$\times$32, 64$\times$64, and 256$\times$256) on Transformer-based and UNet-based DMs.
We reduced the computational load by approximately $50\%$, and the FID only increased by 0.60$\sim$2.19 (1.22 on average).
The higher the resolution, especially for models of similar sizes, the better the FID effect, such as U-ViT on CIFAR10 32$\times$32 and CelebA 64$\times$64.
This way, our method is more suitable for model acceleration in high-resolution and high-fidelity image generation.

In addition to changes in FID, it is also important to intuitively evaluate the data generated by the sparse acceleration model.
From the generated images in Fig. LABEL:fig:samples, it can be seen that there is almost no difference between the images generated by our accelerated model and the images generated by the original model due to a slight change in FID.

In addition to the theoretical acceleration results of $50\%$, Table 3, we show our hardware test of the model in the basic setting. The acceleration is defined as the ratio between the running of a dense model and a sparse model in the same setting. The acceleration ratio should almost be the same across all datasets while the weights also have nothing to do with this.

In the testing code we provide, we set hyper-parameters head_dimension=1024, num_head=8, depth=1 to get the following result. We do the experiment on the two significant hyper-parameters: mlp_ratio (determines the dimension of the MLP layer according to attention layer dimension) and patch_size (determines the length of sequence for attention computation). Theoretically, only the ratio between sequence length and head dimension will impact the acceleration ratio, so we don’t bother to sweep the head dimension further.

We can see that the acceleration is among 1.1 based on the concrete parameters. It is worth mentioning that we only use the most naive acceleration of the 2:4 sparsity provided in PyTorch. Many other recent works like Hu et al. () provide further acceleration, and all these kinds of methods can be integrated into our models and achieve faster training and inference.

**Table 3: 2:4 ($50\%$ ) sparse model acceleration results tested on 4 A40 GPU and CIFAR10 32$\times$32. The row indices and column indices are the hyper-parameters we choose.**
| Acceleration | patch_size=1 | patch_size=2 | patch_size=4 |
| --- | --- | --- | --- |
| mlp_ratio=1 | 1.02 | 1.01 | 1.01 |
| mlp_ratio=2 | 1.04 | 0.97 | 1.01 |
| mlp_ratio=4 | 1.23 | 1.10 | 0.98 |
| mlp_ratio=8 | 1.22 | 1.17 | 1.08 |

Other Sparsity Results
To better understand the performance of our sparse training method, except for the $50\%$ sparsity ratio of 2:4, we have also conducted experiments on other sparsity ratios. Since the ASP method is a sparsity tool provided by Nvidia for GPU hardware acceleration, no other sparsity ratios are provided. The following is mainly a comparison of different sparsity levels with other methods.

Fig. LABEL:fig:sparse_inference(a) shows the effect of different sparsity ratios on FID. We evaluated ten sparsity ratios with 32:32, 31:32, 15:16, 7:8, 3:4, 2:4, 1:4, 1:8, 1:16, and 1:32.
As shown in this figure, it does not mean that the greater the sparsity, the better the FID. For example, the sparsity of 31:32 ($\frac{31}{32}=0.96875$) and the sparsity of 15:16 ($\frac{15}{16}=0.9375$), which is higher than the FID of 7:8 with a sparsity of ($\frac{7}{8}=0.875$). In addition to GPU hardware acceleration, the sparsity ratio 2:4 also achieves good FID, proving the 2:4 sparse mask training and inference effectiveness.

Except for the 2:4 sparsity ratio, Fig. LABEL:fig:sparse_inference(c) shows that our method achieves significantly lower FID than the Diff-Pruning method at different sparse ratios on CIFAR10 and CelebA64.

### 4.3 Ablation Study

Our method involves performing sparse mask fine-tuning on the existing trained model. Masks need to be trained during the training process. We designed three ablation experiments to demonstrate better the rationality of our method design and the necessity of each step. The first one is the STE-based method. The second one is traditional progressive sparse training of DM. The last one
involves fixing the mask during the training process and not performing sparse training.

STE Sparse Training As shown in Fig. LABEL:fig:sparse_inference(d) and LABEL:fig:sparse_inference(e), on datasets CIFAR10 and CelebA64, whether it is 2:4 sparse or other sparsity ratio, the training method of STE is worse than our method, mainly because our method adds sparse mask information into the backward propagation process.

Traditional Progressive Sparse Training As shown in Fig. LABEL:fig:sparse_inference(b), DM is trained on datasets CIFAR10.
Every 100,000 training steps, the sparsity of masks is increased for progressive sparse training.
Training after switching sparse masks for the first time is almost ineffective.
The traditional progressive sparse training does not work well on the diffusion model, especially when switching the sparsity rate every time. It is equivalent to training stopped.
However, the fixed sparse mask can be continuously trained until convergence.
This comparison result motivates us to propose a new method for progressively sparse training of diffusion models.

Sparse Mask Inference
Unlike our trained masks, we generate untrained sparse masks for sparse pruning diffusion models. Dense model weights are imported into the sparse model, and the sparse model is not trained.
Only use STE-based methods for sparse training.
As shown in Fig. LABEL:fig:sparse_inference(f), on datasets CIFAR10, the result of the untrained mask is the worst because the mask did not participate in training. The STE-trained mask is worse than ours because the mask did not use progressive sparse training.
At the highest sparsity rate, our method achieves an FID 137 lower than that of random masks. Our method achieves 1 less FID at the lowest sparsity rate than the untrained mask.

## 5 Related Work

### 5.1 N: M Sparsity

A neural network with $N$:$M$ sparsity ($M$-weights contains $N$ non-zero values) satisfies that, in each group of $M$ consecutive weights of the network, there are at most $N$ weights have non-zero values Zhou et al. (); Chmiel et al. ().
APEX’s Automatic SParsity (ASP) Pool and Yu (); Mishra et al. ()) is a 2:4 sparse tool provided by NVIDIA.
This tool obtains a 2:4 sparse network and can achieve nearly 2x runtime speedup on NVIDIA Ampere and Hopper architecture GPUs Nvidia ().
A provable and efficient method for finding N: M transposed masks for accelerating sparse neural training Hubara et al. ().
Compared with dense networks, sparse network training has gradient changes, and methods such as STE Bengio et al. () should be used to improve training performance. In the forward stage, sparse weight is obtained by pruning dense weight. In the backward stage, the gradient w.r.t. sparse weight will be directly applied to dense weight.
SR-STE Zhang et al. () enhances the backward pass information by integrating the sparse weight of the forward pass into the backward pass.
DominoSearch Sun et al. () found mixed N: M sparsity schemes from pre-trained dense, deep neural networks to achieve higher accuracy than the uniform-sparsity scheme with equivalent complexity constraints.
The N: M sparse method is also used in SparseGPT Frantar and Alistarh () to accelerate the LLM generation model GPT.
Although many N: M sparse methods exist, none have been proven effective for sparse DM.

### 5.2 Diffusion Model Pruning

Pruning is one of the most used methods to reduce the calculation time of DNN, including DM.
The pruning method was divided into structured pruning and unstructured pruning.
Structural pruning aims to physically remove a group of parameters, thereby reducing the size of neural networks. In contrast, unstructured pruning involves zeroing out specific weights without altering the network structure Fang et al. ().
Sparsity can reduce the memory footprint of DM to fit mobile devices and shorten training time for ever-growing networks Hoefler et al. ().
Pay attention to features and selection of useful features for the target dataset Wang et al. (), which can also reduce computational complexity.
The parameter-grouping patterns vary widely across different models, making architecture-specific pruners, which rely on manually designed grouping schemes, non-generalizable to new architectures. Depgraph Fang et al. () can tackle general structural pruning of arbitrary architecture.
Structured sparsity was also used for large language models Ma et al. (); Frantar and Alistarh (); Guo et al. ().
Because Transformer has been proven to outperform the other networks in multiple applications, including DM.
ToMe Bolya and Hoffman () merge redundant tokens of Transformer.
There are also proposals for early exiting for accelerated inference in DM Moon et al. ().
Although there are many pruning methods for DM, either the sparse model FID is poor or structured 2:4 sparsity is not implemented.

## 6 Conclusion

In this paper, we studied how to improve the efficiency of DM by sparse matrix for 2:4 sparse acceleration GPU.
The existing STE-based methods make it difficult to optimize sparse DM. To address this issue, we improve the STE method and propose a progressive approach to transfer knowledge from dense models to sparse models gradually.
Our method is tested on the latest Transformer-based DM, U-ViT, and UNet-based DM, DDPM.
We trained the 2:4 and other sparse models to perform better than other methods.
Our approach also provides an effective solution for DM deployment on NVIDIA Ampere architecture GPUs, achieving about 1.2x acceleration.

Limitations and future work.
The MACs of the sparse diffusion model are significantly reduced, but the acceleration effect on real processors is limited, especially the 2:4 sparse acceleration that can currently only reach about 1.2 times. Future work will design sparse acceleration schemes for different processors to achieve better processor acceleration.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: Could the authors provide the parameter counts for each layer of the SD model before and after sparse pruning?
soundness: 2
strengths: The writing is very clear, and the main idea is highlighted effectively.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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?
soundness: 3
strengths: 1.	This paper is well-written.
2.	The motivation is clear enough.
3.	The organization of this paper is great.
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.
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.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: Please address the weakness stated above.
soundness: 1
strengths: - The 2:4 sparse model calculation offers practical values for practitioners using NVIDIA Ampere architecture GPUs.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
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.
soundness: 2
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.
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.
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.