arXiv:2412.02852cs.CV2024-12被引 10

无需重训练,可大幅压缩生成模型参数量

Learnable Sparsity for Vision Generative Models

  • 提出可学习掩码实现无重训练的结构化剪枝
  • 在SDXL和FLUX上剪枝20%参数,性能几乎无损
  • 适合追求高效部署的模型优化研究者

扩散模型在视觉任务中取得显著进展,但其性能提升往往依赖于模型规模扩大,导致计算复杂度和内存需求激增,影响部署效率并增加推理成本与环境负担。现有剪枝方法多需大量重训练以维持性能,而大模型重训练成本极高,实用性受限。本文提出一种无需重训练的通用结构剪枝框架,通过可学习的可微掩码对扩散模型进行稀疏化。为保障去噪过程质量,设计贯穿整个扩散流程的端到端剪枝目标。针对端到端剪枝内存占用高的问题,进一步提出时间步梯度检查点技术,显著降低优化阶段内存消耗,使在有限内存下实现端到端剪枝成为可能。在先进U-Net扩散模型SDXL与扩散变换器FLUX上的实验表明,该方法可有效剪枝高达20%参数,且性能下降可忽略不计。此外,方法还可应用于时间步蒸馏后的扩散模型。

原文摘要 · Abstract (English)

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.

扩散模型模型剪枝低资源部署

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。