DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization
Authors: haoweiz
Dehua Tang
Ji Liu
Mingjie Lu
Jintu Zheng
Jinzhang Peng
Dong Li
Yu Wang
Fan Jiang
Lu Tian
Spandan Tiwari
Ashish Sirasao
Jun-Hai Yong
Bin Wang
Emad Barsoum
Venue: NeurIPS 2024 poster (2024)
Keywords: Diffusion
Pruning
Speedup
Gradient optimization
SuperNet
Primary area: generative_models
TL;DR: A diffusion pruner via few-step gradient optimization without retraining the diffusion model.

Abstract: Diffusion models have achieved remarkable progress in the field of image generation due to their outstanding capabilities. However, these models require substantial computing resources because of the multi-step denoising process during inference. While traditional pruning methods have been employed to optimize these models, the retraining process necessitates large-scale training datasets and extensive computational costs to maintain generalization ability, making it neither convenient nor efficient. Recent studies attempt to utilize the similarity of features across adjacent denoising stages to reduce computational costs through simple and static strategies. However, these strategies cannot fully harness the potential of the similar feature patterns across adjacent timesteps. In this work, we propose a novel pruning method that derives an efficient diffusion model via a more intelligent and differentiable pruner. At the core of our approach is casting the model pruning process into a SubNet search process. Specifically, we first introduce a SuperNet based on standard diffusion via adding some backup connections built upon the similar features. We then construct a plugin pruner network and design optimization losses to identify redundant computation. Finally, our method can identify an optimal SubNet through few-step gradient optimization and a simple post-processing procedure. We conduct extensive experiments on various diffusion models including Stable Diffusion series and DiTs. Our DiP-GO approach achieves 4.4 x speedup for SD-1.5 without any loss of accuracy, significantly outperforming the previous state-of-the-art methods.

DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization
Haowei Zhu" 12, Dehua Tang', Ji Liu', Mingjie Lu', Jintu Zheng', Jinzhang Peng', Dong Li Yu Wang', Fan Jiang', Lu Tian', Spandan Tiwari', Ashish Sirasao's Junhai Yong",
Bin Wang!?, Emad Barsoum!
Advanced Micro Devices, Inc. TTinngua University
zhuhw23emails. tsinghua. edu. en;
{dehua. tang, Ji.liu, jinz.peng, peng, d.li, lu. tian, ebarsoum) @amd com;
{yongjh, wangbins]
vangbins)@tsinghuu... @tsinghua edu.
Abstract
Diffusion models have achieved remarkable progress in the field of image genera- tion due to their outstanding capabilities. However, these models require substantial computing resources because of the multi-step denoising process during inference. While traditional pruning methods have been employed to optimize these models, the retraining process necessitates large-scale training datasets and extensive com- putational costs to maintain generalization ability, making it neither convenient efficient. Recent studies attempt to utilize the similarity of features across adjacent denoising stages to reduce computational costs through simple and static strategies. However, these strategies cannot fully harness the potential of the similar feature patterns across adjacent timesteps. In this work, we propose a novel pruning method that derives an efficient diffusion model via more intelligent and differ- entiable pruner. At the core of our approach is casting the model pruning process into SubNet search process. Specifically, on standard diffusion via adding some backup connections built upon the similar features. We then construct a plugin pruner network and design optimization losses to identify redundant computation. Finally, our method can identify an optimal SubNet through few-step gradient optimization and a simple post-process poo--prcessing pro- cedure. We conduct extensive experiments various diffusion models including Stable Diffusion series and DiT's Our DiP-GO approach achieves 4.4x speedup for SD-1.5 without any loss of accuracy, significantly outperforming the previous satt--of-theare methods.
first introduce SuperNet based
Introduction
Diffusion models have undergone significant advancements over the past years due to the outstanding capabilities of diffusion probabilistic models (DPMs) D. DPMs typically consist of processes:
the noise diffusion process and the reverse denoising process. Given their remarkable superiority in content generation, diffusion models have made sigmificant progress in various fields of general image generation, including text-to-image generation (2] layout-to-image generation 4 [.) image editing 6 7]. and image personalization 8 D. Furthermore. diffusion models have contributed to
advancements in autonomous driving, ranging from driving dataset generation 12 to perception model enhancement [13] 14) through diffusion strategies. However, DPMs often incur considerable computational overhead during both the training and inference phases. The high cost of inference, due
Work done during an internship at AMD. Corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

