Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Post-Training Quantization for Diffusion Transformer via Hierarchical Timestep Grouping

Ning Ding, Jing Han, Yuchuan Tian, Chao Xu, Kai Han, Yehui Tang

TL;DR

A timestep-aware shift-and-scale strategy to smooth the activation distribution to reduce the quantization error, and a re-parameterization scheme which absorbs the quantization parameters into nearby module to avoid redundant computations are proposed.

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.

Post-Training Quantization for Diffusion Transformer via Hierarchical Timestep Grouping

TL;DR

A timestep-aware shift-and-scale strategy to smooth the activation distribution to reduce the quantization error, and a re-parameterization scheme which absorbs the quantization parameters into nearby module to avoid redundant computations are proposed.

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.

Paper Structure

This paper contains 21 sections, 17 equations, 4 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Reviewing Quantization Difficulty in DiT
  4. Channel-dependent activation outliers
  5. Timestep-dependent outlier distributions
  6. Hierarchical Timestep Grouping
  7. Temporally-grouped Channel-wise Shifting
  8. Temporally-aggregated Channel-wise Scaling
  9. Re-parameterization for Faster Inference
  10. Experiment
  11. Experimental Setting
  12. Models and Diffusion Setup.
  13. Calibration Dataset
  14. Quantization Configurations.
  15. Baseline Methods.
  16. ...and 6 more sections

Figures (4)

  • Figure 1: (a) Per-channel input activation of blocks[27].mlp.fc1 at step 99, step 63 and step 0. Channels with indexes between 900-970 are displayed. (b) Activations of channel No.910 across different timesteps of 100-step DDPM scheduler.
  • Figure 2: Original activation (left) and shifted activation (right). Activations are taken from blocks[27].mlp.fc1 at step 90.
  • Figure 3: Scaling the shifted activation leads to further less quantization error for non-outliers. Activations are taken from blocks[27].mlp.fc1 at step 90.
  • Figure 4: Schematic diagram of the HTG framework for DiT.