Joint Post-Training Quantization of Vision Transformers with Learned Prompt-Guided Data Generation

Deze paper introduceert een data-vrij raamwerk voor gezamenlijke post-training kwantisatie van Vision Transformers, dat gebruikmaakt van door Stable Diffusion gegenereerde afbeeldingen om state-of-the-art prestaties te behalen op ImageNet zonder gelabelde data.

De kern

Stel je voor dat je een enorme, superintelligente robot hebt die foto's kan herkennen. Deze robot is zo slim dat hij bijna alles kan zien, maar hij is ook zo zwaar en traag dat hij niet in je telefoon past. Hij heeft een enorme krachtbron nodig en neemt te veel ruimte in beslag. Dit is wat er gebeurt met moderne AI-modellen, genaamd Vision Transformers (ViT). Ze zijn fantastisch, maar te groot voor de echte wereld.

De oplossing? Quantisatie. Dat klinkt als een ingewikkeld woord, maar het is eigenlijk als het verkleinen van de robot. Je maakt zijn "gedachten" minder precies (van 32-bits naar bijvoorbeeld 4-bits), zodat hij lichter wordt en sneller draait. Maar hier zit de valkuil: als je te veel precisie weghaalt, wordt de robot dom en kan hij geen foto's meer herkennen.

Deze paper introduceert een slimme manier om deze robot te verkleinen zonder dat hij zijn intelligentie verliest, zelfs zonder dat je duizenden echte foto's hebt om hem te oefenen. Hier is hoe ze het doen, vertaald in alledaagse taal:

1. Het Probleem: De "Losse Onderdelen" Benadering

Tot nu toe probeerden mensen de robot te verkleinen door hem stuk voor stuk te bekijken. Ze namen één onderdeel (bijvoorbeeld de ogen), pasten de precisie aan, en gingen dan naar het volgende.

• De analogie: Stel je voor dat je een orkest repeteert. De oude methode was: "Speel jij je viool stukje perfect, en jij je drumstokje perfect." Maar als je dat allemaal apart doet, klinkt het samen als een chaos. De delen hangen namelijk aan elkaar; wat de viool doet, beïnvloedt de drums.
• De oplossing van deze paper: Ze kijken naar het hele orkest tegelijk. Ze optimaliseren alle onderdelen samen in één keer. Hierdoor kunnen de delen elkaar "redden" als er een foutje optreedt. Het resultaat is dat de robot zelfs met heel weinig precisie (zoals 1,58 bits!) nog steeds bijna net zo slim blijft als het origineel.

2. Het Nieuwe Probleem: Geen Echte Foto's

Normaal gesproken moet je een robot oefenen met duizenden echte foto's van katten, auto's en bomen om te leren hoe hij moet verkleinen. Maar wat als je die foto's niet mag hebben (bijvoorbeeld vanwege privacy) of als je ze niet hebt?

• De oude methode: Je probeerde de robot te oefenen met simpele beschrijvingen, zoals "een foto van een kat".
• Het probleem: De AI die de foto's maakt, wordt dan saai. Hij maakt 100 keer dezelfde grijze kat op een witte achtergrond. Of hij maakt een kat die eruitziet als een hond. Dit helpt de robot niet om de echte wereld te begrijpen.

3. De Oplossing: De "Meester-Regisseur" (Learned Prompts)

Hier komt het creatieve deel van deze paper. In plaats van simpele zinnen te gebruiken, laten ze de computer leren hoe hij moet praten met de beeld-maker.

• De analogie: Stel je voor dat je een regisseur bent voor een film.
  • De oude methode was: "Draai een scène met een kat." De acteur (de AI) denkt dan: "Oké, ik speel een kat," en doet het altijd hetzelfde.
  • De nieuwe methode van deze paper is: Ze laten de regisseur leren om 20 verschillende manieren te bedenken om een kat te beschrijven.
    • Regisseur 1: "Een oranje kat die op een muur zit in de zon."
    • Regisseur 2: "Een zwarte kat die in de regen loopt."
    • Regisseur 3: "Een kat die speelt met een bal van wol."
• Ze gebruiken een slimme truc: ze laten de robot (die de foto's maakt) kijken naar een andere, super-slimme robot (die de foto's herkent). Als de foto die gemaakt wordt, niet herkend wordt als een "kat", zeggen ze: "Nee, probeer het anders!"
• Ze zorgen er ook voor dat de regisseurs niet allemaal hetzelfde zeggen. Ze dwingen ze om divers te zijn. Zo krijgen ze een enorme verzameling van unieke, mooie en verschillende foto's van katten, auto's en bergen.

