MaskInversion: Localized Embeddings via Optimization of Explainability Maps

Probleemstelling

Contrastieve vision-language foundation modellen (zoals CLIP) hebben enorme prestaties geleverd in het globale afstemmen van visuele en taalkundige concepten. Ze trainen echter op het matchen van globale [CLS]-tokens (image class tokens) met tekst, wat resulteert in representaties die de gehele afbeelding samenvatten. Hierdoor worstelen deze modellen met taken die precieze lokalisatie vereisen of het herkennen van specifieke beeldregio's.

Bestaande oplossingen hebben beperkingen:

• Naïeve methoden: Het bijsnijden (cropping) van een regio verwijdert vaak belangrijke context, en het aggregeren van lokale tokens leidt niet altijd tot een goed uitgelijnde representatie.
• Finetuning-benaderingen: Methoden zoals AlphaCLIP vereisen het finetunen van het model met miljoenen mask-text paren, wat rekenkundig duur is en de generalisatie van het oorspronkelijke model kan beïnvloeden.
• Input-modificaties: Methoden die de inputafbeelding manipuleren (bijv. met gekleurde kaders of blur) kunnen leiden tot een domeinverschil (domain gap) met de trainingsdata.

Het doel is om een methode te ontwikkelen die test-tijd (zonder het basismodel aan te passen) een specifieke embedding leert voor een regio die is gespecificeerd door een masker, terwijl de rijke wereldkennis van het gefronste foundation model behouden blijft.

Methodologie: MaskInversion

MaskInversion is een methode die een lokaal georiënteerde embedding (een token) leert voor een specifieke regio in een afbeelding, gedefinieerd door een query-masker. Het proces verloopt als volgt:

1. Initialisatie:

   • De methode start met het initialiseren van een learnable embedding token ($LET_m$) door de globale [CLS]-token van het gefronste vision encoder (bijv. CLIP) te kopiëren.
   • De inputafbeelding wordt slechts één keer door het model gestuurd (forward pass) om de basisactivaties te krijgen.

2. Optimalisatie via Explainability Maps:

   • In plaats van het model te trainen, wordt de embedding token iteratief geoptimaliseerd.
   • Het doel is dat de explainability map (een warmtekaart die aangeeft waar het model naar kijkt) van deze token overeenkomt met het query-masker.
   • De loss functie is gebaseerd op de Dice Loss, die de overlap meet tussen de gegenereerde explainability map en het binair masker.
   • De token wordt bijgewerkt via gradient descent om de discrepantie tussen de map en het masker te minimaliseren.

3. Regularisatie (Optioneel):

   • Om te voorkomen dat de lokale embedding volledig losraakt van de globale context van de afbeelding, wordt een regularisatie-loss toegevoegd. Deze forceert de geleerde token om dicht bij de oorspronkelijke globale [CLS]-token te blijven.
   • Een hyperparameter ($\alpha$) reguleert de balans tussen lokaal detail en globale context.

4. Gradient Decomposition (Efficiëntie):

   • Het berekenen van de gradient voor de explainability map in elke iteratie is rekenkundig duur (vereist tweede-orde afgeleiden).
   • De auteurs stellen een gradient decomposition strategie voor. Omdat de embedding token onafhankelijk is van de modelactivaties tijdens de forward pass, kan de gradient van de score ten opzichte van de activaties ($\frac{\partial \bar{z}}{\partial A}$) één keer worden berekend.
   • Voor elke iteratie wordt de update vervolgens berekend als een eenvoudig dot-product tussen deze vooraf berekende gradient en de huidige token. Dit elimineert de noodzaak voor herhaalde backpropagation en versnelt het proces aanzienlijk, vooral bij meerdere maskers per afbeelding.

Belangrijkste Bijdragen

1. Test-tijd Lokale Embeddings: Een methode om een token te leren die de kenmerken van een specifieke regio vastlegt zonder het foundation model te finetunen. De token fungeert als een "drop-in replacement" voor de [CLS]-token in downstream taken.
2. Gradient Decomposition: Een techniek die de computationele kosten van het optimaliseren van meerdere maskers drastisch verlaagt door de gradient-berekening te decomponeren.
3. Brede Toepasbaarheid: De methode wordt getest op diverse zero-shot taken, waaronder referentiële expressie-comprehensie, open-vocabulary classificatie, lokale beeldbeschrijving (captioning) en lokale beeldgeneratie.

Resultaten

De auteurs evalueren MaskInversion op diverse datasets (PascalVOC, MSCOCO, RefCOCO, RefCOCO+, OpenImagesV7) en vergelijken het met state-of-the-art (SOTA) methoden.

• Referentiële Expressie Retrieval: MaskInversion behaalt SOTA-prestaties op RefCOCO en RefCOCO+, zelfs met kleinere modellen (ViT-B/32), en schaalt uitstekend met grotere backbones (ViT-H/14), waarbij het de beste resultaten behaalt.
• Class Retrieval (Zero-shot): Op semantische segmentatiedatasets (PascalVOC, PascalContext) en instance datasets (COCO, OpenImagesV7) overtreft MaskInversion alle concurrenten, inclusief AlphaCLIP (dat wel gefinetuned is). MaskInversion presteert beter zonder enige training van het basismodel.
• Lokale Captioning: Door de lokale embedding te gebruiken als input voor CLIPCap, verdubbelt de nauwkeurigheid van het genereren van beschrijvingen voor specifieke regio's vergeleken met standaard CLIP, en presteert het aanzienlijk beter dan AlphaCLIP.
• Lokale Diffusie: De gegenereerde embeddings kunnen worden gebruikt in diffusiemodellen om variaties van specifieke objecten in een afbeelding te genereren, waarbij de focus correct ligt op het gemaskerde gebied.
• Efficiëntie: De gradient decomposition strategie reduceert de runtime aanzienlijk bij het verwerken van meerdere maskers (bijv. 50 maskers) in vergelijking met de "vanilla" aanpak.

Betekenis en Impact

MaskInversion biedt een krachtige, training-free oplossing voor het probleem van lokale visuele taalbegrip in foundation modellen.

• Efficiëntie: Het elimineert de noodzaak voor kostbare finetuning of het verzamelen van enorme datasets met mask-text paren.
• Flexibiliteit: Het werkt met elk bestaand gefronst vision-language model en kan direct worden toegepast op diverse downstream taken als een plug-in module.
• Kwaliteit: Het behoudt de sterke wereldkennis van het foundation model terwijl het tegelijkertijd de aandacht van het model richt op specifieke regio's, wat resulteert in superieure prestaties ten opzichte van methoden die de input manipuleren of het model finetunen.

Kortom, MaskInversion demonstreert dat het mogelijk is om geavanceerde, contextbewuste lokale representaties te extraheren uit grote foundation modellen door simpelweg de explainability maps te optimaliseren, wat een nieuwe weg opent voor zero-shot visuele analyse.