HiAP: A Multi-Granular Stochastic Auto-Pruning Framework for Vision Transformers

Dit paper introduceert HiAP, een end-to-end trainingsframework dat Vision Transformers op meerdere granulariteiten automatisch en stochastisch inplant om de reken- en geheugenkosten op randapparaten te verminderen zonder complexe multi-stap pipelines of handmatige heuristieken.

De kern

Stel je voor dat een Vision Transformer (ViT) een gigantische, hyper-intelligente chef-kok is die miljoenen foto's tegelijk moet analyseren. Deze chef is ongelooflijk goed in het herkennen van objecten, maar hij heeft een enorm probleem: hij staat in een keuken die te klein is voor zijn apparatuur. Hij heeft te veel ingrediënten nodig (geheugen) en te veel werkbladruimte (rekenkracht) om zijn gerechten te bereiden. Hierdoor kan hij niet werken in kleine keukens, zoals die in je smartphone of een slimme camera.

De huidige oplossing is vaak om de chef te dwingen om minder te koken, maar de methoden die daarvoor worden gebruikt zijn vaak omslachtig. Ze vragen een menselijke expert om te beslissen welke ingrediënten weg moeten, of ze gebruiken een ingewikkeld, meerstaps proces waarbij de chef eerst een proefmenu maakt, dan wordt geëvalueerd, en pas daarna echt begint met koken.

HiAP (Hierarchical Auto-Pruning) is een nieuwe, slimme manier om deze chef te trainen om zelf te beslissen wat hij nodig heeft. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De Twee-Lagen "Zelfreinigende" Keuken

Stel je voor dat HiAP twee soorten schakelaars heeft die de chef zelf kan bedienen:

• De Macro-schakelaars (De Grote Knoppen): Dit zijn de schakelaars die hele kookstations (zoals de "Aandachtshoofden" of de "Voedingsnetwerken") volledig uitschakelen. Als een station niet nodig is, wordt het gewoon verwijderd. Dit bespaart enorm veel ruimte op het werkblad (geheugenbandbreedte), omdat de chef niet meer hoeft te lopen naar die lege stations.
• De Micro-schakelaars (De Kleine Knoppen): Binnen de stations die wel blijven, kan de chef nu ook de hoeveelheid ingrediënten aanpassen. Hij kan bijvoorbeeld beslissen dat hij in een bepaalde pan niet 64 kruiden gebruikt, maar slechts 32. Dit maakt het gerecht lichter en sneller te bereiden zonder de smaak (de nauwkeurigheid) te veel te veranderen.

Het unieke aan HiAP is dat de chef beide schakelaars tegelijkertijd kan bedienen. Hij kan een heel station sluiten én tegelijkertijd de inhoud van de andere stations verkleinen.

2. De "Gumbel-Sigmoid" Magische Zandloper

Hoe leert de chef dit? In plaats van dat een mens zegt "verwijder dit", gebruikt HiAP een wiskundig trucje dat lijkt op een magische zandloper.

• Aan het begin van de training is de zandloper nog heel losjes. De schakelaars zijn vaag en "zacht". De chef probeert uit wat er gebeurt als hij een station halfslachtig uitschakelt. Hij leert hierdoor welke routes belangrijk zijn.
• Naarmate de tijd vordert, wordt de zandloper steeds strakker. De schakelaars worden harder. De chef begint te zien: "Ah, dit station werkt echt niet, ik doe het helemaal dicht." Of: "Deze kruiden zijn overbodig, ik doe ze weg."
• Aan het einde zijn de schakelaars volledig hard (aan of uit). De chef heeft nu een perfect, compact recept dat hij direct kan gebruiken. Er is geen tweede ronde nodig om het te "finetunen". Hij is klaar om te koken.

3. Waarom is dit zo slim?

De meeste andere methoden zijn als een architect die eerst een plattegrond tekent, dan bouwt, en dan pas ziet dat het te zwaar is, waarna hij moet slopen en herbouwen. Dat kost tijd en energie.

HiAP is als een slimme, zelforganiserende tuin.