4. Het Resultaat: De "Zelfgemaakte" School

Door deze leerzame, diverse foto's te gebruiken, kan de robot worden getraind alsof hij duizenden echte foto's heeft gezien.

• De analogie: Het is alsof je een student voorbereidt op een examen. Normaal doe je dit met echte oefenexamens. Maar als je die niet hebt, maak je zelf oefenexamens. De meeste mensen maken saaie, saaie oefenexamens. Deze paper maakt echter duizenden unieke, creatieve oefenexamens die precies de juiste moeilijkheidsgraad hebben.
• Het resultaat is verbazingwekkend: De robot die getraind is met deze "zelfgemaakte" foto's, werkt bijna net zo goed als diegene die met echte foto's is getraind.

Samenvatting in één zin

De auteurs hebben een manier bedacht om enorme, zware AI-modellen te verkleinen tot een formaat dat op je telefoon past, door ze samen te optimaliseren in plaats van stuk voor stuk, en ze te trainen met kunstmatige, gevarieerde foto's die door een slimme AI-regisseur zijn bedacht, zodat je geen echte foto's meer nodig hebt.

Dit betekent dat we in de toekomst veel slimmere AI-apps kunnen hebben die sneller draaien, minder batterij verbruiken en geen enorme databases met foto's nodig hebben om te werken.

Technische samenvatting

Probleemstelling

Vision Transformers (ViT's) hebben zich bewezen als krachtige modellen voor visuele herkenningsopdrachten, maar ze zijn zeer rekenintensief en geheugenintensief, wat hun implementatie op randapparatuur (edge devices) bemoeilijkt. Quantisatie is een veelgebruikte oplossing om de precisie van gewichten en activaties te verlagen, waardoor het model kleiner en sneller wordt.

Echter, bestaande methoden voor Post-Training Quantization (PTQ) hebben ernstige beperkingen bij ViT's:

1. Inter-block afhankelijkheden: ViT's hebben sterke afhankelijkheden tussen blokken en lagen. Traditionele PTQ-methoden die per laag of per blok werken (zoals bij CNN's), falen omdat ze deze globale correlaties negeren.
2. Niet-uniforme activaties: De activatiedistributies in ViT's zijn vaak niet-Gaussisch en bevatten extreme uitbijters (outliers), wat standaard uniforme quantisatie onstabiel maakt.
3. Data-afhankelijkheid: De meeste geavanceerde PTQ-methoden vereisen gelabelde kalibratiegegevens (zoals ImageNet), wat een probleem is bij privacybeperkingen of data-schaarste.
4. Extreme low-bit quantisatie: Er is tot nu toe geen succesvolle PTQ-methode voor ViT's die werkt op extreem lage bitbreedtes (bijv. ternaire gewichten, W1.58), zonder volledige hertraining (QAT).

Methodologie

De auteurs stellen een end-to-end gezamenlijke PTQ-framework voor dat alle lagen en blokken van een Vision Transformer gelijktijdig optimaliseert zonder gebruik te maken van gelabelde data. Het framework bestaat uit twee hoofdbestanddelen:

1. End-to-End Gezamenlijke Optimalisatie

In plaats van blokken sequentieel te reconstrueren, wordt het hele netwerk in één keer geoptimaliseerd.

• Uniforme Quantisatie: Er wordt gebruikgemaakt van een standaard uniforme quantisator met leerbare stapgrootte ($\Delta$) en nulpunt ($z$).
• Per-kanaal herschaling (Channel-Wise Rescaling): Om de grote variatie in activatiemagnitudes tussen kanalen aan te pakken, wordt een herschaling- en verschuivingsmechanisme ($\alpha$ en $\beta$) geïntroduceerd. Dit gladde de activatiedynamiek over kanalen en verplaatst de quantisatie-moeilijkheid van de activaties naar de gewichten, wat stabieler is.
• Gewichtsverfijning: Een extra term ($W_{refine}$) wordt toegevoegd om de gekwantiseerde gewichten licht aan te passen ten opzichte van de originele full-precision gewichten.
• Verliesfunctie: De optimalisatie gebruikt een distillatie-achtige loss zonder labels:
  • MSE Loss: Herconstructie van tussenliggende features tussen full-precision en gekwantiseerd model.
  • KL-divergentie: Afstemming van de eindlogits (voorspellingen) tussen beide modellen.
  • Regularisatie: L1-regularisatie op de verfijningstermen.
