Adaptive MLP Pruning for Large Vision Transformers

Deze paper introduceert Adaptive MLP Pruning (AMP), een methode die de parameters en FLOPs van grote Vision Transformers met ongeveer 40% vermindert zonder merkbare prestatieverlies door het gebruik van een label-vrije entropie-criterium voor nauwkeurigere neuron-evaluatie en een adaptieve binaire zoekalgoritme voor het bepalen van de pruning-diepte.

De kern

Titel: Slimmer, niet groter: Hoe we enorme beeldherkenningsmodellen op de maat knippen

Stel je voor dat je een gigantische, hypergeavanceerde robot hebt die foto's kan bekijken en begrijpen. Deze robot is zo slim dat hij bijna alles herkent, maar hij is ook zo groot en zwaar dat hij een eigen elektriciteitscentrale nodig heeft om te draaien. Hij neemt veel ruimte in beslag en is traag. Dit is precies het probleem met de huidige "grote visuele transformers" (AI-modellen voor beeldherkenning). Ze werken fantastisch, maar ze zijn te zwaar voor gewone computers of telefoons.

De onderzoekers van dit paper hebben een oplossing bedacht die ze AMP noemen (Adaptive MLP Pruning). Laten we uitleggen hoe dit werkt met een paar simpele vergelijkingen.

1. Het probleem: De overvolle boekenkast

Deze AI-modellen bestaan uit verschillende onderdelen. De onderzoekers ontdekten dat het grootste deel van het gewicht (de "parameters") zit in een specifiek onderdeel dat ze MLP noemen.

• De analogie: Stel je voor dat de AI een enorme bibliotheek is. De onderzoekers merkten op dat 80% van de boeken in één specifieke kamer (de MLP-kamer) staan. Maar als je die kamer bekijkt, zie je dat er veel boeken liggen die niemand leest. Ze zijn er alleen maar voor de show. Ze nemen ruimte in, maar dragen niet echt bij aan de slimheid van de robot.

2. De oplossing: De "Slimme Schaar"

In plaats van willekeurig boeken weg te gooien (wat de robot dom zou maken), gebruiken ze een slimme methode om te bepalen welke boeken echt nodig zijn.

Stap 1: De "Geurtest" (Informatie-entropie)
Oude methoden keken alleen naar het juiste antwoord. Als de robot een foto van een hond zag en de label was "hond", keken ze alleen naar hoe goed hij "hond" zei. Ze negeerden alles wat hij over katten of auto's dacht.

• De nieuwe aanpak: De onderzoekers gebruiken een methode die ze Informatie-Entropie noemen.
• De analogie: Stel je voor dat je een detective bent. De oude methode vraagt alleen: "Is dit de dader?" De nieuwe methode vraagt: "Wat denkt de detective over iedereen in de kamer?" Als de detective twijfelt over iedereen, is hij verward (hoge entropie). Als hij heel zeker is over wie de dader is en wie niet, is hij helder (lage entropie).
• Door te kijken naar alle twijfels en zekerheden van de AI, niet alleen het juiste antwoord, kunnen ze veel nauwkeuriger zien welke "boeken" (neuronen) echt belangrijk zijn en welke alleen maar rommel zijn. Dit werkt zelfs als je niet weet wat de "juiste antwoorden" zijn (geen labels nodig).

Stap 2: De "Op maat gemaakte knipbeurt" (Adaptief knippen)
Vaak zeggen mensen: "Knip gewoon 40% van de boeken weg." Maar dat is dom, want sommige kamers hebben misschien 60% rommel en andere maar 10%.

• De nieuwe aanpak: Ze gebruiken een binaire zoekmethode (zoals het raden van een getal tussen 1 en 100).
• De analogie: Je wilt de bibliotheek kleiner maken, maar je wilt niet dat de robot stopt met werken. Je begint met een beetje knippen. Als de robot nog steeds perfect werkt, knip je nog meer. Als hij begint te struikelen, stop je en knip je iets minder. Je doet dit steeds sneller (halverend) totdat je precies de juiste maat hebt gevonden.
• Dit betekent dat ze niet vooraf zeggen "we halen 40% weg", maar ze zoeken automatisch de perfecte balans voor elk onderdeel van de robot.

