Gradient flow-based iterative pruning for efficient and high-quality lightweight diffusion models

Diffusion Models (DMs) stand out among generative models for their impressive capabilities, but their application is often constrained by slower inference speeds and higher computational costs. Previous works have attempted to address these issues through one-shot structural pruning, deriving lightweight DMs from pre-trained models. However, this approach frequently leads to a substantial degradation in generation quality and risks the removal of critical parameters. To overcome these limitations, we propose an iterative pruning method based on gradient flow. Our approach integrates a gradient flow pruning process with a gradient flow pruning criterion, employing a progressive soft pruning strategy to ensure the continuity of the mask throughout the pruning process. Specifically, the gradient flow pruning criterion identifies and prunes parameters whose removal would increase the gradient norm of the loss function, thus facilitating faster convergence during the iterative pruning stage. By guiding the mask along the gradient flow in the sparse space, our method mitigates the abrupt information loss typically associated with one-shot pruning and achieves a more optimized mask compared to vanilla iterative pruning. Additionally, we design a coarse-grained pruning framework that achieves comparable results while significantly reducing time consumption. Extensive experiments conducted on widely used datasets demonstrate that our method not only delivers superior performance but also enhances efficiency and preserves greater consistency with the original pre-trained models.