FA-Seg: A Fast and Accurate Diffusion-Based Method for Open-Vocabulary Segmentation

1. Enunciato del Problema

La Segmentazione Semantica a Vocabolario Aperto (OVSS) mira a segmentare gli oggetti nelle immagini basandosi su categorie testuali arbitrarie, senza richiedere annotazioni dense a livello di pixel per ogni possibile classe.

• Limiti dei Metodi Esistenti:
  • Modelli Vision-Language (es. CLIP): Sebbene efficaci per il riconoscimento zero-shot, si basano su embedding globali immagine-testo, risultando spesso in una scarsa precisione spaziale e confini sfocati a livello di pixel.
  • Modelli Basati su Diffusione: Sebbene codifichino naturalmente caratteristiche spaziali fini tramite meccanismi di attenzione, i metodi esistenti soffrono spesso di elevati costi computazionali. Molti richiedono molteplici passaggi di inferenza per classe, ottimizzazioni estese o generazione di dati sintetici, creando un compromesso tra qualità della segmentazione ed efficienza di inferenza.
• Il Divario: C'è un bisogno critico di un metodo OVSS senza addestramento che raggiunga un'alta accuratezza di segmentazione (precisione spaziale fine) mantenendo al contempo un'efficienza computazionale adatta al deployment nel mondo reale.

2. Metodologia: FA-Seg

FA-Seg è un framework senza addestramento che sfrutta modelli di diffusione pre-addestrati da testo a immagine (specificamente Stable Diffusion) per eseguire la segmentazione in un singolo passaggio di inferenza. La pipeline consta di tre fasi principali:

A. Costruzione del Dual-Prompt

Per abilitare sia una fedele ricostruzione dell'immagine che l'estrazione di attenzione specifica per la classe, FA-Seg costruisce due prompt complementari:

