Effiziente Post-Training Vektorquantisierung für Diffusion Transformer Modelle

Einführung

Die Diffusion Transformer Modelle (DiTs) haben sich als leistungsfähige Werkzeuge für die Bildgenerierung etabliert und dabei die Netzwerkarchitektur von traditionellen UNets zu Transformern verändert. Trotz ihrer herausragenden Fähigkeiten in der Bild- und Videogenerierung stellen die großen Parametergrößen von DiTs eine Herausforderung bei der Inferenz auf Edge-Geräten dar. Die Vektorquantisierung (VQ) ermöglicht die Zerlegung von Modellgewichten in ein Codebuch und Zuordnungen, was eine extreme Gewichtsquantisierung und eine signifikante Reduzierung des Speicherbedarfs ermöglicht. In diesem Artikel stellen wir VQ4DiT vor, eine schnelle Post-Training Vektorquantisierungsmethode für DiTs. Traditionelle VQ-Methoden kalibrieren nur das Codebuch, ohne die Zuordnungen anzupassen, was zu suboptimalen Ergebnissen führt. VQ4DiT berechnet das Kandidatenzuordnungset für jeden Gewichtssubvektor basierend auf der euklidischen Distanz und rekonstruiert den Subvektor basierend auf dem gewichteten Durchschnitt. Anschließend wird die optimale Zuordnung effizient ausgewählt und das Codebuch kalibriert.

Hintergrund und verwandte Arbeiten

Diffusion Transformer Modelle

Diffusion Transformer Modelle (DiTs) haben sich aufgrund ihrer überlegenen Leistung in der Bildgenerierung viel Aufmerksamkeit erarbeitet. DiTs bestehen aus mehreren in Reihe geschalteten Transformer-Blöcken und nutzen die Skalierungseigenschaften der Transformer, was eine flexible Parametererweiterung zur Leistungssteigerung ermöglicht. Im Vergleich zu anderen UNet-basierten Diffusionsmodellen haben DiTs gezeigt, dass sie qualitativ hochwertigere Bilder generieren können.

Herausforderungen der Vektorquantisierung für DiTs

Die bestehenden Quantisierungsmethoden sind nicht direkt auf DiTs anwendbar, da diese verschiedene Netzwerkstrukturen und algorithmische Konzepte im Vergleich zu UNet-basierten Diffusionsmodellen aufweisen. Zudem führen Post-Training Quantisierungsmethoden (PTQ) zu erheblichen Genauigkeitsverlusten, wenn Gewichte auf extrem niedrige Bitbreiten quantisiert werden. Traditionelle VQ-Methoden kalibrieren nur das Codebuch und passen die Zuordnungen nicht an, was zu inkonsistenten Gradienten führt und suboptimale Ergebnisse liefert.

VQ4DiT: Eine neue Methode zur Vektorquantisierung

VQ4DiT führt eine neuartige Post-Training Vektorquantisierungstechnik für die Quantisierung von DiTs mit extrem niedriger Bitbreite ein. Diese Methode umfasst folgende Schritte: - Mapping der Gewichtssubvektoren jeder Schicht auf ein Codebuch mittels K-Means-Algorithmus. - Bestimmung eines Kandidatenzuordnungsets für jeden Gewichtssubvektor basierend auf euklidischer Distanz und Rekonstruktion des Subvektors basierend auf dem gewichteten Durchschnitt. - Kalibrierung des Codebuchs und der Zuordnungen jeder Schicht gleichzeitig mittels Zero-Data und blockweiser Kalibrierungsmethode. VQ4DiT gewährleistet, dass das quantisierte Modell Ergebnisse erzielt, die mit denen des Modells in Gleitkommapräsision vergleichbar sind.

Experimente und Ergebnisse

VQ4DiT wurde auf einem DiT XL/2 Modell mit einer einzigen NVIDIA A100 GPU quantifiziert, wobei die Quantisierungszeiten zwischen 20 Minuten und 5 Stunden variieren, abhängig von den unterschiedlichen Quantisierungseinstellungen. Die Experimente zeigen, dass VQ4DiT neue Maßstäbe im Hinblick auf das Verhältnis von Modellgröße und Leistung setzt und Gewichte auf 2-Bit-Präzision quantisiert, während eine akzeptable Bildgenerierungsqualität beibehalten wird.

Parameter und Kalibrierung

In jeder DiT-Schicht werden die Gewichtssubvektoren mittels K-Means-Algorithmus auf ein Codebuch gemappt. Dabei wird ein Kandidatenzuordnungset für jeden Gewichtssubvektor bestimmt und der Subvektor basierend auf dem gewichteten Durchschnitt rekonstruiert. Die endgültige Zuordnung wird effizient ausgewählt und das Codebuch kalibriert, um die besten Ergebnisse zu erzielen.

Fazit

VQ4DiT stellt eine effiziente Methode zur Post-Training Vektorquantisierung für Diffusion Transformer Modelle dar. Indem sowohl das Codebuch als auch die Zuordnungen kalibriert werden, erreicht VQ4DiT Ergebnisse, die mit denen von Modellen in Gleitkommapräzision vergleichbar sind. Diese Methode ermöglicht die Quantisierung von Modellen auf extrem niedrige Bitbreiten und eröffnet somit neue Möglichkeiten für den Einsatz von DiTs auf Edge-Geräten mit begrenzten Ressourcen.

Bibliographie

https://arxiv.org/html/2408.17131v1 https://arxiv.org/pdf/2408.17131 https://paperreading.club/page?id=249044 https://paperswithcode.com/paper/ptq4dit-post-training-quantization-for https://www.researchgate.net/publication/381703958_Q-DiT_Accurate_Post-Training_Quantization_for_Diffusion_Transformers https://openaccess.thecvf.com/content/CVPR2024/papers/Wang_Towards_Accurate_Post-training_Quantization_for_Diffusion_Models_CVPR_2024_paper.pdf https://researchmgt.monash.edu/ws/portalfiles/portal/605910764/592678525_oa.pdf https://proceedings.neurips.cc/paper/2021/file/ec8956637a99787bd197eacd77acce5e-Paper.pdf https://openreview.net/forum?id=Y3g1PV5R9l¬eId=Wg1KTrcQfC