Entropy-Guided Automated Progressive Pruning for Diffusion and Flow Models

Changlin Li, Jiawei Zhang, Zeyi Shi, Zongxin Yang, Xiaojun Chang

Keywords: Pruning; Diffusion models; Flow models

TL;DR: We propose EntPruner, an entropy-guided automatic progressive pruning framework for diffusion and flow models, achieving up to 2.22× inference speedup while maintaining competitive generation quality on ImageNet and three downstream datasets.

Abstract: Large-scale vision generative models, including diffusion and flow models, have demonstrated remarkable performance in visual generation tasks. However, transferring these pre-trained models to downstream tasks often results in significant parameter redundancy. In this paper, we propose EntPruner, an entropy-guided automatic progressive pruning framework for diffusion and flow models. First, we introduce entropy-guided pruning, a block-level importance assessment strategy tailored for transformer-based diffusion and flow models. As the importance of each module can vary significantly across downstream tasks, EntPruner prioritizes pruning of less important blocks using data-dependent transfer entropy as a guiding metric. Second, leveraging the entropy ranking, we propose a zero-shot Neural Architecture Search (NAS) framework during training to automatically determine when and how much to prune. This dynamic strategy avoids the pitfalls of one-shot pruning, mitigating mode collapse, and preserving model performance. Extensive experiments on DiT and SiT models demonstrate the effectiveness of EntPruner, achieving up to 2.22× inference speedup while maintaining competitive generation quality on ImageNet and three downstream datasets.

Primary Area: generative models

Submission Number: 4560