文章链接：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可以看作是一堆残差层、Transformer层和几个上/下采样的卷积层的堆叠，省略了跳跃连接。这些层中的大部分在它们自己的输入和输出形状上保持一致。因此，这些层可以直接移除，而不会引起形状不匹配的问题。具体来说，除了涉及输入和输出通道数量变化的几个残差层外，SDM的U-Net的所有残差层和Transformer层都可以剪枝，如下图2所示。可剪枝层占整个SDM-v1.5和SDXL的整个U-Net的总参数的88%和98%。将所有的n个可剪枝层的集合表示为，其中是单个残差层或Transformer层。
先前的方法重新训练剪枝的SDM U-Net使用了三个目标，即任务损失、对数蒸馏（输出蒸馏）损失和特征蒸馏损失：
其中，?、?和表示潜在表示、对应的条件和对应的时间步。~表示真值噪声。和 分别是学生和老师的输出，其中学生是被剪枝的U-Net，而老师是冻结的原始U-Net。分别表示老师和学生在第?个阶段结束时的特征图。将 这样的不同?的项称为本文中的特征损失项。和是知识蒸馏目标。这三个目标也在下图2中说明。总体目标是。其中和是超参数，默认情况下都设置为1。
一次性层剪枝
层剪枝的目标是从U-Net中删除一组可剪枝层，以满足给定的剪枝比率，同时导致网络性能的最小下降。选择最小化原始网络和剪枝网络之间的均方误差损失，给定一个剪枝比率。还有其他可能的目标，我们将在本节末尾讨论。设表示原始网络的输出，表示已移除的第?个层，表示删除了层的网络的输出。注意，?是一个未确定的变量。优化问题的公式如下：
其中，表示层的参数总数，?是要剪枝的参数数量，由原始网络的总参数数乘以剪枝比率计算得到。直接解决这个优化问题需要在的所有组合上进行搜索，这是NP难的。因此，我们试图找到一个替代目标。
目标（2）的上界。根据三角不等式，我们得到了目标（2）的上界。
因此，我们为目标（2）找到了一个替代目标。
因此，解决优化问题（4）仍然是NP难的。然而，这引导我们找到了另一个更好的替代目标。
上界（4）的近似。我们注意到公式（4）的每一项都是一个网络（剪枝或未剪枝）与删除一个额外层的网络之间的均方误差损失，例如， 是层被移除的网络与层,被移除的网络之间的均方误差损失。我们猜想，删除相同层??导致的均方误差损失在不同的网络中是相似的。该假设被表述如下：
最终，借助假设（5），我们可以通过优化其近似的目标来优化目标（4）。
使用目标（6）来代替原始目标（2）意味着一种有趣的属性，称为可加性，即网络受到多个扰动引起的输出失真大约等于每个单一扰动引起的输出失真之和。在第4.3节中进行的进一步实验验证了我们方法在SDMs上的可加性属性，显示出最终的目标（6）是对原始目标（2）的良好近似，同时，假设（5）也得到了很好的支持。
在这里，项充当剪枝标准。我们只需要计算对于每个层，即原始网络和仅移除层的网络之间的输出损失，总时间复杂度为?(?)，然后优化问题（6）变成了0-1背包问题的一个变种。类似于原始的0-1背包问题，这个变种问题也可以通过动态规划算法求解，或者可以通过贪心算法近似求解。两种伪代码都在补充材料中给出。当使用贪心算法时，解会退化为，与先前方法采用的相同方案。贪心算法更简单易行，但动态规划算法始终更好。我们在工作中使用了动态规划算法。最后，一次性层剪枝方法总结如下。
讨论。在本节中，将一次性层剪枝问题形式化为一个组合优化问题，并推导出一个剪枝标准，即输出损失。此外，利用类似的推导形式，从公式（2）到公式（6），我们可以构建一些其他可能的剪枝标准，例如，Δ任务损失、ΔCLIP分数，分别由和计算，其中L是任务损失，C是CLIP分数。这些标准的推导在补充材料中。
实验和消融研究表明，我们的输出损失剪枝标准显著满足可加性属性，并在不同SDM模型上实现了最佳性能。此外，对于我们构建的三个剪枝标准，即输出损失、Δ任务损失和ΔCLIP分数，对于每个模型，具有更强可加性属性的标准实现了更好的剪枝性能。这个观察是合理的，因为一次性层剪枝直接优化了替代目标，如（6），而不是原始目标，如（2），优化那些更好地近似原始目标的替代目标通常可以实现更好的剪枝性能。基于以上观察和讨论，我们可以得出结论，我们的输出损失标准的有效性来自其良好的可加性属性。
规范化特征蒸馏
在实践中，我们发现重新训练的关键目标是特征损失。然而，进一步的检查揭示了前面在重新训练阶段采用的特征蒸馏方法存在不平衡问题。如下图3所示，不同阶段末端的特征图的L2-范数和不同特征损失项的值变化显著。在图3中显示的相同现象在使用我们的层剪枝或手工层删除的不同剪枝设置下，无论是在SDXL还是SDM-v1.5上都观察到。在整个蒸馏过程中，最高的特征损失项约比最低的特征损失项大10000倍，并产生大约1000倍更大的梯度，稀释了数值上不显著的特征损失项的梯度。
由于具有更大 L2-范数的特征图自然倾向于产生更大的特征损失项，不同特征损失项的显著幅度差异归因于两个因素，即在不同阶段末端，教师和学生之间特征图的固有差异不同，以及特征图的 L2-范数不同。第二个因素在很大程度上加剧了不平衡问题。
基于以上观察，简单地将所有特征损失项相加将导致少数特征损失项严重占优，阻碍数值上不显著的特征损失项的减少，导致性能下降。为此，我们提出了一种简单而有效的特征蒸馏重新加权策略，以消除第二个因素对不平衡问题的影响。我们选择利用教师特征图的 L2-范数重新加权特征损失项，并将特征蒸馏适应于我们的剪枝方案。规范化特征损失的制定如下:
其中，? 是剪枝后仍保留有残差层或transformer层的阶段集合，|? | 是集合 ? 的大小。因此，整体重新训练目标是:
在没有进行任何超参数调整的情况下，将 和 设置为 1。
讨论。进一步研究了我们的方法对重新训练过程中不同 项的影响。如图4所示，通过应用规范化特征蒸馏，观察到某些阶段的项显著下降，而其他地方略有增加，证明了一些特征损失项对其他特征损失项的支配程度得到了缓解。进一步的消融研究表明，我们的方法最终取得了显著的性能改进。
实验
实现
在本小节中，主要详细阐述了基本的实现设置。更多细节请参见补充材料。
模型选择。在后面实验中，采用 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) 的上界的近似，有必要检验近似的精度。通过选择不同的 ? 和 ，覆盖从 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 实现了最先进的性能压缩。
参考文献
[1] LAPTOP-Diff: Layer Pruning and Normalized Distillation for Compressing Diffusion Models
更多精彩内容，请关注公众号：AI生成未来
欢迎加群交流AIGC技术，添加小助手