Thicker and Quicker: A Jumbo Token for Fast Plain Vision Transformers

Probleemstelling

Vision Transformers (ViTs) zijn momenteel de standaardarchitectuur voor veel computerzichttaken vanwege hun veelzijdigheid en nauwkeurigheid. Ze worden gebruikt in fundamentele modellen (zoals DINOv2), taal-gealigneerde modellen (CLIP) en segmentatiemodellen (SAM). Echter, standaard "plain" ViTs (die alleen gebruikmaken van attention en geen hiërarchie hebben) zijn vaak traag en niet altijd praktisch wanneer efficiëntie cruciaal is.

Bestaande oplossingen voor snellere ViTs hebben twee grote nadelen:

1. Hybride architecturen: Ze combineren ViTs met convoluties of hiërarchische structuren (zoals EfficientViT of MobileNetV4). Dit verliest de generieke voordelen van ViTs, zoals compatibiliteit met state-of-the-art (SOTA) zelftoezichtende leermethodes (SSL), flexibele verwerking van verschillende invoervormen (bijv. tijdreeksen) en token-dropping.
2. Token-verkleining: Het verkleinen van tokens om rekentijd te besparen gaat ten koste van de nauwkeurigheid.

De kernuitdaging is dus: hoe kan men ViTs sneller maken zonder de unieke architecturale eigenschappen (alleen attention, niet-hiërarchisch) te verliezen die ze zo flexibel en krachtig maken?

Methodologie: De Jumbo Token

De auteurs stellen een nieuwe architectuur voor genaamd Jumbo, die de snelheid van ViTs verbetert door de breedte (width) asymmetrisch te schalen in plaats van uniform.

Kernprincipes:

1. Asymmetrische Token-breedte: In plaats van alle tokens even breed te maken, introduceert Jumbo een speciale "Jumbo token" die $J$ keer breder is dan de standaard patch-tokens. Deze vervangt de traditionele CLS-token.
2. Split-Attention Mechanisme:
   • Voor de self-attention laag wordt de brede Jumbo token opgesplitst in $J$ kleinere tokens (elk met dezelfde breedte als de patch-tokens).
   • Deze gesplitste tokens worden samen met de patch-tokens door de Multi-Head Self-Attention (MHSA) laag gevoerd.
   • Na de attention worden de gesplitste Jumbo tokens weer samengevoegd (geconcateneerd) tot één brede token.
3. Specifieke FFN (Feed-Forward Network):
   • De samengevoegde Jumbo token wordt verwerkt door een eigen, breder FFN.
   • Cruciaal voor efficiëntie: Dit Jumbo FFN wordt niet per laag herhaald met nieuwe parameters. De parameters van het Jumbo FFN worden gedeeld over alle lagen van het netwerk. Dit vermindert het geheugengebruik aanzienlijk.
   • De patch-tokens worden verwerkt door het standaard, gedeelde FFN van de ViT.
4. Behoud van "Plain ViT" Eigenschappen:
   • De architectuur blijft "attention-only" en "non-hierarchical".
   • Dit betekent dat Jumbo volledig compatibel blijft met bestaande ViT-methoden zoals token dropping (voor efficiëntie), masking (voor MAE pre-training), en diverse downstream heads (segmentatie, detectie).

Belangrijkste Bijdragen

• Nieuwe Architectuur: Introductie van de Jumbo token die de globale representatiecapaciteit verhoogt zonder de snelheid te vertragen, dankzij het splitsen voor attention en het delen van FFN-parameters.
• Efficiëntie zonder Verlies van Generaliteit: Jumbo is de eerste architecturen die de snelheid van gespecialiseerde, efficiënte modellen (zoals EfficientViT) benadert of overtreft, maar wel de volledige compatibiliteit behoudt met het "plain ViT" ecosysteem.
• Uitgebreide Validatie: Het paper toont aan dat Jumbo werkt over verschillende schalen (Nano tot Large), verschillende taken (classificatie, segmentatie, tijdreeksen) en verschillende leerparadigma's (supervised, self-supervised, test-time adaptation).

Resultaten

De experimenten tonen aan dat Jumbo de Pareto-grens (de optimale balans tussen snelheid en nauwkeurigheid) verbetert ten opzichte van bestaande methoden:

1. Beeldclassificatie (ImageNet-1K & 21K):

   • Jumbo presteert beter dan ViT+Registers (de huidige SOTA voor plain ViTs) met een verbetering van 0,1% tot 13% op ImageNet-1K, afhankelijk van het modelgrootte.
   • Op de uitdagendere ImageNet-21K dataset is de verbetering 1,2% tot 3,1%.
   • Jumbo is 1,9x sneller dan ViT+Registers bij gelijke nauwkeurigheid.
   • Het presteert beter dan gespecialiseerde efficiënte modellen zoals EfficientViT, SHViT en MobileNetV4, terwijl het de ViT-voordelen behoudt.

2. Semantische Segmentatie (ADE20K):

   • Jumbo behaalt een verbetering van 1,9% tot 3,1% in mIoU ten opzichte van ViT+Registers.

3. Zelftoezichtend Leren (Self-Supervised Learning - MAE):

   • Bij pre-training met MAE op ImageNet-1K verbetert Jumbo de linear probing nauwkeurigheid met 4,9%.
   • Een ViT-Base+Jumbo model bereikt dezelfde nauwkeurigheid als een ViT-Large model, maar met 2,3x minder parameters, 3,5x minder FLOPs en 3,1x hogere doorvoer.

4. Robuustheid en Test-Time Adaptation (TTA):

   • Jumbo is robuuster tegen corrupties (ImageNet-C) en presteert beter bij test-time adaptatie (met SAR), met een verbetering van 5,2%.

5. Tijdreeksen en Taal:

   • Jumbo generaliseert succesvol naar tijdreeksdata (UCR/UEA benchmarks), waar het de beste rangschikking behaalt.
   • Preliminaire resultaten op taal-taken (MLM en image-caption retrieval) tonen ook verbeteringen.

Significantie

De betekenis van dit werk ligt in het oplossen van het fundamentele compromis tussen snelheid en flexibiliteit in Vision Transformers.

• Ecosysteem Compatibiliteit: Veel efficiënte architecturen zijn te gespecialiseerd om te profiteren van de snelle innovaties in het ViT-ecosysteem (zoals nieuwe SSL-methodes of token-dropping). Jumbo is de eerste die deze efficiëntie biedt zonder die compatibiliteit op te offeren.
• Schaalbaarheid: De methode werkt effectief van zeer kleine modellen (Nano) tot grote modellen (Large), wat het ideaal maakt voor diverse toepassingen, van mobiele apparaten tot datacenters.
• Eenvoud: De architectuur is opvallend simpel (geen convoluties, geen hiërarchie), wat het implementeren en aanpassen voor nieuwe data-types (zoals 3D of video) zeer eenvoudig maakt.

Kortom, Jumbo bewijst dat het mogelijk is om ViTs sneller en nauwkeuriger te maken door slimme, asymmetrische schaling van capaciteit, in plaats van het introduceren van complexe hybride structuren.