arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2501.09464cs.LG
收藏

基于梯度流的稀疏扩散模型剪枝

Pruning for Sparse Diffusion Models based on Gradient Flow

Shanghai Jiao Tong University(上海交通大学);Fudan University(复旦大学)

Ben Wan, Tianyi Zheng, Zhaoyu Chen, Yuxiao Wang, Jia Wang

AI总结:

针对扩散模型推理慢、计算成本高及一次性结构剪枝损害生成质量的问题,提出基于梯度流的迭代软剪枝方法与准则,实验表明其在效率和模型一致性上性能更优。

摘要:

扩散模型(Diffusion Models,DMs)在生成模型中具有出色的能力,但受限于较慢的推理速度和较高的计算成本。以往工作采用一次性结构剪枝,从预训练模型中获得轻量化 DMs,但这种方法通常会导致生成质量显著下降,并可能移除关键权重。因此,我们提出一种基于梯度流的迭代剪枝方法,包括梯度流剪枝过程和梯度流剪枝准则。我们采用渐进式软剪枝策略,以保持掩码矩阵的连续性,并依据稀疏空间中的剪枝准则,引导掩码矩阵沿能量函数的梯度流变化,从而避免一次性剪枝通常造成的突发性信息损失。基于梯度流的准则会剪除那些移除后会使损失函数梯度范数增大的参数,并能使被剪枝模型在迭代剪枝阶段快速收敛。我们在广泛使用的数据集上进行了大量实验,结果表明,该方法在效率以及与预训练模型的一致性方面均取得了优越性能。

英文摘要:

Diffusion Models (DMs) have impressive capabilities among generation models, but are limited to slower inference speeds and higher computational costs. Previous works utilize one-shot structure pruning to derive lightweight DMs from pre-trained ones, but this approach often leads to a significant drop in generation quality and may result in the removal of crucial weights. Thus we propose a iterative pruning method based on gradient flow, including the gradient flow pruning process and the gradient flow pruning criterion. We employ a progressive soft pruning strategy to maintain the continuity of the mask matrix and guide it along the gradient flow of the energy function based on the pruning criterion in sparse space, thereby avoiding the sudden information loss typically caused by one-shot pruning. Gradient-flow based criterion prune parameters whose removal increases the gradient norm of loss function and can enable fast convergence for a pruned model in iterative pruning stage. Our extensive experiments on widely used datasets demonstrate that our method achieves superior performance in efficiency and consistency with pre-trained models.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