arXiv:2503.06930cs.CV2025-03被引 3

针对扩散Transformer的量化难题,提出分层时间步分组方法提升推理效率。

Post-Training Quantization for Diffusion Transformer via Hierarchical Timestep Grouping

  • 按时间步相似性分组激活值,降低量化误差。
  • 实现W8A8量化,生成质量达当前最优FiD分数。
  • 支持更激进的W4A8量化,不损失图像生成质量。

扩散Transformer(DiT)因其强大的生成能力成为图像生成模型的首选。与以往基于卷积的UNet模型不同,DiT完全由堆叠的Transformer模块构成,具备类似大语言模型的可扩展性。然而,模型规模增长和多步采样机制给部署与推理带来巨大压力。本文提出一种专为扩散Transformer设计的后训练量化框架。首先发现量化困难主要源于时序依赖的通道特异性异常值。提出一种时间步感知的移位缩放策略,平滑激活分布以减少量化误差。其次,基于相邻时间步激活分布相似的观察,采用分层聚类将去噪时间步划分为多个组,并设计重参数化方案,将量化参数融入邻近模块,避免冗余计算。大量实验表明,本方法成功将DiT模型量化至8位权重和8位激活(W8A8),达到当前最优的FiD得分;还可进一步实现4位权重、8位激活(W4A8)量化,且不牺牲生成质量。

原文摘要 · Abstract (English)

Diffusion Transformer (DiT) has now become the preferred choice for building image generation models due to its great generation capability. Unlike previous convolution-based UNet models, DiT is purely composed of a stack of transformer blocks, which renders DiT excellent in scalability like large language models. However, the growing model size and multi-step sampling paradigm bring about considerable pressure on deployment and inference. In this work, we propose a post-training quantization framework tailored for Diffusion Transforms to tackle these challenges. We firstly locate that the quantization difficulty of DiT mainly originates from the time-dependent channel-specific outliers. We propose a timestep-aware shift-and-scale strategy to smooth the activation distribution to reduce the quantization error. Secondly, based on the observation that activations of adjacent timesteps have similar distributions, we utilize a hierarchical clustering scheme to divide the denoising timesteps into multiple groups. We further design a re-parameterization scheme which absorbs the quantization parameters into nearby module to avoid redundant computations. Comprehensive experiments demonstrate that out PTQ method successfully quantize the Diffusion Transformer into 8-bit weight and 8-bit activation (W8A8) with state-of-the-art FiD score. And our method can further quantize DiT model into 4-bit weight and 8-bit activation (W4A8) without sacrificing generation quality.

扩散模型量化Transformer图像生成

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