VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion
  Transformers
Authors: name: Juncan Deng; name: Shuaiting Li; name: Zeyu Wang; name: Hong Gu; name: Kedong Xu; name: Kejie Huang
Published: 2024-08-30 09:15:54

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.