TaQ-DiT: Time-aware Quantization for Diffusion Transformers

TL;DR

This paper introduces TaQ-DiT, a novel time-aware quantization method for diffusion transformers that improves model compression and inference speed by addressing layer sensitivity and reconstruction issues.

Contribution

It proposes joint weight-activation reconstruction and time-variance-aware transformations to enhance quantization performance of diffusion transformers.

Findings

Outperforms previous quantization methods at 4-bit weights and 8-bit activations.

Addresses non-convergence in separate weight and activation reconstruction.

Effectively handles sensitivity of Post-GELU activations during quantization.

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.