文章链接：
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的总参数的
和
。将所有的
个可剪枝层的集合表示为
, 其中
是单个残差层或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 背包问题的一个变种。类似于原始的
背包问题, 这个变种问题也可以通过动态规划算法求解, 或者可以通过贪心算法近似求解。两种伪代码都在补充材料中给出。当使用贪心算法时, 解会退化为
, 与先前方法采用的相同方案。贪心算法更简单易行, 但动态规划算法始终更好。我们在工作中使用了动态规划算法。最后, 一次性层剪枝方法总结如下。
讨论
。在本节中, 将一次性层剪枝问题形式化为一个组合优化问题, 并推导出一个剪枝标准, 即输出损失
。此外, 利用类似的推导形式, 从公式 (2) 到公式 (6), 我们可以构建一些其他可能的前枝标准, 例如,
任务损失、
CLIP分数, 分别由
和
计算, 其中L是任务损失, C是CLIP分数。这些标准的推导在补充材料中。
实验和消融研究表明, 我们的输出损失剪枝标准显著满足可加性属性, 并在不同SDM模型上实现了最佳性能。此外, 对于我们构建的三个剪枝标准, 即输出损失、
任务损失和
CLIP分数, 对于每个模型, 具有更强可加性属性的标准实现了更好的剪枝性能。这个观察是合理的, 因为一次性层剪枝直接优化了替代目标, 如 (6), 而不是原始目标, 如 (2), 优化那些更好地近似原始目标的替代目标通常可以实现更好的剪枝性能。基于以上观察和讨论, 我们可以得出结论, 我们的输出损失标准的有效性来自其良好的可加性属性。
规范化特征蒸馏
在实践中, 我们发现重新训练的关键目标是特征损失
