ScalePredictor: Instance-aware Scale Learning for Accurate Quantization of Vision Transformers

Technische Samenvatting: ScalePredictor

Probleemstelling

Vision Transformers (ViTs) hebben de state-of-the-art prestaties bereikt in computer vision, maar worden geconfronteerd met aanzienlijke implementatie-uitdagingen op hulpbron-beperkte edge-apparaten vanwege hun computationele intensiteit. Post-Training Quantization (PTQ) biedt een oplossing door modellen te comprimeren naar low-bit formaten (bijv. 8-bit of lager) zonder hertraining. Bestaande PTQ-methoden vertrouwen echter voornamelijk op statische quantisatie, waarbij een enkele set quantisatieparameters (scales) uniform op alle input-instanties wordt toegepast.

Dit statische paradigma is inherent suboptimaal voor ViTs omdat natuurlijke afbeeldingen een enorme diversiteit vertonen, wat leidt tot aanzienlijke variaties in de activatieverdelingen over verschillende samples heen. Deze heterogeniteit is bijzonder schadelijk in ultra-low-bit instellingen (bijv. W3/A3), waar vaste scales er niet in slagen om instantie-specifieke dynamiek te vangen. Omgekeerd vereisen dynamische quantisatie strategieën die parameters per sample aanpassen vaak runtime full-tensor reducties (min/max berekeningen) voor elke laag. Deze operaties introduceren ernstige synchronisatiebarrières en geheugen-gebonden overheads op moderne hardware, wat de efficiëntievoordelen van quantisatie drastisch ondermijnt en de inferentie-snelheid negatief beïnvloedt.

Methodologie: ScalePredictor

De auteurs stellen ScalePredictor voor, een dynamisch quantisatieframework dat instantie-bewuste scale-learning bereikt met minimale computationele overhead. De kerninzicht die de methode drijft, is een ontdekte correlatie: de optimale quantisatie-scales voor diepere lagen zijn sterk afhankelijk van de activatie-amplitude die in ondiepere lagen wordt waargenomen.

Het framework bestaat uit drie hoofdonderdelen:

1. Robuuste Range Extractie:
   In plaats van globale min/max waarden te berekenen (gevoelig voor uitschieters) of dure percentiel-sortering uit te voeren, extraheert ScalePredictor een "robuuste range" ($r(x)$) uit de Patch Embedding output (de ondiepste laag).

   • De geflatterde tensor wordt gepartitioneerd in $K$ niet-overlappende chunks.
   • Lokale min/max waarden worden berekend voor elke chunk.
   • Deze lokale extremen worden gemiddeld om een globale robuuste range af te leiden.
   • Deze aanpak fungeert als een grofmazige percentiel-benadering, waardoor de invloed van uitschieters effectief wordt onderdrukt terwijl de $O(N \cdot C)$ complexiteit behouden blijft en sortering wordt vermeden.

2. Taylor-Gemotiveerde Polinomiale Predictor:
   Een leerbare polynomiale module brengt de geëxtraheerde robuuste range $r(x)$ in kaart naar de optimale quantisatie-scales voor alle daaropvolgende lagen.

   • Op basis van de observatie dat de relatie tussen ondiepe statistieken en diepe optimale scales een glad, niet-lineair trend volgt, modelleert de methode de scale $s_l(x)$ voor laag $l$ als een polynomiale functie van $r(x)$:
     $$s_l(x) = \sum_{k=0}^{d} c_{l,k} \cdot (r(x))^k$$
   • De coëfficiënten $c_{l,k}$ worden tijdens een kalibratiefase geleerd met behulp van een block-wise reconstructieverlies.
   • Dit maakt het mogelijk om de instantie-specifieke scales voor alle lagen in één enkele pass te voorspellen, zonder toegang te krijgen tot diepere activatie-tensors.

3. Kalibratie en Inferentie:

   • Kalibratie: De polynomiale coëfficiënten worden geoptimaliseerd met behulp van een kleine kalibratieset. De constante term ($c_{l,0}$) wordt geïnitialiseerd met de optimale statische scale om stabiliteit te garanderen, terwijl hogere-orde termen worden geleerd om instantie-variaties te vangen.
   • Inferentie: Voor een nieuwe input wordt de robuuste range één keer berekend bij de Patch Embedding laag. Alle daaropvolgende laag-scales worden gegenereerd via polynomiale evaluatie ($O(1)$ FLOPs), waardoor de noodzaak voor per-laag runtime reducties wordt geëlimineerd.

