Learnable Sparsity for Vision Generative Models

TL;DR

This paper introduces a low-cost, retraining-free structural pruning method for diffusion models in vision tasks, which learns differentiable masks and employs memory-efficient techniques to prune up to 20% of parameters with minimal performance loss.

Contribution

It presents a novel end-to-end, model-agnostic pruning framework for diffusion models that does not require retraining and incorporates memory-saving gradient checkpointing.

Findings

Prunes up to 20% of parameters with minimal quality loss.

Effective pruning achieved on state-of-the-art diffusion models SDXL and FLUX.

Maintains performance even when applied to time step distilled diffusion models.

Abstract

Diffusion models have achieved impressive advancements in various vision tasks. However, these gains often rely on increasing model size, which escalates computational complexity and memory demands, complicating deployment, raising inference costs, and causing environmental impact. While some studies have explored pruning techniques to improve the memory efficiency of diffusion models, most existing methods require extensive retraining to retain the model performance. Retraining a modern large diffusion model is extremely costly and resource-intensive, which limits the practicality of these methods. In this work, we achieve low-cost diffusion pruning without retraining by proposing a model-agnostic structural pruning framework for diffusion models that learns a differentiable mask to sparsify the model. To ensure effective pruning that preserves the quality of the final denoised latent, we design a novel end-to-end pruning objective that spans the entire diffusion process. As end-to-end pruning is memory-intensive, we further propose time step gradient checkpointing, a technique that significantly reduces memory usage during optimization, enabling end-to-end pruning within a limited memory budget. Results on state-of-the-art U-Net diffusion models SDXL and diffusion transformers (FLUX) demonstrate that our method can effectively prune up to 20% parameters with minimal perceptible performance degradation, and notably, without the need for model retraining. We also showcase that our method can still prune on top of time step distilled diffusion models.