Stap 3: De "Leermeester" (Kennisoverdracht)
Na het knippen is de robot misschien een beetje verward omdat hij zijn boekenkast heeft opgeruimd.

• De oplossing: Ze laten de oorspronkelijke, gigantische robot (de meester) naar de nieuwe, kleinere robot (de leerling) kijken.
• De analogie: De grote robot zegt: "Kijk, ik zie een hond. Jij moet ook denken aan een hond." De kleine robot leert zo snel mogelijk van de grote robot hoe hij de wereld moet zien, zonder dat hij zelf opnieuw hoeft te leren. Hierdoor wordt de kleine robot bijna net zo slim als de grote, maar hij is veel lichter.

Wat levert dit op?

De resultaten zijn indrukwekkend:

• Gewicht: De modellen worden ongeveer 40% lichter.
• Snelheid: Ze zijn 1,5 keer sneller in het verwerken van beelden.
• Kwaliteit: Na het "leren" van de grote robot, werkt de kleine robot bijna even goed als de originele. In sommige gevallen is hij zelfs nog iets beter!
• Flexibiliteit: Het werkt zelfs voor modellen waar de onderzoekers niet alle geheimen van kennen (zoals de DINOv2-modellen), omdat ze geen "geheime antwoorden" nodig hebben om te knippen.

Kortom:
De onderzoekers hebben een manier gevonden om de "rommel" uit de grootste en zwaarste AI-modellen te halen, zonder dat ze hun intelligentie verliezen. Het is alsof je een zware winterjas omtovert tot een lichte, warme jas die precies past, in plaats van een zware, onhandige mantel. Hierdoor kunnen deze slimme modellen straks op veel meer apparaten draaien, van telefoons tot drones.

Technische samenvatting

Probleemstelling

