一次性剪枝扩散模型,提速同时保持画质
OBS-Diff: Accurate Pruning For Diffusion Models in One-Shot
- 重用经典剪枝法,适配扩散模型复杂结构与多种稀疏方式
- 按时间步加权构建海森矩阵,重点保留早期去噪关键参数
- 分组渐进剪枝降低计算开销,适合部署压缩场景
大规模文本生成图像的扩散模型虽强大,但计算成本高昂。现有一次性网络剪枝方法难以直接应用于这类模型,因其迭代去噪特性。本文提出OBS-Diff,一种新型一次性剪枝框架,实现大模型的精准、免训练压缩。具体而言:(i) 重激活经典最优脑外科(OBS),适配现代扩散模型复杂架构,支持无结构、N:M半结构化及结构化(MHA头和前馈神经元)等多种稀疏粒度;(ii) 从误差累积角度出发,提出新型时间步感知海森矩阵构造方法,采用对数递减权重方案,赋予早期时间步更高重要性,缓解误差累积风险;(iii) 提出计算高效的分组渐进剪枝策略,摊销昂贵校准过程。大量实验表明,OBS-Diff在扩散模型一次性剪枝中达到当前最优性能,实现显著推理加速,视觉质量损失极小。
原文摘要 · Abstract (English)
Large-scale text-to-image diffusion models, while powerful, suffer from prohibitive computational cost. Existing one-shot network pruning methods can hardly be directly applied to them due to the iterative denoising nature of diffusion models. To bridge the gap, this paper presents OBS-Diff, a novel one-shot pruning framework that enables accurate and training-free compression of large-scale text-to-image diffusion models. Specifically, (i) OBS-Diff revitalizes the classic Optimal Brain Surgeon (OBS), adapting it to the complex architectures of modern diffusion models and supporting diverse pruning granularity, including unstructured, N:M semi-structured, and structured (MHA heads and FFN neurons) sparsity; (ii) To align the pruning criteria with the iterative dynamics of the diffusion process, by examining the problem from an error-accumulation perspective, we propose a novel timestep-aware Hessian construction that incorporates a logarithmic-decrease weighting scheme, assigning greater importance to earlier timesteps to mitigate potential error accumulation; (iii) Furthermore, a computationally efficient group-wise sequential pruning strategy is proposed to amortize the expensive calibration process. Extensive experiments show that OBS-Diff achieves state-of-the-art one-shot pruning for diffusion models, delivering inference acceleration with minimal degradation in visual quality.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。