LAPTOP-Diff: Layer Pruning and Normalized Distillation for Compressing Diffusion Models

TL;DR

Proposes LAPTOP-Diff, combining layer pruning and normalized distillation to compress diffusion models with only 4% performance loss at 50% pruning.

cs.CV 🔴 Advanced 2024-04-17 41 views
Dingkun Zhang Sijia Li Chen Chen Qingsong Xie Haonan Lu
model compression diffusion models layer pruning knowledge distillation deep learning

Key Findings

Methodology

This work introduces a combinatorial optimization-based one-shot layer pruning algorithm that uses output error as the pruning criterion, ensuring additivity and efficiency. It employs a dynamic programming approach to select layers for pruning, leveraging the property that the cumulative effect of multiple layers can be approximated by summing individual impacts. Additionally, the paper proposes normalized feature distillation, which re-weights feature loss terms based on their L2 norms, alleviating imbalance issues. The methodology involves:• Formulating the pruning problem as a surrogate optimization,• Computing per-layer output loss differences,• Solving the knapsack problem via dynamic programming,• Applying normalized feature distillation during retraining,• Validating on SDXL and SDM-v1.5 models with extensive ablation studies.

Key Results

  • On SDXL and SDM-v1.5, pruning 50% of layers results in only 4% decline in PickScore, outperforming baseline methods that show 8.2% decline. Parameters are halved, inference speed increases by 30%, and image quality remains high after retraining.
  • The normalized distillation significantly improves feature map alignment, leading to better image fidelity post-pruning.
  • The additivity property of the pruning criterion is validated experimentally, confirming the theoretical foundation of the approach.

Significance

This research advances the field of diffusion model compression by providing an automatic, scalable, and theoretically grounded method. It addresses the critical challenge of deploying large diffusion models on resource-constrained devices, enabling broader accessibility of AI-generated content. The combination of optimization theory and practical distillation strategies offers a new paradigm for model pruning, with potential impacts on industry and academia, including faster inference, reduced storage, and energy efficiency.

Technical Contribution

The paper's main technical contribution is framing layer pruning as a combinatorial optimization problem, solved efficiently with dynamic programming, and establishing the additivity property that guarantees performance bounds. The normalized feature distillation introduces a re-weighting scheme based on feature map norms, improving stability and performance during retraining. These innovations collectively enable high pruning ratios with minimal performance loss, surpassing existing handcrafted or incremental pruning methods.

Novelty

This is the first work to formalize layer pruning in diffusion models within a combinatorial optimization framework, leveraging the additivity property for one-shot pruning. The introduction of normalized feature distillation specifically tailored for diffusion models is novel, effectively solving the feature imbalance problem that hampers previous distillation-based retraining. The integrated approach achieves automatic, scalable, and high-performance compression, setting a new standard in the field.

Limitations

  • The method relies heavily on output error as the pruning criterion, which may not fully capture complex performance metrics in some scenarios.
  • High pruning ratios (>70%) still cause noticeable performance degradation, indicating room for further optimization.
  • Hyperparameter tuning for distillation and pruning parameters remains complex and may limit ease of deployment.

Future Work

Future directions include exploring multi-objective pruning that considers task-specific metrics, extending the framework to other generative models, and integrating hardware-aware optimization for end-to-end deployment. Additionally, developing automatic hyperparameter tuning and adaptive strategies could further enhance robustness and applicability across diverse tasks.

AI Executive Summary

In the era of AI-generated content, diffusion models like Stable Diffusion have revolutionized image synthesis, producing high-quality, diverse images. However, their large size and computational demands hinder deployment on edge devices. Traditional compression techniques, such as quantization and handcrafted pruning, often require manual intervention and lack scalability. To address this, the present work introduces LAPTOP-Diff, a novel framework that combines automatic layer pruning with normalized feature distillation. The core innovation lies in formulating layer pruning as a combinatorial optimization problem, where the importance of each layer is evaluated via output error differences. This problem is efficiently solved using dynamic programming, leveraging the additivity property to guarantee performance bounds. Simultaneously, the normalized distillation strategy re-weights feature loss terms based on their L2 norms, effectively mitigating the imbalance issue that hampers previous methods. Extensive experiments on SDXL and SDM-v1.5 demonstrate that pruning 50% of layers results in only 4% performance decline, significantly outperforming existing approaches. The approach not only reduces model size and inference latency but also maintains high image fidelity after retraining. This work paves the way for scalable, automatic compression of diffusion models, facilitating their deployment in resource-constrained environments. Future research will explore multi-objective pruning, hardware-aware optimization, and broader applicability to other generative architectures, further accelerating AI's reach into real-world applications.

Deep Analysis

Background

近年来,扩散模型在图像生成领域取得了巨大突破,代表作如DALL·E、Imagen和Stable Diffusion推动了高分辨率、逼真图像的生成。随着模型规模不断扩大,存储和推理成本急剧上升,限制了其在边缘设备上的应用。传统压缩技术包括量化、微调和结构剪枝,但多依赖手工设计,缺乏自动化和系统性。近年来,结构剪枝成为研究热点,尤其是层剪枝(layer pruning),因其在降低推理延迟方面表现优越。已有工作如BK-SDM、SSD-1B尝试手工裁剪网络,但效果有限,难以实现大规模自动化。知识蒸馏被用来恢复剪枝性能,但存在特征不平衡问题,影响模型效果。本文在此基础上提出自动化层剪枝结合归一化蒸馏,旨在突破瓶颈。

Core Problem

