[Paper Review] DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization

Haowei Zhu, Dehua Tang+13 more|arXiv (Cornell University)|Oct 22, 2024

Optical Systems and Laser Technology4 citations

TL;DR

DiP-GO proposes a novel, retraining-free pruning method for diffusion models by formulating model pruning as a differentiable SubNet search problem. It introduces a SuperNet with backup connections based on feature similarity across denoising steps, uses a plugin pruner with optimized loss functions, and identifies an optimal SubNet via few-step gradient optimization, achieving 4.4× inference speedup on Stable Diffusion 1.5 without accuracy loss.

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.

Motivation & Objective

• To address the high computational cost of diffusion model inference, especially due to multi-step denoising.
• To overcome the limitations of static, non-adaptive pruning strategies that fail to exploit dynamic feature similarities across timesteps.
• To enable efficient, high-accuracy model pruning without retraining the pretrained diffusion model.
• To develop a differentiable, gradient-optimized pruning framework that identifies optimal subnetworks (SubNets) from a SuperNet.
• To achieve significant speedup while preserving image generation quality across diverse diffusion architectures.

Proposed method

• Construct a SuperNet by augmenting standard diffusion models with backup connections based on feature similarity across adjacent denoising steps.
• Formulate the pruning process as a SubNet search problem, where the optimal SubNet is derived from the SuperNet.
• Design a plugin pruner network that learns to identify redundant computation paths using task-specific optimization losses.
• Apply few-step gradient optimization to efficiently search for the optimal SubNet, reducing computational cost compared to full retraining.
• Integrate a lightweight post-processing step to ensure the final SubNet meets user-specified pruning ratios and structural constraints.
• Leverage gradient checkpointing to mitigate memory issues during backward passes over expanded denoising timesteps.

Figure 1: Overview of the SuperNet and SubNet. Standard diffusion models execute the full inference path step by step. In our framework, we propose a SuperNet based on the original flow and integrate backup connections to facilitate block removal. This allows the partial inference SubNet to efficien

Experimental results

Research questions

• RQ1 Can we formulate diffusion model pruning as a differentiable SubNet search problem to avoid retraining?
• RQ2 How can we effectively exploit feature similarities across adjacent denoising steps to reduce redundant computation?
• RQ3 Can a few-step gradient optimization strategy identify high-performing SubNets more efficiently than existing pruning methods?
• RQ4 To what extent can we accelerate inference without compromising generation quality in diffusion models?
• RQ5 Can the proposed method generalize across diverse diffusion architectures like Stable Diffusion and DiT?

Key findings

• DiP-GO achieves a 4.4× speedup in inference time on Stable Diffusion 1.5 without any drop in image quality or need for retraining.
• The method outperforms previous state-of-the-art pruning techniques by effectively leveraging dynamic feature patterns across timesteps.
• The approach generalizes across multiple diffusion models, including Stable Diffusion 1.5, 2.1, XL, and DiT, demonstrating broad applicability.
• The plugin pruner with optimized loss functions successfully identifies redundant components, enabling efficient SubNet discovery.
• The use of few-step gradient optimization enables fast convergence to high-performing SubNets, reducing search cost.
• The post-processing step ensures that the final SubNet adheres to pruning ratio constraints while maintaining model capability.

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\times T$ queries as input and passes them through stacked self-attention (SA) and fully connected (FC) layers to capture the structur