ToaSt: Token Channel Selection and Structured Pruning for Efficient ViT

Each language version is independently generated for its own context, not a direct translation.

Stel je voor dat Vision Transformers (ViT) enorme, superintelligente bibliotheken zijn die beelden "lezen" om ze te begrijpen. Ze zijn ongelooflijk slim, maar ze zijn ook zwaar, traag en verbruiken veel energie. Het is alsof je een vrachtwagen gebruikt om een postzegel te bezorgen: het werkt, maar het is inefficiënt en kost veel brandstof.

De auteurs van dit paper, ToaSt, hebben een slimme oplossing bedacht om deze vrachtwagen om te bouwen tot een snelle, zuinige sportwagen, zonder dat hij minder goed kan rijden. Ze noemen hun methode ToaSt (Token Channel Selection and Structured Pruning).

Hier is hoe het werkt, vertaald naar alledaagse termen:

1. Het Probleem: Twee soorten "rommel"

Deze AI-modellen hebben twee grote gebieden waar veel onnodige energie wordt verbruikt:

De "Aandacht" (MHSA): Dit is het gedeelte waar het model kijkt naar verschillende delen van een foto tegelijk. Het is alsof een detective elke foto in duizenden stukjes snijdt en elk stukje met elk ander stukje vergelijkt. Dit kost enorm veel tijd.
De "Feed-Forward" (FFN): Dit is het gedeelte waar het model de informatie verwerkt en "nadenkt". Dit is eigenlijk het zwaarste deel van de motor (meer dan 60% van het werk!). Hier zitten vaak veel dubbelzinnige of nutteloze kanalen die alleen maar ruis maken.

2. De Oplossing: ToaSt in twee stappen

ToaSt pakt deze twee problemen aan met een heel slimme, gescheiden aanpak.

Stap 1: De "Aandacht" opschonen (Structuur-Pruning)

Stel je voor dat de detective in de eerste stap duizenden collega's heeft die allemaal hetzelfde verhaal vertellen. In plaats van ze allemaal te ontslaan (wat het verhaal zou verstoren), kiest ToaSt ervoor om de kwaliteit van de gesprekken te verbeteren.

De analogie: Stel je een koor voor. Als één zanger een noot mist, moet de hele groep dat doen, anders klinkt het niet goed. ToaSt kijkt naar de "zangers" (de gewichten in de computer) en zegt: "Jullie zingen allemaal hetzelfde liedje, maar sommigen zingen het iets minder goed."
De truc: Ze verwijderen de slechtste zangers, maar doen dit gekoppeld. Als ze een zanger uit de "Q"-groep verwijderen, verwijderen ze direct de corresponderende zanger uit de "K"-groep. Zo blijft de harmonie perfect.
Het resultaat: De detective heeft minder mensen nodig om te praten, maar de boodschap blijft even helder. En het beste: ze hoeven niet opnieuw te leren zingen (geen dure hertraining).

Stap 2: De "Denk-kanalen" selecteren (Token Channel Selection)

Dit is het meest creatieve deel. Het model heeft een enorme denkruimte (de FFN-laag) die vol zit met "dode" neuronen of ruis.

De analogie: Stel je voor dat je een enorme zaal hebt met 1000 ramen. Je wilt weten of er buiten een auto voorbijrijdt. Je hoeft niet naar alle 1000 ramen te kijken; je kunt gewoon naar de ramen kijken waar het licht is, of waar de beweging is.
De truc: ToaSt kijkt niet naar de hele zaal, maar pakt een klein steekproefje van de ramen (tokens). Op basis van dit steekproefje ziet het model direct welke ramen (kanalen) alleen maar ruis laten zien en welke echt belangrijk zijn.
Het resultaat: Ze sluiten de ramen die alleen maar ruis laten zien. Omdat ze dit doen op basis van wat er nu gebeurt (tijdens het kijken naar de foto), hoeven ze het model niet opnieuw te trainen. Het is alsof je direct de gordijnen dichttrekt bij de ramen waar niemand zit.

3. Waarom is dit zo speciaal?

