Comp-Diff: A Unified Pruning and Distillation Framework for Compressing Diffusion Models

Recently, generative models such as diffusion models (DMs) have gained prominence in various applications, and there is a growing demand for their deployment on resource-constrained devices. Model pruning provides an effective solution by reducing the model redundancy without significantly impacting performance. However, most existing model pruning methods are designed for classification models and often lead to substantial performance degradation when applied to generative models. To address this issue, we propose Comp-Diff, a novel two-stage framework of pruning and knowledge distillation tailored for diffusion models. In the pruning stage, we propose a new structured content-aware pruning (CaP) method within Comp-Diff to identify and preserve informative units (filters/channels) that actually contribute to the generative capability of the model. Specifically, we introduce input perturbations to the pre-trained model and measure each unit’s importance score using gradients induced by these perturbations. Units with higher importance scores are considered more informative and are retained to maintain the model’s generative power. In the fine-tuning stage of Comp-Diff, we propose the distribution-aware knowledge distillation (DaKD) method, which effectively transfers fine-grained knowledge from the original model to the pruned one on both attention and noise distribution levels. In addition, DaKD includes an adversarial loss to improve the quality and diversity of generated outputs. To verify and evaluate our method, we apply the proposed Comp-Diff on three representative tasks: unconditional image generation, conditional image generation, and text-to-image generation. Extensive experiments on both multi-step and one-step diffusion models demonstrate that the proposed framework consistently yields compact models and outperforms existing pruning techniques by a large margin.