• De tuinier (de computer) plant zaden (de volledige structuur).
• De tuinier geeft de planten een signaal: "Jullie moeten binnen een bepaalde omheining blijven."
• De planten groeien en concurreren om ruimte. De zwakke takken (onnodige onderdelen) sterven vanzelf af, terwijl de sterke takken blijven staan.
• Uiteindelijk heb je een prachtige, compacte struik die precies past in de pot, zonder dat de tuinier handmatig elke tak heeft moeten snoeien.

Het Resultaat

Met HiAP kunnen we deze gigantische, zware AI-modellen verkleinen tot iets dat past in een gewone telefoon of drone, zonder dat ze minder goed worden in het herkennen van dingen.

• Snelheid: Omdat de "lege stations" volledig worden verwijderd, hoeft de computer niet meer naar die lege plekken te kijken. Dit bespaart tijd en batterij.
• Eenvoud: Er is geen dure menselijke expert nodig om te beslissen wat weg moet. Het model leert het zelf in één keer.
• Efficiëntie: Het model wordt niet alleen kleiner, maar ook slimmer in hoe het zijn middelen gebruikt. Het houdt de belangrijke dingen vast en gooit alleen het echte rommel weg.

Kortom: HiAP is de slimme assistent die een overvolle, inefficiënte AI omtovert tot een strakke, snelle en efficiënte machine die overal mee naartoe genomen kan worden.

Technische samenvatting

Hieronder volgt een gedetailleerde technische samenvatting van het paper "HiAP: A MULTI-GRANULAR STOCHASTIC AUTO-PRUNING FRAMEWORK FOR VISION TRANSFORMERS" in het Nederlands.

Probleemstelling

Vision Transformers (ViT) zijn dominant geworden in computervisie, maar hun hoge reken- en geheugeneisen maken ze moeilijk inzetbaar op apparaten met beperkte middelen (edge devices). Bestaande methoden voor gestructureerd pruning (het verwijderen van volledige componenten zoals attention heads of FFN-neuronen) hebben twee belangrijke beperkingen:

1. Enkele granulariteit: Methodes richten zich vaak alleen op micro-structuren (binnen-head dimensies) of macro-structuren (hele blokken/heads). Micro-pruning vermindert wel de FLOPs, maar lost het geheugenbandbreedte-probleem (DRAM/SRAM-toegang) niet op omdat alle lagen nog steeds geladen moeten worden. Macro-pruning vermindert geheugenoverhead, maar kan leiden tot significante prestatieverlies door het verlies van representatievermogen.
2. Complexiteit en afhankelijkheid van menselijke input: Veel huidige differentieerbare zoekmethoden vereisen complexe, multi-stadia pijplijnen, handmatige drempelwaarden (thresholding) en expertkennis om de juiste sparsiteit te bereiken. Ze vertrouwen vaak op post-hoc heuristieken in plaats van dat het model zelf leert wat er moet worden verwijderd.

Methodologie: HiAP (Hierarchical Auto-Pruning)

HiAP is een continu relaxatie-framework dat optimale sub-netwerken ontdekt in één enkele, end-to-end trainingsfase zonder vooraf gedefinieerde sparsiteitsdoelen of handmatige heuristieken.

1. Hiërarchisch Stochastisch Gating
Het kernidee is de introductie van stochastische Gumbel-Sigmoid gates op twee niveaus:

• Macro-niveau: Gates ($g_{l,h}$ en $b_l$) bepalen of hele attention heads of volledige FFN-blokken (Feed-Forward Networks) worden behouden of volledig worden omzeild. Dit vermindert de geheugenbandbreedte-overhead aanzienlijk.
• Micro-niveau: Gates ($d_{l,h,j}$ en $c_{l,k}$) werken binnen de actieve macro-structuren om specifieke intra-head dimensies en FFN-neuronen selectief te verwijderen. Dit vermindert de rekenkosten (FLOPs).

2. Differentieerbare Kostenmodellering
Om de zoektocht te sturen, formuleert HiAP een exacte, differentieerbare berekening van de Multiply-Accumulate (MAC) operaties. De kostenfunctie is lineair opgebouwd uit:

• Macro-kosten: De overhead van het laden van gehele attention heads en FFN-blokken.
• Micro-kosten: De kosten van de resterende interne dimensies en neuronale activaties.
  Dit stelt het model in staat om de trade-off tussen geheugenbandbreedte en rekenkracht autonoom te balanceren.

3. Training en Temperatuur Annealing