Geen dure hertraining: Andere methoden moeten het model vaak maandenlang opnieuw laten leren na het verwijderen van onderdelen. ToaSt doet dit bijna direct. Het is alsof je een auto instelt in plaats van hem opnieuw te bouwen.
Grotere modellen profiteren meer: Hoe groter en slimmer het model is (zoals de "Huge" versie), hoe makkelijker het is om de "dode" onderdelen te vinden. Het is alsof een zeer rijke persoon makkelijker 50% van zijn spullen kan wegdoen zonder dat het leven eronder lijdt, omdat hij zoveel overbodige spullen heeft.
Beter dan voorheen: In tests bleek dat ToaSt niet alleen sneller is, maar soms zelfs slimmer wordt dan het origineel. Door de ruis weg te halen, ziet het model de echte details scherper.

Samenvatting

ToaSt is als een slimme renovatie van een zware machine. Ze verwijderen de zware, onnodige onderdelen (de ruis) en synchroniseren de overgebleven onderdelen zodat ze perfect samenwerken. Het resultaat is een model dat:

Sneller is (minder brandstofverbruik).
Lichter is (past op je telefoon).
Slimmer is (door de ruis te verwijderen).

Het is een manier om de kracht van de grootste AI-modellen beschikbaar te maken voor iedereen, zonder dat je een supercomputer nodig hebt.

Each language version is independently generated for its own context, not a direct translation.

Probleemstelling

Vision Transformers (ViTs) hebben indrukwekkende resultaten geboekt op diverse visuele taken, maar hun toepassing in resource-beperkte omgevingen (zoals mobiele apparaten en edge computing) wordt gehinderd door hoge rekenkosten. De complexiteit van ViTs komt voort uit twee bronnen:

Zelf-attention mechanismen: Deze hebben een kwadratische complexiteit $O(N^2)$ met betrekking tot de sequentielengte $N$ .
Feed-Forward Networks (FFN): Deze lagen dragen ongeveer 61% bij aan de totale FLOPs (rekenoperaties) door operaties die schalen met het verborgen dimensie $D^2$ .

Bestaande compressietechnieken hebben beperkingen:

Gestructureerde gewichts- pruning: Vereist vaak langdurige hertraining (retraining) die even duur kan zijn als het oorspronkelijke trainen. Bovendien richten deze methoden zich vaak alleen op de attention-mechanismen en negeren ze de redundante FFN-lagen.
Token-compressie: Verkleint de sequentielengte $N$ , maar lost de redundantie op kanaalniveau binnen MHSA en FFN niet op. Bovendien propageren beslissingen over het verwijderen van tokens globaal door alle volgende lagen, wat de optimalisatie bemoeilijkt.

Methodologie: ToaSt Framework

ToaSt (Token Channel Selection and Structured Pruning) is een gekoppeld (decoupled) framework dat specifieke strategieën toepast op verschillende componenten van de ViT-architectuur om hertraining te minimaliseren en globale propagatie te voorkomen. Het werkt in twee onafhankelijke fasen:

1. Gestructureerde Gekoppelde Gewichts-Pruning voor MHSA

Voor de Multi-Head Self-Attention (MHSA) modules wordt de interne hoofddimensie ( $d_k$ ) gereduceerd via een "coupled weight pruning" strategie.

Koppeling: Om de wiskundige integriteit te behouden, worden de indexen gesynchroniseerd over de gekoppelde gewichtsmatrices ( $W_Q, W_K, W_V, W_{proj}$ ). Als kolom $j$ van $W_Q$ wordt verwijderd, moet kolom $j$ van $W_K$ ook worden verwijderd (Q-K synchronisatie), en evenzo voor $W_V$ en $W_{proj}$ (V-Proj synchronisatie).
Selectie-criterium: De belangrijkheid van dimensies wordt bepaald door de Geometrische Mediaan (Geometric Median - GM) van de vooraf getrainde gewichten. Dimensies die het dichtst bij het centrum van de verdeling liggen (dicht bij de GM) worden als het meest redundant beschouwd en verwijderd.
Strategie: Er wordt een "Head-wise Uniform" strategie toegepast waarbij alle hoofden dezelfde verkleining krijgen. De eerste laag wordt overgeslagen om de interface met patch-embeddings te behouden, terwijl de resterende lagen tot 90% worden gepreund.

2. Token Channel Selection (TCS) voor FFN

Voor de Feed-Forward Networks (waar de meeste FLOPs zitten) wordt een trainingsvrije (training-free) kanaalselectie-methode gebruikt.

