VQ4DiT: Quantizzazione Vettoriale Post-Training Efficient per i Transformer a Diffusione

VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers

August 30, 2024

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

cs.AI

Abstract

I modelli Diffusion Transformers (DiTs) hanno trasformato l'architettura di rete dai tradizionali UNet ai transformer, dimostrando capacità eccezionali nella generazione di immagini. Sebbene i DiTs siano stati ampiamente applicati a compiti di generazione di video ad alta definizione, le loro dimensioni di parametro elevate ostacolano l'inferenza sui dispositivi periferici. La quantizzazione vettoriale (VQ) può decomporre il peso del modello in un codebook e assegnazioni, consentendo una quantizzazione estrema del peso e riducendo significativamente l'utilizzo della memoria. In questo articolo, proponiamo VQ4DiT, un metodo di quantizzazione vettoriale post-training veloce per i DiTs. Abbiamo riscontrato che i metodi tradizionali di VQ calibrano solo il codebook senza calibrare le assegnazioni. Ciò porta a sottovettori di peso assegnati erroneamente alla stessa assegnazione, fornendo gradienti non coerenti al codebook e risultando in un risultato non ottimale. Per affrontare questa sfida, VQ4DiT calcola l'insieme di assegnazioni candidato per ciascun sottovettore di peso basandosi sulla distanza euclidea e ricostruisce il sottovettore basandosi sulla media ponderata. Successivamente, utilizzando il metodo di calibrazione a zero-dati e a livello di blocco, viene selezionata efficientemente l'assegnazione ottimale dall'insieme mentre si calibra il codebook. VQ4DiT quantizza un modello DiT XL/2 su una singola GPU NVIDIA A100 in 20 minuti fino a 5 ore a seconda delle diverse impostazioni di quantizzazione. Gli esperimenti mostrano che VQ4DiT stabilisce un nuovo stato dell'arte nei compromessi tra dimensioni del modello e prestazioni, quantizzando i pesi con una precisione di 2 bit mantenendo una qualità accettabile nella generazione di immagini.

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.