• Efficiëntie: Het proces convergeert binnen 1-2,5 uur op één GPU (voor ViT-small).

2. Data-vrije Kalibratie met Geleerde Prompts

Om de afhankelijkheid van echte kalibratiegegevens (ImageNet) te elimineren, wordt een strategie ontwikkeld die Stable Diffusion Turbo gebruikt om synthetische data te genereren.

• Geleerde Multi-Prompts: In plaats van handmatig ontworpen tekstprompts (bijv. "een foto van een "), worden voor elke klasse meerdere unieke prompt-embeddings geleerd.
• Optimalisatie: Deze embeddings worden geoptimaliseerd door een bevooroordeelde, full-precision ViT als leraar. De doelen zijn:
  1. Classificatie: De gegenereerde afbeeldingen moeten correct worden geclassificeerd als de doelklasse.
  2. Diversiteit: Er worden orthogonale en variantie-losses toegevoegd om te zorgen dat de gegenereerde afbeeldingen variëren in lay-out, achtergrond, stijl en objectconfiguratie.
• Resultaat: Dit creëert een diverse set synthetische afbeeldingen die de activatiedistributie van echte ImageNet-data nauwkeurig benaderen.

Belangrijkste Bijdragen

1. End-to-End PTQ Framework: Een methode die alle quantisatieparameters (stapgrootte, herschaling, gewichtsverfijning) over het hele transformer-netwerk gelijktijdig optimaliseert zonder labels.
2. Data-vrije Strategie: Een innovatieve aanpak die Stable Diffusion Turbo combineert met geleerde multi-mode prompts om diverse, semantisch correcte synthetische kalibratiegegevens te genereren.
3. State-of-the-Art Resultaten: Het bereiken van nieuwe hoogtes in PTQ-prestaties, inclusief de eerste succesvolle resultaten voor ViT's onder extreme low-bit instellingen (W1.58A8).
4. Efficiëntie: Het framework is computatie-efficiënt en werkt binnen een uur op één GPU.

Resultaten

De methode is getest op ViT, DeiT en Swin-T modellen op de ImageNet-dataset.

• Prestaties bij Extreme Low-Bit:
  • Bij W1.58A8 (ternaire gewichten, 8-bit activaties) behaalt de methode aanzienlijk betere resultaten dan bestaande methoden (die vaak faalden of <50% nauwkeurigheid hadden). Bijvoorbeeld: 68.45% voor ViT-S (Real data) en 63.71% (Synthetische data), vergeleken met ~4.8% voor concurrenten.
  • Bij W4A4 en W3A3 worden state-of-the-art resultaten behaald, vaak met een verbetering van enkele procentpunten ten opzichte van RepQ-ViT, FIMA-Q en APHQ-ViT.
• Data-vrije vs. Real Data:
  • De data-vrije variant (met 100k synthetische afbeeldingen) presteert bijna even goed als de variant met echte data (verschil vaak <1-2%).
  • Geleerde prompts presteren aanzienlijk beter dan simpele tekst-sjablonen (bijv. "een foto van een..."), vooral bij lage bitbreedtes.
• Schalbaarheid: De nauwkeurigheid verbetert naarmate de kalibratiedataset groter wordt (tot ca. 10.000 samples), waarna het verzadigt. Het framework profiteert van grotere datasets, in tegenstelling tot blok-gebaseerde methoden.

Betekenis en Impact

Dit paper opent de deur voor de efficiënte implementatie van Vision Transformers op randapparatuur (edge devices) zonder dat er toegang is tot de oorspronkelijke trainingsdata.

• Edge AI: Het maakt het mogelijk om zware transformer-modellen te draaien op apparaten met beperkt geheugen en rekenkracht.
• Privacy: De data-vrije aanpak lost het probleem van data-privacy op, aangezien er geen gevoelige echte afbeeldingen nodig zijn voor kalibratie.
• Technische Doorbraak: Het bewijst dat extreme low-bit quantisatie (ternair) mogelijk is voor complexe transformer-architecturen, wat eerder als onhaalbaar werd beschouwd zonder volledige hertraining.

Samenvattend biedt deze studie een robuust, schaalbaar en data-onafhankelijk framework dat de kloof tussen full-precision Vision Transformers en efficiënte, gekwantiseerde implementaties op edge-apparaten dicht.