MaskInversion: Localized Embeddings via Optimization of Explainability Maps

1. Il Problema

I modelli fondazionali visione-linguaggio (come CLIP) hanno ottenuto risultati eccezionali nell'allineamento globale tra immagini e testo. Tuttavia, questi modelli sono addestrati per allineare l'intera immagine con il suo testo descrittivo (tramite il token [CLS]), catturando solo informazioni globali. Di conseguenza, faticano in compiti che richiedono una localizzazione precisa o il riconoscimento di regioni specifiche di un'immagine.

Le soluzioni naive, come ritagliare la regione di interesse (che perde il contesto) o aggregare i token locali sopra una maschera (che spesso non produce rappresentazioni allineate correttamente), non sono ottimali. Metodi precedenti che tentano di risolvere questo problema spesso richiedono:

• L'addestramento (fine-tuning) del modello fondazionale con dati specifici (es. AlphaCLIP).
• La modifica diretta dell'immagine di input (es. aggiungere cerchi colorati o sfocature) per guidare l'attenzione.

L'obiettivo è quindi generare un embedding localizzato per una regione specifica di un'immagine, definita da una maschera, senza modificare i pesi del modello fondazionale pre-addestrato e mantenendo la capacità di utilizzare il modello in uno scenario zero-shot.

2. Metodologia: MaskInversion

Il metodo proposto, MaskInversion, è un approccio di ottimizzazione a tempo di test (test-time) che apprende un token di embedding localizzato per una data maschera di query.

Processo Principale

1. Inizializzazione: Si parte inizializzando un token di embedding localizzato ($LET_m$) copiando il token [CLS] globale prodotto dal modello fondazionale (es. CLIP) per l'immagine intera.
2. Mappa di Spiegabilità (Explainability Map): Si calcola la mappa di spiegabilità derivata dal modello pre-addestrato per questo token. Questa mappa indica quali regioni dell'immagine influenzano maggiormente l'embedding corrente.
3. Ottimizzazione Iterativa: L'obiettivo è aggiornare il vettore $LET_m$ in modo che la sua mappa di spiegabilità corrisponda il più possibile alla maschera di query fornita.
   • La funzione di perdita principale è una Soft Dice Loss che misura la discrepanza tra la mappa di spiegabilità generata e la maschera binaria di input.
   • Si utilizza la discesa del gradiente per aggiornare solo il vettore $LET_m$, mantenendo il modello fondazionale frozen (congelato).
4. Regolarizzazione (Opzionale): Per compiti che beneficiano del contesto globale (es. riferimenti complessi), viene introdotta una perdita di regolarizzazione ($L_{reg}$) che allinea l'embedding localizzato all'embedding globale dell'immagine, bilanciando informazioni locali e globali tramite un iperparametro $\alpha$.

Ottimizzazione Computazionale: Decomposizione del Gradiente

Calcolare i gradienti per le mappe di spiegabilità ad ogni iterazione è costoso. Gli autori propongono una strategia di decomposizione del gradiente:

• Poiché il token localizzato $LET_m$ è indipendente dalle attivazioni interne del modello durante il calcolo del gradiente della mappa di spiegabilità, il gradiente può essere scomposto.
• Invece di ricalcolare il gradiente completo ad ogni passo, si calcola una volta sola la derivata delle attivazioni rispetto al punteggio ($\frac{\partial \bar{z}}{\partial A}$) e si esegue un semplice prodotto scalare con il token in evoluzione.
• Questo riduce drasticamente il costo computazionale, specialmente quando si devono generare embedding per molte maschere diverse sulla stessa immagine.

3. Contributi Chiave

1. Embedding Localizzato Zero-Shot: Il metodo apprende un singolo token che cattura le caratteristiche di una regione specifica definita da una maschera, senza richiedere alcun fine-tuning del modello base. Questo token può essere usato come sostituto diretto ("drop-in replacement") del token [CLS] in qualsiasi applicazione basata su quel backbone.
2. Efficienza Computazionale: La strategia di decomposizione del gradiente rende il processo scalabile per immagini con molte maschere, eliminando la necessità di backpropagation completa ad ogni iterazione.
3. Versatilità: Il metodo è stato valutato su una vasta gamma di compiti downstream, dimostrando la sua efficacia in scenari diversi senza bisogno di addestramento specifico per ogni task.

4. Risultati Sperimentali

Il metodo è stato valutato su diversi dataset (PascalVOC, MSCOCO, RefCOCO, RefCOCO+, OpenImagesV7) e confrontato con lo stato dell'arte (SOTA).

• Recupero di Espressioni di Riferimento (Referring Expression Retrieval): Su RefCOCO, RefCOCO+ e PhraseCut, MaskInversion ha superato tutti i metodi basati su fine-tuning (come AlphaCLIP) e metodi senza addestramento (come ReCLIP, FGVP), raggiungendo il nuovo SOTA, specialmente con backbones più grandi (ViT-H/14).
• Recupero di Classi (Class Retrieval): Nel task di classificazione zero-shot di regioni specifiche, MaskInversion ha ottenuto prestazioni superiori rispetto a metodi come MaskCLIP, CLIPSurgery e SCLIP su tutti i dataset di segmentazione. Ha dimostrato di poter distinguere istanze multiple della stessa classe.
• Captioning Localizzato: Sostituendo il token [CLS] di CLIPCap con l'embedding di MaskInversion, il modello genera descrizioni testuali focalizzate esclusivamente sulla regione indicata dalla maschera, superando significativamente CLIP standard e AlphaCLIP.
• Generazione di Immagini (Localized Diffusion): Integrato con modelli di diffusione (es. $\lambda$-ECLIPSE), l'embedding localizzato permette di generare variazioni di specifiche parti dell'immagine, confermando che il token cattura correttamente le proprietà semantiche della regione.

Ablation Studies:

• Qualità della Maschera: Il metodo è robusto; maschere approssimative (es. bounding box + SAM) funzionano quasi quanto le maschere di ground-truth.
• Decomposizione del Gradiente: Riduce il tempo di esecuzione da 1.27s a 0.44s per 100 maschere, rendendo il metodo pratico per applicazioni reali.
• Iperparametro $\alpha$: Un valore ottimale di $\alpha \approx 5$ bilancia efficacemente contesto globale e dettaglio locale.

5. Significato e Impatto

MaskInversion rappresenta un passo avanti significativo nell'utilizzo dei modelli fondazionali per compiti di localizzazione fine. La sua principale innovazione risiede nel non modificare il modello pre-addestrato, evitando il costo e la complessità del fine-tuning su grandi dataset di maschere-testo, e nel non alterare l'immagine di input, preservando il contesto naturale.

Grazie alla sua natura "drop-in replacement", MaskInversion può essere integrato facilmente in pipeline esistenti per:

• Migliorare la comprensione delle immagini in sistemi di recupero informazioni.
• Abilitare la generazione di descrizioni o immagini focalizzate su oggetti specifici.
• Fornire rappresentazioni semantiche precise per robotica o analisi medica, dove la localizzazione è critica.

In sintesi, il paper dimostra che è possibile estrarre rappresentazioni locali ricche e allineate da modelli fondazionali globali semplicemente ottimizzando un token vettoriale guidato da mappe di spiegabilità, offrendo una soluzione efficiente, scalabile e versatile.