TaQ-DiT: Time-aware Quantization for Diffusion Transformers

TL;DR

TaQ-DiT presents a DiT-specific post-training quantization framework addressing two core challenges: non-convergence from separate reconstruction of weights and activations, and time-varying quantization sensitivity of Post-GELU activations. It introduces joint reconstruction to align weight and activation quantizations and time-variance-aware static transformations, including momentum-based shifting and reconstruction-driven migration, to stabilize Post-GELU quantization across timesteps and channels. Empirical results on ImageNet with $4$-bit weights and $8$-bit activations (W$4$A$8$) show TaQ-DiT outperforming state-of-the-art DiT quantization baselines, including large improvements in FID and IS, and better visual quality. The approach delivers substantial practical impact by enabling efficient, high-quality diffusion-based image generation with DiTs suitable for real-world deployment.

Abstract

Transformer-based diffusion models, dubbed Diffusion Transformers (DiTs), have achieved state-of-the-art performance in image and video generation tasks. However, their large model size and slow inference speed limit their practical applications, calling for model compression methods such as quantization. Unfortunately, existing DiT quantization methods overlook (1) the impact of reconstruction and (2) the varying quantization sensitivities across different layers, which hinder their achievable performance. To tackle these issues, we propose innovative time-aware quantization for DiTs (TaQ-DiT). Specifically, (1) we observe a non-convergence issue when reconstructing weights and activations separately during quantization and introduce a joint reconstruction method to resolve this problem. (2) We discover that Post-GELU activations are particularly sensitive to quantization due to their significant variability across different denoising steps as well as extreme asymmetries and variations within each step. To address this, we propose time-variance-aware transformations to facilitate more effective quantization. Experimental results show that when quantizing DiTs' weights to 4-bit and activations to 8-bit (W4A8), our method significantly surpasses previous quantization methods.

Figures (4)

• Figure 1: Illustrating (a) tensor-wise, (b) token-wise, and (c) input-channel-wise quantization for activations, and (d) channel-wise quantization for weights. To ensure hardware-efficient computation, we adopt channel-wise quantization for all activations and tensor-wise quantization for all weights.
• Figure 2: (a)-(c) The trajectories of separate and joint reconstruction. (d) Activation ranges of Post-GELU activations across different timesteps.
• Figure 3: Distributions and channel-wise ranges of Post-GELU activations, where (a)(d) original, (b)(e) after shifting, and (c)(f) after splitting/migration.
• Figure 4: Comparisons of generated images between ours and baselines.