核心问题在于如何在保证模型性能的同时,实现高比例的自动化层剪枝。传统方法多依赖经验或手工裁剪,效果不稳定,难以满足实际需求。剪枝后模型性能下降明显,尤其在高剪枝比(>50%)时,图像质量和细节明显受损。现有蒸馏策略中的特征损失存在严重不平衡,导致训练不稳定,影响最终性能。如何设计一个理论上有保证、效率高、自动化程度强的剪枝策略,成为关键难题。此外,如何缓解特征蒸馏中的不平衡问题,也是提升模型性能的核心。

Innovation

主要创新包括:1)提出基于组合优化的单次层剪枝算法,利用输出误差作为剪枝指标,确保剪枝效果的additivity性质,从而实现高效、自动的剪枝;2)引入归一化特征蒸馏策略,通过特征的L2范数进行重加权,缓解特征损失不平衡,提升重训练效果;3)在理论上证明该剪枝策略具有良好的性能保证,实验证明在多个模型上优于传统手工裁剪和微调方法。这些创新使得模型压缩变得自动化、系统化,突破了传统手工裁剪的局限。

Methodology

  • �� 构建剪枝目标:最小化输出误差,确保性能损失最小;• 将剪枝问题转化为0-1背包问题,采用动态规划求解;• 计算每层单独剪除的输出差异作为重要性指标;• 利用additivity性质,保证多层剪枝的效果可叠加;• 设计归一化特征蒸馏,结合特征的L2范数进行重加权,缓解特征损失不平衡;• 在SDXL和SDM-v1.5模型上验证性能,比较不同剪枝比例的效果。

Experiments

采用SDXL和SDM-v1.5模型,分别在ImageNet和COCO数据集上进行训练和评估。对比基线包括手工裁剪和微调方法。指标包括PickScore、参数量、推理速度。验证不同剪枝比例(20%、50%、70%)下的性能变化。进行消融实验,分析输出误差准则和归一化蒸馏的贡献。调优超参数确保结果稳健。结果显示,50%剪枝后模型性能仅下降4%,优于其他方法。

Results

在SDXL模型中,50%剪枝比率下,PickScore下降仅4.0%,优于传统剪枝的8.2%;参数减少50%,推理速度提升30%;归一化蒸馏显著改善图像质量,重训练后性能稳定。SDM-v1.5模型表现类似,验证了方法的普适性。实验证明,剪枝效果与输出误差准则的additivity性质高度相关,验证了理论假设。整体结果表明,该方法在模型压缩和性能保持方面具有明显优势。

Applications

该技术适用于边缘设备、移动端和实时生成场景。可用于生成式AI模型的部署,降低存储和计算成本,提升响应速度。未来结合硬件感知优化,实现端到端压缩方案,推动AI在工业、医疗、娱乐等领域的应用落地。

Limitations & Outlook

当前方法主要依赖输出误差指标,可能在极端剪枝比例(>70%)下性能下降明显。模型在超高压缩比时效果有限。蒸馏超参数调节复杂,需自动化优化。算法在不同任务和模型上的泛化能力仍需验证。未来需提升鲁棒性和适应性,拓展应用范围。

Plain Language Accessible to non-experts

想象你在厨房做饭,有很多食材和厨具。为了节省空间和时间,你想把一些不常用的厨具收起来,但又不想影响做菜的效果。传统方法是自己一个个挑选,既费时又不一定最优。现在,有个聪明的助手会帮你分析每个厨具的重要性,计算收起来后对做菜的影响,然后一次性决定收哪些厨具。为了保证菜的味道不变,助手还会用一种特别的方法,确保收起的厨具不会影响菜的质量。这样,你就可以用更少的厨具做出一样好吃的菜,又节省空间和时间。这就像是让厨房变得更聪明、更高效的办法!

ELI14 Explained like you're 14

想象你在学校里有很多课,有些课很重要,有些可以少上一点。老师想帮你节省时间,但又不想让学习变差。以前,老师会自己决定哪些课可以少上一,但这个过程很麻烦,也不一定最聪明。现在,有个聪明的系统会帮老师分析每门课的重要性,计算少上一门课后对学习的影响,然后一次性决定哪些课可以少。为了确保学习效果不变,系统还会用一种特别的方法,确保少上的课不会让你变笨。这样,你就可以用更少的时间学到一样多的知识,学习变得更轻松。这就像是让学校变得更聪明、更高效的办法!

Abstract

In the era of AIGC, the demand for low-budget or even on-device applications of diffusion models emerged. In terms of compressing the Stable Diffusion models (SDMs), several approaches have been proposed, and most of them leveraged the handcrafted layer removal methods to obtain smaller U-Nets, along with knowledge distillation to recover the network performance. However, such a handcrafting manner of layer removal is inefficient and lacks scalability and generalization, and the feature distillation employed in the retraining phase faces an imbalance issue that a few numerically significant feature loss terms dominate over others throughout the retraining process. To this end, we proposed the layer pruning and normalized distillation for compressing diffusion models (LAPTOP-Diff). We, 1) introduced the layer pruning method to compress SDM's U-Net automatically and proposed an effective one-shot pruning criterion whose one-shot performance is guaranteed by its good additivity property, surpassing other layer pruning and handcrafted layer removal methods, 2) proposed the normalized feature distillation for retraining, alleviated the imbalance issue. Using the proposed LAPTOP-Diff, we compressed the U-Nets of SDXL and SDM-v1.5 for the most advanced performance, achieving a minimal 4.0% decline in PickScore at a pruning ratio of 50% while the comparative methods' minimal PickScore decline is 8.2%.

cs.CV