Structural Pruning for Diffusion Models

Summary

This paper presents Diff-Pruning, an efficient compression method for learning lightweight diffusion models from pre-existing ones. DPMs have shown impressive capabilities in generative modeling, but they often come with significant computational overhead during training and inference. Diff-Pruning addresses this challenge by introducing a Taylor expansion over pruned timesteps, which eliminates non-contributory diffusion steps and combines informative gradients to identify important weights. The authors empirically evaluate the proposed method on four diverse datasets, highlighting two primary benefits: 1) Efficiency, with a 50% reduction in FLOPs at a fraction of the original training expenditure, and 2) Consistency, as the pruned diffusion models retain generative behavior congruent with their pre-trained counterparts.

Strengths

- The paper introduces a novel method, Diff-Pruning, **specifically designed** for compressing diffusion models. - The authors conduct empirical assessments on four diverse datasets, providing a comprehensive analysis of the proposed method's performance. The evaluation demonstrates the effectiveness of Diff-Pruning in terms of efficiency and consistency. - The paper appears to be well-written overall.

Weaknesses

- Provide additional details about the pruning process: Although this paper describes the essence of diffo-pruning as Taylor expansions at pruning time steps, providing more specific details about the pruning process itself would enhance the clarity of the proposed method. Providing step-by-step explanations or pseudocode algorithms would help readers understand and replicate this approach. - In Eq. (4), (5), (6), (10), (11), the symbols $|\cdot|$, $||\cdot||$, and $||\cdot||_0$ appear. Are they representing the same thing? The author needs to carefully check if the symbols in the formulas are correct. Additionally, the notation $\nabla L_t(\mathbf{\theta})(\mathbf{\theta}^\prime - \theta)$ should be reviewed for its validity since $\theta$ is a vector. Furthermore, what does the "$\cdot$" symbol represent in Equation (6)? Is it the dot product? If so, why is "$|\cdot|$" included? - I am unsure how Eq. (6) is derived from Eq. (5). Why can it serve as an importance criterion? Does it have a direct relationship with the model's performance? If so, the author should demonstrate the relationship between this criterion and the performance of the diffusion model, such as through correlation analysis, and so on. - Why does the pruned model exhibit better performance than the pretrained model? Can the author explain this phenomenon? - I would like to know if the FID and SSIM metrics are sensitive to the pruned model. - Can the MACs metric demonstrate the superiority of the proposed method in terms of performance acceleration? I am interested in understanding the effectiveness of the proposed method on training or inference speed under different GPU, e.g. Frame Per Second (FPS). Overall, this article is well-written and easy to follow, but there are some questions that need to be addressed. If the author can address these questions, I would consider increasing the score.

Questions

see Weaknesses.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The paper briefly mentions limitations but does not provide detailed explanations. Expanding on the limitations section will help readers understand the potential constraints or assumptions of the proposed method and its applicability in different scenarios. Additionally, it is necessary to discuss specific technical shortcomings.

Summary

This work introduces a structural pruning method for diffusion models, called Diff-pruning. The authors leverage Taylor expansion on each term of the ELBO loss as the creteria to decide the pruned weights. By calculating the error of the final image induced by the pruning at timestep t and discussing about the effect of converged loss on the higher order terms of the Taylor expansion, this work manually drop the pruning creteria of timesteps near noise by setting a threshold $\mathcal{T}$ related to the relative loss. Experiments compare the performance of the proposed pruning method with other common pruning methods, as well as training from scratch.

Strengths

1. This paper proposes a novel pruning method that utilizes the multi-step property of the diffusion model. 2. The experiment results demonstrate the effectiveness of the proposed method compared with baseline methods.

Weaknesses

1. My major concern is the motivation of the method is not clearly explained: - The authors conclude an important conclusion from eq.9. However, the derivation process of eq.9 has mistakes and typos. There is a neural inference term with $x_t$ as input in eq.3, the final error $\delta_0$ can not be derived in such a simple form. Moreover, the meaning of the symbol "$\delta$" is ambiguous. It seems to represent noise prediction error in line 146, but $\delta_{t-1}$ and $\delta_0$ seems to represent the error of $x_t$ and $x_0$ caused by pruning. And the subscript $s$ of $\alpha_s$ is wrongly written in eq.8. - How could Equ. 9 derive the conclusion "prediction errors occurring at larger t primarily impact the high-level content of generated images, while smaller t values concentrate on refining the images with relatively small modifications"? Relevant paper presents that the main difference between high-level contents and details is that they come from different frequency components of the whole image [1] , while Equ.9 focus on the "amplitude of the error $\delta_0$". It requires further clarification. - The authors claim that the final distortion is progressively magnified according to eq.9, which indicates that the error at timesteps near the noise has a greater impact on the final image. However, the proposed method truncates the pruning criteria at large t, which does not match such observation. 2. Lack of detailed description of the baseline method ToMe. [1] Diffusion probabilistic model made slim. arXiv preprint arXiv:2211.17106, 2022.

Questions

1. How is eq.9 derived and why the choice of $\alpha_t$ according to this equation is opposite to the proposed method (as described in Weaknesses)?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The authors do not include the limitations and potential negative societal impact in the paper.

