Locality-Attending Vision Transformer

Dit artikel introduceert LocAtViT, een eenvoudige toevoeging aan Vision Transformers die door het moduleren van self-attention met een leerbaar Gaussisch kernel de prestaties op segmentatietaken aanzienlijk verbetert zonder de bestaande classificatiecapaciteiten te schaden.

De kern

Locality-Attending Vision Transformer (LocAtViT): Een simpele uitleg

Stel je voor dat je een kunstwerk bekijkt. Een Vision Transformer (ViT) is als een kunstcriticus die heel goed is in het begrijpen van het gehele schilderij. Hij kan zien dat er een bos is, een rivier en een zonsondergang, en hij kan dat perfect benoemen. Hij kijkt naar alles tegelijk, van links naar rechts, van boven naar beneden. Dit maakt hem een meester in het herkennen van wat er op de foto staat (bijvoorbeeld: "dit is een schoolbus").

Maar er is een probleem: als je deze kunstcriticus vraagt om precies te tekenen waar de wielen van de bus zitten, of om de randen van de bomen na te trekken, faalt hij een beetje. Omdat hij zo gefocust is op het "grote plaatje", verliest hij de fijne details. Hij ziet de bus, maar hij ziet niet precies waar de bus eindigt en de weg begint. Dit is een groot probleem voor taken zoals segmentatie (het precies inkleuren van elk object in een foto).

De auteurs van dit paper hebben een slimme oplossing bedacht: LocAtViT. Ze noemen het een "add-on" (een extraatje) dat je op de bestaande kunstcriticus plakt, zonder hem te vervangen.

Hier is hoe het werkt, vertaald naar alledaagse analogieën:

1. De "Gluurder" vs. De "Buurtbewoner" (GAug)

In de originele ViT kijkt elk stukje van de foto (een "patch") naar elk ander stukje in de hele foto. Het is alsof je in een drukke stad staat en naar iedereen in de stad kijkt, van de overkant van de wereld tot de persoon naast je. Dat is handig voor het grote plaatje, maar je mist de details van je directe omgeving.

LocAtViT voegt een Gaussische kern toe.

• De analogie: Stel je voor dat je een magische bril opzet. Met deze bril zie je nog steeds de hele stad, maar je ziet je directe buren (de mensen die naast je staan) heel scherp en helder. De mensen die verder weg staan, worden iets vager.
• In het kort: Het model leert om extra goed te kijken naar de stukjes die direct naast elkaar liggen. Dit zorgt ervoor dat de randen van objecten (zoals de wielen van de bus) scherper worden, zonder dat het model zijn vermogen om de hele scène te begrijpen verliest.

2. De "Groepsleider" vs. De "Persoonlijke Coach" (PRR)

In de originele ViT is er een speciale token (een soort "hoofd") genaamd [CLS]. Dit is de hoofd-leraar die het eindantwoord geeft (bijvoorbeeld: "Dit is een schoolbus"). Alle andere stukjes van de foto (de "patches") werken hard, maar aan het einde krijgt alleen de hoofd-leraar een cijfer. De andere stukjes krijgen geen directe feedback.

• Het probleem: Omdat de andere stukjes nooit een cijfer krijgen, leren ze niet goed hoe ze hun eigen specifieke taak moeten doen. Ze worden allemaal hetzelfde, alsof ze allemaal proberen de "hoofd-leraar" na te bootsen. Voor het precies inkleuren van objecten is dat slecht; je wilt dat elk stukje van de foto zijn eigen identiteit behoudt.

LocAtViT introduceert Patch Representation Refinement (PRR).

• De analogie: Stel je voor dat de hoofd-leraar (de [CLS]) een vergadering houdt met de klas. In de oude versie luistert de leraar alleen naar zichzelf en negeert hij de klas. In de nieuwe versie (LocAtViT) doet de leraar een korte, slimme check-in met de klas voordat hij het antwoord geeft. Hij vraagt: "Hoe ziet het eruit bij jou, patchje linksboven? En jij, rechtsonder?"
• Het resultaat: Hierdoor krijgen alle stukjes van de foto (de patches) weer hun eigen stem en hun eigen identiteit. Ze leren dat ze belangrijk zijn voor de uiteindelijke taak. Dit zorgt ervoor dat het model niet alleen weet wat er op de foto is, maar ook precies waar het zit.