1. Prompt di Didascalia ($\mathcal{P}$): Generato da un modello di descrizione di immagini (es. BLIP) per guidare l'inversione DDIM nella ricostruzione dell'immagine di input.
2. Prompt di Classe ($\mathcal{P}'$): Una lista di categorie di oggetti candidate (generate tramite TagCLIP + BLIP) aggiunta alla didascalia.

• Meccanismo: Il modello di diffusione utilizza $\mathcal{P}$ per la ricostruzione ma sfrutta $\mathcal{P}'$ per generare mappe di attenzione incrociata specificamente per le categorie candidate. Ciò permette al modello di focalizzarsi su concetti semantici specifici senza ri-inferenze iterative.

B. Inversione Efficiente ed Estrazione dell'Attenzione

• Inversione Rapida: Invece della standard inversione DDIM multi-step, FA-Seg impiega 2-Rectified Flow (2-RF). Questo modello distillato permette una ricostruzione accurata dell'immagine in un processo (1+1)-step (un passo in avanti, uno indietro), riducendo drasticamente il tempo di inferenza.
• Estrazione dell'Attenzione: Durante il processo di denoising, il metodo estrae:
  • Mappe di Attenzione Incrociata: Che collegano i token testuali (dal prompt di classe) alle regioni dell'immagine.
  • Mappe di Auto-Attenzione: Che catturano la struttura spaziale interna e le similarità pairwise dell'immagine.
  • Queste vengono estratte attraverso molteplici risoluzioni ($r \in \{8, 16, 32, 64\}$).

C. Metodo di Affinamento Gerarchico dell'Attenzione (HARD)

Per produrre maschere di alta qualità, FA-Seg fonde le mappe di attenzione da diverse risoluzioni:

1. Fusione dell'Attenzione Incrociata: Le mappe di attenzione incrociata da diverse risoluzioni vengono upsampled e aggregate utilizzando una somma pesata. Le risoluzioni inferiori (es. 16) sono pesate più alto per il contesto semantico, mentre le risoluzioni superiori sono pesate meno per ridurre il rumore.
2. Affinamento dell'Auto-Attenzione: La mappa di attenzione incrociata fusa è ulteriormente raffinata utilizzando matrici di affinità di auto-attenzione. Le mappe di auto-attenzione forniscono dettagli strutturali (confini e forme) che l'attenzione incrociata non possiede. Il metodo trasforma i tensori di auto-attenzione per corrispondere alla risoluzione della mappa di attenzione incrociata e li applica come pesi moltiplicativi per affinare i confini degli oggetti.
3. Flipping al Tempo di Test (TTF): Per migliorare la coerenza spaziale, il processo viene ripetuto su una versione orizzontalmente ribaltata dell'immagine di input. Le mappe di attenzione risultanti vengono riallineate e mediate con le mappe originali per ridurre il rumore e migliorare la robustezza.

3. Contributi Chiave

1. Inferenza Unificata in Singolo Passaggio: A differenza dei metodi precedenti che richiedono passaggi di inferenza separati per ogni classe, FA-Seg genera maschere di segmentazione per tutte le classi candidate simultaneamente in un singolo inferenza (1+1)-step.
2. Meccanismo Dual-Prompt: Una strategia innovativa che separa il prompt di ricostruzione dal prompt consapevole della classe, permettendo al modello di diffusione di estrarre mappe di attenzione discriminative per categorie specifiche senza perdere la fedeltà dell'immagine.
3. Affinamento Gerarchico dell'Attenzione (HARD): Una strategia di fusione multi-scala che combina l'attenzione incrociata (per la localizzazione semantica) e l'auto-attenzione (per il raffinamento strutturale) attraverso le risoluzioni per ottenere confini di oggetti precisi.
4. Efficienza Senza Addestramento: Il metodo non richiede fine-tuning, non richiede generazione di dati sintetici e utilizza risorse computazionali minime, raggiungendo prestazioni all'avanguardia con meno di 1 secondo di tempo di inferenza.

4. Risultati Sperimentali

Gli autori hanno valutato FA-Seg su tre benchmark standard: PASCAL VOC, PASCAL Context e COCO Object.

• Prestazioni (mIoU):
  • FA-Seg ha raggiunto un mIoU medio del 43,8% su tutti e tre i dataset.
  • Ha stabilito un nuovo stato dell'arte (SOTA) per i metodi basati su diffusione senza addestramento, superando competitor come InvSeg (42,4%), DiffSegmentor (41,8%) e vari approcci basati su CLIP.
  • Punteggi specifici: 64,0% (PASCAL VOC), 31,3% (PASCAL Context), 36,2% (COCO).
• Efficienza:
  • Tempo di Inferenza: 0,36 secondi per immagine su una GPU RTX 4090. Questo è significativamente più veloce di InvSeg (7,9s su H100) e DiffSegmentor (il tempo scala con il numero di classi).
  • Memoria: Consuma solo 13,4 GB di VRAM, rispetto ai 32,4 GB di InvSeg.
• Studi di Ablazione:
  • La rimozione del raffinamento dell'auto-attenzione ha causato il calo di prestazioni più grande (-15,1%), dimostrando il suo ruolo critico nella precisione spaziale.
  • La fusione pesata delle mappe multi-risoluzione ha superato la semplice fusione media.
  • La combinazione TagCLIP+BLIP per la generazione delle classi candidate si è dimostrata la più robusta.

5. Significato e Impatto

• Colmare il Divario: FA-Seg colma con successo il divario tra qualità della segmentazione ed efficienza di inferenza, dimostrando che la segmentazione a vocabolario aperto ad alta precisione non richiede un pesante addestramento o ottimizzazioni multi-step.
• Deployment Pratico: Riducendo il tempo di inferenza a meno di un secondo e l'uso di memoria a ~13GB, FA-Seg rende la segmentazione a vocabolario aperto in tempo reale fattibile per applicazioni nel mondo reale (es. robotica, guida autonoma, strumenti interattivi).
• Fondamento per Lavori Futuri: Il paper stabilisce una solida baseline per la segmentazione tramite diffusione senza addestramento, suggerendo che i futuri miglioramenti potrebbero focalizzarsi sulla generazione adattiva di candidati e sulla segmentazione a livello di istanza.

In conclusione, FA-Seg rappresenta un avanzamento significativo nel campo della segmentazione a vocabolario aperto sfruttando le capacità spaziali intrinseche dei modelli di diffusione attraverso un framework altamente ottimizzato, senza addestramento e computazionalmente efficiente.