VQ4DiT: Effiziente Post-Training Vektorquantisierung für Diffusionstransformatoren

VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers

August 30, 2024

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

cs.AI

Zusammenfassung

Die Diffusionstransformator-Modelle (DiTs) haben die Netzwerkarchitektur von traditionellen UNets auf Transformer umgestellt und dabei außergewöhnliche Fähigkeiten bei der Bildgenerierung gezeigt. Obwohl DiTs weit verbreitet bei der Generierung von hochauflösenden Videos eingesetzt wurden, behindert ihre große Parametergröße die Inferenz auf Edge-Geräten. Die Vektorquantisierung (VQ) kann das Modellgewicht in einen Codebuch und Zuweisungen zerlegen, was eine extrem starke Gewichtsquantisierung ermöglicht und den Speicherverbrauch signifikant reduziert. In diesem Paper schlagen wir VQ4DiT vor, eine schnelle Post-Training-Vektorquantisierungsmethode für DiTs. Wir haben festgestellt, dass traditionelle VQ-Methoden nur das Codebuch kalibrieren, ohne die Zuweisungen zu kalibrieren. Dies führt dazu, dass Gewichts-Teilvektoren fälschlicherweise derselben Zuweisung zugeordnet werden, was inkonsistente Gradienten zum Codebuch liefert und zu einem suboptimalen Ergebnis führt. Um diese Herausforderung anzugehen, berechnet VQ4DiT für jeden Gewichts-Teilvektor auf der Grundlage des euklidischen Abstands den Kandidaten-Zuweisungssatz und rekonstruiert den Teilvektor auf der Grundlage des gewichteten Durchschnitts. Anschließend wird unter Verwendung der Zero-Data- und blockweisen Kalibrierungsmethode die optimale Zuweisung aus dem Satz effizient ausgewählt, während das Codebuch kalibriert wird. VQ4DiT quantisiert ein DiT XL/2-Modell auf einer einzelnen NVIDIA A100 GPU innerhalb von 20 Minuten bis 5 Stunden, abhängig von den verschiedenen Quantisierungseinstellungen. Experimente zeigen, dass VQ4DiT einen neuen Stand der Technik bei der Modellgröße und der Leistungsbilanz etabliert, indem Gewichte auf eine Präzision von 2 Bit quantisiert werden, während eine akzeptable Bildgenerierungsqualität beibehalten wird.

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.