Empirische Analyse: De auteurs ontdekten drie kenmerken in diepere lagen: hoge lineaire reconstructie-fideliteit ( $R^2 \approx 1.0$ , wat lineaire afhankelijkheid aangeeft), een instortende effectieve rang (effective rank), en toenemende activatie-sparseheid.
Statistische Sampling: In plaats van alle tokens te analyseren (wat duur is), wordt een willekeurige subset van tokens (2-20%) gebruikt om de kanaalbelangrijkheid te schatten.
Unificatie van Belangrijkheid: Een score wordt berekend die rekening houdt met de CLS-token activatie en de attention-weights van patches. Dit prioriteert kanalen die globale semantische informatie bevatten.
Adaptieve Verwijdering:
- FC1 (Expansie): Conservatieve pruning in vroege lagen om diversiteit te behouden.
- FC2 (Reductie): Aggressieve pruning (tot 90%) in diepere lagen waar redundantie hoog is.
Voordeel: Omdat hele kanalen worden verwijderd (in plaats van individuele gewichten), blijft de matrixstructuur dicht (dense), wat directe versnelling op standaard GPU-hardware mogelijk maakt zonder speciale sparse libraries.

Belangrijkste Bijdragen

Gestructureerde Gekoppelde Pruning voor MHSA: Een methode die de interne hoofddimensie verkleint door gesynchroniseerde verwijdering over gewichtsmatrices heen, wat laag-onafhankelijke compressie mogelijk maakt zonder de interface tussen transformer-blokken te wijzigen.
Token Channel Selection (TCS): Een trainingsvrije aanpak voor FFN-lagen die gebaseerd is op een empirische analyse van redundantie (sparsiteit, rang, $R^2$ ). Dit elimineert de noodzaak voor dure hertraining en filtert effectief ruis.
Superieure Trade-offs: Het framework toont aan dat grotere modellen (zoals ViT-MAE-Huge) minder hertraining nodig hebben na compressie, en dat de gecomprimeerde modellen consistent beter presteren dan de baselines.

Resultaten

De methodologie is geëvalueerd op negen modellen (DeiT, ViT-MAE, Swin Transformer) voor beeldclassificatie (ImageNet-1K) en objectdetectie (COCO).

ImageNet-1K Classificatie:
- ViT-MAE-Huge: Bereikt 88,52% Top-1 nauwkeurigheid (+1,64% t.o.v. de baseline) met een reductie van 39,4% in FLOPs.
- Efficiëntie: ToaSt levert een snelheidswinst (throughput) van 1,28x tot 2,07x op een NVIDIA H100 GPU. Bijvoorbeeld, DeiT-Small bereikt 4783 beelden/sec (2,07x sneller) met een nauwkeurigheidswinst van +3,58%.
- Hertraining: ViT-MAE-Huge herstelt zijn prestaties binnen slechts ~15 epochs van fine-tuning, vergeleken met ~290 epochs voor kleinere modellen.
COCO Objectdetectie:
- Gecomprimeerde Swin-Backbones behouden of verbeteren de prestaties. Swin-Small bereikt 52,2 mAP (tegenover 51,9 mAP van de baseline) met een reductie in FLOPs.
Vergelijking met State-of-the-Art: ToaSt presteert consistent beter dan token-compressie methoden (zoals ToMe en DiffRate) bij vergelijkbare FLOPs-budgetten, vaak met 1-4% hogere nauwkeurigheid.

Significantie

ToaSt biedt een doorbraak in de efficiëntie van Vision Transformers door twee kritieke problemen op te lossen: de hoge kosten van hertraining bij gewichts-pruning en de beperkte compressie van token-methoden.

Hardware-vriendelijk: Door gestructureerde pruning toe te passen, blijft de rekenstructuur dicht, wat essentieel is voor snelle inferentie op standaard hardware.
Schaalbaarheid: De bevinding dat grotere foundation modellen minder hertraining nodig hebben, suggereert dat deze techniek zeer goed schaalbaar is voor toekomstige, nog grotere modellen.
Algemene Toepasbaarheid: De methode werkt effectief over verschillende architecturen (Van DeiT tot Swin) en taken (classificatie en detectie), wat aantoont dat de gereduceerde redundantie echt architecturaal is en niet taak-specifiek.

Kortom, ToaSt stelt een nieuwe standaard voor ViT-compressie door een decoupled aanpak te combineren die zowel de attention- als de FFN-componenten optimaliseert zonder de nadelen van bestaande methoden.