to the multi-step denoising computation during the sampling process, can significantly impact their practical application. Many efforts [T5HTT have been made to improve the efficiency of diffusion models, which can be broadly divided into two types of optimization: inference sampling optimization and model structural optimization.
Sampling optimization methods reduce the number of sampling steps for generation without compro- mising image quality. For instance, DDIM [15) reduces these steps by exploring a Markovaann process without requiring model retraining. LCM [8] [9] enable image generation in fewer steps with retraining requirements. Structural optimization methods []
aim to reduce com- putational overhead through efficient model design and model pruning. These methods require retraining the diffusion model, which entails significant computational overhead and large-scale datasets, making them neither convenient nor efficient. DeepCache [22] proposes novel training- free paradigm based on the U-Net architecture in diffusion models, caching and retrieving features across adjacent denoising stages to reduce redundant computation costs. However, DeepCache only reuses the output feature of U-Net block a denoising step via a simple and static strategy. We believe many intermediate features remain untapped, and the simple static strategy cannot fully exploit the potential of similar feature patterns across adjacent timesteps during inference, as observed in recent studies 15 18 22
20 21
To address these challenges, we introduce Diffusion Pruning via Few-step Gradient Optimization (DiP-GO), a method designed to achieve efficient model pruming with enhanced dynamism and intelligence. Our upproach rethinks the diffusion model during inference by proposing a SuperNet based on standard diffusion via adding some backup connections built upon the similar features, conceptualizing the inference process as a specific SubNet derived from our proposed SuperNet. We reformulate the diffusion model pruning into a SubNet search process. By addressing the out- of-memory issue inherent in the backward process during expanded denoising timesteps using the gradient checkpoint [23] method, we introduce a plugin pruner that discovers an optimal SubNet surpassing existing methods through carefully designed optimization losses. Extensive experiments validate the effectivvness of our approach, demonstrating a 4.4x speedup on Stable Diffusion 1.5. Moreover, our method efficiently prunes the DiT model [3] without requiring retraining the diffusion model, achieving significant inference speedup. Our contribution can be summarized as follows: (1) We define a SuperNet based on standard diffusion and show how to obtain a SubNet. This transforms the diffusion optimization problem into an efficient SubNet search process without the need for retraining pretrained diffusion models. (2) We design a plugin pruner tailored specifically for diffusion models. This pruner optimizes pruning constraints while maximizing synthesis capability. Additionally, we introduce a post-processing method for the pruner to ensure that the SubNet meets specified pruning requirements. (3) We conduct extensive pruning experiments across various diffusion models, including Stable Diffusion 1.5, Stable Diffusion 2.1, Stable Diffusion XL, and DiT. Extensive experiments demonstrate the superiority of our method, achieving notable 4.4 x speedup during inference on Stable Diffusion 1.5 without the need for retraining the diffusion model.
2 Related Work
2.1 Efficient Diffusion Models
The diffusion models, celebrated for their iterative demoising process during inference, play a pivotal role in content generation but often hindered by time-consuming operations. To mitigate this challenge, extensive research has focused on accelerating diffusion models. Acceleration efforts typically approach the problem from two primary perspectives:
Efficient Sampling Methods. Recent works focus on reducing the number of denoising steps required for content generation. DDIM [15] achieves this by exploring non-Markovian process related to neural ODEs. Fast high-order solvers 24 25) for diffusion ordinary differential equations also enhance sampling speed. LCMs 18 19 treat the reverse diffusion process as augmented probability flow ODE (PF-ODE) problem, inspired by Consistency Models (CMs) [66,, enabling generation in fewer steps. PNDM (7 emphasizes efficient sampling without retraining diffusion model. Additionally, ADD pretrained diffusion models into high-fidelity image generators using only single sampling steps.
28 combines adversarial training and score distillation to transform
Efficient Structural Methods. Other efforts concentrate on reducing the computational overhead associated with each denoising step. Previous methods [16] 17 22 have typically conducted extensive

---

empirical studies to identify and remove non- non-critionn critical layers from U-Net architectures to achieve faster networks. BK-SDM [16] customizes three efficient U-Nets by strategically removing residual and attention blocks. Derived from BK-SDM, KOALA 17 develops two efficient U-Nets of varying sizes tailored for SI-XII applications Diff-pruning 20 employs Taylor expansion over pruned timesteps to pinpoint essential layer weights, optimizing model efficiency without sacrificing performance. DeepCache [22] enhances inference efficiency by reusing predictions from blocks in previous timesteps within the U-Net architecture. LAPTOP-DiL Difff [21] tackles optimization problems with one-shot pruning approach, incorporating normalized feature distillation to streamline retraining processes. T-GATE [29] not only reduces computation overhead but also marginally lowers FID scores by omitting text conditions during fidelity- improtteemmr stages.
In addition to the two primary acceleration methods, other strategies such as distillation [8 30, B1
early stopping (32,, and quantization 33 are commonly employed to enhance performance and efficiency. However, most of these strategies necessitate retraining pretrained models. Our method falls under the category of efficient structural methods by focusing on reducing inference time at each timestep. Importantly, these efficiency gains are achieved without retraining the diffusion model.
2.2 Model Optimization
Network Pruning. The taxonomy of pruning methodologies typically divides into two main cate- gories: unstructured pruning methods [34/36] and structural pruning methods [37. Unstructured pruning methods involve masking parameters without structural constraints by zeroing them out, often requiring specialized software or hardware accelerators. In contrast, structured pruning methods generally remove regular parameters substructures from networks. Recent works have been inter- ested in accelerating transformers, Dynamic skipping blocks, which involve selectively removing layers while maintaining the overall structure, have emerged paradigm for transformer compres- sion 41444. However, applying structural pruning techniques to diffusion modeling poses unique challenges that necessitate reevaluating conventional pruning methods.
3 Methodology
In this study. we introduce the Diffusion Pruner via Few-step Gradient Optimization (DiP-GO), which utilizes a neural network to predict whether to skip or keep each computational block during inference. Our primary objective is to identify the optimal subset of computational blocks that facilitate denoising with minimal computational overhead. comprises three main components: neural network pruner, optimization losses, and a post-process algorithm to derive the pruned model based on the predictions of pruner. The neural network pruner is designed with learnable queries inspired by DETR 45) to predict the state of each block. Our proposed optimization losses include sparsity and consistency constraints for generation quality, guiding the pruner to accurately assess the importance of each block. In this Section, we first revisit the framework of diffusion models in Section emphasizing their potential for exploring pruned networks. In Section 3.2] introduce a SuperNet based on diffusion models and demonstrate how to derive a SubNet or pruned network from it for inference acceleration, highlighting the challenges method, including the neural network
illustrated in Figure 2 our method
in achieving optimal SubNet. Section 3.3 details pruner, optimization losses, and post-process a algorithm for obtaining a SubNet that meets pruning requirements. Finally, we provide insights into the training and inference processes of our method.
3.1 Preliminary
We begin with a brief introduction to diffusion models. Diffusion models are structured learn a series of sequential state transitions with the goal of iteratively refining random noise sampled from a known prior distribution towards a target distribution zo that matches the data distribution. During the forward diffusion process, the transition from Iz-1 to x, is initially determined by forward transition function, which can be described as follows:
q( ((=))))-1) N(z /1 BLFr-1, B.I)
(1)
where the hyperparameter (8, 1)] increases with each successive time step t.
3

---

