Zero-Shot Quantization via Weight-Space Arithmetic

Probleemstelling

Diepe neurale netwerken worden doorgaans getraind en opgeslagen in hoge precisie (float32/float16), wat leidt tot grote geheugenvoetafdrukken en hoge bandbreedte-eisen tijdens implementatie. Integer-quantization (het vertalen van parameters naar lage-bit gehele getallen) lost dit op, maar brengt een fundamenteel dilemma met zich mee:

• Post-Training Quantization (PTQ): Kan worden toegepast op een reeds getraind model zonder extra trainingsdata of rekenkracht. Echter, bij extreem lage bit-breedtes (bijv. 3-bit of lager) veroorzaakt PTQ vaak ernstige degradatie in prestaties omdat de geoptimaliseerde parameters niet robuust zijn tegen de quantisatiefouten.
• Quantization-Aware Training (QAT): Lost dit op door het model te laten trainen met simulaties van quantisatie, wat resulteert in veel betere prestaties bij lage bit-breedtes. Het nadeel is dat QAT echter extra trainingsdata, rekenkracht en tijd vereist, wat niet altijd beschikbaar is.

De centrale vraag van het artikel is: Kan de robuustheid die tijdens QAT wordt geleerd, worden geïsoleerd en hergebruikt voor andere taken zonder dat het ontvangermodel opnieuw getraind hoeft te worden?

Methodologie: Weight-Space Arithmetic

De auteurs introduceren een methode die Weight-Space Arithmetic (rekenen in de ruimte van modelgewichten) combineert met quantisatie-robustheid. De kern van de methode is het concept van de Quantization Vector (QV).

1. Definitie van de Quantization Vector (QV):
   De auteurs nemen een donor-taak waarvoor twee checkpoints beschikbaar zijn:

   • $\theta_{D}$: Een standaard fijngefineerd model (zonder QAT).
   • $\theta_{D,QAT}$: Een model dat op dezelfde taak is getraind met QAT.
     De QV ($\rho_D$) wordt gedefinieerd als het verschil tussen deze twee checkpoints:
     $$\rho_D = \theta_{D,QAT} - \theta_{D}$$
     Deze vector vertegenwoordigt de verplaatsing in de parameter-ruimte die nodig is om een model robuust te maken tegen quantisatieruis.

2. Zero-Shot Patching:
   Voor een ontvanger-taak (receiver) met een standaard checkpoint $\theta_R$ (zonder QAT), wordt de QV van de donor toegepast om een "gepatcht" model te creëren:
   $$\theta_{R \leftarrow D} = \theta_{R} + \lambda \rho_D$$
   Hierbij is $\lambda$ een schalingscoëfficiënt. Deze operatie verplaatst de parameters van het ontvangermodel in de richting die robuustheid garandeert, zonder dat er trainingdata van de ontvanger-taak nodig is.

3. Aannames:
   De methode gaat uit van een gedeelde pre-trained initialisatie (bijv. een ViT-model dat op ImageNet is voorgeïmplementeerd). Hierdoor bevinden de checkpoints zich in een "compatible basin" in de weight space, waardoor lineaire operaties zinvol zijn.

Belangrijkste Bijdragen

1. Introductie van Quantization Vectors (QV): Het concept dat de verplaatsing tussen een standaard checkpoint en een QAT-checkpoint een lineair samenvoegbare vector is die structurele robuustheid tegen quantisatie encodeert.
2. Zero-Shot Cross-Task Patching Framework: Een protocol dat QAT-robustheid overdraagt van een donor-taak naar een ontvanger-taak zonder enige training van de ontvanger (data-vrij en compute-arm).
3. Empirisch Bewijs van Transferabiliteit: Het aantonen dat QAT-geïnduceerde robuustheid geen lokaal fenomeen is, maar een herbruikbaar kenmerk van de geometrie van de weight space dat over verschillende visuele taken kan worden getransfereerd.

Resultaten

De auteurs hebben hun methode geëvalueerd op Vision Transformer (ViT) modellen (ViT-T/S/B/L) over 22 verschillende beeldclassificatietaken (zoals ImageNet, CIFAR, Stanford Cars, etc.) met 3-bit symmetrische per-kanaal gewichtsquantisatie.

• Prestatieverbetering: Het patchen van een model met een donor-QV kan de robuustheid tegen PTQ-ruis met tot 60% verbeteren ten opzichte van standaard PTQ.
• Richting vs. Grootte:
  • De richting van de QV is universeel bruikbaar. Zelfs zonder tuning van de schaalfactor ($\lambda=1$) tonen veel donor-ontvanger paren verbetering.
  • De grootte ($\lambda$) is cruciaal. Bij een vaste schaal ($\lambda=1$) kan het soms leiden tot destructieve interferentie (prestatiedaling). Door $\lambda$ te optimaliseren (via een kleine validatie-set), verdwijnt bijna alle destructieve interferentie en worden de prestatiesystematisch verbeterd.
• Donor-Selectie: Modellen getraind op grote, diverse datasets (zoals ImageNet en Tiny ImageNet) blijken de sterkste "donors" te zijn, met een 100% positieve transfer-ratio over alle ontvangers wanneer de schaal correct is ingesteld.
• Modelgrootte: De methode werkt voor ViT-T (klein) tot ViT-L (groot). Grotere modellen tonen soms meer polarisatie (sterkere winst maar ook grotere risico's bij verkeerde schaling), maar de fundamentele transferabiliteit blijft bestaan.

Significantie en Conclusie

Dit werk heeft een fundamentele implicatie voor het veld van efficiënte deep learning:

• Robuustheid is een transferable eigenschap: Quantisatie-robustheid is niet slechts een bijproduct van taakspecifieke training, maar een geometrische eigenschap van de weight space die kan worden "gepakt" en hergebruikt.
• Alternatief voor QAT: Het biedt een zero-shot, low-cost alternatief voor QAT. Voor scenario's waar trainingsdata ontbreekt of rekenkracht beperkt is (bijv. edge devices), kan een model nu robuust worden gemaakt voor extreme lage bit-breedtes door simpelweg een vector van een ander model toe te voegen.
• Uitbreiding van Weight-Space Arithmetic: Waar eerdere werken zich richtten op het overdragen van semantische kennis (bijv. nieuwe taken of talen), toont dit werk aan dat weight-space vectoren ook structurele en computationele eigenschappen (zoals overleven bij lage precisie) kunnen coderen.

Samenvattend bewijst de paper dat we de "leerervaring" van QAT kunnen isoleren als een vector en deze als een "patch" kunnen gebruiken om willekeurige modellen robuust te maken voor extreme compressie, zonder de kosten van opnieuw trainen.