Waarom is dit zo cool?

1. Het is een "plug-and-play" oplossing: Je hoeft geen hele nieuwe architectuur te bouwen. Het is als het toevoegen van een nieuwe lens aan een bestaande camera. Je kunt het op bijna elke ViT plakken.
2. Twee vliegen in één klap: Het model wordt beter in het precies inkleuren van objecten (segmentatie), maar het wordt niet slechter in het herkennen van de foto (classificatie). Sterker nog, vaak wordt het zelfs beter in beide!
3. Het werkt ook op de "grote" modellen: De auteurs tonen aan dat dit werkt op de enorme modellen die vandaag de dag gebruikt worden (foundation models), wat betekent dat we in de toekomst betere AI kunnen hebben die niet alleen weet wat er op een foto staat, maar ook precies weet waar de randen zijn.

Samenvattend:
De auteurs hebben een slimme manier gevonden om Vision Transformers te "leren kijken" naar hun directe omgeving, terwijl ze hun vermogen om het grote plaatje te zien behouden. Het is alsof je een kunstcriticus een bril geeft die hem scherp maakt op details, zonder zijn brede blik te verliezen. Hierdoor wordt hij niet alleen een betere kunstkenner, maar ook een meester in het tekenen van de contouren.

Technische samenvatting

Probleemstelling

Vision Transformers (ViT) hebben zich bewezen als krachtige modellen voor beeldclassificatie door gebruik te maken van globale zelf-attention (self-attention) om lange-afstandsafhankelijkheden te vangen. Echter, deze globale focus vormt een uitdaging voor dichte voorspellingstaken (zoals semantische segmentatie), die nauwkeurige lokalisatie en fijnkorrelige ruimtelijke details vereisen.

De auteurs identificeren twee kernproblemen in standaard ViT's:

1. Verlies van lokale structuur: De globale attention-mechanisme kan fijne ruimtelijke details "verduisteren". Empirisch bewijs toont aan dat in een voor ViT getraind model, patch-tokens geleidelijk hun unieke lokale structuur verliezen en steeds meer aligneren met de globale [CLS]-token.
2. Gradient-flow issues: Bij klassieke ViT-training wordt alleen de [CLS]-token gebruikt voor het verlies (loss). Dit betekent dat de output van de ruimtelijke patches geen directe supervisie ontvangt. Voor segmentatie is dit funest, omdat de representaties van deze patches cruciaal zijn. Bovendien leidt het gebruik van Global Average Pooling (GAP) als alternatief tot een uniforme gradient-flow, wat betekent dat achtergrondpatches even zwaar worden gewogen als objectpatches, wat resulteert in suboptimale representaties voor segmentatie.

Methodologie: LocAtViT

Het paper introduceert LocAtViT, een modulaire "add-on" die bestaande ViT-architecturen verbetert zonder de trainingsregime of de basisarchitectuur fundamenteel te veranderen. De methode bestaat uit twee componenten:

1. Gaussian-Augmented (GAug) Attention

Om de aandacht van het model te biasen naar lokale omgevingen, wordt een leerbare Gaussische kern toegevoegd aan de attention-logits.

• Mechanisme: In plaats van alleen te vertrouwen op de standaard attention-matrix ($qk^T/\sqrt{d}$), wordt een supplementaire matrix $S$ toegevoegd.
• Gaussische Kern: Voor elke patch wordt een Gaussische kern berekend die afloopt met de afstand tot de bronpatch. De variantie ($\sigma^2$) van deze kern is niet statisch, maar wordt leerbaar voorspeld op basis van de query-vector van de patch (via een kleine lineaire laag). Dit stelt het model in staat om dynamisch te bepalen hoe "lokaal" of "globaal" de aandacht moet zijn per patch.
• Schaling: Een leerbaar gewicht ($\alpha$) schaal de Gaussische bijdrage, zodat deze in balans blijft met de originele attention-logits. Dit zorgt voor een "soft" locality-bias die data-afhankelijk is, in plaats van een harde constraint.

2. Patch Representation Refinement (PRR)

