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

Area of Science:

• Artificial Intelligence
• Machine Learning
• Computer Vision

Background:

• Diffusion Models (DMs) are powerful generative models but suffer from slow inference and high computational costs.
• Existing lightweight DM approaches using one-shot pruning often degrade generation quality and remove crucial parameters.

Purpose of the Study:

• To develop an efficient and effective method for creating lightweight Diffusion Models.
• To overcome the limitations of one-shot pruning in maintaining generation quality and parameter integrity.

Main Methods:

• Proposed an iterative pruning method guided by gradient flow.
• Integrated a gradient flow pruning criterion to identify and remove parameters that minimally impact the loss function's gradient norm.
• Employed a progressive soft pruning strategy for mask continuity and designed a coarse-grained pruning framework for efficiency.

Main Results:

• The gradient flow pruning method mitigates abrupt information loss, yielding better-optimized masks than vanilla iterative pruning.
• Achieved superior performance, enhanced efficiency, and preserved greater consistency with original pre-trained models.
• The coarse-grained pruning framework demonstrated comparable results with significantly reduced time consumption.

Conclusions:

• The proposed iterative gradient flow pruning method effectively creates lightweight Diffusion Models with high generation quality and efficiency.
• This approach offers a more robust alternative to one-shot pruning for model compression in generative AI.