Accelerating Vision Transformers with Adaptive Patch Sizes

Probleemstelling

Vision Transformers (ViTs) zijn de dominante architectuur voor visuele herkenning, maar ze hebben een fundamenteel schaalprobleem. Traditionele ViTs verdelen een invoerimage in een uniforme raster van patches van vaste grootte (bijv. 16x16 pixels), ongeacht de inhoud van het beeld. Dit leidt tot twee belangrijke nadelen:

1. Quadratische complexiteit: De kosten van self-attention groeien kwadratisch met de lengte van de token-sequentie. Hoogresolutiebeelden genereren dus extreem lange sequenties.
2. Inefficiëntie door redundantie: Veel beelden bevatten grote gebieden met weinig detail (bijv. een blauwe lucht of een witte muur). Het toepassen van dezelfde hoeveelheid rekenkracht op deze homogene gebieden als op complexe gebieden (bijv. een vogelkop met veel textuur) is inefficiënt. Bestaande oplossingen zoals token-pruning of token-merging hebben vaak vaste verhoudingen of introduceren onregelmatige vormen die de snelheid in de praktijk tenietdoen door padding en inefficiënte kernels.

Methodologie: Adaptive Patch Transformer (APT)

De auteurs stellen Adaptive Patch Transformer (APT) voor, een methode die de patch-grootte dynamisch aanpast op basis van de visuele complexiteit van het beeld.

1. Content-Aware Patching (Entropy-based):
In plaats van een vaste patch-grootte te gebruiken, verdeelt APT het beeld in een hiërarchische structuur (vergelijkbaar met een quadtree).

• Het systeem berekent de entropie van patches op meerdere schalen. Entropie dient als maatstaf voor de informatie-inhoud en voorspelbaarheid; lage entropie betekent hoge redundantie (homogeen gebied), hoge entropie betekent complexe details.
• Gebieden met lage entropie (homogeen) krijgen grotere patches toegewezen.
• Gebieden met hoge entropie (complex) krijgen kleinere patches toegewezen.
• Dit resulteert in een aanzienlijke reductie van het totale aantal tokens zonder verlies van belangrijke informatie.

2. Patch Aggregatie en Embedding:
Een uitdaging bij variabele patch-groottes is het projecteren van deze verschillende groottes naar een uniforme embedding-dimensie ($d_{embed}$) voor de transformer. APT lost dit op met een innovatieve aanpak:

• Resizing + Sub-patch Embedding: Grote patches worden geresized naar de basisgrootte ($p \times p$) én tegelijkertijd opgesplitst in hun onderliggende sub-patches.
• Zero-Initialized MLP: De embedding van de grote patch wordt berekend door de geaggregeerde embeddings van de sub-patches (via een convolutielaag) te combineren met de embedding van de geresized patch. Cruciaal is dat deze combinatie een nul-geinitialiseerde MLP (ZeroMLP) gebruikt.
  • Voordeel: Hierdoor kan het model direct worden toegepast op een reeds getrainde ViT zonder dat de prestaties direct instorten. De ZeroMLP zorgt ervoor dat het model geleidelijk de hogeresolutie-details kan integreren tijdens het fine-tunen, wat convergentie in slechts 1 epoch mogelijk maakt.

3. Dynamische Invoer en Sequence Packing:
Omdat het aantal tokens per beeld varieert, gebruikt APT sequence packing. Tokens van meerdere beelden in een batch worden samengevoegd tot één lange sequentie met een blokgewijze masker (block-diagonal mask). Dit zorgt ervoor dat tokens alleen naar tokens van hetzelfde beeld kijken. Deze techniek is compatibel met geoptimaliseerde attention-backends zoals FlashAttention, wat essentieel is voor de snelheidswinst.

4. Toepassing op Dichte Taken:
Voor taken zoals objectdetectie en semantische segmentatie, waar een feature map met dezelfde aspectratio als het beeld nodig is, herhaalt APT de embeddings van grote patches om een volledig differentieerbare feature map te genereren die kan worden opgeschaald.

Belangrijkste Bijdragen

1. Introductie van APT: Een nieuwe architectuur die ViTs versnelt door content-bewuste patch-groottes te gebruiken, wat leidt tot minder tokens en lagere rekenkosten.
2. Snelheid zonder Prestatieverlies: APT bereikt een drastische versnelling in training en inferentie terwijl de nauwkeurigheid behouden blijft, zelfs op hoge resoluties.
3. Snelle Convergentie: Dankzij de Zero-MLP kan APT worden toegepast op bestaande, voorgeprogrammeerde ViT-modellen en convergeren in slechts één epoch van fine-tuning.
4. Brede Toepasbaarheid: Het werkt niet alleen voor beeldclassificatie, maar ook voor complexe taken zoals Visual Question Answering (VQA), objectdetectie en semantische segmentatie.

Resultaten

De auteurs testen APT op verschillende modellen (ViT-B, ViT-L, ViT-H) en taken:

• ImageNet Classificatie:

  • Doorvoer: Een toename van 40% voor ViT-L en 50% voor ViT-H.
  • Trainingstijd: Tot 30% snellere training en inferentie.
  • Nauwkeurigheid: APT behaalt dezelfde nauwkeurigheid als de baseline (bijv. 86.1% voor ViT-L op 336x336) na slechts 1 epoch fine-tuning. Zelfs zonder training is de daling in nauwkeurigheid minimaal.
  • Vergelijking: APT presteert beter dan bestaande token-merging methoden (zoals ToMe, EViT) en is compatibel met FlashAttention, wat deze baselines vaak niet zijn.

• Downstream Taken:

  • Visual QA (LLaVA): Een doorvoerverhoging van 23% met behoud of zelfs verbetering van de prestaties op benchmarks zoals VQAv2 en GQA.
  • Objectdetectie (COCO): Op 1536x1536 resolutie wordt 30% minder tokens gebruikt met een 30% versnelling in training/inferentie, terwijl de mAP (mean Average Precision) gelijk blijft.
  • Semantische Segmentatie (ADE20K): Een reductie van 28-32% in tokens met behoud van pixel-nauwkeurigheid.

Betekenis en Impact

Dit paper biedt een significante doorbraak in de efficiëntie van Vision Transformers. Door de vastheid van patch-groottes los te laten en in te spelen op de inhoud van het beeld, lost APT het schaalprobleem van ViTs op bij hoge resoluties.

De belangrijkste implicaties zijn:

• Kostenefficiëntie: Het verlaagt de rekenkosten (FLOPs) en de wandkloktijd voor het trainen van state-of-the-art modellen aanzienlijk.
• Toegankelijkheid: Het maakt het mogelijk om grotere modellen (zoals ViT-H) op hogere resoluties te trainen en af te leiden op hardware die daarvoor te beperkt was.
• Flexibiliteit: De methode is plug-and-play voor bestaande modellen en vereist geen training vanaf nul, wat de adoptie in de industrie en onderzoeksgemeenschap versnelt.

Samenvattend toont APT aan dat adaptieve tokenisatie, geïnspireerd op taalmodellen (zoals Byte-Pair Encoding), ook voor visuele taken superieur is aan vaste splitsingen, mits correct geïmplementeerd om compatibiliteit met moderne hardware-optimalisaties te behouden.