Belangrijkste Bijdragen

1. Instantie-bewuste PTQ voor ViTs: ScalePredictor wordt gepresenteerd als de eerste post-training quantisatiemethode voor ViTs die gebruikmaakt van dynamische quantisatie om instantie-niveau distributie-heterogeniteit aan te pakken zonder hertraining.
2. Efficiënt Scale Learning Mechanisme: De integratie van een robuust range-extractiemechanisme (chunk-gebaseerde min/max averaging) met een Taylor-geïnspireerde polynomiale projectie maakt nauwkeurige, on-the-fly scale-voorspelling mogelijk met verwaarloosbare inferentie-overhead.
3. Ontdekking van de Shallow-Deep Correlatie: Het werk onthult een consistente functionele afhankelijkheid waarbij de activatie-statistieken van ondiepe lagen (specifiek van de Patch Embeddings) effectief de optimale quantisatie-scales van diepere transformer-blokken kunnen voorspellen over diverse architecturen (ViT, DeiT, Swin).

Experimentele Resultaten

Uitgebreide experimenten zijn uitgevoerd op de ImageNet-dataset over zeven ViT-architecturen (ViT-S/B, DeiT-T/S/B, Swin-S/B) en vier PTQ-baselines (BRECQ, QDrop, AdaLog, I&S-ViT).

• Verbeteringen in Nauwkeurigheid: ScalePredictor presteert consequent beter dan statische PTQ-baselines, met name in low-bit regimes.
  • Opvallende winsten zijn onder meer een verbetering van 10,58% (van 33,57% naar 44,15%) op Swin-Small met behulp van I&S-ViT bij W2/A3.
  • Significante verbeteringen werden ook waargenomen bij W3/A3, zoals een winst van 8,71% op ViT-Small met behulp van BRECQ.
  • Zelfs in gevallen waar statische baselines instortten (bijv. BRECQ op ViT-Small bij W2/A3 met 1,10% nauwkeurigheid), herstelde ScalePredictor de prestaties naar 4,15%.
• Vergelijking met JIT Dynamische Quantisatie: In tegen tegenstelling tot Just-In-Time (JIT) dynamische quantisatie-methoden (bijv. JIT Min/Max), die de nauwkeurigheid ernstig verslechteren door de gekalibreerde gewichts-activatie relatie te verbreken, behoudt ScalePredictor de kalibratie-consistentie terwijl het adaptiviteit mogelijk maakt.
• Efficiëntie:
  • Latency: Op een NVIDIA RTX 3090 introduceert ScalePredictor slechts 0,6% end-to-end latency overhead vergeleken met statische PTQ.
  • Scale Berekening: De methode is 40–52× sneller dan JIT Min/Max en 507–1025× sneller dan JIT Percentile voor het berekenen van quantisatie-scales, omdat het per-laag tensor reducties vermijdt.

Betekenis en Claims

Het artikel claimt dat ScalePredictor er succesvol in is geslaagd de kloof te overbruggen tussen de nauwkeurigheid van dynamische quantisatie en de efficiëntie van statische quantisatie. Door gebruik te maken van de verborgen correlatie tussen de verdelingen in ondiepe lagen en de diepe laag optimale scales, vermijdt de methode de kostbare "just-in-time" kalibratie die vereist is voor traditionele dynamische benaderingen.

De auteurs benadrukken dat ScalePredictor een gunstige nauwkeurigheid-efficiëntie trade-off biedt, wat het een levensvatbare oplossing maakt voor het implementeren van low-bit Vision Transformers op edge-apparaten. De methode is ontworpen om naadloos te worden geïntegreerd in bestaande PTQ-frameworks, waarbij enkel de toevoeging van een lichtgewicht predictor module en een robuust range-extractiestap vereist is, zonder de onderliggende netwerkarchitectuur te wijzigen of volledige hertraining te vereisen.