Self-supervised pretraining for an iterative image size agnostic vision transformer

Probleemstelling

Vision Transformers (ViT) zijn de standaard geworden voor zelftoezichtend leren (Self-Supervised Learning of SSL) en vormen de ruggengraat van moderne fundamentele visiemodellen. Ondanks hun succes vertonen standaard ViT-architecturen twee kritieke beperkingen:

1. Rekenkundige inefficiëntie: De complexiteit van ViT schaalt kwadratisch met het aantal patches. Dit betekent dat de verwerkingstijd en het geheugenverbruik explosief toenemen bij hoge resoluties.
2. Resolutie-afhankelijkheid: Standaard ViT's gebruiken een vaste pixelgrootte voor patches (bijv. 16x16 pixels). Als de invoerresolutie na het trainen verandert (bijv. van 224x224 naar 1024x1024), verandert de semantische inhoud binnen een patch. Dit leidt tot een "distribution shift" (verdelingsschifting) en een ineenstorting van de prestaties, waardoor fundamentele modellen zoals DINO beperkt blijven tot lage resoluties.

Bestaande werkzame oplossingen, zoals de foveale transformer van Prisadnikov et al. [44], lossen dit op door iteratief te werken met "multi-zoom" patches, maar deze werden tot nu toe alleen getraind met toezicht (supervised learning) en hadden nog steeds een rekenkundige bottleneck bij het extraheren van patches.

Methodologie

De auteurs introduceren een nieuw framework dat een dynamische, foveale transformer combineert met zelftoezichtend pretraining (gebaseerd op DINO) zonder gebruik te maken van Backpropagation Through Time (BPTT).

1. Architectuur: Iteratieve Foveale Transformer
Het model werkt niet in één keer over het hele beeld, maar iteratief in stappen:

• State Tokens: Het model onderhoudt een set van leerbare "state tokens" die fungeren als intern geheugen van de tot nu toe verwerkte scène.
• Gaze Policy: Een reinforcement learning (RL) agent bepaalt waar het model als volgende moet kijken (fixatiepunt).
• Multi-zoom Patches: Rond het kijkpunt worden concentrische patches geëxtraheerd. Cruciaal is dat de grootte van deze patches wordt gedefinieerd als een ratio van de totale beeldgrootte (top-down), niet als een vaste pixelgrootte. Dit maakt het model resolutie-onafhankelijk.
• Geen BPTT: De state tokens worden tussen iteraties losgekoppeld van de computationele graaf (stop-gradient), waardoor training efficiënt blijft zonder de kosten van BPTT.

2. Nieuwe Zelftoezichtdoelstelling: Sequential-to-Global Distillation
De auteurs passen het DINO-objectief aan voor deze sequentiële context:

• Leerling (Student): Verwerkt een dynamische reeks van 8 lokale "glimpses" (kijkpunten) en update zijn state tokens iteratief.
• Leraar (Teacher): Verwerkt een statisch, globaal beeld (een grote top-down grid) en genereert doel-embeddings.
• Distillatie: De student wordt getraind om zijn state tokens (na elke stap) af te stemmen op de embeddings van de leraar. Hierdoor leert het model om een globaal begrip van de scène op te bouwen door lokaal, iteratief te kijken.

3. Efficiënte Patch-extractie via Integral Images
Om de rekenkundige bottleneck bij het extraheren van patches op te lossen, gebruiken de auteurs Integral Images (Som-gebiedstabellen).

• In plaats van de afbeelding herhaaldelijk te herschalen of te bijsnijden, wordt één keer een integral image berekend.
• Hieruit kunnen rechthoekige patches van elke grootte en positie in constante tijd $O(1)$ worden geëxtraheerd, ongeacht de oorspronkelijke beeldresolutie.
• Dit maakt het totale computebudget per stap echt onafhankelijk van de invoergrootte.

4. Trainingspipeline
Het trainingsproces verloopt in drie fasen:

1. Zelftoezichtend Pretraining: De transformer wordt getraind op ImageNet-1K zonder labels.
2. Linear Probing: De transformer wordt bevroren en een lineaire classificatielaag wordt getraind (met willekeurige kijkpunten).
3. Policy Learning: De transformer en de taaklaag worden bevroren. De gaze policy wordt getraind met GRPO (Group Relative Policy Optimization) om de beste kijkpunten te vinden op basis van de classificatiezekerheid.

Belangrijkste Bijdragen

1. Eerste Zelftoezichtend Pretraining voor Dynamische Foveale Transformers: Het is de eerste keer dat een dergelijk dynamisch model succesvol wordt voorgeprogrammeerd met SSL (DINO) zonder BPTT.
2. Sequential-to-Global Framework: Een nieuwe distillatie-methode die sequentiële lokale observaties koppelt aan een globaal doel, waardoor het model scene-begrip kan opbouwen.
3. Resolutie-onafhankelijke Extractie: De introductie van een dichte foveale grid en het gebruik van integral images zorgt voor een strikt $O(1)$ computekost per stap, ongeacht de beeldgrootte.
4. Dense Foveal Grid: Een verbetering ten opzichte van eerdere werk, waarbij een dicht raster van hoge-resolutie patches rond het kijkpunt wordt toegevoegd om de zelf-attentie te voeden met voldoende detail.

Resultaten

De auteurs evalueren hun model op ImageNet-1K en fijnkorrelige benchmarks (CUB-200-2011, Oxford 102 Flowers):

• Prestaties: Het model bereikt 76,1% Top-1 nauwkeurigheid op ImageNet-1K in de statische ViT-modus, wat vergelijkbaar is met standaard DINO-baselines.
• Dynamische Modus: In de iteratieve modus (8 stappen) met een geleerde gaze policy bereikt het model 75,0% nauwkeurigheid. Dit is aanzienlijk beter dan met willekeurige kijkpunten (71,7%) en toont aan dat het model actief de meest discriminerende gebieden zoekt.
• Resolutie Robuustheid: In tegenstelling tot standaard ViT's, die instorten bij ultra-hoge resoluties (>1024px), behoudt het voorgestelde model zijn prestaties. Het presteert zelfs 10% beter dan een gesuperviseerde baseline bij extreme resoluties.
• Efficiëntie: De latentie (verwerkingstijd) en het GFLOPs-gebruik blijven constant naarmate de beeldgrootte toeneemt, terwijl dit bij standaard ViT's kwadratisch stijgt.

Betekenis en Impact

Dit werk opent de deur voor fundamentele visiemodellen die echt resolutie-onafhankelijk zijn. Het lost het dilemma op tussen de hoge rekenkosten van hoge-resolutie verwerking en de noodzaak van gedetailleerde visuele analyse.

• Het maakt pretraining op zeer hoge resoluties mogelijk zonder de rekenkosten te laten exploderen.
• Het bewijst dat actieve visie (het selectief kijken) effectief kan worden geleerd via zelftoezicht, wat relevant is voor robotica en systemen met beperkte rekenkracht.
• De methode biedt een schaalbare route voor toekomstige multimodale modellen die omgaan met variabele invoergroottes zonder kwaliteitsverlies.

Samenvattend introduceert deze paper een doorbraak in de efficiëntie en flexibiliteit van Vision Transformers, waardoor ze geschikt worden voor real-world toepassingen met wisselende en hoge resoluties, terwijl ze de state-of-the-art prestaties behouden.