回复
为Stable Diffusion模型瘦身并达到SOTA！LAPTOP-Diff：剪枝蒸馏新高度（哈工大&OPPO） 精华
文章链接：https://arxiv.org/pdf/2404.11098
在AIGC时代，对低成本甚至设备端应用扩散模型的需求日益增加。在压缩Stable Diffusion模型（SDM）方面，提出了几种方法，其中大多数利用手工设计的层移除方法来获得更小的U-Net，同时利用知识蒸馏来恢复网络性能。然而，这种手工设计的层移除方式效率低下，缺乏可扩展性和泛化性，并且在重新训练阶段使用的特征蒸馏面临着不平衡问题，即一些数值显著的特征损失项在整个重新训练过程中占主导地位。
为此，本文提出了用于压缩扩散模型（LAPTOP-Diff）的层剪枝和归一化蒸馏方法。
- 引入了层剪枝方法来自动压缩SDM的U-Net，并提出了一个有效的一次性剪枝准则，其一次性性能由其良好的可加性特性保证，超越了其他层剪枝和手工设计的层移除方法;
- 提出了归一化特征蒸馏用于重新训练，缓解了不平衡问题。
利用所提出的LAPTOP-Diff，压缩了SDXL和SDM-v1.5的U-Net，以获得最先进的性能，即使在剪枝比率为50%时，PickScore仅下降了最小的4.0％，而比较方法的最小PickScore下降为8.2％。
引言
在最近几年，文本到图像（T2I）合成的生成建模经历了快速进展。特别是，扩散模型以其生成高分辨率、照片逼真且多样化的图像的能力而出现。在所有扩散模型中，Stable Diffusion模型（SDM）是最具影响力的一个，在AIGC社区中扮演着至关重要的角色，作为一个开源框架，为广泛的下游应用提供了基础。
然而，SDM卓越的性能伴随着其相当大的内存消耗和延迟，使其在个人计算机甚至移动设备上的部署受到严重限制。此外，最近版本的SDM系列，如SDXL，往往会增加更多的参数，导致内存消耗和延迟进一步增加。
为了减少SDM的推理预算，提出了几种方法，即降噪步骤的减少、高效的架构设计、结构剪枝、量化和硬件优化。这些方法通常是相互正交的。
在这些方法中，高效的架构设计和结构剪枝是被低估的。一方面，先前的高效架构设计方法通常通过大量的经验研究来识别SDM的U-Net中不重要的层，然后移除它们以获得更小更快的网络。这种手工设计的方式通常无法达到最佳性能，并且缺乏可扩展性和泛化性。
注意到，这些手工设计的层移除方法可以被自动方案中的层剪枝方法所替代，以获得更好的可扩展性和性能。另一方面，先前的SDM结构剪枝方法侧重于精细的剪枝，即剪枝参数矩阵的行和列。然而，有研究表明，与更细粒度的结构剪枝相比，粗粒度的层剪枝通常在减少模型延迟方面效率更低，并且有趣的是，层剪枝有可能达到与更细粒度的结构剪枝甚至更好的性能。基于以上两点，层剪枝方法值得研究。
在层移除或剪枝之后，SDMs通常无法直接生成清晰的图像。先前的方法利用知识蒸馏来重新训练被剪枝的网络以恢复其性能。先前的方法通常利用三种类型的目标，即常规训练目标、对数蒸馏（输出蒸馏）目标和特征蒸馏目标。在这三个部分中，特征蒸馏是关键部分。然而，进一步的检查发现，先前基于蒸馏的重新训练方法存在一个不平衡问题，即一些特征损失项在整个重新训练过程中占主导地位，导致性能下降。
本文提出了用于压缩扩散模型的层剪枝和归一化蒸馏方法（LAPTOP-Diff），将SDM的高效架构设计和结构剪枝推向自动化、可扩展性和更大性能的领域。从组合优化问题的高角度来阐述层剪枝问题，并以简单而有效的一次性方式解决它。
受益于这种视角，我们探索了几种其他可能的剪枝标准，并通过消融研究，发现一次性层剪枝标准的有效性来自其良好的可加性属性。此外，确定了先前基于蒸馏的重新训练方法存在一个不平衡问题，即一些特征损失项在整个重新训练过程中占主导地位，通过提出的归一化特征蒸馏来缓解这个问题。贡献总结如下：
- 探索了层剪枝在SDMs上的一个被低估的结构剪枝方法，并提出了一个有效的一次性剪枝标准，其一次性性能由其良好的可加性属性保证，超越了其他层剪枝和手工设计的层移除方法，将先前基于层移除的高效架构设计方法推向自动化、可扩展性和更大性能。
- 通过提出的归一化特征蒸馏方法缓解了先前基于蒸馏的重新训练存在的不平衡问题。
- 提出的LAPTOP-Diff在不同SDMs和剪枝比率下大大超越了基于层移除的高效架构设计方法，在不同SDMs和剪枝比率下具有更好的网络性能。
相关工作
扩散模型
扩散模型是一种利用迭代去噪过程来合成数据的生成模型。在文本到图像（T2I）合成领域，诸如DALL·E、Imagen、Deepfloyd IF和Stable Diffusion等扩散模型展示了它们生成高分辨率、照片逼真且多样化图像的显著能力。在各种扩散模型中，Stable Diffusion在学术界和工业界都是最具影响力的一个。Stable Diffusion模型（SDM）是一种在低维潜在空间执行迭代去噪过程，然后通过VAE解码器将潜在表示转换为像素空间图像的潜在扩散模型。SDM系列还有一个更新版本，即SDXL，在更高分辨率的1024×1024下展示出卓越的图像生成质量。
然而，扩散模型令人印象深刻的性能伴随着相当大的内存消耗和延迟。为了减少SDM的模型预算，已经探索了几种方法，例如，减少去噪步骤、量化、硬件优化、高效的架构设计和结构剪枝。
SDM的高效架构设计
与许多其他减少SDM模型预算的方法正交，高效的架构设计主要旨在设计原始SDM U-Net的有效子结构，因为大部分SDM的内存消耗和延迟来自其U-Net。这种类型的先前方法通常通过大量的经验研究来识别SDM U-Net中不重要的层，并移除它们以获得更小更快的网络。
例如，BK-SDM通过层移除为SDM-v1或SDM-v2手工设计了3个不同大小的高效U-Net，部分遵循了对BERT模型进行压缩的经验结论。SSD-1B和Segmind-Vega（在本文的其余部分将它们简称为SSD和Vega）通过人类评估识别不重要的层，然后施加层移除，为SDXL手工设计了2个不同大小的高效U-Net。KOALA是从BK-SDM衍生出来的，也通过层移除为SDXL手工设计了2个不同大小的高效U-Net。
这种手工设计的方式通常无法达到最佳性能，并且缺乏可扩展性和泛化性。我们注意到这些基于层移除的方法可以归类为手工设计的层剪枝。因此，我们提出这些手工设计的层移除方法可以被层剪枝方法在自动方案中替代，以获得更好的可扩展性和性能。
层剪枝
层剪枝，也称为深度剪枝，是一种旨在自动评估和移除不重要层的结构剪枝方法。与其他结构剪枝方法不同，由于其粗粒度的特性，层剪枝受到的关注较少。相比于那些细粒度的结构剪枝方法，层剪枝通常被认为效果较差。然而，有研究表明，与细粒度的结构剪枝方法相比，粗粒度的层剪枝通常更有效地减少模型的延迟，并且有可能实现相同甚至更好的性能。
虽然已经提出了几种层剪枝方法，但是没有一种方法是从组合优化问题的角度来看待层剪枝。基于幅度和基于泰勒展开的剪枝是常见的基准线。基于幅度的层剪枝使用层中参数的幅度总和作为层重要性标准，基于泰勒展开的层剪枝使用损失函数的一阶泰勒展开作为层重要性标准。在通过不同的重要性标准评估每个层的重要性后，先前的层剪枝方法选择修剪最不重要的层。
本文从更高的角度将层剪枝问题描述为组合优化问题，并以简单而有效的一次性方式解决了它，形成了一个有效的一次性剪枝标准，超越了其他层剪枝和手工设计的层移除方法。此外，通过这样的视角，能够确定一次性层剪枝标准的有效性来自于其良好的可加性属性。
基于蒸馏的重新训练
在层移除或剪枝之后，SDMs通常无法直接生成清晰的图像。先前的方法通过利用知识蒸馏来重新训练被剪枝的SDMs以恢复其性能。在重新训练阶段使用的知识蒸馏通常包括三部分，即常规训练目标、对数蒸馏（输出蒸馏）目标和特征蒸馏目标。在这三个部分中，特征蒸馏是最有效的。然而，在实践中，发现了基于蒸馏的重新训练过程中存在一个不平衡问题。为此提出了一种简单而有效的重新加权策略来缓解这个问题。
方法论
初步
先前的方法重新训练剪枝的SDM U-Net使用了三个目标，即任务损失、对数蒸馏（输出蒸馏）损失和特征蒸馏损失：
一次性层剪枝
目标（2）的上界。根据三角不等式，我们得到了目标（2）的上界。
因此，我们为目标（2）找到了一个替代目标。
最终，借助假设（5），我们可以通过优化其近似的目标来优化目标（4）。
使用目标（6）来代替原始目标（2）意味着一种有趣的属性，称为可加性，即网络受到多个扰动引起的输出失真大约等于每个单一扰动引起的输出失真之和。在第4.3节中进行的进一步实验验证了我们方法在SDMs上的可加性属性，显示出最终的目标（6）是对原始目标（2）的良好近似，同时，假设（5）也得到了很好的支持。
实验和消融研究表明，我们的输出损失剪枝标准显著满足可加性属性，并在不同SDM模型上实现了最佳性能。此外，对于我们构建的三个剪枝标准，即输出损失、Δ任务损失和ΔCLIP分数，对于每个模型，具有更强可加性属性的标准实现了更好的剪枝性能。这个观察是合理的，因为一次性层剪枝直接优化了替代目标，如（6），而不是原始目标，如（2），优化那些更好地近似原始目标的替代目标通常可以实现更好的剪枝性能。基于以上观察和讨论，我们可以得出结论，我们的输出损失标准的有效性来自其良好的可加性属性。
规范化特征蒸馏
在实践中，我们发现重新训练的关键目标是特征损失LFeatKD。然而，进一步的检查揭示了前面在重新训练阶段采用的特征蒸馏方法存在不平衡问题。如下图3所示，不同阶段末端的特征图的L2-范数和不同特征损失项的值变化显著。在图3中显示的相同现象在使用我们的层剪枝或手工层删除的不同剪枝设置下，无论是在SDXL还是SDM-v1.5上都观察到。在整个蒸馏过程中，最高的特征损失项约比最低的特征损失项大10000倍，并产生大约1000倍更大的梯度，稀释了数值上不显著的特征损失项的梯度。
由于具有更大 L2-范数的特征图自然倾向于产生更大的特征损失项，不同特征损失项的显著幅度差异归因于两个因素，即在不同阶段末端，教师和学生之间特征图的固有差异不同，以及特征图的 L2-范数不同。第二个因素在很大程度上加剧了不平衡问题。
基于以上观察，简单地将所有特征损失项相加将导致少数特征损失项严重占优，阻碍数值上不显著的特征损失项的减少，导致性能下降。为此，我们提出了一种简单而有效的特征蒸馏重新加权策略，以消除第二个因素对不平衡问题的影响。我们选择利用教师特征图的 L2-范数重新加权特征损失项，并将特征蒸馏适应于我们的剪枝方案。规范化特征损失的制定如下:
其中，𝑉 是剪枝后仍保留有残差层或transformer层的阶段集合，|𝑉 | 是集合 𝑉 的大小。因此，整体重新训练目标是:
实验
实现
在本小节中，主要详细阐述了基本的实现设置。更多细节请参见补充材料。
模型选择。在后面实验中，采用 512×512 分辨率以进行快速验证，对 SDXL 模型使用 ProtoVisionXL-v6.2.0，因为官方发布的 SDXL-Base-1.0在 512×512 分辨率附近生成的图像异常，而对于 SDM-v1.5 模型，我们使用了官方发布的 stable-diffusion-v1-5。实验，我们使用了每种比较方法的相同教师模型。
数据集。对于用于剪枝的校准数据集，使用 LAION-2B的随机抽样的 1K 子集。在消融研究中，使用了 LAION-2B的随机抽样的 0.34M 子集。为了与其他方法进行比较，使用了相同的数据集或选择了质量和数量较低的数据集，如果其他方法报告的数据集难以复现，会选择质量和数量较低的数据集。数据集的选择详见补充材料。
评估指标。尽管评估生成式 T2I 模型常用的指标是 FID和 CLIP 分数，但最近的研究表明这两个指标与视觉美学和人类偏好之间的相关性较低。因此，
采用了 3 个高级指标来评估模型的综合性能。使用 HPS v2、PickScore和 ImageReward来评估生成图像的视觉质量和文本-图像一致性。在它们各自的基准数据集上计算这 3 个指标，即 HPS v2 在其 3.2K 基准数据集上计算，ImageReward 在 ImageRewardDB的随机抽样的 3K 子集上计算，而 PickScore 在 Pick-a-Pic v1的随机抽样的 3K 子集上计算。
与其他方法的比较
到目前为止，对于 SDM 模型有三种手工制作的层移除方法。KOALA、SSD 和 Vega适用于 SDXL，而 BK-SDM可应用于 SDM-v1 或 SDM-v2。将 LAPTOP-Diff 与这三种方法在相应的模型上进行比较。使用 SDXL-Base-1.0与 KOALA 1.16B 和 782M 进行比较，使用 RealisticVision-v4.0与更先进的 BK-SDM实现（由 Segmind）进行比较，即 579M small-sd和 323M tiny-sd。对于与 1.3B SSD和 745M Vega进行比较，仅使用它们在多教师蒸馏策略中使用的三个教师模型之一，即 ZavychromaXL-v1.0。对于所有比较方法的评估，我们使用它们发布的模型权重。
表1 显示了与最先进的压缩 SDM 模型的比较结果。提出的 LAPTOP-Diff 实现了最先进的性能。与 SSD和 Vega的视觉比较如图5 所示，其他方法的更多视觉比较见补充材料。可以观察到，与其他方法相比，我们的方法在不同提示下实现了更好的视觉效果。值得注意的是，我们的大约 50% 压缩的 SDXL 模型几乎可以达到原始模型的相同视觉质量。
添加性质的验证
由于最终的替代目标 (6) 是原始目标 (2) 的上界的近似，有必要检验近似的精度。通过选择不同的 𝑚 和 Rm，覆盖从 0% 到 90% 的剪枝比率，模拟了近似准则 (6) 和真实准则 (2) 的值。我们还对剪枝准则 Δtask loss 和 ΔCLIP score 进行了相同的实验。如图6 所示，我们的输出损失准则在不同模型下显著满足添加性质，即使在极端的剪枝比率下也是如此，几乎所有观察点都位于恒等线附近。因此，我们最终的替代目标 (6) 是原始目标 (2) 的出色近似。另一方面，其他准则未能始终满足添加性质。虽然我们的输出损失准则在不同模型之间具有最强的添加性质，但 ΔCLIP score 在 SDXL 上具有第二强的添加性质，而 Δtask loss 在 SDM-v1.5 上具有第二强的添加性质。
剪枝准则消融实验
实验表明，我们的输出损失准则在不同模型上显著并一致地满足添加性质，而其他准则未能达到这一点。此外，我们评估了使用不同剪枝准则的剪枝性能，包括我们构建的剪枝准则和两个基线剪枝准则，即幅度度量和损失函数的一阶泰勒展开。我们首先使用不同的层剪枝准则对 SDXL 和 SDM-v1.5 进行剪枝，然后通过普通蒸馏对修剪后的网络进行相同次数的重新训练。如表2 所示，在不同模型上，输出损失准则在剪枝准则中取得了最高的剪枝性能。
除了不同的剪枝准则，还将层剪枝方法与手工制作的层移除方法进行了比较。对于 SDXL 和 SDM-v1.5 的实验，分别使用 SSD和 BK-SDM的相同手工制作层移除设置。表2 中的结果显示，除了基于幅度的层剪枝被 SDXL 上的手工制作方法 SSD 打败外，几乎所有的层剪枝方法通常都优于手工制作方法。
值得注意的是，在我们构建的三种剪枝准则中，即输出损失、Δtask loss 和 ΔCLIP score，虽然输出损失在不同模型上取得了最佳性能，但 ΔCLIP score 在 SDXL 上取得了第二好的性能，而 Δtask loss 在 SDM-v1.5 上取得了第二好的性能，与它们添加性质的排名相一致。
知识蒸馏消融实验
通过消融研究验证了我们的归一化特征蒸馏的有效性。使用我们的一次性层剪枝方法对 SDXL 和 SDM-v1.5 进行剪枝，然后使用普通蒸馏或我们的归一化蒸馏对修剪后的网络进行相同次数的重新训练。如表3 所示，我们的方法在 SDXL 和 SDM-v1.5 上都取得了很好的性能改进。
剪枝分析
以剪枝比例为 50% 的 SDXL 的 U-Net 和剪枝比例为 33% 的 SDM-v1.5 的 U-Net 为例进行了剪枝架构的分析。如表4 所示，对于 SDXL 和 SDM-v1.5，许多位于中间阶段及其附近的层被认为不太重要，因此被剪枝了。这一观察结果与先前的手工制作层移除方法一致。然而，我们观察到的是，与之前的手工制作层移除方法不同，更多的层在 Dn 阶段被剪枝，而不是在 Up 阶段。对于 SDXL，我们观察到在 Dn 阶段有 30 层中有 18 层被剪枝，而在 Up 阶段只有 10 层中有 10 层被剪枝。同样的现象也在 SDM-v1.5 上观察到，在 Dn 阶段有 14 层中有 4 层被剪枝，而在 Up 阶段只有 21 层中有 1 层被剪枝。这一观察结果与其他研究一致，即 Dn 阶段，即 U-Net 的编码器，不如 U-Net 的其他部分重要。
结论
本项工作提出了用于压缩扩散模型的层剪枝和归一化蒸馏（LAPTOP-Diff）。引入了层剪枝方法来实现自动化、可扩展性和更好的性能，并提出了一种有效的一次性剪枝准则，即输出损失准则，其有效性是通过其良好的添加性质来保证的。进一步通过提出的归一化特征蒸馏来缓解先前基于蒸馏的重训练中的不平衡问题。使用提出的 LAPTOP-Diff，为 SDMs 实现了最先进的性能压缩。
本文转自 AI生成未来 ，作者：Dingkun Zhang等
赞
收藏
回复
相关推荐