Deze component lost het probleem van de gradient-flow en de kwaliteit van de patch-representaties op, specifiek voor de laag vlak voor de classificatie-head.

• Probleem: Standaard ViT's aggregeren patches vaak via Global Average Pooling (GAP) of gebruiken alleen de [CLS]-token, wat de unieke bijdrage van individuele patches onderdrukt.
• Oplossing: PRR introduceert een parameterloze multi-head self-attention laag direct voor de classificatie-head. Deze laag herschikt de tokens en voert een attention-operatie uit waarbij elke patch informatie verzamelt van alle andere patches op een niet-uniforme manier.
• Effect: Dit zorgt ervoor dat supervisie en gradients effectief worden doorgegeven aan de ruimtelijke patch-outputs, waardoor deze representaties "betekenisvol" worden voor dichte taken, zonder extra parameters toe te voegen.

Belangrijkste Bijdragen

1. Modulaire Architectuur: LocAtViT is een lichtgewicht toevoeging die compatibel is met bestaande ViT-varianten (zoals ViT, Swin, RegViT) en foundation modellen.
2. Dual-Optimaliteit: Het model verbetert de prestaties voor segmentatie aanzienlijk zonder de classificatieprestaties te offeren; in veel gevallen verbetert het zelfs de classificatie-accuraatheid.
3. Gradient-Flow Innovatie: Het introduceert PRR als een oplossing voor het vaak over het hoofd geziene probleem van gradient-flow naar ruimtelijke tokens in ViT's.
4. Data-Dependent Locality: In tegenstelling tot veel andere methoden die vaste vensters of statische posities gebruiken, gebruikt GAug een dynamische, query-afhankelijke Gaussische kern.

Resultaten

De auteurs evalueren LocAtViT op drie segmentatie-benchmarks (ADE20K, PASCAL Context, COCO Stuff) en diverse classificatiedatasets (ImageNet-1K, mini-ImageNet, CIFAR-100).

• Segmentatie: Er worden aanzienlijke verbeteringen geboekt. Bijvoorbeeld, voor de ViT Tiny op ADE20K stijgt de mIoU met +6,17% (van 17,30% naar 23,47%). Voor ViT Base is de stijging +4,24%.
• Classificatie: De prestaties op ImageNet-1K blijven gelijk of verbeteren licht (bijv. +1,55% voor ViT Tiny).
• Robuustheid: De methode werkt ook goed op andere architecturen zoals Swin Transformer en RegViT, hoewel de winst kleiner is bij modellen die al beperkte attention-mechanismen hebben (zoals windowed attention in Swin).
• Foundation Models: De methode verbetert ook de prestaties van zelf-supervisie modellen (DINO) en toont aan dat de representaties beter geschikt zijn voor "in-context" scene understanding (geëvalueerd via Hummingbird retrieval).
• Kwalitatieve Analyse: Attention maps tonen aan dat LocAtViT meer geconcentreerde en coherente activering heeft op objectkenmerken (bijv. een schoolbus) vergeleken met de verspreide aandacht van standaard ViT.

Betekenis en Impact

Dit werk is significant omdat het een brug slaat tussen de kracht van globale context in ViT's en de noodzaak van lokale details voor pixel-accurate taken.

• Efficiëntie: Het lost het probleem op zonder complexe nieuwe architecturen te bouwen of de trainingsregime te veranderen.
• Toekomstgericht: Gezien de wijdverbreide adoptie van ViT als backbone voor foundation modellen (zoals CLIP, DINO), biedt LocAtViT een eenvoudige, "plug-and-play" upgrade om deze modellen direct bruikbaarder te maken voor dichte voorspellingstaken zonder extra fine-tuning of decoder-lagen.
• Filosofie: Het paper pleit voor een "segmentation-in-mind" pretraining, waarbij ViT's van nature zowel globale als lokale informatie behouden, in plaats van dat dit later moet worden geforceerd door complexe aanpassingen.

Kortom, LocAtViT demonstreert dat minimale, weloverwogen wijzigingen in de attention-mechanismen en de token-aggregatie de ViT-architectuur aanzienlijk kunnen verbeteren voor een breder scala aan computer visie-taken.