Summary

The paper introduces a pruning-based method aimed at addressing the efficiency challenges of diffusion models. Unlike existing approaches that focus on accelerating sampling or enhancing architectures, it specifically focus on the time cost of compression and the consistency in generated images. The authors propose a novel variant of Taylor approximation as the importance function for pruning, which effectively preserves both high-level content and low-level details of generated images, while minimizing the impact of noisy steps. Experimental results demonstrate that the compressed diffusion model successfully generates similar images to the pre-trained model.

Strengths

I agree that efficiency and consistency are indeed crucial considerations in compressing diffusion models, which have received limited attention in prior works. This work contributes to this field by building several initial experiments on datasets such as Cifar, CelebA, and LSUN, effectively showcasing the advantages of the method. This work establishes a solid baseline for future explorations into efficient diffusion models. Moreover, the impressive ability of the pruned model to generate consistent results is particularly noteworthy, which is user-friendly in deployment.

Weaknesses

It is observed that in certain cases, the generated images may still exhibit some distortions or changes, such as the watermarks shown in Figure 2. It is worth noting that some of these elements may contain important information that should ideally be preserved. To address this, it would be beneficial to explore methods that allow for controllable preservation of such information while pruning the models. This ability would provide the flexibility to focus on specific parts of the generated images based on different scenarios.

Questions

It would be valuable for the author to include discussions regarding the controllability of the observed distortions or changes, as mentioned in the weaknesses.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

N/A

Summary

This paper presents Diff-Pruning, a novel structural compression technique for learning efficient diffusion models from pre-trained ones. The fundamental idea behind Diff-Pruning is the utilization of Taylor expansion on pruned timesteps, which effectively combines informative and clean gradients to estimate the importance of weights. The authors demonstrate that Diff-Pruning not only achieves compression of pre-trained models within a few training iterations but also maintains the original generation capabilities intact.

Strengths

- The experimental results demonstrate that structural pruning can serve as a powerful and efficient compressor for diffusion models. An important advantage is that the pruned model inherently retains the generative behavior of the pre-trained model. Considering the time-consuming nature of training new diffusion models, this work proves valuable for various downstream applications. - The proposed method, which employs the Taylor expansion over pruned timesteps, is both well-motivated and practical. The idea of balancing the contribution of different stages with binary weighting is interesting and easy to practice. - The authors provide extensive experimental evidence that substantiates their claims regarding the efficiency and consistency of Diff-Pruning.

Weaknesses

- Extra comparison of the memory requirement for different pruning methods might provide a more complete assessment of the proposed methods. - The experiments reveal that structural pruning can potentially have a negative impact on the performance of pre-trained models. This observation deviates slightly from the results seen in discriminative tasks such as classification, where lossless compression can be achieved for certain networks like ResNets. It would be beneficial for the authors to provide further clarification on this phenomenon

Questions

Please refer to the weaknesses above.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

This work has no negative societal impact.

Summary

The paper proposed a method to prune diffusion models to achieve 50% FLOPS reduction at 10% to 20% of original training budget. The authors show that prune and fine-tune strategy based on the random selection, magnitude based selection, Taylor expansion based selection all lead to non-ideal / suboptimal performance. They develops a method based on the modification of Taylor expansion based selection and show it achieves good performance-quality tradeoff. They quantitatively evaluated on standard datasets used for diffusion models at small and medium resolution (from 32^2 to 256^2).

Strengths

The paper works on a popular task in recent literature of generative models and presented a practice to prune diffusion models while preserves its generative capability. They demonstrated empirical results on 256^2 resolution unconditional image generation, which seems promising.

Weaknesses

The lack of objective quantitative metrics in image generation literature makes the evaluation of pruned model very difficult. In addition to FID score, the authors also use SSIM score to compare the image generated from the full model and the pruned model under the same random seed is used. However, this is still hard for me to evaluate the quality of empirical results given limited examples presented in the paper. It would be otherwise more informative, if the authors can demonstrate their method being valid for conditional diffusion models, or even latent diffusion models. The technical contribution is also limited, as the method authors used already exist in the literature but has not been adopted for diffusion models.

Questions

I don't have other questions

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

I have a feeling that authors tried very hard to prove their technique working to some extent. It would be helpful as a research publication if the authors can present more results telling when and where their method fails or explore more on the boundary of quality and performance trade-offs.

Thanks for the authors' response

The authors' response addresses most of my concerns. Therefore, I'm increasing my score. BTW, I still have several questions. * A2 (a): According to my understanding, eq 9 assumes that the shift of $x_{t-1}$ (caused by pruning) does not affect functional shift at timesteps smaller than $t$, so that this error can be passed to $x_0$ in the way of sequenced multiplication. However, since neural network is a very complicated non-linear term that may have a significant effect on the final error. So, will this assumption be too strong to be true? What're the authors' opinions? * A3: My major concern is that the "high-level" details and "low-level" contents are directly related to the **frequency** of the image, not the amplitude of the error. So maybe it's inappropriate to derive the conclusion at Line 150-152 from eq 9.