To generate samples from learned diffusion model, it involves series of reverse state transitions from 2т - Zp to denoise random noise 2y (V into the clean data point zg. At each timestep, the denoised output -;-1 1 is predicted by approximating the noise prediction network, which is conditioned on the time embedding & and the previous data point æy:
3c v1- - T
po (0 = (d
(ot Va
20 (Ti, 1)); (8I)
where the covariance constant 33 = at, āe = IT ag and 2e (x t) are the parameterized deep neural networks. With the reverse Markov chain, we can iteratively sample from the learnable transition kernel Iz-1 pe (o( until =
Diffusion modes typically require multi-step conditional sampling to gradually obtain the target sample point Zg. However, recent studies [15 18 22 have highlighted that multi-step inference pro- cesses involve substantial redundant feature computations, particularly in noise prediction networks like UNet and Transformer. For example, in Stable Diffusion 1.4 models with 25 steps, Multiply- Accumulate Operations (MACs) UNet can comprise up 87.2% of the total computational load 16. This underscores significant potential for accelerating inference by effectively eliminating these redundancies. In this work, we propose accelerating the diffusion model by integrating differentiable pruning network designed to identify and remove these redundant computations.
3.2 SuperNet and SubNet of Diffusion Model
Our goal is to identify and remove unimportant blocks during inference to accelerate the process. To achieve this, we introduce a SuperNet based on the diffusion model. This SuperNet is designed to facilitate block removal while ensuring the pruned model maintains inference capability through additional connections. Our approach effectively eliminates unimportant blocks during inference, essentially deriving SubNet from the SuperNet by skipping these unnecessary components. Thus, the pruning process can be conceptualized as a SubNet search within the SuperNet framework.
have observed that diffusion models often exhibit similar feature patterns across adjacent timesteps during inference. Building on this insight, enhance the standard diffusion model's inference phase by introducing additional connections from the current timestep to the previous one. These connections serve as backups for blocks that may be removed, ensuring each block retains valid inputs even if its dependent blocks eliminated for acceleration. Specifically, for all inputs of each block across all timesteps except the inital step during inference, establish backup input connection the corresponding block in the previous timestep, as illustrated in Figure]
How to Construct SuperNet. Recent studies 15
18 22
24-
x
E
Pruned Blocks
X Removed Degendences
Original Connections ...
Potential Backup C Connections
Searched Backup C Connections
a) Standard Diffusion Model
b) SuperNet
c) SubNet
Figure 1: Overview of the SuperNet and SubNet. Standard diffusion models execute the full inference path step by step. In our framework, we propose SuperNet based on the original flow and integrate backup connections to facilitate block removal. This allows the partial inference SubNet to efficiently eliminate redundant computational costs.
How to Obtain a SubNet. To construct the SuperNet for the standard diffusion model, we introduce additional connections that ensure valid SubNet selects either the original input connection or backup input connection, but not both simultaneously. This design principle mandates that dependent block is pruned, its original input connection is also eliminated to reflect the block's removal. Conversely, if the dependent block retained, the backup input connection is removed to maintain efficient inference, depicted in Figure [

---

Pruner Training Processing
Layers
Block Importance Soores (T
X
x N3
Xy
Prodution Hoel
D
LIx- Optimization
Prun
()
S0 PC
SP
DV
D
A
M
Furward
Bacloward
traized
Praning Post Processing
Hiney sorch Intballcatoom
Dowmaunple Blcks
Citaaio cament thrchekii,
2 Upoodo Seles of Caledeeeggnee
Upsample
Pruned Bircks
Expected Prening Ratio
4 jugge d com 5 Updates, kndapp ceokop
Remove Degendences
D
Ignared Poths
Dinzry south aad
Figure 2: Overview of our diffusion pruner. a) DiP-GO employs a pruner network to learn the importance scores of blocks in the diffusion sampling process. It takes N x T queries as input and passes them through stacked self-attention (SA) and fully connected (FC) layers to capture the structural information existing diffusion models. The network predicts the partial inference paths based the N x T importance scores and is optimized by consistent and sparse loss. b) Once trained, the pruner network discarded. We can infer the optimal partial inference path with expected computational costs via post-processing based the predicted importance scores.
We draw inspiration from the Lottery Ticket Hypothesis (LTH) [46], which posits the existence of sub-network capable of achieving comparable performance to the original over over-parameterized network for given task, but with fewer unnecessary weights. Moreover, prior work [22] has explored manually removing redundant computations by caching features across adjacent steps. Thus, our approach seeks to identify an optimal SubNet from the SuperNet, maximizing diffusion model acceleration while minimizing any loss in generation quality.
Hard to Obtain an Optimal SubNet. The challenge of obtaining an optimal SubNet is compounded by the large number of blocks expanded during inference. In a diffusion pipeline with N x T blocks (where N the number of blocks per timestep and T the number of timesteps), each block's decision to be kept or removed results in 2 xxT possible configurations. For instance, a 50-step PLMS setup (27), considering 9 blocks the U-Net, yields 2450 choices (> 135). Traditional search methods like random search and genetic algorithms [4 often struggle in such vast search spaces. Gradient-based optimization offers a promising approach to tackle this challenge. However, there are significant hurdles to overcome. First, effectively modeling discrete block states (kept or removed) with parametric methods poses difficulties. Second, training the entire model, comprising both the parametric model and the expanded diffusion model with denoising timesteps, risks encountering out-of-memory (OOM) issues.
3.3 Our DiP-GO Approach
In this study, we introduce diffusion pruner network designed to predict importance scores for all blocks during reverse sampling as depicted in Figure To optimize the pruner network effectively, employ two key optimization losses: consistency and sparsity losses, leveraging few-step gradient optimization. Addressing the OOM issue inherent in such computations, we implement gradient checkpointing and half-precision floating-po poingg-pontttto representation techniques, enabling efficient search processes on a single GPU. Once the pruner network trained, we extract predicted importance scores for all blocks. Subsequently, we devise post-pree pesc-prcessing algorithm utilize these scores, generating pruned SubNets of diffusion models that satisfy specific pruning criteria.
Pruner Network. Our pruner network comprises three main components: N x T learnable queries, a query encoder, and a prediction head. We design the learnable queries to match the number of all blocks during inference. These queries are optimized with sparsity and consistency loss constraints to
5

---

learn the contextual information necessary for predicting the importance score of each block. For the query encoder, we provide two options: simple version with several stacked linear layers, and more complex version with several stacked self-attention layers to facilitate interaction among the learnable queries. Our experiments demonstrate that both versions can effectively obtain optimal SubNets in various diffusion models under different pruning requirements. The prediction head consists of N X T simple branches, each containing two stacked linear layers followed by a softmax operation. The final linear layer has a dimension of 2, and the softmax output represents the importance scores of a block. During training or inference, the query embeddings are transformed into output embeddings via the query encoder. These embeddings then independently decoded into binary vectors by the multi-layer prediction head, resulting in N X importance scores for all blocks.
Optimization Losses. The k-th predicted binary vector of importance score, denoted as s², represents the likelihood of its corresponding block being removed or kept in the denoising process. A gate {0, 1) 1)TN is derived based on 8, where g = 0 or g computation block, respectively. Only the blocks that are kept according to g will be calculated the denoising process. However, directly converting predicted probabilities s into discrete gates g with arg max is non-differentiable. To address this issue, we utilize the Fraiigh-Thhruuge. (ST) Estimator 48 to approximate the real gradient Veg with the gradient of the soft prediction V8s. To encourage both high-fidelity predictions and minimal computation block usage, we design our training objective function combination of consistent loss Le and sparse loss Lg, formulated follows:
indicate removing or keeping the k-th
L(xT; W) = L. + a.f. = { [(al= 20 + NT [( 200))
Σx k
NT gh if sparsity <T if sparsily 2
(3)
Here, ag represents hyperparameter used to balance the consistent and sparse losses. and W denote the pruner network and pretrained diffusion model, respectively. f(-) denotes a distance function that evaluates the consistency between the generated clean data point x from partial inference of the pruned SubNet and the x8² from full inference. This function can be any distance measure, and in this work, we utilize a negative SSIM loss [49] The sparse loss encourages minimal computational usage and is weighted by the computational flops proportion 7* of the k-th block, thereby imposing greater penalty on heavier blocks. The calculation of y takes into account the cascading relationships between blocks. Specifically, when a block is pruned, the associated dependent blocks will also pruned. Therefore, the flops reduction from pruning a block includes the block itself and its dependent blocks. We denote the flops reduction ratio after pruning the k-th block as X*= The flops ratio y is in the range [0. 1]. The sparse loss is only introduced when the sparsity (pruning ratio) is below a certain threshold r. This compound loss controls the trade-off between efficiency (block usage) and accuracy (generation quality).
Post- post-Processing Algorithm. After training the pruner network, our diffusion pruner is able to predict which computation blocks during inference contribute less to generation quality based on the importance scores for all the blocks. As the importance scores are continuous values in inference phase, they can not be utilized directly to identify which blocks should be removed to meet given pruning requirements. Therefore, we present a post-process algorithm obtain an appropriate threshold for these importance score to meet the pruning requirements as shown in Algorithm [ in Appendix B Considering the required pruning sparsity, we use bisection lookup to select the appropriate threshold value to identify which blocks should be removed to meet the pruning ratio. Specifically, the blocks whose important scores below the threshold should be removed and the kept blocks should update their input connections as mentioned in Section 3.2 to maintain the pruned model inference. Thus a pruned model met the pruning ratio has been obtained.
Training and Inference Details. In the training phase, the prompt inputs are fed into the diffusion model obtain two kinds of outputs, one generated by the baseline diffusion model and the other is generated by the pruned model obtained via the current predictions of the pruner network. Then our proposed losses are utilized to optimize the pruner network to enable distinguishing the less important blocks. In the pruner's 's network, we initialize the weight of the last linear layer's output channel to 0 and its bias to 1. This setup ensures that at the beginning of training, the consistency loss is O and the sparsity loss is 1, facilitating smooth training. As training progresses, the sparsity loss gradually decreases while the number of pruned blocks increases, causing the consistency loss to rise. To maintain network fidelity after pruning, we switch to training only with the consistency loss once the sparsity loss reaches 0.2, continuing until training is complete. Once the pruner well trained, we can obtain pruned models to meet the pruning requirements via our post-process algorithm.

---

Experiments
4.1 Experimental Setup
Pre-trained Model and Datasets. We select four official pretrained Diffusion Models (i.e., SD-1.5 (2). SDD221 [2], SD-XL
(50] and DiT [3] to evaluate our approach. The SD series models are constructed on the U-Net 51 and the DiT is constructed on the transformer 52 We utilize a subset of the DiffusionDB 5 dataset comprising 1000 samples to train our pruner network, utilizing only
textual prompts. Following previous works (29 222 we evaluate the DiP-GO on three public datasets,
i.e. PartiPrompts [54), MS-COCO 2017
155 and ImageNet [56].
Evaluation Metries. We employ the Fréchet Inception Distance (FID) metrics assess the quality of images created by the generative models. FID quantifies the dissimilarity between the Gaussian distributions of synthetic and real images. A lower FID score indicates a closer resemblance to real images in the generative model. Additionally, we utilize the CLIP Score (58) (ViT-g/14) to evaluate the relational compatibility between images and text.
Implementation Details. For Stable Diffusion models, we utilize the SGD optimizer with a cosine learning schedule for 1000 steps of training. The batch size, learning rate, and weight decay are to 1. 0.1, and x 10-4, respectively. The hyperparameters 0t., and the query embedding dimension D, along with the encoder layer number L, are to 1, 0.2, 512, and 1, respectively. For the Diffusion Transformer model, we use the same experimental configuration as for the stable diffusion model, except that the learning rate set 10. To evaluate the inference efficiency, we evaluate the Multiply Accumulate Calculation (MACs). Parameters (Params), and Speedup for all models with batch size of in the PyTorch 2.1 environment on the AMD MI250 platform. Besides, we report MACs in those tables, which refer to the totals MACs for all steps.
Table I: Comparison with PLMS, BK-SDM and DeepCache SD-1.5. We utilize prompts in PartiPrompt and COC02017 validation set.
PartiPrompts Speedup
COC02017
Method
Pruning Type
MACs
CLIP Score
MACs
Speedup 1
CLIP Score
PLMS 50 steps
Baseline
16.94T
1.00x
29.51
16,94T
1.00x
30.30
BK-SDM Base
Structured
1.49X
PLMS 25 steps
Fast Sampler
8.47T
2.04x
29.33
8,47T
1.91
29.99
PLMS Skip- Skkip-ntter Interval=2
Structured
8,47T
2.04x
19.74
8.47T
1.1x
16.78
DeepCache
Structured
6.52T
2.15x
29.46
6,52T
2.11x
30.23
Ours (w/ Pruned-0.80)
Structured
3.38T
4,4xx
29.51
3.38T
4.40x
30.29
BK-SDM- - Small
Structured
10.88T
L.75x
27.94
10.88T
1.68x
27.96
PLMS 15 steps
Sampler
5.08T
2.89x
28.58
5,08T
2.59x
29.39
Ours (w/ Pruned-0.85)
Structured
2.54T
5,52x
29.07
2.54T
5.46x
29.84
Table 2: Comparison of computational complexity, inference speed, CL.P Score and FID on the
MS-COCO 2017 validation set on SD-2.11
Inference Method
MACs,
Speedupt
CLIP Score
FID-5K +
SD-2,1-50 steps 2
38.04T
1.00x
31.55
27.29
SD-2.1-20 steps
15.21T
2.49x
31.53
27.83
Ours (w/ Pruned-0.7)
11.42T
3.02x
31.50
25.98
Ours (w/ Pruned-0.8)
7.61T
3.81x
30.92
27.69
4.2 Comparison with State-of-the-Art Methods on Different Base Models
Stable Diffusion on PartiPrompt and COCO2017. We compare our method with the state-of-the art (SOTA) compression methods on Stable Diffusion 1.5 (SD-1.5), and the results are summarized in Table Compared the SOTA DeepCache (22), our approach demonstrates significant performance improvements, achieving nearly 2x fewer MACs while maintaining better CLIP Scores. Our method can achieve 4.4x speedup compared to the baseline model. Furthermore, our method does not require training the diffusion model, which preserves the pre-trained knowledge the diffusion model. Also, apply method the SD-2.1 model to verify the effectiveness, as shown in Table 2 our method achieves significant acceleration while maintaining generation quality, demonstrating its superiority.
Diffusion Transformers on ImageNet. To the best of our knowledge, we are the first to apply pruning to DiT 3 model. Therefore, we have replicated a training-free acceleration method, DeepCache
7

---

Table 3: Comparison of pruning type, computational complexity, FID and inference speed on the
ImageNet validation datasets on DIT. * denotes the results reproduced with diffusers [59]
Method
Pruning Type
MACs
FID:S0K
Speedup
DrT-XI/22250 steps
29.66T
2.27
1.00x
DIT-XL/2*-250 steps
Bascline
29.66T
2.97
1.00x
DIT-XL/29-110 steps
Sampler
13.05T
3.06
DiT-XI 224-104 steps
Fast Sampler
11.86T
3.17
2.46x
DeepCache(DIT-XX/2P>---
Structured Pruning
15.88T
3.07
1.76x
Ours (DiT-XL/2" Pruned-0.6)
Structured Pruning
11.86T
3.01
2.43x
D steps DeepCache(DTTxXX/25)))) Ours (DiT-X11/2* w/ Pruned-0.75)
Fast Sumpler
8.307
3.35
3.49x
Structured Pruning
6.77T
3.20
3.44x
Strittred Pruning
7.40I
3.14
3.60x
with intervals = and 5, DiT for comparison. The results in Table show that method can speed up the original DiT model by a factor of 2.4 with minimal performance loss, while DeepCache has lower speedup ratio when applied to the DiT model. This can be attributed to DeepCache's 's overreliance on pre-defined structures, whereas our method can automatically learn the optimal pruning strategy for the given model, thereby achieving superior performance.
4.3 Compatibility with Fast Sampler
We investigate the compatibility of DiP-GO with methods that prioritize reducing sampling steps using faster samplers: DDIM [5]. DPM-Solver 225, and LCM [8]. As shown in Table it indicates that our method further improves computational efficiency on existing fast samplers. Specifically, we reduce MACs by a factor of 5 on the SD-1.5 with DDIM sampler and by 3.36 on the SD-2.1 with DPM-Solver Our method achieves nearly unchanged performance with significant acceleration. Additionally, our method benefits from information redundancy in multi-step optimization processes, showing relatively limited acceleration performance on fewer-step LCM due to its low redundancy in features across adjacent timesteps.
Table 4: Comparison with PLMS, SSIM, and LCM samplers. We evaluate the effectivvness of our methods on COCO2017 validation set.
Sampler
Base Model
Ours
MAC' L
CLIP Score
MACs
CLIP Score
DDIM (SD-1.5 w/ 50 steps)
16.94T
30.30
3.38T
30.29
DPM (SD-2,1 w/ 50 steps)
38.04T
31.55
11.42T
31.50
DPM (SD-2.1 w/ 25 steps)
19.02 T
31.59
9.51T
31.52
LCM (SD-XL 4 steps)
11.5TT
31.92
11.58T
31.30
4.4 Ablation Study
Compared with Different Consistent Constraints. We further compare other alternatives explored for our consistent loss designs, we further scrutinize additional options, including L1, L2, SSIM, and L1+SSIM losses, as depicted in Table 5 The results demonstrate that SSIM emerges as the most effective choice, boasting the highest CLIP-Score. In contrast, the L1 loss function often results in image blurring or distortion due to its sensitivity to pixel-level differences, the L2 loss may yield overly smoothed images by penalizing squared differences between pixels. Conversely, the combination of LI+SSIM loss attempts to address these limitations but can complicate the training process and suffer from trade-offs. Therefore, SSIM emerges the preferred choice in our consistent loss designs, offering superior accuracy and stability while preserving image quality.
Table 5: Comparison with different consistent loss types. Here we conduct pruning experiments with 80% sparsity on COC02017 validation using SD-115.
Loss Type
L1
L2
SSIM
LI +
SSIM
CI.IP Scoret
29.94
29.71
30.29
29.77
Effect of Gradient Optimization. the traditional search algorithm can also obtain SubNets from proposed SuperNet. It is crucial to validate whether traditional search-based algorithms yield
8

---

positive effectiveness. We assess two search algorithms: random search and genetic algorithm- search [47] Table We have iterated the search 1000 times using the first 500 images of the test set as calibration dataset. Remarkably, we observe that the search time of traditional search algorithms significantly longer than the training time of our method due large number of evaluations. Moreover, due to the vast search space, traditional search algorithms struggle to achieve satisfactory results. Additionally, traditional search algorithms lack the "once-for-all characteristic, requiring re-execution when faced with deployment scenarios demanding different computational resources. In contrast, leveraging the parametric pruner network, our method achieves superior performance with reduced running time and is more adaptable to diverse development scenarios.
Table 6: Comparison of cost time, computational complexity and CLIP-Score between Random Search and GA search strategies on Stable Diffusion 1.5.
Method
Cost GPU Hours |
Pruning Ratio
MACs 4
CLIP Score 1
PLMS 50 steps
16,94T
30,30
Random Scarch
25
0.80
2.96T
28.73
Seurch
25
0.80
3.34T
29.37
Ours
2.3
0.80
3.38T
30,29
Random Search
24
0.85
2.90T
27.22
GA Search
24
0.85
2.73T
28.61
Ours
2.2
0.85
2.54T
29.84
Random Search
23
0.90
1.4T
24.07
Search
23
2.04T
25.14
Ours
2.2
0,90
1.69T
28.72
Qualitative Analysis of Increased Prune Ratio. In Figure we visualize the generated images we increase the pruning ratio. With the increase in pruning ratio, the model's inference speed significantly improves, allowing us to achieve up to a fourfold increase in inference speed. However, the pruning ratio increases, some patterns in the image content deviate from those in the original images. Nevertheless, our main objects in the figures consistently adhere to the textual conditions. Subtle changes in background details typically do not compromise image quality, quantitatively analyzed in T
Original
Prupe 0.3
Prune 0.4
Prune 0.5
Prune 0.6
Pruns 0.7
Prime 0.8
A tran traveling down tracks next lights
castle and the big ben clocktower next to a river
A old television set displaying an old computer in front of two bookshelves
Time: 5.01s
Time: 3.59s
Time: 3.13s
Time: 2.65s
Time: 2.18s
Time: 1.70s
Time: 1.23s
Figure 3: Visualization of generated images. shows evolving patterns as pruning ratios increase. Despite these changes, main objects in the images remain consistent with the textual conditions.
Conclusion
This work explores resolving diffusion accelerating tasks by reducing redundant feature calculations across adjacent timesteps. We present a novel diffusion pruning framework and cast the model pruning process as a SubNet search problem. Our approach introduces a plugin pruner network that identifies an optimal SubNet through few-step gradient optimization. Results on wide range of Stable Diffusion (SD) and DiT series models verify the effectiveness of our method. We achieve a 4.4x speedup Stable Diffusion 1.5 and effectively prune the DiT model with few step optimizations.

---

References
[1] Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. NewrIPS, 33:6840-6851, 2020.
[2] Robin Rombach, Andreas Blatmmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In CVPR, pages 10684-10695, 2022.
[3] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4 4195-4205, 2023.
[4] Lymin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control to text-to- ext-to-image diffusion models. In ICCV, pages 3836-3888 -847, 2023.
[5] Chengyou Jia, Minnan Luo, Zhuohang Dang, Guang Dai, Xiaojun Chang, Mengmeng Wang, and Jingdong Wang. Ssmg: Spatial-semantic map guided diffusion model for free-form layout-to-image generation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 2480-2488 2024.
[6] Xi Chen, Lianghua Huang, Yu Liu, Yujun Shen, Deli Zhao, and Hengshuang Zhao. Anydoor: Zero-shot object-level image customization. arXiv preprint arXiv:2307. 09777, 2023.
[7] Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang, Tali Dekel, Inbar Mosseri, and Michal Irani. Imagic: Text-based real image editing with diffusion models. In Proceedings of the JEEF/CVE Conference on Computer Vision and Pattern Recognition, pages 6007-6017, 2023.
[8] Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1-500-22511 2023.
[9] Nupur Kumari, Bingliang Zhang, Richard Zhang, Eli Shechtman, and Jun- Yan Zhu. Multi- concept customization of text-to-image diffusion. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1931-1941, 2023.
[10] Xiaofan Li, Yifu Zhang, and Xiaoqing Ye. Drivingdiifusion: Layout-guided multi-view driving scene video generation with latent diffusion model. arXiv preprint arXiv:2310.07771, 2023.
[11] Ruiyuan Gao, Kai Chen, Enze Xie, Lanqing Hong, Zhenguo Li, Dit-Yan Yeung, and Qiang Xu. Magicdrive: Street view generation with diverse 3d geometry control. arXiv preprint arXiv:23 2023.
[12] Kai Chen, Enze Xie, Zhe Chen, Lanqing Hong. Zhenguo Li, and Dit-Yan Yeung. Integrating geometric control into text-to-in to-iimge diffusion models for high-quality detection data generation via text prompt. arXiv: 2306.04601 2023.
[13] Jiacheng Chen, Ruizhi Deng, and Yasutaka Furukawa. Polydiffuse: Polygonal shape reconstruc- tion via guided set diffusion models. ArXiv, abs/2306.01 01611, 2023.
[14] Zhendong Wang, Jianmin Bao, Wengang Zhou, Weilun Wang, Hezhen Hu, Hong Chen, and Houqiang Li. Dire for diffusion-generatee image detection. arXiv preprint arXiv:2303.091 09295, 2023.
[15] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500, 2020.
[16] Bo-Kyeong Kim. Hyoung-Ky Song. Thibault Castells, and Shinkook Choi. Bk-sdm: Architec- turally compressed stable diffusion for efficient text-to-image generation. ICML Workshop on Efficient Systems for Foundation Models (ES- (ES-FoMo), 2023.
[17] Youngwan Lee, Kwanyong Park, Yoohrim Cho, Yong Ju Lee, and Sung Ju Hwang. Koala: Self-attention matters in knowledge distillation of latent diffusion models for memory-efficient and fast image synthesis, arXiv preprint ar arXiv:2312.0400 2023.
10

---

[18] Simian Luo, Yigin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthesizing high-resolution images with few-step inference, 2023.
[19] Simian Luo, Yiqin Tan, Suraj Patil, Daniel Gu, Patrick von Platen, Apolinário Passos, Longbo Huang, Jian Li, and Hang Zhao. Lem-lora: universal stable-diffusion acceleration module, 2023.
[20] Gongfan Fang, Xinyin Ma, and Xinchao Wang. Structural pruning for diffusion models. In Advances in Neural Information Processing Systems, 2023.
[21] Dingkun Zhang, Sijia Li, Chen Chen, Qingsong Xie, and Haonan Lu. Laptop-diff: Layer pruning and normalized distillation for compressing diffusion models, 2024.
[22] Xinyin Ma, Gongfan Fang, and Xinchao Wang. Deepcache: Accelerating diffusion models for free. In The JEEE/CVF Conference on Computer Vision and Pattern Recognition, 2024.
[23] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv: 1ee.:10044600... 06744, 2016.
[24] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
[25] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models, 2023.
[26] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models, arXiv preprint arXiv:2303.01460 2023.
[27] Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09775 2022.
[28] Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion distillation. arXiv preprint arXiv:2311.17042 2023.
[29] Wentian Zhang, Haozhe Liu, Jinheng Xie, Francesco Faccio, Mike Zheng Shou, and Jüüren Schmidhuber. Cross-attention makes inference cumbersome text-to-image diffusion models. arXiv preprint arXiv:2404.02747, 2024.
[30] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00511 2022.
[31] David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap, Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbott, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv:2303.04244: 2023.
[32] Zhaoyang Lyu, Xudong Xu, Ceyuan Yang. Dahua Lin, and Bo Dai. Accelerating diffusion models via early stop of the diffusion process. arXiv preprint arXiv:2205. /2524, 2022.
[33] Yefei He, Luping Liu, Jing Liu, Weijia Wu, Hong Zhou, and Bohan Zhuang. Ptqd: Accurate post-training quantization for diffusion models, 2023.
[34] Sejun Park, Jaeho Lee, Sangwoo Mo, and Jinwoo Shin. Lookahead: far-sighted alternative of magnitude-based pruning, arXiv preprint arXiv:2002.04800 2020.
[35] Xin Dong, Shangyu Chen, and Sinno Pan. Learning to prune deep neural networks via layer-wise optimal brain surgeon. Advances in neural information processing systems, 30, 2017.
[36] Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine-tuning. Advances in neural information processing systems, 33:20378-20388 2020.
[37] Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient convnets, 2017.

---

[38] Sara Elkerdawy, Mostafa Elhoushi, Abhineet Singh, Hong Zhang, and Nilanjan Ray. To filter prune, or to layer prune, that is the question. In Proceedings of the Asian Conference on Computer Vision, 2020.
[39] Xiaohan Ding, Guiguang Ding, Yuchen Guo, and Jungong Han. Centripetal sgd for pruning very deep convolutional networks with complicated structure. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages
2019.
[40] Liyang Liu, Shilong Zhang, Zhanghui Kuang, Aojun Zhou, Jing-Hao Xue, Xinjiang Wang, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. Group fisher pruning for practical network compression. In International Conference on Machine Learning, pages 7021-7032. PMLR, 2021.
[41] Liu, Dehua Tang, Yuanxian Huang, Li Zhang, Xiaocheng Zeng, Dong Li, Mingjie Lu, Jinzhang Peng. Yu Wang, Fan Jiang, Lu Tian, and Ashish Sirasao. Updp: A unified progressive depth pruner for cnn and vision transformer, 2024.
[42] Minjia Zhang and Yuxiong He. Accelerating training of transformer based language models with progressive layer dropping, Advances in Neural Information Processing Systems, 33:1401 1- 14023, 2020.
[43] Yihe Dong, Jean-Boy Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 2793-2803. PMLR, 2021.
[44] Lu Hou, Zhiqi Huang, Lifeng Shang. Xin Jiang, Xiao Chen, and Qun Liu. Dynabert: Dynamic bert with adaptive width and depth. Advances Neural Information Processing Systems, 33:9782-977 9793, 2020.
[45] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference computer vision, pages 213-229. Springer, 2020.
[46] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint trXiv:1803.03635, 2018.
[47] John Holland. Genetic algorithms. Scientific american, 267())666733, 1992.
[48] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizztion with gumbel-softmax. arXiv preprini arXiv: 1.0::11/111010 144, 2016.
[49] Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE transactions on image processing, (3(4):600- 612, 2004.
[50] Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952, 2023.
[51] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention-MICCA. 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234-241. Springer. 2015.
[52] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[53] Zijie J. Wang, Evan Montoya, David Munechika, Haoyang Yang, Benjamin Hoover, and Duen Horng Chau. DiffusionDD: A large-scale prompt gallery dataset for text-to-image generative models. MxXiv2221014886 [cs], 2022.
12

---

[54] Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Bureu Karagol Ayan, et Scaling autoregressive models for content-rich text-to-imaae generation. arXiv preprint arXiv:2206. 10789, 2(3):5, 2022.
[55] Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pages 740-755. Springer, 2014.
[56] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. lece, 2009.
[57] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[58] Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference-free evaluation metric for image captioning, arXiv preprint arXiv:2104.087/6 2021.
[59] Patrick Platen, Suraj Patil, Anton Lozhkov, Pedro Cuenca, Nathan Lambert, Kashif Rasul, Mishig Davaadorj, Dhruv Nair, Sayak Paul, Steven Liu, William Berman, Yiyi Xu, and Thomas Wolf. Diffusers: State-of-the-art diffusion models, 2023.
[60] Junsong Chen, Jincheng Yu. Chongjian Ge, Lewei Yao. Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, et al. Pixart-a: Fast training of diffusion transformer for photorealistic text-to-image synthesis. arXiv preprint arXiv:2310.00426, 2023.
[61] Jaemin Cho, Abhay Zala, and Mohit Bansal. Dall-eval: Probing the reasoning skills and social biases of text-to-image generation models. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 3043-3054, 2023.
[62] Ranjita Naik and Besmira Nushi. Social biases through the text-to-image generation lens. In Proceedings of the 2023 AAAT/ACM Conference Al, Ethics, and Society, pages 786-808, 2023.
[63] Candace Ross, Boris Katz, and Andrei Barbu. Measuring social biases in grounded vision and language embeddings. arXiv preprini arXiv:2002.08911, 2020.
[64] Siwei Lyu. Deepfake detection: Current challenges and next steps. pages 1-6, 2020.
13

---

A Memory Optimization Details
Gradient Checkpointing. Due to the multi-step Markovian nature sampling in diffusion models, updating the entire sampling process using gradient accumulation incurs significant memory costs, even with a batch size of To mitigate this issue, we employ gradient checkpointing and half- precision floating- point training to reduce memory consumption. The core idea behind gradient checkpointing is to selectively preserve portion of activation values during forward propagation, discarding the rest. During backpropagation, the gradients of the discarded activation values are computed using the saved gradients of the preserved nodes, effectively reducing memory usage. Additionally,
gradient accumulation, wherein gradients computed over multiple iterations are accumulated and then backpropagated in a single batch for parameter updates, thus allowing for larger batch sizes under limited memory usage.
B
Pseudo Code
Here, we show the details our proposed post-process algorithm via pseudo code as followings.
Algorithm Diffusion Pruner
Input: A pretrained diffusion model M, importance scores S, a pruning ratio p Output: The pruned diffusion model M
left 0.0
2: right + 1.0
3: while True do
4: current + (left + right)/2 S for / in 1..1,,.... 7 do
5:
6:
7:
for each block score in S* do if's < current then
8:
9:
8 8 0
10:
end if end for end for
11:
12:
13:
update scores of blocks (S*) remove dependent blocks to set them zeros. p¹, M t prune diffusion model(S". .M) // obtain the pruned ratio and the pruned model.
14:
15:
if afabspp - p) < 0.0125 then
16:
break
17:
else if p" < then left current
19: else
right current 21: end if
20:
T T2//2
23:
end while
return M
14

---

C
Additional Experiment Results
Original
Ours (Pruned 0.6)
Original
Ours (Prumed-0.6) 0.6
Figure 4: Visualization of DiT model generated images: samples using DDIM-250 steps (uplink) and pruned 60% MACs (downlink). The speedup ratio here is 2.4x.
C.1 More Qualitative Results
Comparison with DiT Baselines. We provide the original unpruned DiT model and a version pruned by 0.6 ratio to generate comparison images in Figure It can be observed that the plots generated by the pruned model are almost identical those produced by the original model. Although there are slight differences in details, such as the appearance of the dog's eyes, these do not significantly affect the overall image quality.
Comparison with SD Baselines. We provide qualitative comparisons with the SD baseline and DeepCache, shown in Figure Our method demonstrates superior image-text consistency and image quality compared to existing methods.
Pruning Gate Visualization. Our method exhibits specific pattern of pruning ratios with respect to the timesteps. As shown in Figure [6] fewer blocks are pruned during the middle denoising stage (approximately between steps 65 and 1,50), 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.
Feature Similarity Analysis. Recent studies have confirmed feature similarity across adjacent time steps 22 29 We also conducted analysis of feature similarity between adjacent steps in fast samplers. Specifically, we sampled 200 images from the COC02017 validation set and calculated the average cosine similarity between the features of the penultimate upsampling block across all steps for two typical fast samplers, resulting in a similarity matrix, shown in Figure | The heatmap in Figure highlights the high degree of similarity between features at consecutive time steps.
C.2 More Quantitative Results
More Ablations. We conducted an ablation study on a and present the results in Table Our method achieves the best performance when 1.0. we also conducted an ablation study on Without 7 pruning 80% on SD1.5 resulted in a CLIP score of 29.50 (w/ ry: 30.29).
More Baseline Comparison. We also evaluated our method on Pix-rt-a 60 achieving excellent pruning results, as shown in Table below. Our method exhibits minimal performance loss on PixArt-a with 0.4 pruning ratio.
15

---

SD1.5
DeepCache
Ours
A red and white bicycle next to a building and door.
tall church tower sitting on the side of road.
A close- of a plate of food containing meat and beans.
Figure 5: A qualitative comparison with existing methods provided. We compare our method (prune 0.75) with DeepCache (N=4).
Pruner
Keep
Rumbvi
10
20
310 Tirestaps
200
Figure 6: The visualization results of the pruning gates for DiT-XL/2 at 250 steps with pruning ratio of 0.75.
D Limitations
A limitation of our method arises from its training process of the pruner network. Our method necessitates tuning an additional pruner network for the pre-trained diffusion model. This may entail users investing additional time when adapting method to specific diffusion models. For example, we train DiP-GO for SD-1.5 1.5 on a single AMD Instinct MI250 GPU for ~ 2.5 hours. However, we note that the introduced time is small compared to training lightweight diffusion model. Besides, same as existing work, our method struggles to maintain performance with extremely high pruning ratios, presenting challenge for deploying diffusion models in scenarios with severely limited computational resources.
E
Social Impact
Generative models have demonstrated promising results content generation [50, 60
However, due to the high inference costs, current methods struggle to achieve rapid application and deployment. Our approach introduces an efficient acceleration method for diffusion models, enabling nearly lossless speedup. Moreover, our method does not require retraining of the pretrained models and is compatible with various diffusion models, making highly generalizable. This makes suitable for rapid deployment of generative models on mobile and edge devices.
16

---

1.0
1.0
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
(a) DDIM (SD-1.3 1.5 w/ 50 steps)
(b) DPM SSD-2.1 w/ 25 steps)
Figure 7: Feature similarity across adjacent time steps in fast samplers.
Table 7: Comparison of different values. Pruning experiments with 80% pruning ratio
conducted on COCO2017 validation using SD-1.5.
0.1
0.5
1.0
2.0
CLIP Scoret
29.77
29.93
30.29
30.17
Table 8: Comparison with 20-step DPM-Solver sampler for diffusion transformer model. We
evaluate the effectiveness of our methods on COCO2017 validation set.
Method +
MACs
Speedup f CLIP Score t
PixArt-a: w/ 20-step DPM
85.65 T
1.0
30.43
Pruned-0.4 (Ours)
51.39
1.6
30.41
Nevertheless, since generative models are pretrained on large-scale internet datasets, the data they generate may contain inherent social biases and stereotypes of misuse, such in the creation of DeepFakes [64], which could pose significant social harm. While reducing the usage cost, it is crucial to prevent the low-cost generative models from being misused, leading negative societal impacts. Therefore, it is necessary establish relevant laws and regulations, create a well-regulated community environment, and provide guidelines to ensure responsible dissemination and use of generative models.
Additionally, there is a risk
17

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract provides a concise summary of the main findings and contributions of the paper, while the introduction elaborates the problem statement and research objectives, thereby clarifying the contributions.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitation Section in Appendix D we expound upon the limitations the work conducted and provide brief discussion thereof.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
18

---

Answer: [No]
Justification: None.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: In Section 4. introduced the details of experimental setup and model training to ensure reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
19

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification: The code will be released after it successfully passes company's internal review.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nnis.cc
mublic//uudees/Codesbbbass caasssionoolyyyy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ wttps/cccppplcccuua.eeaeaecccuea
Stleee//CoSSSSbmsssS.S.. www.ssssiicggcooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section training and testing.
we introduced the details of experimental setup and model
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: The experiments conducted in our paper do not involve the use of error bars or statistical significance analysis, thus this aspect not applicable to our study.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
20

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: For our experiments, we furnished detailed specifications of the GPU models used along with their corresponding tasks. Furthermore, included specifie information regarding the model training batch size and the number of training iterations.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: principles.
We have carefully reviewed the NeurIPS Code of Ethics and adhere to its
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss both potential positive societal impacts and negative societal impacts of the work performed in Appendix.
21

---

Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., denloymemt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Our paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The creators original owners of assets, such code, data, models, used in the paper, are properly credited. Additionally, the license and terms of use associated with these assets explicitly mentioned and respected in accordance with ethical and legal standards.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
22

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperssi thoode com/ com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's ' creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper does not release new assets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the datasct/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing research with human subjects.
Guidelines:
23

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
24


Reviews:
[{"review_id": "uVf3FWFzcE", "reviewer": "Reviewer_1Yei", "summary": "The paper introduces DiP-GO, a novel pruning method for diffusion models. Unlike the majority of existing methods that require extensive retraining with large datasets, DiP-GO employs (1) a SuperNet with backup connections, and (2) a plugin pruner network to identify redundant computations. The authors formulate a subnetwork searching problem which requires few-step gradient optimization instead of expensive retraining. Extensive experiments demonstrate DiP-GO’s effectiveness, outperforming state-of-the-art methods. The paper also explores compatibility with fast samplers and presents a fair amount of ablation studies.", "questions": "- The covariance constant $\\beta$ in Equation (2) should be defined.\n- In Equation (3), does the flops ratio $\\gamma$ in [0,1]?\n- In line 250, how is the threshold for sparsity loss set to 0.2? Is this merely a hyperparameter?", "limitations": "The paper provided limitations in Appendix D.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper tackles a timely and practically-relevant problem supported by a fair amount of experiments. Diffusion model pruning is an area with limited prior research, making this work particularly valuable.\n- The proposed method demonstrates superior performance compared to the baselines, and the paper provides a comprehensive review of relevant previous works.\n- Overall, the paper is well-written and easy to follow.", "weaknesses": "- I wonder why the pruner network takes $T \\times N$ random learnable queries and prediction heads. If the goal is to leverage the similarity between feature representations from adjacent timesteps, rather than distant relationships like between timestep 0 and 1, a more intuitive design might involve using $2N$ learnable queries and a timestep embedding vector (e.g., say $t_{emb}$). The final score could then be averaged or normalized as there will be duplicated score outputs, e.g., $s_t$ can be calculated from outputs $(s_{t-1}, s_t)$ and $(s_t, s_{t+1})$. This approach could result in a smaller network size and faster training for the pruner network. Please correct me for any misunderstanding.\n- Qualitative comparison with respect to baselines such as Diff-Pruning may help to improve the paper. Currently, there are only two figures (Figures 3 and 4).\n- I wonder if there is any pattern of pruning ratio with respect to the time-steps. For instance, DiP-GO may aggressively prune blocks near $t=0$, or there may exist a repeating pruned block pattern across time-steps. This leads to the question of design of $\\gamma$ in Equation (3). Did the authors ablate results concerning $\\gamma$?\n- As the proposed method offers 2-4X speedup, is this method better than naively skipping time-steps, say for every two time-steps?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nBelow are our detailed responses to the weaknesses and questions raised:\n\n**Re: Weakness #1**\n\nWe sincerely appreciate your valuable feedback. We believe your concerns primarily stem from two aspects:\n\n1. 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.\n\n2. 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$\n 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.\n\n**Re: Weakness #2**\n\nThank 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.\n \n\n\n**Re: Weakness #3**\n\n1. 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.\n\n2. 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**).\n\n**Re: Weakness #4**\n\nWe 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.\n\n\n**Re: Question #1**\n\nThank you for pointing this out. We will refine the definition of the covariance constant β in the final version.\n\n\n**Re: Question #2**\n\nYes. The flops ratio $\\gamma$ is in the range [0, 1].\n\n\n**Re: Question #3**\n\nThe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hIFv2eYnUk", "reviewer": "Reviewer_zGQU", "summary": "This paper proposes DiP-GO, a novel pruning method to address the computational challenge of Diffusion model during inference. The key innovation lies in the creation of a SuperNet, which includes backup connections based on similar features across adjacent time steps, and a plugin pruner network optimized through gradient optimization to identity whether to keep or remove a computational block. This approach eliminates the need to retrain the diffusion model, significantly reducing computational costs. Extensive experiments on various diffusion model, including Stable Diffusion series and DiTs, demonstrate that DiP-GO achieves substantial speedups without compromising accuracy.", "questions": "- Feature similarity in fast samplers: This approach relies heavily on the similarity of features across adjacent time steps. For diffusion model with fast sampler, only a few steps are required which might result in dissimilar features across adjacent time. A more strategic demonstration of this similarity could make statement more convincing.\n- How does different choice of alpha in equation 3 impact the overall performance?", "limitations": "The DiP-GO is limited in training the pruner and maintaining performance with extremely high pruning ratios. The authors address the pruner training limitation as the training time is still relatively short compared to retraining a model. For extremely high pruning ratio cases, it is possible that the pruned model is not over-parameterized. A theoretical work about the upper bound of the pruning ratio would be helpful.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The creation of SuperNet and transforming network pruning into a SubNet search, despite its simplicity, significantly improves performance in terms of efficiency and evaluation metrics such as FID and CLIP score.\n- This method avoids retraining diffusion models, saving substantial computational resource and time\n- This method has been applied to both traditional U-Net based diffusion models and transformer-based DiTs, validating its effectiveness across different architectures.", "weaknesses": "- As noted by the authors, memory overhead is a potential issue in this work. Although techniques like gradient checkpointing and half-precision floating-point representation are used to mitigate this, they might sacrifice the performance", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nWe address the raised concerns as follows:\n\n**Re Weaknesses #1:**\n\nGradient 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.\n\nAdditionally, 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.\n\n**Re Questions #1:**\n\nFeature 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.\n\n**Re Questions #2:**\n\nWe 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.\n\n| α   | 0.1  | 0.5  | 1.0  | 2.0  |\n|-----|------|------|------|------|\n| CLIP-Score | 29.77 | 29.93 | 30.29 | 30.17 |\n\n*Table 1:* Comparison of different α values. Pruning experiments with 80% pruning ratio were conducted on COCO2017 validation using SD-1.5.\n\n[1] DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters, KDD 2020, Tutorial.\n\n[2] ZeRO: memory optimizations toward training trillion parameter models, SC 2020.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8w2SzKb9d5", "reviewer": "Reviewer_A73E", "summary": "This paper proposes a novel differentiable pruner for diffusion models. The core of the approach involves transforming the model pruning process into a SubNet search process.", "questions": "Please refer to the weaknesses part.", "limitations": "Yes. The authors have addressed the limitations and social impacts.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The main idea is to transfer the model pruning process into a SubNet search process, eliminating the need to retrain pretrained diffusion models.\n2. Compared to traditional search methods, this differentiable approach is much more efficient.\n3. Extensive experiments demonstrate the superiority of this method.", "weaknesses": "1. In Figure 2(a), the dimension of the prune queries is  T \\times N \\times D . The meaning of dimension  D  is not discussed earlier in the paper. Could you clarify what  D  represents?\n2. DiP-GO is tested with 50 steps for SD and 250 steps for DiT, while 20 or 25 steps are more common nowadays. How does your algorithm perform under these more typical step scenarios?\n3. The paper only provides theoretical speedup measurements in terms of MACs or Speedup. What is the actual inference latency?\n4. In Table 3, the faster sampler method is compared under pruned-0.75 with 70 steps. However, under pruned-0.6, the results for the fast sampler method are not shown. Could you explain this omission?\n5. The usual step count for DPM-Solver is 20 or 25 [1,2,3]. Why did the author choose 50 steps in Table 4?\n6. According to Table 4, DiP-GO does not seem to be friendly to LCM. \n\n[1]. Cross-Attention Makes Inference Cumbersome in Text-to-Image Diffusion Models, arXiv.\n[2]. PixArt-$\\alpha$: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis, ICLR24.\n[3]. Your Student is Better Than Expected: Adaptive Teacher-Student Collaboration for Text-Conditional Diffusion Models, arXiv.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our sincere gratitude for the detailed and professional attention you have given to our work during the review process.\n\n**Re: Weakness #1**\n\n    $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.\n\n**Re: Weakness #2**\n\n    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.\n\n| Model | MACs  | Speedup  | CLIP Score |\n|-----------------------------|------|---|------------|\n| SD-2.1 with 25-step DPM     | 19.02 T | 1.0$\\times$    | 31.59      |\n| SD-2.1 Pruned-0.5  (Ours)       | 9.51 T |1.8$\\times$    | 31.52      |\n\n*Table 1:* Comparison with a 25-step DPM-Solver sampler for SD model. We evaluate the effectiveness of our methods on COCO2017 validation set.\n\n| Model | MACs  | Speedup  | CLIP Score |\n|-----------------------------|------|---|------------|\n| PixArt-$\\alpha$ with 20-step DPM     | 85.65 T | 1.0$\\times$    | 30.43      |\n| PixArt-$\\alpha$ Pruned-0.4  (Ours)       | 51.39 T |1.6$\\times$    | 30.41      |\n\n*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.\n\n\n**Re: Weakness #3**\n\n   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.\n\n**Re: Weakness #4**\n\nWe 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.\n\n| Method | Pruning Type | MACs  | FID-50K  | Speedup  |\n| --- | --- | --- | --- | --- |\n| DiT-XL/2*-250 steps | Baseline | 29.66T | 2.97 | 1.00 $\\times$ |\n| DiT-XL/2*-110 steps | Fast Sampler | 13.05T | 3.06 | 2.13 $\\times$ |\n| DiT-XL/2*-100 steps | Fast Sampler | 11.86T | 3.17 | 2.46 $\\times$ |\n| **Ours (DiT-XL/2\\* w/ Pruned-0.6)** | Structured Pruning | 11.86T | **3.01** | **2.43 $\\times$**  |\n\n*Table 3:*  Comparison of fast sampler methods under pruned-0.6.\n\n**Re: Weakness #5**\n\n    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].\n\nFurther, 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.\n\n\n**Re: Weakness #6**\n\n   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.\n\n[1] Deepcache: Accelerating Diffusion Models for Free, CVPR 2024.\n\n[2] Scalable Diffusion Models with Transformers, ICCV 2023.\n\n[3] Cross-Attention Makes Inference Cumbersome in Text-to-Image Diffusion Models, arXiv.\n\n[4] PixArt-$\\alpha$: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis, ICLR24.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ParMTSbYmT", "reviewer": "Reviewer_JBdU", "summary": "The paper introduces Diffusion Pruning via Few-step Gradient Optimization (DiP-GO), which is a new diffusion model pruning method. The method addresses the high computational cost of diffusion models' multi-step denoising process, which hinders their practical use. Traditional pruning methods require resource-intensive, inefficient retraining with large datasets. DiP-GO proposes intelligent, dynamic pruning without retraining.\nThe paper proposes a SuperNet using standard diffusion models with backup connections for similar features. These turn pruning into SubNet search, eliminating retraining.\n\nPlugin pruner design optimizes pruning constraints and synthesis. The pruner finds an optimal SubNet by identifying redundant computation with carefully designed optimization losses.\n\nA post-processing method ensures the pruned SubNet meets specific requirements, improving pruning efficiency and effectiveness.", "questions": "1.\tHow exactly are the additional backup connections in the SuperNet constructed? Are they predefined or learned during the training of the pruner network?\n2.\tCan you provide detailed computational resources required to train the pruner network? How does this compare to the computational resources required for retraining the diffusion models in traditional pruning methods?\n3.\tThe threshold $\\tau$ vary among different datasets or not?", "limitations": "The authors adequately addressed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tDiP-GO efficiently predicts the importance scores of computational blocks, allowing for dynamic and intelligent pruning without the need for retraining the diffusion models.\n2.\tThe pruner network employs consistency and sparsity optimization losses to ensure generation quality while minimizing computational usage. This balance between accuracy and efficiency is crucial for effective model pruning.\n3.\tThe authors validated DiP-GO across different diffusion models, including Stable Diffusion series and DiTs, showing its versatility and robustness. The extensive experiments show that DiP-GO achieves great speedup (up to 4.4× on Stable Diffusion 1.5) without loss of accuracy, outperforming previous state-of-the-art methods.", "weaknesses": "1.\tThe vast search space (due to the large number of blocks and timesteps) may pose challenges, as the method must efficiently navigate and optimize within this space to identify the optimal SubNet.\n2.\tAt higher pruning ratios, there is a risk of performance degradation where the quality of generated images might be compromised. This is particularly concerning for applications requiring high fidelity and detail. As the pruning ratio increases, some patterns in the image content may deviate from those in the original images. Although the main objects typically adhere to the textual conditions, subtle changes in background details can lead to noticeable artifacts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our sincere gratitude for the detailed and professional attention you have given to our work during the review process.\n\n**Re: Weakness#1**\n\n    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.\n\n**Re: Weakness #2**\n\n    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.\n\n**Re: Question #1**\n\n    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.\n\n**Re: Question #2**\n\n   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.  \n\n\n\n**Re: Question #3**\n\n    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.\n\n[1] Structural Pruning for Diffusion Models. NeurIPS 2023. \n\n[2] Stable-diffusion-v1-5 in hugging face.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Thank you to all the reviewers. We mainly upload the images needed for the rebuttal here. Detailed rebuttal responses and tables have already been sent to each reviewer separately.

Meta-review: This paper proposes DiP-GO, a novel differentiable pruner for diffusion models. DiP-GO involves transforming the model pruning process into a search process. Specifically, DiP-GO employs a SuperNet with backup connections, and a pruner network to identify redundant computations. The subnetwork searching problem involves few-step gradient optimization instead of expensive retraining. Extensive experiments demonstrate the effectiveness of DiP-GO. It tackles a practically-important problem. The descriptions of many parts of the paper are not clear and need more details, as suggested by several reviewers. Some of the information from the rebuttal should be included in the final version of the paper.
Decision: Accept (poster)