VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers

TL;DR

VQ4DiT introduces a fast post-training vector quantization method for diffusion transformer models, significantly reducing model size while maintaining image generation quality, enabling deployment on edge devices.

Contribution

The paper proposes VQ4DiT, a novel vector quantization approach that calibrates both codebook and assignments, improving quantization efficiency and performance for diffusion transformers.

Findings

Quantizes DiT XL/2 model to 2-bit precision.

Achieves state-of-the-art size-performance trade-offs.

Completes quantization within 20 minutes to 5 hours.

Abstract

The Diffusion Transformers Models (DiTs) have transitioned the network architecture from traditional UNets to transformers, demonstrating exceptional capabilities in image generation. Although DiTs have been widely applied to high-definition video generation tasks, their large parameter size hinders inference on edge devices. Vector quantization (VQ) can decompose model weight into a codebook and assignments, allowing extreme weight quantization and significantly reducing memory usage. In this paper, we propose VQ4DiT, a fast post-training vector quantization method for DiTs. We found that traditional VQ methods calibrate only the codebook without calibrating the assignments. This leads to weight sub-vectors being incorrectly assigned to the same assignment, providing inconsistent gradients to the codebook and resulting in a suboptimal result. To address this challenge, VQ4DiT calculates the candidate assignment set for each weight sub-vector based on Euclidean distance and reconstructs the sub-vector based on the weighted average. Then, using the zero-data and block-wise calibration method, the optimal assignment from the set is efficiently selected while calibrating the codebook. VQ4DiT quantizes a DiT XL/2 model on a single NVIDIA A100 GPU within 20 minutes to 5 hours depending on the different quantization settings. Experiments show that VQ4DiT establishes a new state-of-the-art in model size and performance trade-offs, quantizing weights to 2-bit precision while retaining acceptable image generation quality.