arXiv:2412.16700cs.CV2024-12AAAI被引 7

提出自适应量化方法,让扩散模型推理更快更省显存。

TCAQ-DM: Timestep-Channel Adaptive Quantization for Diffusion Models

  • 按时间步与通道动态调整量化策略,提升精度
  • 在W6A6下FID接近全精度模型,在W4A4可生成可用图像
  • 适合部署扩散模型的轻量级设备使用

扩散模型在图像和视频生成任务中表现卓越,但推理时需大量内存和计算开销,源于复杂的网络结构和多阶段迭代过程。近期后训练量化(PTQ)技术通过将浮点运算转为低比特运算,被证明是降低推理成本的有效途径。然而,现有方法难以应对不同通道与时间步间激活分布的巨大差异,以及量化与推理阶段输入不一致的问题,仍有较大改进空间。为此,本文提出针对扩散模型的时间-通道自适应量化方法(TCAQ-DM)。首先设计时间-通道联合重参数化(TCR)模块,平衡各时间步与通道的激活范围,促进后续重构;其次引入动态自适应量化(DAQ)模块,根据每层后Softmax的分布特性选择最优量化器,缓解量化误差;最后提出渐进式对齐重构(PAR)策略,减少因输入不匹配带来的偏差。在多个基准和不同扩散模型上的实验表明,该方法显著优于当前最优方案,在CIFAR-10上于W6A6设置下达到与全精度模型相当的FID指标,且在W4A4设置下仍能生成可用图像。

原文摘要 · Abstract (English)

Diffusion models have achieved remarkable success in the image and video generation tasks. Nevertheless, they often require a large amount of memory and time overhead during inference, due to the complex network architecture and considerable number of timesteps for iterative diffusion. Recently, the post-training quantization (PTQ) technique has proved a promising way to reduce the inference cost by quantizing the float-point operations to low-bit ones. However, most of them fail to tackle with the large variations in the distribution of activations across distinct channels and timesteps, as well as the inconsistent of input between quantization and inference on diffusion models, thus leaving much room for improvement. To address the above issues, we propose a novel method dubbed Timestep-Channel Adaptive Quantization for Diffusion Models (TCAQ-DM). Specifically, we develop a timestep-channel joint reparameterization (TCR) module to balance the activation range along both the timesteps and channels, facilitating the successive reconstruction procedure. Subsequently, we employ a dynamically adaptive quantization (DAQ) module that mitigate the quantization error by selecting an optimal quantizer for each post-Softmax layers according to their specific types of distributions. Moreover, we present a progressively aligned reconstruction (PAR) strategy to mitigate the bias caused by the input mismatch. Extensive experiments on various benchmarks and distinct diffusion models demonstrate that the proposed method substantially outperforms the state-of-the-art approaches in most cases, especially yielding comparable FID metrics to the full precision model on CIFAR-10 in the W6A6 setting, while enabling generating available images in the W4A4 settings.

扩散模型量化推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。