基于梯度流的迭代剪枝让扩散模型更轻更快,且生成质量损失小。
Pruning for Sparse Diffusion Models based on Gradient Flow
- 通过梯度流引导渐进式剪枝,保持掩码连续性。
- 剪掉使损失梯度范数增大的参数,提升剪枝后收敛速度。
- 适合追求高效推理又不想牺牲生成质量的研究者。
扩散模型(DMs)在生成任务中表现优异,但存在推理速度慢、计算开销高的问题。已有工作采用一次性结构剪枝从预训练模型中提取轻量版本,但常导致生成质量显著下降,可能误删关键权重。为此,本文提出一种基于梯度流的迭代剪枝方法,包含梯度流剪枝过程与剪枝判据。采用渐进式软剪枝策略,使掩码矩阵沿能量函数的梯度流方向演变,避免一次性剪枝带来的信息突失。基于梯度流的判据会剪除那些移除后会使损失函数梯度范数增加的参数,有助于剪枝后模型在迭代阶段快速收敛。大量实验在常用数据集上表明,该方法在效率与与预训练模型的一致性方面均优于现有方法。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。