VQ4DiT : Quantification vectorielle efficace après l'entraînement pour les transformateurs de diffusion

VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers

August 30, 2024

Auteurs: Juncan Deng, Shuaiting Li, Zeyu Wang, Hong Gu, Kedong Xu, Kejie Huang

cs.AI

Résumé

Les modèles de transformateurs de diffusion (DiTs) ont fait passer l'architecture réseau des UNets traditionnels aux transformateurs, démontrant des capacités exceptionnelles en génération d'images. Bien que les DiTs aient été largement appliqués aux tâches de génération de vidéos haute définition, leur grande taille de paramètres entrave l'inférence sur les appareils périphériques. La quantification vectorielle (VQ) peut décomposer le poids du modèle en un codebook et des affectations, permettant une quantification de poids extrême et réduisant significativement l'utilisation de la mémoire. Dans cet article, nous proposons VQ4DiT, une méthode de quantification vectorielle post-entraînement rapide pour les DiTs. Nous avons constaté que les méthodes VQ traditionnelles ne calibrent que le codebook sans calibrer les affectations. Cela conduit à ce que des sous-vecteurs de poids soient incorrectement attribués à la même affectation, fournissant des gradients incohérents au codebook et donnant un résultat sous-optimal. Pour relever ce défi, VQ4DiT calcule l'ensemble d'affectations candidat pour chaque sous-vecteur de poids en fonction de la distance euclidienne et reconstruit le sous-vecteur en fonction de la moyenne pondérée. Ensuite, en utilisant la méthode de calibration zéro-donnée et par bloc, l'affectation optimale de l'ensemble est sélectionnée efficacement tout en calibrant le codebook. VQ4DiT quantifie un modèle DiT XL/2 sur un seul GPU NVIDIA A100 en 20 minutes à 5 heures selon les différents réglages de quantification. Les expériences montrent que VQ4DiT établit un nouvel état de l'art en termes de compromis entre la taille du modèle et les performances, quantifiant les poids avec une précision de 2 bits tout en conservant une qualité de génération d'images acceptable.

English

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.