Pruning for Sparse Diffusion Models based on Gradient Flow

TL;DR

This paper introduces an iterative gradient flow-based pruning method for diffusion models that improves efficiency and maintains generation quality better than traditional one-shot pruning techniques.

Contribution

It proposes a novel gradient flow pruning approach with a progressive soft pruning strategy to preserve model performance during sparsification.

Findings

Achieves higher efficiency in diffusion models.

Maintains better generation quality after pruning.

Enables faster convergence in sparse models.

Abstract

Diffusion Models (DMs) have impressive capabilities among generation models, but are limited to slower inference speeds and higher computational costs. Previous works utilize one-shot structure pruning to derive lightweight DMs from pre-trained ones, but this approach often leads to a significant drop in generation quality and may result in the removal of crucial weights. Thus we propose a iterative pruning method based on gradient flow, including the gradient flow pruning process and the gradient flow pruning criterion. We employ a progressive soft pruning strategy to maintain the continuity of the mask matrix and guide it along the gradient flow of the energy function based on the pruning criterion in sparse space, thereby avoiding the sudden information loss typically caused by one-shot pruning. Gradient-flow based criterion prune parameters whose removal increases the gradient norm of loss function and can enable fast convergence for a pruned model in iterative pruning stage. Our extensive experiments on widely used datasets demonstrate that our method achieves superior performance in efficiency and consistency with pre-trained models.