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

EDA-DM:扩散模型训练后量化的增强分布对齐

EDA-DM: Enhanced Distribution Alignment for Post-Training Quantization of Diffusion Models

Institute of Automation, Chinese Academy of Sciences(中国科学院自动化研究所);School of Artificial Intelligence, University of Chinese Academy of Sciences(中国科学院大学人工智能学院)

Xuewen Liu, Zhikai Li, Junrui Xiao, Mengjuan Chen, Jianquan Li, Qingyi Gu

AI总结:

针对扩散模型训练后量化中激活动态导致的分布不匹配,提出EDA-DM,通过密度多样性引导校准样本选择及Hessian损失优化块重建,实现高效对齐,显著提升性能并加速压缩。

摘要:

扩散模型在图像生成任务中取得了巨大成功。然而,冗长的去噪过程和复杂的神经网络阻碍了其在现实场景中的低延迟应用。量化能有效降低模型复杂度,而无需微调的训练后量化(PTQ)在压缩和加速扩散模型方面极具前景。不幸的是,我们发现由于高度动态的激活值,现有PTQ方法在校准样本层面和重建输出层面均存在分布不匹配问题,导致性能远非理想。本文提出EDA-DM,一种标准化PTQ方法,有效解决上述问题。具体而言,在校准样本层面,我们从潜在空间特征图的密度和多样性中提取信息,指导校准样本的选择以对齐整体样本分布;在重建输出层面,我们从理论上分析了先前重建失败的原因,并基于此洞察,利用层的Hessian损失优化块重建,在不同网络粒度上对齐量化模型和全精度模型的输出。大量实验表明,EDA-DM在多种模型和数据集上显著优于现有PTQ方法。我们的方法在MS-COCO上对流行的Stable-Diffusion实现了1.83倍加速和4倍压缩,CLIP分数仅损失0.05。代码可在http://github.com/BienLuky/EDA-DM获取。

英文摘要:

Diffusion models have achieved great success in image generation tasks. However, the lengthy denoising process and complex neural networks hinder their low-latency applications in real-world scenarios. Quantization can effectively reduce model complexity, and post-training quantization (PTQ), which does not require fine-tuning, is highly promising for compressing and accelerating diffusion models. Unfortunately, we find that due to the highly dynamic activations, existing PTQ methods suffer from distribution mismatch issues at both calibration sample level and reconstruction output level, which makes the performance far from satisfactory. In this paper, we propose EDA-DM, a standardized PTQ method that efficiently addresses the above issues. Specifically, at the calibration sample level, we extract information from the density and diversity of latent space feature maps, which guides the selection of calibration samples to align with the overall sample distribution; and at the reconstruction output level, we theoretically analyze the reasons for previous reconstruction failures and, based on this insight, optimize block reconstruction using the Hessian loss of layers, aligning the outputs of quantized model and full-precision model at different network granularity. Extensive experiments demonstrate that EDA-DM significantly outperforms the existing PTQ methods across various models and datasets. Our method achieves a 1.83 times speedup and 4 times compression for the popular Stable-Diffusion on MS-COCO, with only a 0.05 loss in CLIP score. Code is available at http://github.com/BienLuky/EDA-DM .

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

↑