[论文解读] DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization

Haowei Zhu, Dehua Tang+13 more|arXiv (Cornell University)|Oct 22, 2024

Optical Systems and Laser Technology被引用 4

一句话总结

DiP-GO 提出了一种新颖的、无需微调的扩散模型剪枝方法，通过将模型剪枝建模为可微分子网络（SubNet）搜索问题。该方法构建了一个基于去噪步骤间特征相似性的备份连接超网络（SuperNet），采用优化损失函数的插件式剪枝器，并通过少量梯度优化步骤识别最优子网络，从而在稳定扩散模型 1.5 上实现 4.4 倍的推理加速，且不损失生成精度。

ABSTRACT

Diffusion models have achieved remarkable progress in the field of image generation due to their outstanding capabilities. However, these models require substantial computing resources because of the multi-step denoising process during inference. While traditional pruning methods have been employed to optimize these models, the retraining process necessitates large-scale training datasets and extensive computational costs to maintain generalization ability, making it neither convenient nor efficient. Recent studies attempt to utilize the similarity of features across adjacent denoising stages to reduce computational costs through simple and static strategies. However, these strategies cannot fully harness the potential of the similar feature patterns across adjacent timesteps. In this work, we propose a novel pruning method that derives an efficient diffusion model via a more intelligent and differentiable pruner. At the core of our approach is casting the model pruning process into a SubNet search process. Specifically, we first introduce a SuperNet based on standard diffusion via adding some backup connections built upon the similar features. We then construct a plugin pruner network and design optimization losses to identify redundant computation. Finally, our method can identify an optimal SubNet through few-step gradient optimization and a simple post-processing procedure. We conduct extensive experiments on various diffusion models including Stable Diffusion series and DiTs. Our DiP-GO approach achieves 4.4 x speedup for SD-1.5 without any loss of accuracy, significantly outperforming the previous state-of-the-art methods.

研究动机与目标

• 为解决扩散模型推理的高计算成本，尤其是多步去噪带来的开销。
• 克服静态、非自适应剪枝策略的局限性，这些策略无法有效利用时间步之间的动态特征相似性。
• 在不微调预训练扩散模型的前提下，实现高效且高精度的模型剪枝。
• 开发一种可微分的、基于梯度优化的剪枝框架，从超网络中识别最优子网络（SubNets）。
• 在多种扩散架构上实现显著加速，同时保持图像生成质量。

提出的方法

• 通过在标准扩散模型中引入基于相邻去噪步骤间特征相似性的备份连接，构建超网络（SuperNet）。
• 将剪枝过程建模为子网络（SubNet）搜索问题，其中最优子网络由超网络推导得出。
• 设计一个插件式剪枝器网络，通过任务特定的优化损失函数学习识别冗余计算路径。
• 应用少量梯度优化步骤，高效搜索最优子网络，相比完整微调显著降低计算成本。
• 集成轻量级后处理步骤，确保最终子网络满足用户指定的剪枝比例和结构约束。
• 利用梯度检查点技术，缓解在扩展去噪时间步上的反向传播过程中的内存问题。

Figure 1: Overview of the SuperNet and SubNet. Standard diffusion models execute the full inference path step by step. In our framework, we propose a SuperNet based on the original flow and integrate backup connections to facilitate block removal. This allows the partial inference SubNet to efficien

实验结果

研究问题

• RQ1 我们能否将扩散模型剪枝建模为可微分子网络搜索问题，从而避免微调？
• RQ2 如何有效利用相邻去噪步骤间的特征相似性，以减少冗余计算？
• RQ3 少量梯度优化策略是否能比现有剪枝方法更高效地识别高性能子网络？
• RQ4 在不损害生成质量的前提下，我们能在多大程度上加速扩散模型推理？
• RQ5 所提方法是否能在如稳定扩散和 DiT 等多种扩散架构上实现泛化？

主要发现

• DiP-GO 在稳定扩散模型 1.5 上实现了 4.4 倍的推理速度提升，且未造成图像质量下降，也无需微调。
• 该方法通过有效利用时间步间的动态特征模式，优于以往最先进剪枝技术。
• 该方法在多个扩散模型上具有泛化能力，包括稳定扩散 1.5、2.1、XL 和 DiT，展现出广泛的适用性。
• 配备优化损失函数的插件式剪枝器成功识别出冗余组件，实现了高效的子网络发现。
• 少量梯度优化的使用使模型能快速收敛至高性能子网络，显著降低搜索成本。
• 后处理步骤确保最终子网络在满足剪枝比例约束的同时，仍保持模型能力。

Figure 2: Overview of our diffusion pruner. a) DiP-GO employs a pruner network to learn the importance scores of blocks in the diffusion sampling process. It takes $N\times T$ queries as input and passes them through stacked self-attention (SA) and fully connected (FC) layers to capture the structur