Separators in Enhancing Autoregressive Pretraining for Vision Mamba

Probleemstelling

Hoewel Mamba (een State Space Model) recentelijk veel aandacht heeft getrokken in het domein van computer vision vanwege zijn efficiëntie bij het verwerken van lange sequenties, blijft het potentieel binnen zelftoezichthoudend leren (self-supervised learning) beperkt.

• Bestaande beperking: Huidige methoden voor autoregressieve pretraining (zoals ARM en MAP) zijn beperkt tot het modelleren van afzonderlijke afbeeldingen als korte sequenties.
• Mismatch: Mamba is bij uitstek geschikt voor lange sequenties, maar de huidige pretraining-tasks benutten deze kracht niet. Er is een disbalans tussen de modelcapaciteit (lange sequenties) en de taakontwerp (korte, enkelvoudige afbeeldingen).
• Uitdaging: Het modelleren van 2D-beelddata als een lange, continue sequentie is complex, vooral omdat er geen natuurlijke volgorde bestaat zoals bij tekst, en het samenvoegen van meerdere afbeeldingen zonder verlies van context of structuur lastig is.

Methodologie: STAR (SeparaTors for AutoRegressive pretraining)

De auteurs introduceren een nieuw raamwerk genaamd STAR, dat de autoregressieve pretraining voor Vision Mamba transformeert naar een lange-sequentie taak door meerdere ongerelateerde afbeeldingen te combineren.

Kerncomponenten:

1. Cluster-gebaseerde Scanning:
   • Afbeeldingen worden opgedeeld in niet-overlappende patches.
   • Ruimtelijk aangrenzende patches worden gegroepeerd tot clusters (bijv. 4x4 patches). Deze clusters vormen de basis-eenheid voor voorspelling, in plaats van individuele pixels of patches.
2. De Separator (S):
   • Om meerdere afbeeldingen in één lange sequentie te plaatsen zonder dat ze door elkaar lopen, wordt een unieke separator-cluster toegevoegd aan het begin van elke afbeelding.
   • Structuur van de separator: Het is een cluster van dezelfde grootte als de beeldclusters. De diagonale patches bevatten vectoren met waarden 1, terwijl alle andere posities 0 zijn (een "Identity" separator).
   • Doel: De separator markeert het begin van een nieuwe afbeelding en zorgt ervoor dat het model de onafhankelijkheid van de afbeeldingen behoudt terwijl het toch een lange context kan verwerken.
3. Architectuur:
   • Encoder: Een unidirectionele Vision Mamba (MambaMLP) die de lange sequentie (separator + clusters van afbeelding 1 + separator + clusters van afbeelding 2, etc.) verwerkt.
   • Decoder: Een lichtgewicht decoder met causale attention. De attention-map is bidirectioneel binnen een cluster, maar unidirectioneel tussen clusters.
   • Voorspelling: Het model voorspelt sequentieel de volgende cluster. De voorspelling voor de laatste cluster van een afbeelding is de separator-cluster van de volgende afbeelding.
4. Aanpassing van het Class Token:
   • In tegenstelling tot eerdere Vision Mamba-modellen (zoals Vim) die het class-token in het midden van de sequentie plaatsen, plaatst STAR het class-token aan het einde van elke afbeelding. Dit optimaliseert de prestaties voor lange-sequentie autoregressieve taken.

Belangrijkste Bijdragen

1. Nieuw Framework (STAR): Het introduceren van separators om meerdere ongerelateerde afbeeldingen te modelleren als één lange sequentie. Dit overwint de beperkingen van single-image modeling en benut volledig de kracht van Mamba voor lange sequenties.
2. Optimalisatie van de Separator: Een grondige studie naar het type, de waarde, de positie en het aantal separators. De auteurs concluderen dat een Identity-separator (diagonaal 1, anders 0) geplaatst voor de clusters (SC-configuratie) de beste resultaten oplevert.
3. Architecturale Aanpassing: Het verplaatsen van het class-token naar het einde van de sequentie, wat leidt tot een efficiëntere sequentiemodellering en betere prestaties na lange pretraining.
4. Efficiëntie: Het bereiken van state-of-the-art resultaten met een aanzienlijk lagere rekencomplexiteit en trainingstijd in vergelijking met bestaande methoden zoals MAE of Contrastive Learning.

Resultaten

De methode is getest op de ImageNet-1k dataset.

• Prestaties:
  • Het STAR-Base model bereikte na 1600 epochs autoregressieve pretraining een nauwkeurigheid van 83.5% (Top-1).
  • Dit is vergelijkbaar met de prestaties van VMamba (een geavanceerd Vision Mamba model) en bijna gelijk aan MAE op ViT-B (83.6%), maar met een veel efficiëntere training.
• Vergelijking met Bestaande Methoden:
  • Tegenover ARM: STAR is nauwkeuriger (83.5% vs 83.2% na 1600 epochs) en simpeler, met een inputlengte die 8x zo lang is als die van ARM.
  • Tegenover Contrastive Learning & MAE: STAR is respectievelijk 6,6x en 1,4x sneller in trainingstijd, terwijl het betere of vergelijkbare nauwkeurigheid behaalt.
  • Lichtgewicht Modellen: STAR presteert zeer competitief tegenover andere lichtgewicht modellen zoals RegNetY-16G en DeiT-B.
• Ablatie Studies:
  • Het gebruik van een separator is cruciaal; zonder separator daalt de prestatie aanzienlijk (van ~82% naar ~78%).
  • De optimale inputlengte bleek rond de 640 tokens (4 afbeeldingen) te liggen voor het gebruikte model, hoewel het model tot 2560 tokens kon verwerken.

Significantie

Dit paper markeert een belangrijke stap in de evolutie van Vision Mamba-modellen:

• Paradigmaverschuiving: Het bewijst dat autoregressieve pretraining niet beperkt hoeft te zijn tot korte, enkelvoudige afbeeldingen, maar kan worden uitgebreid naar lange sequenties van meerdere afbeeldingen.
• Efficiëntie: Het biedt een route naar hoogwaardige prestaties met minder rekenkracht en tijd dan traditionele self-supervised methoden (zoals MAE).
• Toekomstperspectief: De methode biedt een veelbelovende richting voor het ontwerp en de training van lichtgewicht visuele modellen, waarbij de schaalbaarheid van Mamba voor lange contexten volledig wordt benut. De resultaten suggereren dat er nog ruimte is voor verbetering door het verlengen van de pretraining-epochs.