• Gumbel-Sigmoid Relaxatie: De binaire gates worden tijdens het trainen relaxeerd naar continue waarden tussen 0 en 1, waardoor gradients kunnen stromen.
• Temperatuur Annealing: De temperatuur ($\tau$) van de Gumbel-Sigmoid wordt geleidelijk verlaagd (van 2.0 naar 0.1). In het begin gedragen de gates zich als een zachte, continue regularisator (stochastische dropout). Naarmate $\tau$ daalt, worden de verdelingen scherper en convergeren ze naar deterministische binaire beslissingen (0 of 1).
• Structuurbeperkingen: Om te voorkomen dat het netwerk hele lagen verwijdert voordat de gewichten zich kunnen aanpassen (structuurinstorting), worden haalbaarheidsstraffen (feasibility penalties) toegevoegd die een minimumaantal behouden heads en dimensies garanderen.
• Single-Phase Discovery: In tegenstelling tot eerdere methoden die een aparte fine-tuning fase vereisen, wordt HiAP getraind in één fase. De gewichten co-adaptëren zich met de geleidelijk "hard wordende" structuur, wat leidt tot een direct inzetbaar, fysiek gecomprimeerd sub-netwerk zonder extra fine-tuning.

Belangrijkste Bijdragen

1. Unificatie van Granulariteit: HiAP is het eerste framework dat macro-pruning (blokken/heads) en micro-pruning (neuronen/dimensies) combineert in één differentieerbaar systeem, waardoor zowel geheugen- als rekenkosten worden geoptimaliseerd.
2. Budget-bewust Zonder Menselijke Interventie: Het model ontdekt zijn eigen optimale sub-architectuur end-to-end. Er zijn geen handmatige drempelwaarden, proxy-rangschikkingen of dure secundaire fine-tuning-fases nodig.
3. Fysieke Extractie: Het resultaat is een fysiek gecomprimeerd ViT-model dat direct op standaard hardware kan worden uitgevoerd, zonder afhankelijkheid van gespecialiseerde sparse convolutie-engines.

Resultaten

De methode is gevalideerd op CIFAR-10 (met ViT-Tiny) en ImageNet-1K (met DeiT-Small).

• ImageNet-1K (DeiT-Small):
  • HiAP reduceerde de rekenkosten van 4.6G naar 3.1G MACs (een reductie van ~33%) met een Top-1 nauwkeurigheid van 79.1%.
  • Dit presteert vergelijkbaar met of beter dan geavanceerde methoden zoals GOHSP en ViT-Slim, maar dan met een veel eenvoudigere pijplijn (één trainingsfase vs. multi-stadia).
  • Bij een nog lagere budget (2.5G MACs) behaalde HiAP 77.95% nauwkeurigheid.
• CIFAR-10:
  • HiAP overtrof handmatige heuristieken (zoals $\ell_1$-norm ranking en uniforme ratio's) consistent in nauwkeurigheid bij zowel gematigde (33%) als agressieve (50%) compressie.
  • Hardware-efficiëntie: Op een GPU daalde de inferentielatentie van 5.57 ms naar 3.86 ms (een snelheidswinst van ~1.44x), wat bevestigt dat de gevonden sub-netwerken daadwerkelijk sneller zijn op standaard hardware.
• Structuuranalyse: Het model leerde om in de vroege trainingsfase hele FFN-blokken en attention heads te verwijderen (macro), en later de interne dimensies te verfijnen (micro). Het identificeerde bijvoorbeeld dat de laatste FFN-blok volledig overbodig was.

Betekenis

HiAP markeert een verschuiving in het veld van modelcompressie voor Vision Transformers. Het demonstreert dat complexe, handmatig ontworpen pruning-strategieën kunnen worden vervangen door een geautomatiseerd, end-to-end leerproces dat zowel de geheugenbandbreedte als de rekenkosten optimaliseert. Door het elimineren van de noodzaak voor multi-stadia pijplijnen en post-hoc tuning, maakt HiAP het veel eenvoudiger en praktischer om efficiënte ViT-modellen te deployen op edge-apparaten. Het biedt een nieuwe staat van de kunst voor de nauwkeurigheids-efficiëntie Pareto-grens, waarbij modellen compacter worden zonder in te leveren op prestaties.