Diffusion Model as a Generalist Segmentation Learner

1. Enunciato del Problema

I sistemi di segmentazione attuali sono altamente frammentati. I modelli sono tipicamente adattati a compiti specifici (semantica, istanza, panottica), a vocabolari specifici (insieme chiuso vs. vocabolario aperto) o a domini specifici (naturale, medico, telerilevamento). Il passaggio tra queste impostazioni richiede spesso la riprogettazione di architetture o il riaddestramento da zero.

Sebbene la ricerca recente abbia tentato di riutilizzare modelli di diffusione testo-immagine preaddestrati per la segmentazione estraendo mappe di attenzione, questi approcci soffrono di limitazioni significative:

• Rumore e Incoerenza: Le mappe di attenzione grezze sono spesso a bassa risoluzione, rumorose e incoerenti tra i vari livelli.
• Dipendenza dal Post-processing: Richiedono frequentemente un pesante post-processing euristico per produrre maschere utilizzabili.
• Mancanza di Generalizzazione: I metodi esistenti si concentrano spesso su un singolo compito e non dimostrano una generalizzazione robusta attraverso domini diversi e query a vocabolario aperto senza personalizzazioni architetturali.

La domanda centrale affrontata da questo paper è: Un singolo modello unificato basato sulla diffusione può fungere da apprendista di segmentazione generalista capace di gestire compiti semantici, a vocabolario aperto e cross-dominio senza modifiche architetturali specifiche per il dominio?

2. Metodologia: Framework DiGSeg

Gli autori propongono DiGSeg, un framework che trasforma un modello di diffusione preaddestrato (nello specifico Stable Diffusion v2) da un prior generativo a un apprendista di segmentazione esplicito tramite un protocollo di fine-tuning unificato.

Architettura Core

Il framework è composto da tre componenti principali:

1. Percorso Latente Visivo:

   • Invece di trattare l'immagine come pixel grezzi, l'immagine RGB di input ($x$) e la maschera di segmentazione ground-truth ($y$) sono codificate in uno spazio latente compatto utilizzando il codificatore VAE congelato di Stable Diffusion.
   • Per mantenere la compatibilità con il VAE a 3 canali, la maschera a singolo canale viene replicata su tre canali.
   • Questo preserva la struttura spaziale consentendo al contempo un apprendimento efficiente di corrispondenze fini-grana nello spazio latente.

2. Condizionamento Testuale Allineato a CLIP:

   • Per abilitare capacità a vocabolario aperto, un codificatore di testo CLIP congelato genera embedding ($t_{clip}$) per nomi di classi o descrizioni in linguaggio naturale.
   • Questi embedding vengono iniettati nell'U-Net di diffusione tramite cross-attention a diverse scale di denoising. Ciò permette al modello di allineare le query testuali con le rappresentazioni visive in evoluzione, sia globalmente che localmente.

3. U-Net di Denoising Consistente con la Segmentazione:

   • Obiettivo di Addestramento: Il modello è addestrato a prevedere il rumore aggiunto al latente di segmentazione ($z_y$), condizionato dal latente dell'immagine ($z_x$) e dall'embedding testuale ($t_{clip}$).
   • Formulazione Generativa: A differenza di lavori precedenti che utilizzano la diffusione come estrattore di caratteristiche, DiGSeg tratta la segmentazione come un compito di generazione di denoising condizionata: $p_\theta(y|x)$. L'U-Net impara a denoisare da rumore casuale a una rappresentazione latente consistente con la segmentazione.
   • Strategia di Rumore: Viene impiegato un nuovo programma di rumore multi-scala annealed. Le perturbazioni ad alta frequenza dominano i primi passaggi per apprendere strutture fini, mentre i componenti a bassa frequenza diventano influenti in seguito per garantire la coerenza semantica.

Processo di Inferenza

