0. 论文信息
标题：DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization
作者：Haowei Zhu, Dehua Tang, Ji Liu, Mingjie Lu, Jintu Zheng, Jinzhang Peng, Dong Li, Yu Wang, Fan Jiang, Lu Tian, Spandan Tiwari, Ashish Sirasao, Jun-Hai Yong, Bin Wang, Emad Barsoum
机构：Advanced Micro Devices, Inc.、Tsinghua University
原文链接：https://arxiv.org/abs/2410.16942
1. 导读
扩散模型由于其突出的能力在图像生成领域取得了显著的进展。然而，由于推断过程中的多步去噪过程，这些模型需要大量的计算资源。虽然传统的剪枝方法已经被用来优化这些模型，但重新训练过程需要大规模的训练数据集和大量的计算成本来保持泛化能力，这使得它既不方便也不高效。最近的研究试图利用相邻去噪阶段的特征的相似性，通过简单和静态的策略来降低计算成本。然而，这些策略不能完全利用相邻时间步中相似特征模式的潜力。在这项工作中，我们提出了一种新的剪枝方法，通过一个更加智能和可微分的剪枝器得到一个有效的扩散模型。我们方法的核心是将模型修剪过程转换为子网搜索过程。具体来说，我们首先通过添加一些建立在相似特征上的备份连接来引入基于标准扩散的超网。然后，我们构建一个插件剪枝网络，并设计优化损失，以确定冗余计算。最后，我们的方法可以通过几个步骤的梯度优化和一个简单的后处理过程来确定一个最佳子网。我们对各种扩散模型进行了大量的实验，包括稳定扩散级数和DiTs。我们的DiP-GO方法实现了SD-1.5的4.4倍加速，而没有任何准确性损失，显著优于以前的最先进方法。
2. 引言
近年来，扩散概率模型（DPMs）的杰出能力使得扩散模型取得了显著进展。DPMs通常包含两个过程：噪声扩散过程和反向去噪过程。鉴于其在内容生成方面的显著优势，扩散模型在通用图像生成领域的多个方面取得了重大进展，包括文本到图像的生成、布局到图像的生成、图像编辑和图像个性化。此外，扩散模型还推动了自动驾驶领域的发展，从驾驶数据集的生成到通过扩散策略增强感知模型。然而，DPMs在训练和推理阶段往往会产生相当大的计算开销。由于采样过程中的多步去噪计算，推理的高成本可能会对其实际应用产生显著影响。为了提高扩散模型的效率，已经做出了许多努力，这些努力大致可以分为两类优化：推理采样优化和模型结构优化。
采样优化方法旨在减少生成所需的采样步骤，同时不牺牲图像质量。例如，DDIM通过探索非马尔可夫过程减少了这些步骤，而无需重新训练模型。LCM则通过重新训练要求，在更少的步骤中实现图像生成。结构优化方法旨在通过高效的模型设计和模型剪枝来减少计算开销。这些方法需要重新训练扩散模型，这涉及大量的计算开销和大规模数据集，因此既不方便也不高效。DeepCache提出了一种基于扩散模型中U-Net架构的新型无训练范式，通过在相邻的去噪阶段缓存和检索特征来减少冗余计算成本。然而，DeepCache仅通过简单且静态的策略在去噪步骤中重用U-Net块的输出特征。我们认为，许多中间特征尚未被开发，且如近期研究所示，简单的静态策略无法充分利用相邻时间步长中相似特征模式的潜力来进行推理。
为了解决这些挑战，我们引入了通过少步梯度优化进行扩散剪枝（DiP-GO）的方法，该方法旨在实现具有增强动态性和智能性的高效模型剪枝。我们的方法通过提出一个基于标准扩散并添加了一些基于相似特征的备用连接的SuperNet，重新思考了推理过程中的扩散模型，将推理过程概念化为从我们提出的SuperNet中派生的特定SubNet。我们将扩散模型剪枝重新定义为SubNet搜索过程。通过使用梯度检查点方法解决扩展去噪时间步长中反向过程固有的内存溢出问题，我们引入了一个插件剪枝器，该剪枝器通过精心设计的优化损失发现了超越现有方法的最佳SubNet。大量实验验证了我们的方法的有效性，在Stable Diffusion 1.5上实现了4.4倍的速度提升。此外，我们的方法无需重新训练扩散模型即可有效剪枝DiT模型，实现了显著的推理加速。
3. 主要贡献
我们的贡献可以总结如下：
（1）我们基于标准扩散定义了一个SuperNet，并展示了如何获得一个SubNet。这将扩散优化问题转化为一个高效的SubNet搜索过程，而无需重新训练预训练的扩散模型。
（2）我们设计了一个专门为扩散模型定制的插件剪枝器。该剪枝器在优化剪枝约束的同时最大化合成能力。此外，我们还为剪枝器引入了一种后处理方法，以确保SubNet满足指定的剪枝要求。
（3）我们在各种扩散模型上进行了广泛的剪枝实验，包括Stable Diffusion 1.5、Stable Diffusion 2.1、Stable Diffusion XL和DiT。大量实验证明了我们的方法的优越性，在Stable Diffusion 1.5的推理过程中实现了显著的4.4倍速度提升，而无需重新训练扩散模型。
推荐课程：
扩散模型入门教程：数学原理、方法与应用
。
4. 方法
我们的目标是在推理过程中识别和移除不重要的块，以加速该过程。为实现这一目标，我们基于扩散模型引入了一个SuperNet。该SuperNet旨在促进块的移除，同时通过额外的连接确保修剪后的模型保持推理能力。本文方法在推理过程中有效地消除了不重要的块，通过跳过这些不必要的组件，基本上从SuperNet中派生出一个SubNet。因此，修剪过程可以概念化为在SuperNet框架内进行SubNet搜索。
如何构建SuperNet。近期研究观察到，在推理过程中，扩散模型通常在相邻时间步上展现出相似的特征模式。基于这一见解，我们通过在当前时间步与前一时间步之间引入额外的连接，增强了标准扩散模型的推理阶段。这些连接作为可能被移除的块的备份，确保即使为加速而消除了其依赖块，每个块也能保留有效的输入。具体而言，对于推理过程中除初始步之外的所有时间步的每个块的输入，我们都建立了到前一时间步中对应块的备份输入连接，如图1所示。
在本研究中，我们介绍了通过少步梯度优化进行扩散剪枝（DiP-GO）的方法，该方法利用神经网络来预测在推理过程中是否跳过或保留每个计算块。我们的主要目标是确定计算块的最优子集，以便以最小的计算开销促进去噪。如图2所示，我们的方法包含三个主要组件：神经网络剪枝器、优化损失和后处理算法，以基于剪枝器的预测导出剪枝后的模型。神经网络剪枝器受DETR启发，设计了可学习的查询来预测每个块的状态。我们提出的优化损失包括用于生成质量的稀疏性和一致性约束，引导剪枝器准确评估每个块的重要性。我们首先回顾了扩散模型的框架，强调了其探索剪枝网络的潜力。之后，我们介绍了一个基于扩散模型的SuperNet，并展示了如何从其中推导出用于推理加速的SubNet或剪枝网络，同时强调了实现最优SubNet的挑战。最后详细介绍了我们的方法，包括神经网络剪枝器、优化损失和后处理算法，以获得满足剪枝要求的SubNet。最后，我们提供了关于我们方法的训练和推理过程的见解。
推荐课程：
扩散模型入门教程：数学原理、方法与应用
。
5. 实验结果
在PartiPrompt和COCO2017数据集上的稳定扩散研究。我们将本文方法与稳定扩散1.5（SD-1.5）上的最先进（SOTA）压缩方法进行了比较，结果总结在表1中。与SOTA方法DeepCache相比，本文方法展现出显著的性能提升，在保持更高CLIP分数的同时，实现了近2倍的MAC（乘加运算次数）减少。与基线模型相比，本文方法可实现4.4倍的速度提升。此外，本文方法无需训练扩散模型，从而保留了扩散模型的预训练知识。同时，我们在SD-2.1模型上应用了本文方法以验证其有效性，如表2所示，本文方法在保持生成质量的同时实现了显著加速，证明了其优越性。
6. 总结 & 未来工作
本文工作探索了通过减少相邻时间步之间的冗余特征计算来解决扩散加速任务。我们提出了一种新颖的扩散修剪框架，并将模型修剪过程视为SubNet搜索问题。本文方法引入了一个插件修剪器网络，通过几步梯度优化来识别最优SubNet。在广泛的稳定扩散（SD）和DiT系列模型上的结果表明了本文方法的有效性。我们在稳定扩散1.5上实现了4.4倍的速度提升，并有效地通过几步优化修剪了DiT模型。
对更多实验结果和文章细节感兴趣的读者，可以阅读一下论文原文~
这里给大家推荐一门我们最新的课程
《
扩散模型入门教程：数学原理、方法与应用
》
：
