SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Considering the significant impact of larger timesteps, it is necessary to incorporate them for importance estimation. To address this problem, Equation 9 naturally provides a simple and practical thresholding strategy for pruning. To achieve this, we introduce a threshold parameter  $\mathcal{T}$  based on the relative loss  $\frac{\mathcal{L}_t}{\mathcal{L}_{max}}$ . Those timesteps with a relative loss below this threshold, i.e.,  $\frac{\mathcal{L}_t}{\mathcal{L}_{max}} < \mathcal{T}$ , are considered uninformative and are disregarded by setting  $\alpha_{t} = 0$ , which yields the finalized importance score:
$$
\mathcal {I} \left(\boldsymbol {\theta} _ {i}, \mathbf {x}\right) = \sum_ {k} \left| \boldsymbol {\theta} _ {i k} \cdot \sum_ {\{t | \frac {\mathcal {L} _ {t}}{\mathcal {L} _ {\text {m a x}}} > \mathcal {T} \}} \nabla_ {\boldsymbol {\theta} _ {i k}} \mathcal {L} _ {t} (\boldsymbol {\theta}, \mathbf {x}) \right| \tag {10}
$$
In practice, we need to select an appropriately large value for  $\mathcal{T}$  to strike a well-balanced preservation of details and content, while also avoiding uninformative gradients from noisy loss terms. The full algorithm is summarized in Alg. 1.
Datasets and Models The efficacy of Diff-Pruning is empirically validated across six diverse datasets, including CIFAR-10  $(32\times 32)[21]$ , CelebA-HQ  $(64\times 64)[32]$ , LSUN Church  $(256\times 256)$ , LSUN Bedroom  $(256\times 256)$  [57] and ImageNet-1K  $(256\times 256)$ . We focus on two popular DPMs in our experiments, i.e., Denoising Diffusion Probability Models (DDPMs) [18] and Latent Diffusion Models (LDMs) [41]. For the sake of reproducibility, we utilize off-the-shelf DPMs from [18] and [41] as pre-trained models and prune these models in a one-shot fashion[23].
Evaluation Metrics In this paper, we concentrate primarily on three types of metrics: 1) Efficiency metrics, which include the number of parameters (#Params) and Multiply-Add Accumulation (MACs); 2) Quality metric, namely the Frechet Inception Distance (FID) [17]; and 3) Consistency metric, represented by Structural Similarity (SSIM) [50]. Unlike previous generative tasks that lacked reference images, we employ the SSIM index to evaluate the similarity between images generated by pre-trained models and pruned models, given identical noise inputs. We deploy a 250-step DDIM sampler [46] for ImageNet and a 100-step DDIM sampler for other experiments.
5.2 An Simple Benchmark for Diffusion Pruning
Scratch Training v.s. Pruning. Table 1 shows our results on CIFAR-10 and CelebA-HQ. The first baseline method that piques our interest is scratch training. Numerous studies on network pruning [10] suggest that training a compact network from scratch can be a formidable competitor. To ensure a fair comparison, we create randomly initialized networks with the same architecture as the pruned ones for scratch training. Our results reveal that scratch training demands relatively more steps to reach convergence. This suggests that training lightweight models from scratch may not be an efficient and economical approach, given its training cost is comparable to that of pre-trained models. Conversely, we observe that all pruning methods are able to converge within approximately 100K steps and outperform scratch training in terms of FID and SSIM scores. Thus, pruning emerges as a potent technique for compressing pre-trained Diffusion Models.
Pruning Criteria. A significant aspect of network pruning is the formulation of pruning criteria, which serve to identify superfluous parameters within networks. Due to the absence of dedicated work on Diffusion model pruning, we adapted three basic pruning methods from discriminative tasks: random pruning, magnitude-based pruning [16], and Taylor-based pruning [36], which we refer to as Random, Magnitude, and Taylor respectively in subsequent sections. For a given parameter  $\theta$ , Random assigns importance scores derived from a uniform distribution to each  $\theta_{i}$  randomly, denoted as  $\mathcal{I}(\theta)\sim \mathcal{U}(0,1)$ . This results in a straightforward baseline devoid of any prior or bias, and has been shown to be a competitive baseline for pruning [28]. Magnitude subscribes to the "smaller-normless-informative" hypothesis [23, 55], modelling the weight importance as  $\mathcal{I}(\pmb {\theta}) = |\pmb {\theta}|$ . In contrast, Taylor is a data-driven criterion that measures importance as  $\mathcal{I}(\pmb {\theta},x) = |\pmb {\theta}\cdot \nabla_{\pmb{\theta}}\mathcal{L}(\pmb {x},\pmb {\theta})|$ , which aims to minimize loss change as discussed in our method. As shown in 1, an intriguing phenomenon is that these three baseline methods do not maintain a consistent ranking on these two datasets. For
Method	CIFAR-10 32 × 32 (100 DDIM steps)
#Params ↓	MACs ↓	FID ↓	SSIM ↑	Train Steps ↓
Pretrained	35.7M	6.1G	4.19	1.000	800K
Scratch Training	9.88	0.887	100K
Scratch Training	5.68	0.905	500K
Scratch Training	19.8M	3.4G	5.39	0.905	800K
Random Pruning	5.62	0.926	100K
Magnitude Pruning	5.48	0.929	100K
Taylor Pruning	5.56	0.928	100K
Ours (T = 0.00)	5.49	0.932	100K
Ours (T = 0.02)	19.8M	3.4G	5.44	0.931	100K
Ours (T = 0.05)	5.29	0.932	100K
Method	CelebA-HQ 64 × 64 (100 DDIM steps)
#Params	MACs	FID	SSIM	Train Steps
Pretrained	78.7M	23.9G	6.48	1.000	500K
Scratch Training	7.08	0.833	100K
Scratch Training	6.73	0.867	300K
Scratch Training	43.7M	13.3G	6.71	0.869	500K
Random Pruning	6.70	0.874	100K
Magnitude Pruning	7.08	0.870	100K
Taylor Pruning	6.64	0.880	100K
Ours (T = 0.00)	6.24	0.885	100K
Ours (T = 0.02)	43.7M	13.3G	6.45	0.878	100K
Ours (T = 0.05)	6.52	0.878	100K
Table 1: Diffusion pruning on CIFAR-10 and CelebA. We leverage Frechet Inception Distance (FID) and Structural Similarity (SSIM) to estimate the quality and similarity of generated samples under the same random seed. A larger SSIM score means more consistent generation.
instance, while Magnitude achieves the best FID performance among the three on CIFAR-10, it performs poorly on CelebA datasets. In contrast, our method delivers stable improvements over baseline methods, demonstrating superior performance on both datasets. Remarkably, our method even surpasses the pre-trained model on CelebA-HQ, with only 100K optimizations. Nonetheless, performance degradation is observed on CIFAR-10, which can be attributed to its more complex scene and a larger number of categories.
5.3 Pruning at Higher Resolutions
DDPMs on LSUN To further validate the efficiency and effectiveness of our proposed Diff-Pruning, we perform pruning experiments on two  $256 \times 256$  scene datasets, LSUN Church, and LSUN Bedroom [57]. The pre-trained models from [18] require around 2.4M and 4.4M training steps, which can be quite time-consuming in practice. We demonstrate that Diff-Pruning can compress these pre-existing models using only  $10\%$  of the standard training resources. We report the number of parameters, MACs, and FID scores in Table 2, and compare the pruned methods to the pre-trained ones as well as those trained from scratch. Results show that the pruned model converges with a passable FID score in  $10\%$  of the standard steps, while a model trained from scratch is still severely under-fitted. Nevertheless, we also discover that compressing a model trained on large-scale datasets, such as LSUN Bedroom, which contains 300K images, proves to be quite challenging with a very limited number of training steps. We show that, in the supplementary materials, the FID scores can be further improved with more training steps. Moreover, we also visualize the generated images in Figure 2 and report the single-image SSIM score to measure the similarity of generated images. By nature, the pruned models can preserve similar generation capabilities since they inherit most parameters from the pre-trained models.
Conditional LDMs on ImageNet Table 3 and Figure 3 illustrate the pruning results of LDM pre-trained on ImageNet-1K. An LDM consists of an encoder, a decoder, and a U-Net model. Around
$\mathrm{SSIM} = 0.792$
$\mathrm{SSIM} = 0.821$
$\mathrm{SSIM} = 0.817$
$\mathrm{SSIM} = 0.952$
$\mathrm{SSIM} = 0.878$
$\mathrm{SSIM} = 0.876$
$\mathrm{SSIM} = 0.907$
$\mathrm{SSIM} = 0.912$
Figure 2: Generated images of the pre-trained models [18] (left) and the pruned models (right) on LSUN Church and LSUN Bedroom. SSIM measures the similarity between generated images.
$\mathrm{SSIM} = 0.881$
LSUN-Church 256 × 256 (DDIM 100 Steps)	LSUN-Bedroom 256 × 256 (DDIM 100 Steps)
Method	#Params	MACs	FID	Steps	Method	#Params	MACs	FID	Steps
Pretrained	113.7M	248.7G	10.6	4.4M	Pretrained	113.7M	248.7G	6.9	2.4M
Scratch Training	63.2M	138.8G	40.2	0.5M	Scratch Training	63.2M	138.8G	50.3	0.2M
Ours (T=0.01)	63.2M	138.8G	13.9	0.5M	Ours (T=0.01)	63.2M	138.8G	18.6	0.2M
Table 2: Pruning diffusion models on LSUN Church and LSUN Bedroom.
400M parameters come from the U-Net architecture and only 55M from the autoencoder. Therefore, we mainly focus on the pruning of the U-Net model. We used the threshold  $\mathcal{T} = 0.1$  to ignore those converged layers and make the pruning process more efficient. With  $T = 0.1$ , only 534 steps participate in the pruning process. After importance estimation, we apply a pre-defined channel sparsity of  $30\%$  to all layers, leading to a lightweight U-Net with 189.43M parameters. Finally, we finetune the pruned model for only 4 epochs with the official training script, with a scaled learning rate of  $0.1 \times lr_{\mathrm{base}}$ .
Pruned Timesteps. First, we conduct an empirical study evaluating the partial Taylor expansion over pruned timesteps. This approach prioritizes steps with larger gradients and strives to preserve as much content and detail as possible, thereby enabling more accurate and efficient pruning. The impacts of timestep pruning are demonstrated in Figure 5. We seek to prune a pre-trained diffusion model over a range of steps, spanning from 50 to 1000, after which we utilize the SSIM metric to gauge the distortion induced by pruning. In diffusion models, earlier steps  $(t\to 0)$  usually present larger gradients compared to the later ones  $(t\rightarrow T)$  [41]. This inherently leads to gradients that have reached a convergence when  $t$  is large. In the CIFAR-10 dataset, we find that the optimal SSIM score
Figure 3: Images sampled from the pruned conditional LDM on ImageNet-1K-256
Method	#Params ↓	MACs ↓	FID ↓	IS ↑	Train Steps ↓
Pretrained LDM	400.92M	99.80G	3.60	247.67	2000K
Scratch Training	189.43M	52.71G	51.45	25.69	100K
Taylor Pruning	11.18	138.97	100K
Ours (T=0.1)	9.16	201.81	100K
Table 3: Compressing conditional Latent Diffusion Models on ImageNet-1K (256 × 256)
Ratio	Pruning Ratios
#Params	MACs	FID ↓	SSIM ↑
0%	35.7M	6.1G	4.19	1.000
16%	27.5M	5.1G	4.62	0.942
44%	19.8M	3.4G	5.29	0.932
56%	14.3M	2.7G	6.36	0.922
70%	8.6M	1.5G	9.33	0.909
Table 4: Pruning with different ratios
Threshold	Thresholding	SSIM ↑
Steps	FID ↓
T = 0.00	1000	5.49	0.932
T = 0.01	707	5.41	0.932
T = 0.02	433	5.44	0.931
T = 0.05	244	5.29	0.932
T = 0.10	127	5.31	0.931
Table 5: Pruning with different threshold  $\mathcal{T}$
can be attained at around 250 steps, and adding more steps can slightly deteriorate the quality of the synthetic images. This primarily stems from the inaccuracy of the first-order Taylor expansion at converged points, where the gradient no longer provides useful information and can even distort informative gradients through accumulation. However, we observe that the situation differs slightly with the CelebA dataset, where more steps can be beneficial for importance estimation.
Pruning Ratios. Table 4 presents the #Params, MACs, FID, and SSIM scores of models subjected to various pruning ratios based on MACs. Notably, our findings reveal that, unlike CNNs employed in discriminative models, diffusion models exhibit a significant sensitivity to changes in model size. Even a modest pruning ratio of  $16\%$  leads to a noticeable degradation in FID score  $(4.19 \rightarrow 4.62)$ . In classification tasks, a perturbation in loss does not necessarily impact the final accuracy; it may only undermine prediction confidence while leaving classification accuracy unaffected. However, in generative models, the FID score is very sensitive, making it more susceptible to domain shift.
Thresholding. In addition, we conducted experiments to investigate the impact of the thresholding parameter  $\mathcal{T}$ . Setting  $\mathcal{T} = 0$  corresponds to a full Taylor expansion at all steps, while  $\mathcal{T} > 0$  denotes pruning of certain timesteps during importance estimation. The quantitative findings presented in Table 5 align with the SSIM results depicted in Figure 5. Notably, Diff-Pruning attains optimal performance when the quality of generated images reaches its peak. For datasets such as CIFAR-10, we observed that a 200-step Taylor expansion is sufficient to achieve satisfactory results. Besides, using a full Taylor expansion, in this case, can be detrimental, as it accumulates noisy gradients over approximately 700 steps, which obscures the correct gradient information from earlier steps.
Figure 4: Generated images of  $5\%$  -pruned models using different important criteria. We report the SSIM of batched images without post-training.
Figure 5: The SSIM of models pruned with different numbers of timesteps. For CIFAR-10, most of the late timesteps can be pruned safely. For CelebA-HQ, using more steps is consistently beneficial.
Visualization of Different Importance Criteria. Figure 4 visualizes the images generated by pruned models using different pruning criteria, including the proposed method with  $\mathcal{T} = 0$  (w/o timestep pruning) and  $\mathcal{T} > 0$ . The SSIM scores of the generated samples are reported for a quantitative comparison. The Diff-Pruning method with  $\mathcal{T} > 0$  achieves superior visual quality, with an SSIM score of 0.905 after pruning. It is observed that employing more timesteps in our method could have a negative impact, leading to greater distortion in both textures and contents.
This work introduces Diff-Pruning, a dedicated method for compressing diffusion models. It utilizes Taylor expansion over pruned timesteps to identify and remove non-critical parameters. The proposed approach is capable of crafting lightweight yet consistent models from pre-trained ones, incurring only about  $10\%$  to  $20\%$  of the cost compared to pre-training. This work may set an initial baseline for future research that aims at improving both the generation quality and the consistency of pruned diffusion models.
This research is supported by the Ministry of Education, Singapore, under its Academic Research Fund Tier 2 (Award Number: MOE-T2EP20122-0006).
[48] Patrick von Platen, Suraj Patil, Anton Lozhkov, Pedro Cuenca, Nathan Lambert, Kashif Rasul, Mishig Davaadorj, and Thomas Wolf. Diffusers: State-of-the-art diffusion models. https://github.com/huggingface/diffusers, 2022.
[49] Zhendong Wang, Yifan Jiang, Huangjie Zheng, Peihao Wang, Pengcheng He, Zhangyang Wang, Weizhu Chen, and Mingyuan Zhou. Patch diffusion: Faster and more data-efficient training of diffusion models. arXiv preprint arXiv:2304.12526, 2023.
[50] Zhou Wang, A.C. Bovik, H.R. Sheikh, and E.P. Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE Transactions on Image Processing, 13(4):600-612, 2004.
[51] Xingyi Yang, Jingwen Ye, and Xinchao Wang. Factorizing knowledge in neural networks. In European Conference on Computer Vision, 2022.
[52] Xingyi Yang, Daquan Zhou, Jiashi Feng, and Xinchao Wang. Diffusion probabilistic model made slim. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
[53] Xingyi Yang, Daquan Zhou, Songhua Liu, Jingwen Ye, and Xinchao Wang. Deep model reassembly. In Advances in Neural Information Processing Systems, 2022.
[54] Lewei Yao, Renjie Pi, Hang Xu, Wei Zhang, Zhenguo Li, and Tong Zhang. Joint-detnas: upgrade your detector with nas, pruning and dynamic distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10175-10184, 2021.
[55] Jianbo Ye, Xin Lu, Zhe Lin, and James Z Wang. Rethinking the smaller-norm-less-informative assumption in channel pruning of convolution layers. arXiv preprint arXiv:1802.00124, 2018.
[56] Zhonghui You, Kun Yan, Jinmian Ye, Meng Ma, and Ping Wang. Gate decorator: Global filter pruning method for accelerating deep convolutional neural networks. Advances in neural information processing systems, 32, 2019.
[57] Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv:1506.03365, 2015.
[58] Lvmin Zhang and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. arXiv preprint arXiv:2302.05543, 2023.
[59] Xunyu Zhu, Jian Li, Yong Liu, Can Ma, and Weiping Wang. A survey on model compression for large language models. arXiv preprint arXiv:2308.07633, 2023.