1. L'immagine di input viene codificata in una rappresentazione latente.
2. Un latente di segmentazione viene inizializzato come rumore gaussiano.
3. Sotto condizionamento testuale, l'U-Net denoisizza iterativamente il latente (utilizzando un campionatore DDIM con una strategia di timestep "trailing" per l'efficienza).
4. Il latente pulito finale viene decodificato nello spazio dei pixel utilizzando il decodificatore VAE congelato per produrre la maschera di segmentazione.
5. Ensemble al Tempo di Test: Più passaggi di inferenza con diversi semi di rumore vengono mediati nello spazio latente per migliorare la coerenza spaziale e ridurre gli artefatti.

3. Contributi Chiave

• Framework Unificato: DiGSeg è il primo framework a riutilizzare con successo un backbone di diffusione preaddestrato in un apprendista di segmentazione generalista che gestisce compiti a vocabolario chiuso, vocabolario aperto e cross-dominio senza modifiche architetturali specifiche per il compito.
• Apprendimento Latente Diretto: Invece di affidarsi a mappe di attenzione rumorose, il metodo affina l'U-Net di denoising per generare direttamente latenti consistenti con la segmentazione, producendo maschere di alta qualità senza pesante post-processing.
• Condizionamento Testuale Multi-Scala: Introduce un meccanismo per iniettare caratteristiche linguistiche attraverso diverse scale di denoising, abilitando un allineamento robusto tra prompt testuali arbitrari e caratteristiche visive.
• Strategia di Rumore Innovativa: Propone un programma di rumore multi-scala annealed che bilancia l'apprendimento di confini fine-grana con la coerenza semantica globale.

4. Risultati Sperimentali

Gli autori hanno valutato DiGSeg su benchmark standard e domini specializzati, dimostrando prestazioni all'avanguardia (SOTA).

• Segmentazione a Vocabolario Aperto:
  • Ha superato 13 baseline (inclusi ODISE, OVSeg, CAT-Seg) su A-847, PC-459, A-150, PC-59 e Cityscapes.
  • Ha ottenuto guadagni significativi (ad esempio, +2.8 mIoU su PC-59, +2.2 su PC-459) rispetto ai metodi secondi classificati.
• Segmentazione Semantica Standard:
  • Ha raggiunto il SOTA su COCO (50.8 mIoU) e ADE20K (58.6 mIoU), superando modelli basati su transformer robusti come Mask2Former e OneFormer.
• Trasferimento Cross-Dominio:
  • Ha dimostrato un forte trasferimento zero-shot su dataset medici (REFUGE-2), agricoli (PhenoBench) e di telerilevamento (DeepGlobe) senza modifiche architetturali.
  • Nota: Le prestazioni nei domini medici sono state inferiori rispetto alle scene naturali, attribuite al divario di dominio nei dati di preaddestramento di CLIP, ma rimangono competitive.
• Efficienza e Utilizzo dei Dati:
  • Il modello è altamente efficiente nei dati, ottenendo prestazioni quasi identiche con solo il 50% dei dati di addestramento.
  • La velocità di inferenza è migliorata tramite DDIM-trailing (raggiungendo la saturazione in 1 passo), sebbene rimanga più lenta rispetto ai modelli feed-forward a causa della natura della diffusione.

5. Significato

Questo lavoro colma fondamentalmente il divario tra generazione visiva e comprensione visiva.

• Cambio di Paradigma: Mette in discussione la nozione secondo cui i modelli di diffusione siano esclusivamente per la sintesi, dimostrando che possono fungere da potenti apprendisti discriminativi generalisti se opportunamente condizionati.
• Generalizzazione: Dimostra che i ricchi prior spaziali e semantici codificati nei modelli generativi su larga scala sono trasferibili a compiti di predizione densa attraverso domini diversi (naturale, medico, agricolo).
• Semplicità: Evitando pipeline multi-stadio complesse o teste specializzate, DiGSeg offre un'interfaccia snella e unificata per la segmentazione, suggerendo un futuro in cui un singolo backbone generativo può gestire un'ampia gamma di compiti di visione artificiale.

In conclusione, DiGSeg stabilisce che i moderni backbone di diffusione, quando affinati con un protocollo condizionato al latente e allineamento testuale, possono superare i modelli di segmentazione specializzati offrendo al contempo una flessibilità superiore negli scenari a vocabolario aperto e cross-dominio.