Post-Training Quantization for Diffusion Transformer via Hierarchical Timestep Grouping
Journal:
arXiv
Published Date:
Mar 10, 2025
Abstract
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.