Grote Vision Transformers (ViT's) tonen uitstekende schaalbaarheid; hun prestaties verbeteren aanzienlijk naarmate het model groter wordt. Echter, deze prestaties gaan gepaard met een enorme hoeveelheid parameters, wat leidt tot exorbitante reken- en geheugeneisen. Dit maakt de kosteneffectieve implementatie in diverse toepassingen moeilijk.
De auteurs analyseren de architectuur van ViT's en ontdekken dat de Multilayer Perceptron (MLP) modules verantwoordelijk zijn voor het grootste deel van de parameters (bijvoorbeeld 81,1% in EVA-CLIP-E). Bestaande pruning-methoden (zoals Taylor-gebaseerde methoden) gebruiken vaak een "one-hot cross-entropy" verliesfunctie om de belangrijkheid van neuronen te beoordelen. Dit heeft twee grote nadelen:

1. Het negeert potentiële voorspellingen voor andere categorieën dan de juiste label, wat leidt tot minder nauwkeurige belangrijkheidsscores.
2. Het vereist toegang tot de oorspronkelijke verliesfunctie of extra modules (zoals de DINO-head bij DINOv2), wat niet altijd beschikbaar is bij gesloten of gedeeltelijk gepubliceerde modellen.

Methodologie: Adaptive MLP Pruning (AMP)

De auteurs stellen een nieuwe methode voor, genaamd Adaptive MLP Pruning (AMP), die bestaat uit drie hoofdfasen:

1. Evaluatie van Neuronbelangrijkheid met Informatie-entropie
In plaats van de traditionele one-hot cross-entropy, introduceren de auteurs een label-vrije informatie-entropie-criterium.

• Principe: De methode gebruikt Taylor-expansie om de gevoeligheid van het model te meten bij het verwijderen van een neuron.
• Innovatie: Om de voorspellingsverdeling volledig te modelleren zonder labels of de oorspronkelijke loss-functie, berekenen ze een "inter-instance similarity matrix" (gelijkheid tussen afbeeldingen in een batch). Op basis van deze gelijkenis wordt een voorspellingskansmatrix gegenereerd via een softmax-operatie.
• Voordeel: De entropie van deze verdeling dient als criterium. Dit maakt het mogelijk om modellen te comprimeren waarvan de trainingsdetails of specifieke head-modules (zoals bij DINOv2 of CLIP) niet volledig openbaar zijn.

2. Adaptieve Pruning via Binaire Zoek
Na het ranken van de verborgen neuronen op basis van hun belangrijkheidsscores, wordt geen vooraf vastgestelde compressie-ratio gebruikt.

• Algoritme: Er wordt een binaire zoekalgoritme toegepast om het optimale aantal te verwijderen neuronen per MLP-module te bepalen.
• Logica: Het algoritme zoekt in een bereik van [0, oorspronkelijke grootte]. Bij elke stap wordt de entropie van het gepreunde model geëvalueerd. Als de toename in entropie (verlies aan onzekerheid/informatie) een bepaalde drempel ($\Delta E$) overschrijdt, wordt de pruning teruggedraaid. Anders wordt er verder gepreund.
• Resultaat: Dit zorgt voor een adaptieve compressie die rekening houdt met de specifieke redundantie van elke module, in plaats van een starre verhouding toe te passen.

3. Kennisdistillatie voor Prestatieherstel
Om de prestaties van het gepreunde model te herstellen, wordt kennisdistillatie toegepast.

• De oorspronkelijke (grote) model fungeert als "leraar" en het gepreunde model als "student".
• Omdat alleen de verborgen lagen van de MLP worden verwijderd en de outputdimensies gelijk blijven, kan kennis direct worden overgedragen via Mean Squared Error (MSE) verlies op de class-token en patch-token embeddings, zonder extra aligneringsmodules.

Belangrijkste Bijdragen

1. Label-vrij Entropie-criterium: Een nieuwe methode voor het beoordelen van neuronbelangrijkheid die volledig gebruikmaakt van de voorspellingsverdeling van het model. Dit elimineert de afhankelijkheid van labels of de oorspronkelijke trainingsloss, waardoor compressie mogelijk is voor modellen met beperkte openbaarmaking (zoals DINOv2).
2. Adaptieve Pruning-strategie: Een methode die binaire zoektoepassing gebruikt om de optimale grootte van MLP-modules dynamisch te bepalen, waardoor vooraf gedefinieerde compressieratio's worden vermeden en redundantie efficiënter wordt verwijderd.
3. Near Lossless Compressie: De methode bereikt een aanzienlijke reductie in parameters en FLOPs met minimaal verlies aan prestaties, zelfs zonder fijnafstelling (finetuning) in sommige scenario's.

Resultaten

De methode is getest op state-of-the-art grote Vision Transformers, waaronder CLIP (OpenCLIP-g, OpenCLIP-G) en DINOv2 (EVA-CLIP-E, EVA-CLIP-8B).

• Compressie: De methode bereikt een reductie van ongeveer 40% in parameters en FLOPs.
• Snelheid: Er is een versnelling van de inferentie van ongeveer 1,5x (gemeten in afbeeldingen per seconde).
• Prestaties:
  • Zelfs zonder finetuning (alleen pruning) behoudt het model veel betere prestaties dan andere pruning-methoden (bijv. +42,7% betere zero-shot classificatie op OpenCLIP-g vergeleken met andere methoden).
  • Na kennisdistillatie (finetuning) herstelt het gepreunde model de prestaties van het originele model volledig, en in sommige gevallen (zoals EVA-CLIP-E) zelfs lichtjes overtreft het de originele prestaties.
  • De resultaten zijn consistent over verschillende benchmarks (ImageNet-1K, ObjectNet, Flickr30K, COCO) en evaluatieprotocollen (zero-shot en kNN).

Betekenis en Toekomstperspectief

Dit paper biedt een krachtige oplossing voor de schaalproblematiek van grote Vision Transformers. Door zich te richten op de MLP-modules en een adaptieve, label-vrije aanpak te gebruiken, maken de auteurs het mogelijk om zeer grote modellen efficiënter te maken zonder de noodzaak van volledige toegang tot de trainingsdata of -architectuur van het originele model.

De auteurs geven aan dat toekomstig werk gericht zal zijn op het uitbreiden van deze adaptieve reductie naar multi-head self-attention modules en het toepassen van de methode op Large Language Models (LLM's).