Which Layer Causes Distribution Deviation? Entropy-Guided Adaptive Pruning for Diffusion and Flow Models

Large-scale vision generative models, including diffusion andflow models, have demonstrated remarkable performance in visual gen-eration tasks. However, transferring these pre-trained models to down-stream tasks often results in significant parameter redundancy. In thispaper, we propose EntPruner, an entropy-guided automatic progressivepruning framework for diffusion and flow models. First, we introduceentropy-guided pruning, a block-level importance assessment strategyspecifically designed for generative models. Unlike discriminative models,generative models require preserving the diversity and condition-fidelityof the output distribution. As the importance of each module can varysignificantly across downstream tasks, EntPruner prioritizes pruning ofless important blocks using data-dependent Conditional Entropy Devia-tion (CED) metric. CED quantifies how much the distribution divergesfrom the learned conditional data distribution after removing a block.Second, we propose a zero-shot adaptive pruning framework to auto-matically determine when and how much to prune during training. Thisdynamic strategy avoids the pitfalls of one-shot pruning, mitigating modecollapse, and preserving model performance. Extensive experiments onDiT and SiT models demonstrate the effectiveness of EntPruner, achiev-ing up to 2.22× inference speedup while maintaining competitive gener-ation quality on ImageNet and three downstream datasets.