VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Deng, Juncan, Li, Shuaiting, Wang, Zeyu, Gu, Hong, Xu, Kedong, Huang, Kejie
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866929479370145792
author Deng, Juncan
Li, Shuaiting
Wang, Zeyu
Gu, Hong
Xu, Kedong
Huang, Kejie
author_facet Deng, Juncan
Li, Shuaiting
Wang, Zeyu
Gu, Hong
Xu, Kedong
Huang, Kejie
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2408_17131
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers
Deng, Juncan
Li, Shuaiting
Wang, Zeyu
Gu, Hong
Xu, Kedong
Huang, Kejie
Computer Vision and Pattern Recognition
Artificial Intelligence
I.2; I.4
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.
title VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers
topic Computer Vision and Pattern Recognition
Artificial Intelligence
I.2; I.4
url https://arxiv.org/abs/2408.17131