VQ4DiT: Efficiënte Post-Training Vector Kwantisatie voor Diffusion Transformers

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

Samenvatting

De Diffusion Transformers Models (DiTs) hebben de netwerkarchitectuur getransformeerd van traditionele UNets naar transformers, waarbij ze uitzonderlijke mogelijkheden hebben getoond in beeldgeneratie. Hoewel DiTs veelvuldig zijn toegepast op taken voor het genereren van hoogwaardige video's, belemmert hun grote parameteromvang de inferentie op edge-apparaten. Vectorquantisatie (VQ) kan modelgewichten ontbinden in een codeboek en toewijzingen, waardoor extreme gewichtsquantisatie mogelijk wordt en het geheugengebruik aanzienlijk wordt verminderd. In dit artikel stellen we VQ4DiT voor, een snelle post-training vectorquantisatiemethode voor DiTs. We ontdekten dat traditionele VQ-methoden alleen het codeboek kalibreren zonder de toewijzingen te kalibreren. Dit leidt ertoe dat gewichtssubvectoren onjuist worden toegewezen aan dezelfde toewijzing, wat inconsistente gradienten aan het codeboek oplevert en resulteert in een suboptimaal resultaat. Om deze uitdaging aan te pakken, berekent VQ4DiT de kandidaat-toewijzingsset voor elke gewichtssubvector op basis van de Euclidische afstand en reconstrueert de subvector op basis van het gewogen gemiddelde. Vervolgens wordt, met behulp van de zero-data en block-wise kalibratiemethode, de optimale toewijzing uit de set efficiënt geselecteerd terwijl het codeboek wordt gekalibreerd. VQ4DiT quantiseert een DiT XL/2-model op een enkele NVIDIA A100 GPU binnen 20 minuten tot 5 uur, afhankelijk van de verschillende quantisatie-instellingen. Experimenten tonen aan dat VQ4DiT een nieuwe state-of-the-art bereikt in de afweging tussen modelgrootte en prestaties, waarbij gewichten worden gequantiseerd tot 2-bit precisie terwijl een acceptabele beeldgeneratiekwaliteit behouden blijft.

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.